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..2ca21ab8f3404 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,35 @@ 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 { - if (Operation *op = getOwnerOfValue(value)) - if (auto bufferizableOp = getOptions().dynCastBufferizableOp(op)) - return bufferizableOp.getAliasingOpOperands(value, *this); + // 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; + }; - // The op is not bufferizable. - return detail::unknownGetAliasingOpOperands(value); + // 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; + } + + // Cache the result + AliasingOpOperandList result = computeAliasingOpOperands(); + aliasingOpOperandsCache[value] = result; + return result; } /// Determine which Values will alias with `opOperand` if the op is bufferized