[MLIR][Bufferization] Bail on automatic deallocation to enable reentrant behaviour

[rengolin]

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.

[llvmbot]

@llvm/pr-subscribers-mlir-bufferization

@llvm/pr-subscribers-mlir

Author: Renato Golin (rengolin)

Changes

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 --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>

[rengolin]

See libxsmm/tpp-mlir#769 (comment)

[rengolin]

@adam-smnk @nhasabni

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[matthias-springer]

This change can lead to double deallocs, e.g.:

func.func private @foo(%c: i1) {
  %0 = memref.alloc() : memref<5xf32>
  %1 = memref.alloc() : memref<5xf32>
  %2 = arith.select %c, %0, %1 : memref<5xf32>
  memref.dealloc %2 : memref<5xf32>
  // additional deallocs will be added for %0 and %1
  return
}

In the above case, we may be able to add a workaround in the buffer deallocation pass to fix the double dealloc. The real problem are deallocs within nested regions. In such cases, the block that contains the dealloc must take ownership of the memref. Otherwise, another block (the one that has ownership instead) will deallocate the memref. There is currently no way to force ownership without redesigning parts of the algorithm. (I think it gets even trickier with loops where we do not know if it has 0 or more iterations.)

Ideally, the buffer deallocation pass should run a single time towards the end of the compilation process. (Note, we have a similar situation with One-Shot Bufferize: if you do partial bufferization, you cannot run One-Shot Bufferize a second time and bufferize the rest. There are certain to_tensor ops that it will reject.)

If you have pre-existing deallocs in your IR, there are two options:

1. Remove all deallocs and let the buffer deallocation pass re-insert them. We have a pass that does that for memref.realloc: --expand-realloc=emit-deallocs=false. Deallocs may not be inserted at the same position. E.g., deallocs will always be inserted in the same block as the alloc.
2. Annotate memref.alloc+memref.dealloc with bufferization.manual_deallocation. (You could do that after running the buffer deallocation pass.) Such ops will be ignored by the buffer deallocation pass. Note, this would not work with the example shown above if only one alloc has that attribute because the dealloc may receive a buffer that is not under manual deallocation.

[matthias-springer]

@maerhart (author of the ownership-based buffer deallocation pass) to make sure that what I said is accurate

[maerhart]

I agree with what @matthias-springer said.

This PR tries to add support for existing deallocation ops without considering all necessary cases. If we want to add support, we need to do it in a way that doesn't break the existing specification of the pass (no double free, etc.). But this is hard to do properly, as Matthias already pointed out.

To add to the second point, we could add the attribute in the deallocation pass itself to those allocs it considered and the deallocs it inserts. It won't be possible to pass it IR as Matthias showed, but it would allow the deallocation pipeline to be run multiple times in sequence with the same configuration. Maybe we'd want a separate attribute, though, to differentiate manual from already automatically deallocated.

[rengolin]

This PR tries to add support for existing deallocation ops without considering all necessary cases. If we want to add support, we need to do it in a way that doesn't break the existing specification of the pass (no double free, etc.).

Absolutely agree. My PR was more asking the question than proposing a solution (hence some comments were incomplete). The answer is clear: not enough. :)

If you have pre-existing deallocs in your IR, there are two options:

1. Remove all deallocs and let the buffer deallocation pass re-insert them. We have a pass that does that for memref.realloc: --expand-realloc=emit-deallocs=false. Deallocs may not be inserted at the same position. E.g., deallocs will always be inserted in the same block as the alloc.
2. Annotate memref.alloc+memref.dealloc with bufferization.manual_deallocation. (You could do that after running the buffer deallocation pass.) Such ops will be ignored by the buffer deallocation pass. Note, this would not work with the example shown above if only one alloc has that attribute because the dealloc may receive a buffer that is not under manual deallocation.

If these pre-existing deallocs are added by the user or another pass, then these could be reasonable steps. But since they're added by the ownership pass itself, it wouldn't be self-consistent.

Maybe we'd want a separate attribute, though, to differentiate manual from already automatically deallocated.

Perhaps. I assumed (from the field comment) that if the dealloc wasn't marked as kManualDeallocation, it was considered "automatically managed" without ownership. However, introducing another attribute will probably cascade checks through the entire pass.

Reading it again after your input, I think my understanding of the pass itself was mistaken. I'll have to re-evaluate the strategy, thanks for the feedback!