From aae42457501c44b6c21303cc014b849b84dc0b5b Mon Sep 17 00:00:00 2001 From: mingzheTerapines Date: Thu, 25 Sep 2025 15:12:26 +0800 Subject: [PATCH 1/2] [mlir][Bufferization] Accelerate bufferization pass Accelerate bufferization pass by caching the result of getAliasingOpOperands function. --- .../IR/BufferizableOpInterface.h | 3 +++ .../IR/BufferizableOpInterface.cpp | 23 +++++++++++++++---- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h index f3b34f9fded7f..12ea3c81615a1 100644 --- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h +++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h @@ -577,6 +577,9 @@ class AnalysisState { /// regions. DenseMap, bool> insideMutuallyExclusiveRegionsCache; + + /// Cache for getAliasingOpOperands results to avoid expensive recomputation. + mutable DenseMap aliasingOpOperandsCache; }; /// BufferizationState provides information about the state of the IR during the diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp index f7b0b87085f3d..0d8f3c331410d 100644 --- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp +++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp @@ -119,6 +119,7 @@ bool AnalysisState::insideMutuallyExclusiveRegions(Operation *op0, void AnalysisState::resetCache() { enclosingRepetitiveRegionCache.clear(); insideMutuallyExclusiveRegionsCache.clear(); + aliasingOpOperandsCache.clear(); } SymbolTableCollection &BufferizationState::getSymbolTables() { @@ -413,12 +414,26 @@ static void setInsertionPointAfter(OpBuilder &b, Value value) { /// Determine which OpOperand* will alias with `value` if the op is bufferized /// in place. Return all tensor OpOperand* if the op is not bufferizable. AliasingOpOperandList AnalysisState::getAliasingOpOperands(Value value) const { + // Check cache first + auto it = aliasingOpOperandsCache.find(value); + if (it != aliasingOpOperandsCache.end()) { + return it->second; + } + + AliasingOpOperandList result; if (Operation *op = getOwnerOfValue(value)) if (auto bufferizableOp = getOptions().dynCastBufferizableOp(op)) - return bufferizableOp.getAliasingOpOperands(value, *this); - - // The op is not bufferizable. - return detail::unknownGetAliasingOpOperands(value); + result = bufferizableOp.getAliasingOpOperands(value, *this); + else + // The op is not bufferizable. + result = detail::unknownGetAliasingOpOperands(value); + else + // The op is not bufferizable. + result = detail::unknownGetAliasingOpOperands(value); + + // Cache the result + aliasingOpOperandsCache[value] = result; + return result; } /// Determine which Values will alias with `opOperand` if the op is bufferized From 7869225e0bc7fbc50d30db6415364c98599fbeee Mon Sep 17 00:00:00 2001 From: mingzheTerapines Date: Fri, 26 Sep 2025 13:19:19 +0800 Subject: [PATCH 2/2] put code to lambda and check in debug build --- .../IR/BufferizableOpInterface.cpp | 31 ++++++++++++------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp index 0d8f3c331410d..2ca21ab8f3404 100644 --- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp +++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp @@ -414,24 +414,33 @@ static void setInsertionPointAfter(OpBuilder &b, Value value) { /// Determine which OpOperand* will alias with `value` if the op is bufferized /// in place. Return all tensor OpOperand* if the op is not bufferizable. AliasingOpOperandList AnalysisState::getAliasingOpOperands(Value value) const { + // Lambda to compute aliasing operands + auto computeAliasingOpOperands = [&]() -> AliasingOpOperandList { + AliasingOpOperandList result; + if (Operation *op = getOwnerOfValue(value)) + if (auto bufferizableOp = getOptions().dynCastBufferizableOp(op)) + result = bufferizableOp.getAliasingOpOperands(value, *this); + else + // The op is not bufferizable. + result = detail::unknownGetAliasingOpOperands(value); + else + // The op is not bufferizable. + result = detail::unknownGetAliasingOpOperands(value); + return result; + }; + // Check cache first auto it = aliasingOpOperandsCache.find(value); if (it != aliasingOpOperandsCache.end()) { +#ifndef NDEBUG + assert(it->second == computeAliasingOpOperands() && + "inconsistent cache result"); +#endif // NDEBUG return it->second; } - AliasingOpOperandList result; - if (Operation *op = getOwnerOfValue(value)) - if (auto bufferizableOp = getOptions().dynCastBufferizableOp(op)) - result = bufferizableOp.getAliasingOpOperands(value, *this); - else - // The op is not bufferizable. - result = detail::unknownGetAliasingOpOperands(value); - else - // The op is not bufferizable. - result = detail::unknownGetAliasingOpOperands(value); - // Cache the result + AliasingOpOperandList result = computeAliasingOpOperands(); aliasingOpOperandsCache[value] = result; return result; }