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

[llvmlistbot at llvm.org]

Author: Morten Borup Petersen
Date: 2023-10-25T09:38:02+02:00
New Revision: bf3a981a7f5ee9074d974fa518461bdb2869710f

URL: https://github.com/llvm/llvm-project/commit/bf3a981a7f5ee9074d974fa518461bdb2869710f
DIFF: https://github.com/llvm/llvm-project/commit/bf3a981a7f5ee9074d974fa518461bdb2869710f.diff

LOG: [MLIR] Properly add operations to blocks during `createOrFold` (#70010)

Fixes #68884.

Added: 
    

Modified: 
    mlir/include/mlir/IR/Builders.h
    mlir/test/lib/Dialect/Test/TestDialect.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/Builders.h b/mlir/include/mlir/IR/Builders.h
index 5e54d4ea49e8251..13fbc3fb928c399 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 erase 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..deef4a9683e7421 100644
--- a/mlir/test/lib/Dialect/Test/TestDialect.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp
@@ -529,6 +529,11 @@ LogicalResult TestOpWithVariadicResultsAndFolder::fold(
 }
 
 OpFoldResult TestOpInPlaceFold::fold(FoldAdaptor adaptor) {
+  // 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.
     getProperties().attr = dyn_cast_or_null<IntegerAttr>(adaptor.getOp());