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

mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp

@@ -861,9 +861,16 @@ 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!");
+      // If we encounter an automatic allocation, some other pass (or this one
+      // in a previous invocation) has added it and we shouldn't try to guess it
+      // if was right or wrong, but we can't do anything, so we just ignore this
+      // operand.
+      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

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

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>