[Mlir-commits] [mlir] 3b74aac - [mlir][bufferize] Do not run the buffer deallocation pass if no allocs escape block boundaries

[Matthias Springer]

Author: Matthias Springer
Date: 2022-03-23T21:07:35+09:00
New Revision: 3b74aac29c97b587c30d940e00b338af384ede59

URL: https://github.com/llvm/llvm-project/commit/3b74aac29c97b587c30d940e00b338af384ede59
DIFF: https://github.com/llvm/llvm-project/commit/3b74aac29c97b587c30d940e00b338af384ede59.diff

LOG: [mlir][bufferize] Do not run the buffer deallocation pass if no allocs escape block boundaries

This fixes a bufferization issue with ops that are not supported by the buffer deallocation pass when `allow-return-allocs=0`.

Differential Revision: https://reviews.llvm.org/D122304

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
    mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
    mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index 861c92d0d7686..cbd83b85a2787 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
@@ -522,10 +522,13 @@ LogicalResult hoistBufferAllocations(Operation *op,
 
 /// Create alloc/dealloc ops as specified in the bufferization options. If
 /// `onlyLeakingAlloc`, only those buffer allocations are processed for which no
-/// buffer deallocation can be created.
+/// buffer deallocation can be created. `changed` is set to `true` if the IR was
+/// modified.
 LogicalResult createAllocDeallocOps(Operation *op,
                                     const BufferizationOptions &options,
-                                    bool onlyLeakingAllocs = false);
+                                    bool onlyLeakingAllocs = false,
+                                    bool *changed = nullptr);
+
 } // namespace bufferization
 } // namespace mlir
 

diff  --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index ee3ea01083986..8c68db885de63 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -512,8 +512,10 @@ LogicalResult bufferization::createMemCpy(OpBuilder &b, Location loc,
 LogicalResult
 bufferization::createAllocDeallocOps(Operation *op,
                                      const BufferizationOptions &options,
-                                     bool onlyLeakingAllocs) {
+                                     bool onlyLeakingAllocs, bool *changed) {
   IRRewriter rewriter(op->getContext());
+  if (changed)
+    *changed = false;
 
   // Bufferization creates memref.alloca ops. After bufferization, these must be
   // rewritten to alloc/dealloc ops as specified in the bufferization options.
@@ -536,6 +538,8 @@ bufferization::createAllocDeallocOps(Operation *op,
     if (failed(alloc))
       return WalkResult::interrupt();
     rewriter.replaceOp(allocaOp, *alloc);
+    if (changed)
+      *changed = true;
 
     // Stop here if the buffer should not be deallocated.
     if (skipDealloc)

diff  --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
index 6a07a40302c0d..48d67fe16b0ed 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
@@ -304,12 +304,21 @@ checkBufferizationResult(Operation *op, const BufferizationOptions &options) {
 LogicalResult
 bufferization::finalizeBuffers(Operation *op,
                                const BufferizationOptions &options) {
+  // Hoist buffers.
   if (failed(hoistBufferAllocations(op, options)))
     return failure();
-  if (failed(createAllocDeallocOps(op, options, /*onlyLeakingAllocs=*/true)))
+
+  // Deallocate buffers that escape block boundaries ("leaking buffers") with
+  // the buffer deallocation pass.
+  bool hasLeakingAlloc = false;
+  if (failed(createAllocDeallocOps(op, options, /*onlyLeakingAllocs=*/true,
+                                   &hasLeakingAlloc)))
     return failure();
-  if (options.createDeallocs && failed(deallocateBuffers(op)))
+  if (options.createDeallocs && hasLeakingAlloc &&
+      failed(deallocateBuffers(op)))
     return failure();
+
+  // Deallocate all remaining buffers at the end of the block.
   if (failed(createAllocDeallocOps(op, options)))
     return failure();