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

[Renato Golin]

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

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.

>From fc47f547c243ddf1f92f0014924827ea81cfd885 Mon Sep 17 00:00:00 2001
From: Renato Golin <rengolin at systemcall.eu>
Date: Tue, 14 Nov 2023 15:52:36 +0000
Subject: [PATCH] [MLIR] Warning on automatic deallocation to enable reentrant
 behaviour

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.
---
 .../OwnershipBasedBufferDeallocation.cpp      | 12 ++++++++---
 ...invalid-buffer-deallocation-reentrant.mlir | 20 +++++++++++++++++++
 .../invalid-buffer-deallocation.mlir          | 18 -----------------
 3 files changed, 29 insertions(+), 21 deletions(-)
 create mode 100644 mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/invalid-buffer-deallocation-reentrant.mlir

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>