[Mlir-commits] [mlir] [MLIR] Properly add operations to blocks during `createOrFold` (PR #70010)

[Morten Borup Petersen]

https://github.com/mortbopet updated https://github.com/llvm/llvm-project/pull/70010

>From 371cb826857a5d2f094e357fba7fda17289d49e7 Mon Sep 17 00:00:00 2001
From: Morten Borup Petersen <morten_bp at live.dk>
Date: Tue, 24 Oct 2023 07:38:51 +0000
Subject: [PATCH 1/3] [MLIR] Properly add operations to blocks during
 `createOrFold`

Fixes #68884.
---
 mlir/include/mlir/IR/Builders.h            | 14 ++++++++------
 mlir/test/lib/Dialect/Test/TestDialect.cpp |  4 ++++
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/mlir/include/mlir/IR/Builders.h b/mlir/include/mlir/IR/Builders.h
index 5e54d4ea49e8251..11c89273e70ea31 100644
--- a/mlir/include/mlir/IR/Builders.h
+++ b/mlir/include/mlir/IR/Builders.h
@@ -504,18 +504,20 @@ class OpBuilder : public Builder {
   template <typename OpTy, typename... Args>
   void createOrFold(SmallVectorImpl<Value> &results, Location location,
                     Args &&...args) {
-    // Create the operation without using 'create' as we don't want to
-    // insert it yet.
+    // Create the operation without using 'create' as we want to control when
+    // the listener is notified.
     OperationState state(location,
                          getCheckRegisteredInfo<OpTy>(location.getContext()));
     OpTy::build(*this, state, std::forward<Args>(args)...);
     Operation *op = Operation::create(state);
+    if (block)
+      block->getOperations().insert(insertPoint, op);
 
-    // Fold the operation. If successful destroy it, otherwise insert it.
+    // Fold the operation. If successful destroy it, otherwise notify.
     if (succeeded(tryFold(op, results)))
-      op->destroy();
-    else
-      insert(op);
+      op->erase();
+    else if (listener)
+      listener->notifyOperationInserted(op);
   }
 
   /// Overload to create or fold a single result operation.
diff --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp
index 8aa7774d2bf008f..fff2b15ca641327 100644
--- a/mlir/test/lib/Dialect/Test/TestDialect.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp
@@ -529,6 +529,10 @@ LogicalResult TestOpWithVariadicResultsAndFolder::fold(
 }
 
 OpFoldResult TestOpInPlaceFold::fold(FoldAdaptor adaptor) {
+  // Excercise the fact that an operation created during createOrFold should be
+  // allowed to enable its block.
+  assert(getBlock() && "expected block to be assigned");
+
   if (adaptor.getOp() && !getProperties().attr) {
     // The folder adds "attr" if not present.
     getProperties().attr = dyn_cast_or_null<IntegerAttr>(adaptor.getOp());

>From 7ad7dcf31bc7a661eb72284875816aa12bcba592 Mon Sep 17 00:00:00 2001
From: Morten Borup Petersen <morten_bp at live.dk>
Date: Tue, 24 Oct 2023 07:42:02 +0000
Subject: [PATCH 2/3] typo

---
 mlir/test/lib/Dialect/Test/TestDialect.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp
index fff2b15ca641327..86b5dc0f4b07049 100644
--- a/mlir/test/lib/Dialect/Test/TestDialect.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp
@@ -529,7 +529,7 @@ LogicalResult TestOpWithVariadicResultsAndFolder::fold(
 }
 
 OpFoldResult TestOpInPlaceFold::fold(FoldAdaptor adaptor) {
-  // Excercise the fact that an operation created during createOrFold should be
+  // Excercise the fact that an operation created with createOrFold should be
   // allowed to enable its block.
   assert(getBlock() && "expected block to be assigned");
 

>From 29277c2a491fd992d32d69e5450016509def9572 Mon Sep 17 00:00:00 2001
From: Morten Borup Petersen <morten_bp at live.dk>
Date: Tue, 24 Oct 2023 08:06:19 +0000
Subject: [PATCH 3/3] review

---
 mlir/include/mlir/IR/Builders.h            | 2 +-
 mlir/test/lib/Dialect/Test/TestDialect.cpp | 7 ++++---
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/IR/Builders.h b/mlir/include/mlir/IR/Builders.h
index 11c89273e70ea31..13fbc3fb928c399 100644
--- a/mlir/include/mlir/IR/Builders.h
+++ b/mlir/include/mlir/IR/Builders.h
@@ -513,7 +513,7 @@ class OpBuilder : public Builder {
     if (block)
       block->getOperations().insert(insertPoint, op);
 
-    // Fold the operation. If successful destroy it, otherwise notify.
+    // Fold the operation. If successful erase it, otherwise notify.
     if (succeeded(tryFold(op, results)))
       op->erase();
     else if (listener)
diff --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp
index 86b5dc0f4b07049..deef4a9683e7421 100644
--- a/mlir/test/lib/Dialect/Test/TestDialect.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp
@@ -529,9 +529,10 @@ LogicalResult TestOpWithVariadicResultsAndFolder::fold(
 }
 
 OpFoldResult TestOpInPlaceFold::fold(FoldAdaptor adaptor) {
-  // Excercise the fact that an operation created with createOrFold should be
-  // allowed to enable its block.
-  assert(getBlock() && "expected block to be assigned");
+  // Exercise the fact that an operation created with createOrFold should be
+  // allowed to access its parent block.
+  assert(getOperation()->getBlock() &&
+         "expected that operation is not unlinked");
 
   if (adaptor.getOp() && !getProperties().attr) {
     // The folder adds "attr" if not present.