[mlir][bufferization] Ownership-based deallocation: Allow manual (de)allocs

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationBase.td

@@ -60,6 +60,17 @@ def Bufferization_Dialect : Dialect {
     /// arguments during One-Shot Module Bufferize.
     constexpr const static ::llvm::StringLiteral
         kBufferLayoutAttrName = "bufferization.buffer_layout";
+
+    /// An attribute that can be attached to ops with an allocation and/or
+    /// deallocation side effect. It indicates that the op is under a "manual
+    /// deallocation" scheme. In the case of an allocation op, the returned
+    /// value is *not* an automatically managed allocation and assigned an
+    /// ownership of "false". Furthermore, only deallocation ops that are
+    /// guaranteed to deallocate a buffer under "manual deallocation" are
+    /// allowed to have this attribute. (Deallocation ops without this
+    /// attribute are rejected by the ownership-based buffer deallocation pass.)
+    constexpr const static ::llvm::StringLiteral
+        kManualDeallocation = "bufferization.manual_deallocation";
   }];
   let hasOperationAttrVerify = 1;
 }

mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp

@@ -28,6 +28,16 @@ constexpr const ::llvm::StringLiteral BufferizationDialect::kWritableAttrName;
 constexpr const ::llvm::StringLiteral
     BufferizationDialect::kBufferLayoutAttrName;
 
+/// An attribute that can be attached to ops with an allocation and/or
+/// deallocation side effect. It indicates that the op is under a "manual
+/// deallocation" scheme. In the case of an allocation op, the returned
+/// value is *not* an automatically managed allocation and assigned an
+/// ownership of "false". Furthermore, only deallocation ops that are
+/// guaranteed to deallocate a buffer under "manual deallocation" are
+/// allowed to have this attribute. (Deallocation ops without this
+/// attribute are rejected by the ownership-based buffer deallocation pass.)
+constexpr const ::llvm::StringLiteral BufferizationDialect::kManualDeallocation;
+
 //===----------------------------------------------------------------------===//
 // Bufferization Dialect Interfaces
 //===----------------------------------------------------------------------===//
