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

[llvmlistbot at llvm.org]

Author: jeanPerier
Date: 2023-12-01T10:03:02+01:00
New Revision: 5a4ca51a91ff28b1d6bdde5403144c29b86e4b54

URL: https://github.com/llvm/llvm-project/commit/5a4ca51a91ff28b1d6bdde5403144c29b86e4b54
DIFF: https://github.com/llvm/llvm-project/commit/5a4ca51a91ff28b1d6bdde5403144c29b86e4b54.diff

LOG: [mlir] notify insertion of parent op first when cloning (#73806)

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).

Added: 
    

Modified: 
    mlir/lib/IR/Builders.cpp
    mlir/test/IR/test-clone.mlir
    mlir/test/lib/IR/TestClone.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/Builders.cpp b/mlir/lib/IR/Builders.cpp
index ab20f4863e11c23..2cabfcd24d35598 100644
--- a/mlir/lib/IR/Builders.cpp
+++ b/mlir/lib/IR/Builders.cpp
@@ -527,7 +527,8 @@ LogicalResult OpBuilder::tryFold(Operation *op,
 
 Operation *OpBuilder::clone(Operation &op, IRMapping &mapper) {
   Operation *newOp = op.clone(mapper);
-  // The `insert` call below handles the notification for inserting `newOp`
+  newOp = insert(newOp);
+  // The `insert` call above 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.
   if (listener) {
@@ -535,9 +536,9 @@ Operation *OpBuilder::clone(Operation &op, IRMapping &mapper) {
       listener->notifyOperationInserted(walkedOp);
     };
     for (Region &region : newOp->getRegions())
-      region.walk(walkFn);
+      region.walk<WalkOrder::PreOrder>(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..0c07593aef32d9d 100644
--- a/mlir/test/IR/test-clone.mlir
+++ b/mlir/test/IR/test-clone.mlir
@@ -1,20 +1,35 @@
-// 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 {
     %r = "test.use"(%arg1) ({
-       "test.yield"(%arg1) : (i32) -> ()
+      %r2 = "test.use2"(%arg1) ({
+         "test.yield2"(%arg1) : (i32) -> ()
+      }) : (i32) -> i32
+      "test.yield"(%r2) : (i32) -> ()
     }) : (i32) -> i32
     return %r : i32
   }
 }
 
+// CHECK: notifyOperationInserted: test.use
+// CHECK-NEXT: notifyOperationInserted: test.use2
+// CHECK-NEXT: notifyOperationInserted: test.yield2
+// 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) -> ()
+// CHECK-NEXT:       %[[r2:.+]] = "test.use2"(%[[arg0]]) ({
+// CHECK-NEXT:         "test.yield2"(%[[arg0]]) : (i32) -> ()
+// CHECK-NEXT:       }) : (i32) -> i32
+// CHECK-NEXT:       "test.yield"(%[[r2]]) : (i32) -> ()
 // CHECK-NEXT:     }) : (i32) -> i32
 // CHECK-NEXT:     %[[i1:.+]] = "test.use"(%[[i0]]) ({
-// CHECK-NEXT:       "test.yield"(%[[i0]]) : (i32) -> ()
+// CHECK-NEXT:       %[[r2:.+]] = "test.use2"(%[[i0]]) ({
+// CHECK-NEXT:         "test.yield2"(%[[i0]]) : (i32) -> ()
+// CHECK-NEXT:       }) : (i32) -> i32
+// CHECK-NEXT:       "test.yield"(%[[r2]]) : (i32) -> ()
 // CHECK-NEXT:     }) : (i32) -> i32
 // CHECK-NEXT:     return %[[i1]] : i32
 // CHECK-NEXT:   }

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)