[mlir] notify insertion of parent op first when cloning

mlir/lib/IR/Builders.cpp

@@ -527,17 +527,18 @@ 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) {
     auto walkFn = [&](Operation *walkedOp) {
       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) {

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:   }

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)