From 3b74aac29c97b587c30d940e00b338af384ede59 Mon Sep 17 00:00:00 2001 From: Matthias Springer Date: Wed, 23 Mar 2022 20:07:23 +0900 Subject: [PATCH] [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 --- .../mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h | 7 +++++-- .../Dialect/Bufferization/IR/BufferizableOpInterface.cpp | 6 +++++- mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp | 13 +++++++++++-- 3 files changed, 21 insertions(+), 5 deletions(-) diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h index 861c92d..cbd83b8 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 ee3ea01..8c68db8 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 6a07a40..48d67fe 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(); -- 2.7.4