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

[llvmlistbot at llvm.org]

https://github.com/jeanPerier created https://github.com/llvm/llvm-project/pull/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).

>From 64dfc9bb7165325465b8e8dbc290bf413a7ca6ab Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier at nvidia.com>
Date: Wed, 29 Nov 2023 07:18:06 -0800
Subject: [PATCH] [mlir] notify insertion of parent op first when cloning

When cloning and operation with a region, the builder was currently
notifying about the insertion of the cloned operations inside the
region before the 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 because the patterns are attempted in order of
notifications for the cloned operations.
---
 mlir/lib/IR/Builders.cpp       | 3 ++-
 mlir/test/IR/test-clone.mlir   | 6 +++++-
 mlir/test/lib/IR/TestClone.cpp | 8 ++++++++
 3 files changed, 15 insertions(+), 2 deletions(-)

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)