[Mlir-commits] [mlir] [MLIR][Bufferization] Bail on automatic deallocation to enable reentrant behaviour (PR #72289)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-bufferization

@llvm/pr-subscribers-mlir

Author: Renato Golin (rengolin)

<details>
<summary>Changes</summary>

Currently, the ownership based dealloc pass fails hard on existing deallocations, but it introduces its own deallocs, and if we run the pass twice (ex. when piping the IR through different tools with their own pipelines), we get a hard crash.

However, having the deallocs is not an IR error or verification failure, it's just a pattern that the pass itself cannot handle, and therefore the correct behaviour is to bail, not crash.

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


3 Files Affected:

- (modified) mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp (+9-3) 
- (added) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/invalid-buffer-deallocation-reentrant.mlir (+20) 
- (modified) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/invalid-buffer-deallocation.mlir (-18) 


``````````diff
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
index 9459cc43547fafc..73a88362756bda7 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
@@ -861,9 +861,15 @@ BufferDeallocation::handleInterface(MemoryEffectOpInterface op) {
 
   for (auto operand : llvm::make_filter_range(op->getOperands(), isMemref)) {
     if (op.getEffectOnValue<MemoryEffects::Free>(operand).has_value()) {
-      if (!op->hasAttr(BufferizationDialect::kManualDeallocation))
-        return op->emitError(
-            "memory free side-effect on MemRef value not supported!");
+      // This should not be an error because the ownership based buffer
+      // deallocation introduces deallocs itself, so running it twice over (say
+      // when piping IR over different tools with their own pipelines) crashes
+      // the compiler on perfectly valid code.
+      if (!op->hasAttr(BufferizationDialect::kManualDeallocation)) {
+        state.resetOwnerships(operand, block);
+        state.updateOwnership(operand, buildBoolValue(builder, op.getLoc(), false));
+        continue;
+      }
 
       // Buffers that were allocated under "manual deallocation" may be
       // manually deallocated. We insert a runtime assertion to cover certain
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/invalid-buffer-deallocation-reentrant.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/invalid-buffer-deallocation-reentrant.mlir
new file mode 100644
index 000000000000000..89271d3422bb028
--- /dev/null
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/invalid-buffer-deallocation-reentrant.mlir
@@ -0,0 +1,20 @@
+// RUN: mlir-opt -ownership-based-buffer-deallocation -split-input-file %s | \
+// RUN: mlir-opt -ownership-based-buffer-deallocation -split-input-file %s
+
+// This should not be an error because the ownership based buffer deallocation introduces
+// deallocs itself, so running it twice over (say when piping IR over different tools with
+// their own pipelines) crashes the compiler on perfectly valid code.
+
+func.func @free_effect() {
+  %alloc = memref.alloc() : memref<2xi32>
+  %new_alloc = memref.realloc %alloc : memref<2xi32> to memref<4xi32>
+  return
+}
+
+// -----
+
+func.func @free_effect() {
+  %alloc = memref.alloc() : memref<2xi32>
+  memref.dealloc %alloc : memref<2xi32>
+  return
+}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/invalid-buffer-deallocation.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/invalid-buffer-deallocation.mlir
index c623891e48362fa..113980d02f92a2b 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/invalid-buffer-deallocation.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/invalid-buffer-deallocation.mlir
@@ -66,24 +66,6 @@ func.func @do_loop_alloc(
 
 // -----
 
-func.func @free_effect() {
-  %alloc = memref.alloc() : memref<2xi32>
-  // expected-error @below {{memory free side-effect on MemRef value not supported!}}
-  %new_alloc = memref.realloc %alloc : memref<2xi32> to memref<4xi32>
-  return
-}
-
-// -----
-
-func.func @free_effect() {
-  %alloc = memref.alloc() : memref<2xi32>
-  // expected-error @below {{memory free side-effect on MemRef value not supported!}}
-  memref.dealloc %alloc : memref<2xi32>
-  return
-}
-
-// -----
-
 func.func @free_effect() {
   %true = arith.constant true
   %alloc = memref.alloc() : memref<2xi32>

``````````

</details>


https://github.com/llvm/llvm-project/pull/72289