[Mlir-commits] [mlir] [mlir] notify insertion of parent op first when cloning (PR #73806)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: None (jeanPerier)

<details>
<summary>Changes</summary>

When cloning an operation with a region, the builder was currently notifying about the insertion of the cloned operations inside the region before the cloned operation itself.

When using cloning inside rewrite pass, this could cause issues if a pattern is expected to be applied on a cloned parent operation before trying to apply patterns on the cloned operations it contains (the patterns are attempted in order of notifications for the cloned operations).

---
Full diff: https://github.com/llvm/llvm-project/pull/73806.diff


3 Files Affected:

- (modified) mlir/lib/IR/Builders.cpp (+2-1) 
- (modified) mlir/test/IR/test-clone.mlir (+5-1) 
- (modified) mlir/test/lib/IR/TestClone.cpp (+8) 


``````````diff
diff --git a/mlir/lib/IR/Builders.cpp b/mlir/lib/IR/Builders.cpp
index ab20f4863e11c23..7bb7358426af5a4 100644
--- a/mlir/lib/IR/Builders.cpp
+++ b/mlir/lib/IR/Builders.cpp
@@ -527,6 +527,7 @@ LogicalResult OpBuilder::tryFold(Operation *op,
 
 Operation *OpBuilder::clone(Operation &op, IRMapping &mapper) {
   Operation *newOp = op.clone(mapper);
+  newOp = insert(newOp);
   // The `insert` call below handles the notification for inserting `newOp`
   // itself. But if `newOp` has any regions, we need to notify the listener
   // about any ops that got inserted inside those regions as part of cloning.
@@ -537,7 +538,7 @@ Operation *OpBuilder::clone(Operation &op, IRMapping &mapper) {
     for (Region &region : newOp->getRegions())
       region.walk(walkFn);
   }
-  return insert(newOp);
+  return newOp;
 }
 
 Operation *OpBuilder::clone(Operation &op) {
diff --git a/mlir/test/IR/test-clone.mlir b/mlir/test/IR/test-clone.mlir
index 575098b642e8ea2..f7a29569d67d625 100644
--- a/mlir/test/IR/test-clone.mlir
+++ b/mlir/test/IR/test-clone.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(test-clone))" -split-input-file
+// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(test-clone))" | FileCheck %s
 
 module {
   func.func @fixpoint(%arg1 : i32) -> i32 {
@@ -9,6 +9,10 @@ module {
   }
 }
 
+// CHECK: notifyOperationInserted: test.use
+// CHECK-NEXT: notifyOperationInserted: test.yield
+// CHECK-NEXT: notifyOperationInserted: func.return
+
 // CHECK:   func @fixpoint(%[[arg0:.+]]: i32) -> i32 {
 // CHECK-NEXT:     %[[i0:.+]] = "test.use"(%[[arg0]]) ({
 // CHECK-NEXT:       "test.yield"(%arg0) : (i32) -> ()
diff --git a/mlir/test/lib/IR/TestClone.cpp b/mlir/test/lib/IR/TestClone.cpp
index 70238608a67c2b5..13a0cfeb402a9cd 100644
--- a/mlir/test/lib/IR/TestClone.cpp
+++ b/mlir/test/lib/IR/TestClone.cpp
@@ -14,6 +14,12 @@ using namespace mlir;
 
 namespace {
 
+struct DumpNotifications : public OpBuilder::Listener {
+  void notifyOperationInserted(Operation *op) override {
+    llvm::outs() << "notifyOperationInserted: " << op->getName() << "\n";
+  }
+};
+
 /// This is a test pass which clones the body of a function. Specifically
 /// this pass replaces f(x) to instead return f(f(x)) in which the cloned body
 /// takes the result of the first operation return as an input.
@@ -50,6 +56,8 @@ struct ClonePass
     }
 
     OpBuilder builder(op->getContext());
+    DumpNotifications dumpNotifications;
+    builder.setListener(&dumpNotifications);
     builder.setInsertionPointToEnd(&regionEntry);
     SmallVector<Operation *> toClone;
     for (Operation &inst : regionEntry)

``````````

</details>


https://github.com/llvm/llvm-project/pull/73806