[mlir][bufferization] Use a cache when checking whether ops are in mutually exclusive regions

[christopherbate]

When profiling one-shot-bufferization over large programs, I found that
the analysis would spend a large amount of time checking whether
two operations are "inside mutually exclusive regions". This change
adds a cache for that information, which can result in a noticeable speedup
depending on program structure.

[llvmbot]

@llvm/pr-subscribers-mlir-bufferization

@llvm/pr-subscribers-mlir

Author: Christopher Bate (christopherbate)

Changes

When profiling one-shot-bufferization over large programs, I found that
the analysis would spend a large amount of time checking whether
two operations are "inside mutually exclusive regions". This change
adds a cache for that information, which can result in a noticeable speedup
depending on program structure.

---

Full diff: https://github.com/llvm/llvm-project/pull/123516.diff

3 Files Affected:

• (modified) mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h (+7)
• (modified) mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp (+17-1)
• (modified) mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp (+2-1)

diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index d1a102e2a6e4e8..6eb2fcb5d64179 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
@@ -563,6 +563,8 @@ class AnalysisState {
 
   virtual void resetCache();
 
+  bool checkInMutuallyExclusiveRegions(Operation *op0, Operation *op1);
+
 protected:
   AnalysisState(const BufferizationOptions &options, TypeID type);
 
@@ -576,6 +578,11 @@ class AnalysisState {
   /// Cache containing closest ancestor repetitive Region.
   DenseMap<std::variant<Operation *, Block *, Region *, Value>, Region *>
       enclosingRepetitiveRegionCache;
+
+  /// Cache that specifies whether the two operations are in mutually exclusive
+  /// regions.
+  DenseMap<std::pair<Operation *, Operation *>, 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 1eb27e44810b0d..4cebfba0d4879c 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::checkInMutuallyExclusiveRegions(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 = 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 fc1b221b4f0369..528e3edfeb4ddf 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
@@ -705,7 +705,8 @@ hasReadAfterWriteInterference(const DenseSet<OpOperand *> &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.checkInMutuallyExclusiveRegions(readingOp,
+                                                  conflictingWritingOp)) {
           LLVM_DEBUG(llvm::dbgs() << "  no conflict: read and write are in "
                                      "mutually exclusive regions\n");
           continue;

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.