[mlir] check whether region and block visitors are interrupted

[ashay-github]

Summary

The visitor functions for Region and Block types did not always
check the value returned by recursive calls. This caused the top-level
visitor invocation to return WalkResult::advance() even if one or more
recursive invocations returned WalkResult::interrupt(). This patch
fixes the problem by check if any recursive call is interrupted, and if
so, return WalkResult::interrupt().

[ashay-github]

It seems Phabricator didn't send out an email when I created the patch. Pinging folks to see if it helps.

@jurahul @rriddle

[dcaballe]

Thanks for the contribution! It looks good to me.

This patch doesn't add unit tests since our test framework
(TestGenericIRVisitorInterruptPass) invokes visitor functions for
operations but not for regions and blocks.

I think it's important that we test this. Would it possible to add a new visitor for regions and blocks?

[ashay-github]

Added tests for region and block visitors.

[ashay-github]

I think it's important that we test this. Would it possible to add a new visitor for regions and blocks?

Thanks, that makes sense. I've added visitor functions for regions and blocks. However, BlockArguments cannot have attributes (which we use to decide whether to interrupt the walk), so the visitor function for blocks skips iterating over its argument.

[dcaballe]

LGTM, thanks!

[rriddle]

mlir/test/lib/IR/TestVisitorsGeneric.cpp

134–135

Same comment here as below.

166–167

168–169

This part doesn't really make sense, result.getDefiningOp() is just going to give you op again.

[ashay-github]

@rriddle My apologies! The fix is in this patch: https://reviews.llvm.org/D129913.

mlir/lib/IR/Visitors.cpp

if (order == WalkOrder::PreOrder) {
       WalkResult result = callback(&region);
       if (result.wasSkipped())
         continue;
       if (result.wasInterrupted())
         return WalkResult::interrupt();
     }
     for (auto &block : region) {
       for (auto &nestedOp : block)
-        walk(&nestedOp, callback, order);
+        if (walk(&nestedOp, callback, order).wasInterrupted())
+          return WalkResult::interrupt();
     }
     if (order == WalkOrder::PostOrder) {
       if (callback(&region).wasInterrupted())
         return WalkResult::interrupt();
       // We don't check if this region was skipped because its walk already
       // finished and the walk will continue with the next region.
     }
   }
for (auto &block : llvm::make_early_inc_range(region)) {
       if (order == WalkOrder::PreOrder) {
         WalkResult result = callback(&block);
         if (result.wasSkipped())
           continue;
         if (result.wasInterrupted())
           return WalkResult::interrupt();
       }
       for (auto &nestedOp : block)
-        walk(&nestedOp, callback, order);
+        if (walk(&nestedOp, callback, order).wasInterrupted())
+          return WalkResult::interrupt();
       if (order == WalkOrder::PostOrder) {
         if (callback(&block).wasInterrupted())
           return WalkResult::interrupt();
         // We don't check if this block was skipped because its walk already
         // finished and the walk will continue with the next block.
       }
     }
   }

mlir/test/IR/generic-block-visitors-interrupt.mlir

+// RUN: mlir-opt -test-generic-ir-block-visitors-interrupt -allow-unregistered-dialect -split-input-file %s | FileCheck %s
+
+func.func @main(%arg0: f32) -> f32 {
+  %v1 = "foo"() {interrupt = true} : () -> f32
+  %v2 = arith.addf %v1, %arg0 : f32
+  return %v2 : f32
+}
+
+// CHECK: step 0 walk was interrupted

mlir/test/IR/generic-region-visitors-interrupt.mlir

+// RUN: mlir-opt -test-generic-ir-region-visitors-interrupt -allow-unregistered-dialect -split-input-file %s | FileCheck %s
+
+func.func @main(%arg0: f32) -> f32 {
+  %v1 = "foo"() {interrupt = true} : () -> f32
+  %v2 = arith.addf %v1, %arg0 : f32
+  return %v2 : f32
+}
+
+// CHECK: step 0 walk was interrupted

mlir/test/lib/IR/TestVisitorsGeneric.cpp

result = outerOp->walk([&](test::TwoRegionOp op, const WalkStage &stage) {
       return walker(op, stage);
     });
 
     if (result.wasInterrupted())
       llvm::outs() << "step " << stepNo++ << " walk was interrupted\n";
   }
 };
 
+struct TestGenericIRBlockVisitorInterruptPass
+    : public PassWrapper<TestGenericIRBlockVisitorInterruptPass,
+                         OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+      TestGenericIRBlockVisitorInterruptPass)
+
+  StringRef getArgument() const final {
+    return "test-generic-ir-block-visitors-interrupt";
+  }
+  StringRef getDescription() const final {
+    return "Test generic IR visitors with interrupts, starting with Blocks.";
+  }
+
+  void runOnOperation() override {
+    int stepNo = 0;
+
+    auto walker = [&](Block *block) {
+      for (Operation &op : *block)
+        for (OpResult result : op.getResults())
+          if (Operation *definingOp = result.getDefiningOp())
+            if (definingOp->getAttrOfType<BoolAttr>("interrupt"))
+              return WalkResult::interrupt();
+
+      llvm::outs() << "step " << stepNo++ << "\n";
+      return WalkResult::advance();
+    };
+
+    auto result = getOperation()->walk(walker);
+    if (result.wasInterrupted())
+      llvm::outs() << "step " << stepNo++ << " walk was interrupted\n";
+  }
+};
+
+struct TestGenericIRRegionVisitorInterruptPass
+    : public PassWrapper<TestGenericIRRegionVisitorInterruptPass,
+                         OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+      TestGenericIRRegionVisitorInterruptPass)
+
+  StringRef getArgument() const final {
+    return "test-generic-ir-region-visitors-interrupt";
+  }
+  StringRef getDescription() const final {
+    return "Test generic IR visitors with interrupts, starting with Regions.";
+  }
+
+  void runOnOperation() override {
+    int stepNo = 0;
+
+    auto walker = [&](Region *region) {
+      for (Block &block : *region)
+        for (Operation &op : block)
+          for (OpResult result : op.getResults())
+            if (Operation *definingOp = result.getDefiningOp())
+              if (definingOp->getAttrOfType<BoolAttr>("interrupt"))
+                return WalkResult::interrupt();
+
+      llvm::outs() << "step " << stepNo++ << "\n";
+      return WalkResult::advance();
+    };
+
+    auto result = getOperation()->walk(walker);
+    if (result.wasInterrupted())
+      llvm::outs() << "step " << stepNo++ << " walk was interrupted\n";
+  }
+};
+
 } // namespace
 
 namespace mlir {
 namespace test {
 void registerTestGenericIRVisitorsPass() {
   PassRegistration<TestGenericIRVisitorPass>();
   PassRegistration<TestGenericIRVisitorInterruptPass>();
+  PassRegistration<TestGenericIRBlockVisitorInterruptPass>();
+  PassRegistration<TestGenericIRRegionVisitorInterruptPass>();
 }
 
 } // namespace test
 } // namespace mlir