[Mlir-commits] [mlir] [mlir][bufferization] Use a cache when checking whether ops are in mutually exclusive regions (PR #123516)

[Christopher Bate]

https://github.com/christopherbate created https://github.com/llvm/llvm-project/pull/123516

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.


>From d4f80862fa3659929c6bda6df0712b8232081cbb Mon Sep 17 00:00:00 2001
From: Christopher Bate <cbate at nvidia.com>
Date: Sun, 19 Jan 2025 15:30:56 +0000
Subject: [PATCH] [mlir][bufferization] Use a cache when checking whether ops
 are in mutually exclusive regions

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.
---
 .../Bufferization/IR/BufferizableOpInterface.h |  7 +++++++
 .../IR/BufferizableOpInterface.cpp             | 18 +++++++++++++++++-
 .../Transforms/OneShotAnalysis.cpp             |  3 ++-
 3 files changed, 26 insertions(+), 2 deletions(-)

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;