[MLIR] Properly add operations to blocks during createOrFold

[mortbopet]

Fixes #68884.

[llvmbot]

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: Morten Borup Petersen (mortbopet)

Changes

Fixes #68884.

---

Full diff: https://github.com/llvm/llvm-project/pull/70010.diff

2 Files Affected:

• (modified) mlir/include/mlir/IR/Builders.h (+8-6)
• (modified) mlir/test/lib/Dialect/Test/TestDialect.cpp (+4)

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

[matthias-springer]

What does it mean to "enable a block"?

[mortbopet]

Typo, meant 'access' its block.