[mlir] Add RewriterBase::replaceAllUsesWith for Blocks.

[ingomueller-net]

Summary

When changing IR in a RewriterPattern, all changes must go through the
rewriter. There are several convenience functions in RewriterBase that
help with high-level modifications, such as replaceAllUsesWith for
Values, but there is currently none to do the same task for Blocks.

[ingomueller-net]

Please review carefully; this is just copy'n'pasting from similar places and I am not sure whether this actually does the notification that is required.

Also, the template may be over-engineering. I have seen IRObjectWithUseList only being inherited from Block and Value, so I might just add this:

cpp
void replaceAllUsesWith(Block *from, Block *to) {
  for (BlockOperand &operand : llvm::make_early_inc_range(from->getUses())) {
    Operation *op = operand.getOwner();
    updateRootInPlace(op, [&]() { operand.set(to); });
  }
}
`

Finally, note that this code has no tests yet. I have not found similar code being tested other than indirectly through more high-level tests and I have not found any place that could use the new function now. Suggestions welcome.

[ingomueller-net]

Ping. (Also added more reviewers.)

[rriddle]

Can you replace the original Value one with the new method? Missing testing right now.

[ingomueller-net]

@rriddle: Thanks a lot for the suggestions!

The new revision implements replaceAllUsesWith(Value, Value) in terms of the new template. I think I can't just replace it with the templated function because Value doesn't itself inherit from IRObjectWithUseList (instead, it contains a ValueImpl object which does), so existing usages of the function would break.

The new revision also implements a test via -test-strict-pattern-driver. (Thanks @springerm for the brainstorming!) I guess this way of testing is somewhat contreived but I did not find examples of similar code being tested. If somebody has a better idea, please let me know.

[ingomueller-net]

Rebasing to latest main...

[ingomueller-net]

Ping...

[guraypp]

I added my comment about modification on replaceAllUsesWith.

I don't have much experience with replacing block.

mlir/include/mlir/IR/PatternMatch.h

509

Why did we move this function here?

511

The template is not needed unless I am missing something. It will increase compilation time here.
All from are Value anyway and they all implement getUses(), and operand.set takes a Value.

[ingomueller-net]

Thanks for the comment @guraypp! I hope my answers inline are helpful?

mlir/include/mlir/IR/PatternMatch.h

509

Both Block and ValueImpl inherit from IRObjectWithUseList, so if we implement the logic for the latter, we don't have to do it again for the two subclasses. For Value, we need an overload that extracts the ValueImpl and calls the function on that.

511

IRObjectWithUseList is also subclassed by Block, and adding that is the whole point of this patch.

operand is a IROperand, which is a template and the parameter to set is of template type IRValueT. IRObjectWithUseList needs to specify the operand type; Block this is BlockOperand, for ValueImpl it is OpOperand, which are two instances of IROperand. The set function thus accepts a Value if it is called for the ValueImpl instance of replaceAllUsesWith and a Block if it is called for the Block instance.

[guraypp]

Thanks for clarification. I am not familiar this part of MLIR, approving this change is not my call.

Just out of curiosity - when you replace two Blocks, does it change the tree and dominance information?

mlir/include/mlir/IR/PatternMatch.h

511

Ah I see, you are calling this function in your pass. There is still an extra function generation due to the use of template. I have slight preference for explicit implementation, but it's up to you.

[ingomueller-net]

In D142525#4122995, @guraypp wrote:

Just out of curiosity - when you replace two Blocks, does it change the tree and dominance information?

I have no idea, TBH! I haven't worked with tree/dominance information yet...

mlir/include/mlir/IR/PatternMatch.h

511

I have thought about that possibility and mentioned it in my initial comment. I don't mind either way but River suggested the template solution (possibly assuming that the Value overload would disappear) so I went with that...

[mehdi_amini]

LG, but I'm unsure I totally understand the test setup.

mlir/test/Transforms/test-strict-pattern-driver.mlir

84

Why is this op not on the worklist initially?

[ingomueller-net]

In D142525#4123430, @mehdi_amini wrote:

LG, but I'm unsure I totally understand the test setup.

Thanks! See my reply inline and let me know if you have suggestions/requests.

mlir/test/Transforms/test-strict-pattern-driver.mlir

84

That's an important point and the fact that it isn't clear indicates that improvements/changes to the test might be required.

That op isn't on the worklist initially, but it gets put into the list when change_block_op modifies its usage of bb1 through replaceAllUsesWith. That call notifies the rewriter about the change of bb1, which then puts all users of it into its worklist. By observing that implicit_change_op is rewritten in the final output, we hence see that replaceAllUsesWith has signaled the updated correctly.

The large comment block below attempts to explain this mechanism. Please let me know if you think it needs clarification.

Also, it is not too late to come up with a better/simpler/clearer way of testing this ;)

[mehdi_amini]

mlir/test/Transforms/test-strict-pattern-driver.mlir

84

I understood we wanted to see if it gets back on the worklist, I just didn't look carefully enough at how the initial worklist was built:

getOperation()->walk([&](Operation *op) {
  StringRef opName = op->getName().getStringRef();
  if (opName == "test.insert_same_op" || opName == "test.change_block_op" ||
      opName == "test.replace_with_new_op" || opName == "test.erase_op") {
    ops.push_back(op);
  }
});

It's all fine :)

mlir/include/mlir/IR/PatternMatch.h

void updateRootInPlace(Operation *root, CallableT &&callable) {
     startRootUpdate(root);
     callable();
     finalizeRootUpdate(root);
   }
 
   /// Find uses of `from` and replace them with `to`. It also marks every
   /// modified uses and notifies the rewriter that an in-place operation
   /// modification is about to happen.
-  void replaceAllUsesWith(Value from, Value to);
+  void replaceAllUsesWith(Value from, Value to) {
+    return replaceAllUsesWith(from.getImpl(), to);
+  }
+  template <typename OperandType, typename ValueT>
+  void replaceAllUsesWith(IRObjectWithUseList<OperandType> *from, ValueT &&to) {
+    for (OperandType &operand : llvm::make_early_inc_range(from->getUses())) {
+      Operation *op = operand.getOwner();
+      updateRootInPlace(op, [&]() { operand.set(to); });
+    }
+  }
 
   /// Find uses of `from` and replace them with `to` if the `functor` returns
   /// true. It also marks every modified uses and notifies the rewriter that an
   /// in-place operation modification is about to happen.
   void replaceUseIf(Value from, Value to,
                     llvm::unique_function<bool(OpOperand &) const> functor);
 
   /// Find uses of `from` and replace them with `to` except if the user is

mlir/lib/IR/PatternMatch.cpp

void RewriterBase::mergeBlocks(Block *source, Block *dest,
 
   // Splice the operations of the 'source' block into the 'dest' block and erase
   // it.
   dest->getOperations().splice(dest->end(), source->getOperations());
   source->dropAllUses();
   source->erase();
 }
 
-/// Find uses of `from` and replace it with `to`
-void RewriterBase::replaceAllUsesWith(Value from, Value to) {
-  for (OpOperand &operand : llvm::make_early_inc_range(from.getUses())) {
-    Operation *op = operand.getOwner();
-    updateRootInPlace(op, [&]() { operand.set(to); });
-  }
-}
-
 /// Find uses of `from` and replace them with `to` if the `functor` returns
 /// true. It also marks every modified uses and notifies the rewriter that an
 /// in-place operation modification is about to happen.
 void RewriterBase::replaceUseIf(
     Value from, Value to,
     llvm::unique_function<bool(OpOperand &) const> functor) {
   for (OpOperand &operand : llvm::make_early_inc_range(from.getUses())) {
     if (functor(operand))

mlir/test/Transforms/test-strict-pattern-driver.mlir

 // RUN: mlir-opt \
+// RUN:     -test-strict-pattern-driver="strictness=AnyOp" \
+// RUN:     --split-input-file %s | FileCheck %s --check-prefix=CHECK-AN
+
+// RUN: mlir-opt \
 // RUN:     -test-strict-pattern-driver="strictness=ExistingAndNewOps" \
 // RUN:     --split-input-file %s | FileCheck %s --check-prefix=CHECK-EN
 
 // RUN: mlir-opt \
 // RUN:     -test-strict-pattern-driver="strictness=ExistingOps" \
 // RUN:     --split-input-file %s | FileCheck %s --check-prefix=CHECK-EX
 
 // CHECK-EN-LABEL: func @test_erase

 // CHECK-EX-LABEL: func @test_replace_with_erase_op
 //  CHECK-EX-SAME:     {pattern_driver_all_erased = true, pattern_driver_changed = true}
 //   CHECK-EX-NOT:   test.replace_with_new_op
 //       CHECK-EX:   test.erase_op
 func.func @test_replace_with_erase_op() {
   "test.replace_with_new_op"() {create_erase_op} : () -> ()
   return
 }
+
+// -----
+
+// CHECK-AN-LABEL: func @test_trigger_rewrite_through_block
+//       CHECK-AN: "test.change_block_op"()[^[[BB0:.*]], ^[[BB0]]]
+//       CHECK-AN: return
+//       CHECK-AN: ^[[BB1:[^:]*]]:
+//       CHECK-AN: "test.implicit_change_op"()[^[[BB1]]]
+func.func @test_trigger_rewrite_through_block() {
+  return
+^bb1:
+  // Uses bb1. ChangeBlockOp replaces that and all other usages of bb1 with bb2.
+  "test.change_block_op"() [^bb1, ^bb2] : () -> ()
+^bb2:
+  return
+^bb3:
+  // Also uses bb1. ChangeBlockOp replaces that usage with bb2. This triggers
+  // this op being put on the worklist, which triggers ImplicitChangeOp, which,
+  // in turn, replaces the successor with bb3.
+  "test.implicit_change_op"() [^bb1] : () -> ()
+}

mlir/test/lib/Dialect/Test/TestPatterns.cpp

public:
   StringRef getArgument() const final { return "test-strict-pattern-driver"; }
   StringRef getDescription() const final {
     return "Test strict mode of pattern driver";
   }
 
   void runOnOperation() override {
     MLIRContext *ctx = &getContext();
     mlir::RewritePatternSet patterns(ctx);
-    patterns.add<InsertSameOp, ReplaceWithNewOp, EraseOp>(ctx);
+    patterns.add<
+        // clang-format off
+        InsertSameOp,
+        ReplaceWithNewOp,
+        EraseOp,
+        ChangeBlockOp,
+        ImplicitChangeOp
+        // clang-format on
+        >(ctx);
     SmallVector<Operation *> ops;
     getOperation()->walk([&](Operation *op) {
       StringRef opName = op->getName().getStringRef();
-      if (opName == "test.insert_same_op" ||
+      if (opName == "test.insert_same_op" || opName == "test.change_block_op" ||
           opName == "test.replace_with_new_op" || opName == "test.erase_op") {
         ops.push_back(op);
       }
     });
 
     GreedyRewriteConfig config;
     if (strictMode == "AnyOp") {
       config.strictMode = GreedyRewriteStrictness::AnyOp;
LogicalResult matchAndRewrite(Operation *op,
             OperationName("test.new_op", op->getContext()).getIdentifier(),
             op->getOperands(), op->getResultTypes());
       }
       rewriter.replaceOp(op, newOp->getResults());
       return success();
     }
   };
 
-  // Remove an operation may introduce the re-visiting of its opreands.
+  // Remove an operation may introduce the re-visiting of its operands.
   class EraseOp : public RewritePattern {
   public:
     EraseOp(MLIRContext *context)
         : RewritePattern("test.erase_op", /*benefit=*/1, context) {}
     LogicalResult matchAndRewrite(Operation *op,
                                   PatternRewriter &rewriter) const override {
       rewriter.eraseOp(op);
       return success();
     }
   };
+
+  // The following two patterns test RewriterBase::replaceAllUsesWith.
+  //
+  // That function replaces all usages of a Block (or a Value) with another one
+  // *and tracks these changes in the rewriter.* The GreedyPatternRewriteDriver
+  // with GreedyRewriteStrictness::AnyOp uses that tracking to construct its
+  // worklist: when an op is modified, it is added to the worklist. The two
+  // patterns below make the tracking observable: ChangeBlockOp replaces all
+  // usages of a block and that pattern is applied because the corresponding ops
+  // are put on the initial worklist (see above). ImplicitChangeOp does an
+  // unrelated change but ops of the corresponding type are *not* on the initial
+  // worklist, so the effect of the second pattern is only visible if the
+  // tracking and subsequent adding to the worklist actually works.
+
+  // Replace all usages of the first successor with the second successor.
+  class ChangeBlockOp : public RewritePattern {
+  public:
+    ChangeBlockOp(MLIRContext *context)
+        : RewritePattern("test.change_block_op", /*benefit=*/1, context) {}
+    LogicalResult matchAndRewrite(Operation *op,
+                                  PatternRewriter &rewriter) const override {
+      if (op->getNumSuccessors() < 2)
+        return failure();
+      Block *firstSuccessor = op->getSuccessor(0);
+      Block *secondSuccessor = op->getSuccessor(1);
+      if (firstSuccessor == secondSuccessor)
+        return failure();
+      // This is the function being tested:
+      rewriter.replaceAllUsesWith(firstSuccessor, secondSuccessor);
+      // Using the following line instead would make the test fail:
+      // firstSuccessor->replaceAllUsesWith(secondSuccessor);
+      return success();
+    }
+  };
+
+  // Changes the successor to the parent block.
+  class ImplicitChangeOp : public RewritePattern {
+  public:
+    ImplicitChangeOp(MLIRContext *context)
+        : RewritePattern("test.implicit_change_op", /*benefit=*/1, context) {}
+    LogicalResult matchAndRewrite(Operation *op,
+                                  PatternRewriter &rewriter) const override {
+      if (op->getNumSuccessors() < 1 || op->getSuccessor(0) == op->getBlock())
+        return failure();
+      rewriter.updateRootInPlace(
+          op, [&]() { op->setSuccessor(op->getBlock(), 0); });
+      return success();
+    }
+  };
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // ReturnType Driver.
 //===----------------------------------------------------------------------===//