@@ -105,6 +115,16 @@ BufferizationDialect::verifyOperationAttribute(Operation *op,
                                                NamedAttribute attr) {
   using bufferization::BufferizableOpInterface;
 
+  if (attr.getName() == kManualDeallocation) {
+    if (!mlir::hasEffect<MemoryEffects::Allocate>(op) &&
+        !mlir::hasEffect<MemoryEffects::Free>(op))
+      return op->emitOpError("attribute '")
+             << kManualDeallocation
+             << "' can be used only on ops that have an allocation and/or free "
+                "side effect";
+    return success();
+  }
+
   return op->emitError()
          << "attribute '" << attr.getName()
          << "' not supported as an op attribute by the bufferization dialect";

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

@@ -21,6 +21,7 @@
 #include "mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
@@ -856,13 +857,32 @@ FailureOr<Operation *> BufferDeallocation::handleInterface(CallOpInterface op) {
 FailureOr<Operation *>
 BufferDeallocation::handleInterface(MemoryEffectOpInterface op) {
   auto *block = op->getBlock();
+  OpBuilder builder = OpBuilder::atBlockBegin(block);
 
-  for (auto operand : llvm::make_filter_range(op->getOperands(), isMemref))
-    if (op.getEffectOnValue<MemoryEffects::Free>(operand).has_value())
-      return op->emitError(
-          "memory free side-effect on MemRef value not supported!");
+  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!");
+
+      // Buffers that were allocated under "manual deallocation" may be
+      // manually deallocated. We insert a runtime assertion to cover certain
+      // cases of invalid IR where an automatically managed buffer allocation
+      // is manually deallocated. This is not a bulletproof check!
+      OpBuilder::InsertionGuard g(builder);
+      builder.setInsertionPoint(op);
+      Ownership ownership = state.getOwnership(operand, block);
+      if (ownership.isUnique()) {
+        Value ownershipInverted = builder.create<arith::XOrIOp>(
+            op.getLoc(), ownership.getIndicator(),
+            buildBoolValue(builder, op.getLoc(), true));
+        builder.create<cf::AssertOp>(
+            op.getLoc(), ownershipInverted,
+            "expected that the block does not have ownership");
+      }
+    }
+  }
 
-  OpBuilder builder = OpBuilder::atBlockBegin(block);
   for (auto res : llvm::make_filter_range(op->getResults(), isMemref)) {
     auto allocEffect = op.getEffectOnValue<MemoryEffects::Allocate>(res);
     if (allocEffect.has_value()) {
@@ -880,6 +900,15 @@ BufferDeallocation::handleInterface(MemoryEffectOpInterface op) {
         continue;
       }
 
+      if (op->hasAttr(BufferizationDialect::kManualDeallocation)) {
+        // This allocation will be deallocated manually. Assign an ownership of
+        // "false", so that it will never be deallocated by the buffer
+        // deallocation pass.
+        state.resetOwnerships(res, block);
+        state.updateOwnership(res, buildBoolValue(builder, op.getLoc(), false));
+        continue;
+      }
+
       state.updateOwnership(res, buildBoolValue(builder, op.getLoc(), true));
       state.addMemrefToDeallocate(res, block);
     }

mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-memoryeffect-interface.mlir

@@ -124,3 +124,35 @@ func.func @op_without_aliasing_and_allocation() -> memref<4xf32> {
 //       CHECK:     [[CLONE:%.+]] = bufferization.clone [[GLOBAL]]
 //       CHECK:     scf.yield [[CLONE]] :
 //       CHECK:   return [[RES]] :
+
+// -----
+
+// Allocations with "bufferization.manual_deallocation" are assigned an
+// ownership of "false".
+
+func.func @manual_deallocation(%c: i1, %f: f32, %idx: index) -> f32 {
+  %0 = memref.alloc() {bufferization.manual_deallocation} : memref<5xf32>
+  linalg.fill ins(%f : f32) outs(%0 : memref<5xf32>)
+  %1 = memref.alloc() : memref<5xf32>
+  linalg.fill ins(%f : f32) outs(%1 : memref<5xf32>)
+  %2 = arith.select %c, %0, %1 : memref<5xf32>
+  %3 = memref.load %2[%idx] : memref<5xf32>
+
+  // Only buffers that are under "manual deallocation" are allowed to be
+  // deallocated with memref.dealloc. For consistency reasons, the
+  // manual_deallocation attribute must also be specified. A runtime insertion
+  // is inserted to ensure that we do not have ownership. (This is not a
+  // bulletproof check, but covers some cases of invalid IR.)
+  memref.dealloc %0 {bufferization.manual_deallocation} : memref<5xf32>
+
+  return %3 : f32
+}
+
+// CHECK-LABEL: func @manual_deallocation(
+//       CHECK:   %[[true:.*]] = arith.constant true
+//       CHECK:   %[[manual_alloc:.*]] = memref.alloc() {bufferization.manual_deallocation} : memref<5xf32>
+//       CHECK:   %[[managed_alloc:.*]] = memref.alloc() : memref<5xf32>
+//       CHECK:   %[[selected:.*]] = arith.select
+//       CHECK:   cf.assert %[[true]], "expected that the block does not have ownership"
+//       CHECK:   memref.dealloc %[[manual_alloc]]
+//       CHECK:   bufferization.dealloc (%[[managed_alloc]] : memref<5xf32>) if (%[[true]])

mlir/test/Dialect/Bufferization/invalid.mlir

@@ -143,3 +143,10 @@ func.func @invalid_dealloc_wrong_number_of_results(%arg0: memref<2xf32>, %arg1:
   %0:3 = "bufferization.dealloc"(%arg0, %arg1, %arg2, %arg2, %arg1) <{operandSegmentSizes = array<i32: 2, 2, 1>}> : (memref<2xf32>, memref<4xi32>, i1, i1, memref<4xi32>) -> (i1, i1, i1)
   return %0#0 : i1
 }
+
+// -----
+
+func.func @invalid_manual_deallocation() {
+  // expected-error @below{{op attribute 'bufferization.manual_deallocation' can be used only on ops that have an allocation and/or free side effect}}
+  arith.constant {bufferization.manual_deallocation} 0  : index
+}