[mlir][bufferization] Remove deallocationFn from BufferizationOptions

[matthias-springer]

One-Shot Bufferize no longer deallocates buffers, so deallocationFn can be removed.

Note: There is a bufferization.dealloc_tensor op that now always bufferizes to memref.dealloc. This op will be phased out soon.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-bufferization

Changes

One-Shot Bufferize no longer deallocates buffers, so deallocationFn can be removed.

Note: There is a bufferization.dealloc_tensor op that now always bufferizes to memref.dealloc. This op will be phased out soon.

---

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

3 Files Affected:

• (modified) mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h (+1-11)
• (modified) mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp (+1-16)
• (modified) mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp (+1-2)

diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index 1c715f8b9a53ef3..ea61ac86b9e08a5 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
@@ -243,10 +243,6 @@ struct BufferizationOptions {
   /// dynamic extents and alignment.
   using AllocationFn = std::function<FailureOr<Value>(
       OpBuilder &, Location, MemRefType, ValueRange, unsigned int)>;
-  /// Deallocator function: Deallocate a buffer that was allocated with
-  /// AllocatorFn.
-  using DeallocationFn =
-      std::function<LogicalResult(OpBuilder &, Location, Value)>;
   /// Memcpy function: Generate a memcpy between two buffers.
   using MemCpyFn =
       std::function<LogicalResult(OpBuilder &, Location, Value, Value)>;
@@ -279,20 +275,14 @@ struct BufferizationOptions {
   /// Return `true` if the given op should be bufferized.
   bool isOpAllowed(Operation *op) const;
 
-  /// Helper functions for allocation, deallocation, memory copying.
+  /// Helper functions for allocation and memory copying.
   std::optional<AllocationFn> allocationFn;
-  std::optional<DeallocationFn> deallocationFn;
   std::optional<MemCpyFn> memCpyFn;
 
   /// Create a memref allocation with the given type and dynamic extents.
   FailureOr<Value> createAlloc(OpBuilder &b, Location loc, MemRefType type,
                                ValueRange dynShape) const;
 
-  /// Creates a memref deallocation. The given memref buffer must have been
-  /// allocated using `createAlloc`.
-  LogicalResult createDealloc(OpBuilder &b, Location loc,
-                              Value allocatedBuffer) const;
-
   /// Creates a memcpy between two given buffers.
   LogicalResult createMemCpy(OpBuilder &b, Location loc, Value from,
                              Value to) const;
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index 57cd303d2076e73..c73f7067e5d6ef5 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -220,9 +220,6 @@ LogicalResult BufferizableOpInterface::resolveTensorOpOperandConflicts(
       return op->emitError("copying of unranked tensors is not implemented");
 
     AliasingValueList aliasingValues = state.getAliasingValues(opOperand);
-    // Is the result yielded from a block? Or are deallocations turned off
-    // entirely? In either case, mark the allocation as "escaping", so that it
-    // will not be deallocated.
     if (aliasingValues.getNumAliases() == 1 &&
         isa<OpResult>(aliasingValues.getAliases()[0].value) &&
         !state.bufferizesToMemoryWrite(opOperand) &&
@@ -770,7 +767,7 @@ void bufferization::replaceOpWithBufferizedValues(RewriterBase &rewriter,
 }
 
 //===----------------------------------------------------------------------===//
-// Bufferization-specific scoped alloc/dealloc insertion support.
+// Bufferization-specific scoped alloc insertion support.
 //===----------------------------------------------------------------------===//
 
 /// Create a memref allocation with the given type and dynamic extents.
@@ -789,18 +786,6 @@ FailureOr<Value> BufferizationOptions::createAlloc(OpBuilder &b, Location loc,
   return b.create<memref::AllocOp>(loc, type, dynShape).getResult();
 }
 
-/// Creates a memref deallocation. The given memref buffer must have been
-/// allocated using `createAlloc`.
-LogicalResult BufferizationOptions::createDealloc(OpBuilder &b, Location loc,
-                                                  Value allocatedBuffer) const {
-  if (deallocationFn)
-    return (*deallocationFn)(b, loc, allocatedBuffer);
-
-  // Default buffer deallocation via DeallocOp.
-  b.create<memref::DeallocOp>(loc, allocatedBuffer);
-  return success();
-}
-
 /// Create a memory copy between two memref buffers.
 LogicalResult BufferizationOptions::createMemCpy(OpBuilder &b, Location loc,
                                                  Value from, Value to) const {
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index 745333de65815ad..7c6c1be351cced1 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
@@ -526,8 +526,7 @@ LogicalResult DeallocTensorOp::bufferize(RewriterBase &rewriter,
   FailureOr<Value> buffer = getBuffer(rewriter, getTensor(), options);
   if (failed(buffer))
     return failure();
-  if (failed(options.createDealloc(rewriter, getLoc(), *buffer)))
-    return failure();
+  rewriter.create<memref::DeallocOp>(getLoc(), *buffer);
   rewriter.eraseOp(getOperation());
   return success();
 }

[qcolombet]

LGTM.

Maybe in the commit message reminds what people should use for deallocation now.