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

[Morten Borup Petersen]

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

Fixes #68884.

>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] [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());