diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h index d1a102e2a6e4e..ea7af3018bda8 100644 --- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h +++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h @@ -563,6 +563,11 @@ class AnalysisState { virtual void resetCache(); + /// Checks whether `op0` and `op1` are inside mutually exclusive regions. + /// The logic defers to `mlir::insideMutuallyExclusiveRegions`, but the + /// result is cached. + bool insideMutuallyExclusiveRegions(Operation *op0, Operation *op1); + protected: AnalysisState(const BufferizationOptions &options, TypeID type); @@ -576,6 +581,11 @@ class AnalysisState { /// Cache containing closest ancestor repetitive Region. DenseMap, Region *> enclosingRepetitiveRegionCache; + + /// Cache that specifies whether the two operations are in mutually exclusive + /// regions. + DenseMap, bool> + insideMutuallyExclusiveRegionsCache; }; /// Create an AllocTensorOp for the given shaped value (memref or tensor). diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp index 1eb27e44810b0..99ffa62c41a4d 100644 --- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp +++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp @@ -107,7 +107,23 @@ Region *AnalysisState::getEnclosingRepetitiveRegion( return region; } -void AnalysisState::resetCache() { enclosingRepetitiveRegionCache.clear(); } +bool AnalysisState::insideMutuallyExclusiveRegions(Operation *op0, + Operation *op1) { + auto key = std::make_pair(op0, op1); + if (auto iter = insideMutuallyExclusiveRegionsCache.find(key); + iter != insideMutuallyExclusiveRegionsCache.end()) + return iter->second; + bool result = ::mlir::insideMutuallyExclusiveRegions(op0, op1); + // Populate results for both orderings of the ops. + insideMutuallyExclusiveRegionsCache[key] = result; + insideMutuallyExclusiveRegionsCache[std::make_pair(op1, op0)] = result; + return result; +} + +void AnalysisState::resetCache() { + enclosingRepetitiveRegionCache.clear(); + insideMutuallyExclusiveRegionsCache.clear(); +} Region *bufferization::getNextEnclosingRepetitiveRegion( Region *region, const BufferizationOptions &options) { diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp index fc1b221b4f036..1eaf999d11c08 100644 --- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp @@ -705,7 +705,8 @@ hasReadAfterWriteInterference(const DenseSet &usesRead, // Note: If ops are executed multiple times (e.g., because they are // inside a loop), mutually exclusive regions may be executed // multiple times. - if (insideMutuallyExclusiveRegions(readingOp, conflictingWritingOp)) { + if (state.insideMutuallyExclusiveRegions(readingOp, + conflictingWritingOp)) { LLVM_DEBUG(llvm::dbgs() << " no conflict: read and write are in " "mutually exclusive regions\n"); continue;