[Mlir-commits] [mlir] [mlir][bufferization] Remove `deallocationFn` from `BufferizationOptions` (PR #67128)

[Matthias Springer]

https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/67128

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.

>From 13790c340cb5d207445834645feb3db4035ed03b Mon Sep 17 00:00:00 2001
From: Matthias Springer <me at m-sp.org>
Date: Fri, 22 Sep 2023 14:49:59 +0200
Subject: [PATCH] [mlir][bufferization][NFC] Remove `deallocationFn` from
 `BufferizationOptions`

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.
---
 .../Bufferization/IR/BufferizableOpInterface.h  | 12 +-----------
 .../IR/BufferizableOpInterface.cpp              | 17 +----------------
 .../Bufferization/IR/BufferizationOps.cpp       |  3 +--
 3 files changed, 3 insertions(+), 29 deletions(-)

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();
 }