[Mlir-commits] [mlir] 47dfcba - [mlir][bufferization] Cache areNonConflictingSubsets results in OneShotAnalysis (#189895)

[llvmlistbot at llvm.org]

Author: Ziliang Zhang
Date: 2026-04-16T10:39:53Z
New Revision: 47dfcba42fe13fa7c9c42d193a66452258990fad

URL: https://github.com/llvm/llvm-project/commit/47dfcba42fe13fa7c9c42d193a66452258990fad
DIFF: https://github.com/llvm/llvm-project/commit/47dfcba42fe13fa7c9c42d193a66452258990fad.diff

LOG: [mlir][bufferization] Cache areNonConflictingSubsets results in OneShotAnalysis (#189895)

The `areNonConflictingSubsets` check in `hasReadAfterWriteInterference`
can be expensive when called repeatedly for the same (uRead,
uConflictingWrite) pairs during bufferization analysis. This patch
caches the results to avoid redundant computation.

Specifically, this adds:
- A private `nonConflictingSubsetCache` DenseMap on
`OneShotAnalysisState` to memoize subset conflict results.
- A public `areNonConflictingSubsetsCached` method that wraps the
lookup-or-compute pattern, following the same convention as the existing
`findDefinitionsCached`.
- Cache invalidation in `resetCache()`.

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
    mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
index 15189d2c1cb87..34ce4fe456199 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
@@ -131,6 +131,11 @@ class OneShotAnalysisState : public AnalysisState {
   /// retrieve them from the cache.
   const SetVector<Value> &findDefinitionsCached(OpOperand *opOperand);
 
+  /// Return whether `uRead` and `uConflictingWrite` are non-conflicting
+  /// subsets, with caching.
+  bool areNonConflictingSubsetsCached(OpOperand *uRead,
+                                      OpOperand *uConflictingWrite);
+
   /// Reset cached data structures.
   void resetCache() override;
 
@@ -232,6 +237,11 @@ class OneShotAnalysisState : public AnalysisState {
   /// Cache definitions of tensor values.
   DenseMap<Value, SetVector<Value>> cachedDefinitions;
 
+  /// Cache results of areNonConflictingSubsets checks. The bool value is `true`
+  /// if the operands are non-conflicting subsets, `false` if they are
+  /// conflicting. The absence of an entry means uncached.
+  DenseMap<std::pair<OpOperand *, OpOperand *>, bool> nonConflictingSubsetCache;
+
   /// Set of all OpResults that were decided to bufferize in-place.
   llvm::DenseSet<OpOperand *> inplaceBufferized;
 

diff  --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
index b57811868a725..4ebb03915348d 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
@@ -732,7 +732,7 @@ hasReadAfterWriteInterference(const DenseSet<OpOperand *> &usesRead,
       }
 
       // No conflict if the operands are non-conflicting subsets.
-      if (areNonConflictingSubsets(uRead, uConflictingWrite, state)) {
+      if (state.areNonConflictingSubsetsCached(uRead, uConflictingWrite)) {
         LDBG() << "  no conflict: non-conflicting subsets";
         continue;
       }
@@ -971,9 +971,19 @@ OneShotAnalysisState::findDefinitionsCached(OpOperand *opOperand) {
   return cachedDefinitions[value];
 }
 
+bool OneShotAnalysisState::areNonConflictingSubsetsCached(
+    OpOperand *uRead, OpOperand *uConflictingWrite) {
+  auto key = std::make_pair(uRead, uConflictingWrite);
+  auto [it, inserted] = nonConflictingSubsetCache.try_emplace(key, false);
+  if (inserted)
+    it->second = areNonConflictingSubsets(uRead, uConflictingWrite, *this);
+  return it->second;
+}
+
 void OneShotAnalysisState::resetCache() {
   AnalysisState::resetCache();
   cachedDefinitions.clear();
+  nonConflictingSubsetCache.clear();
 }
 
 /// Determine if `operand` can be bufferized in-place.