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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Ziliang Zhang (ziliangzl)

<details>
<summary>Changes</summary>

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()`.

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


2 Files Affected:

- (modified) mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h (+8) 
- (modified) mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp (+11-1) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
index 15189d2c1cb87..e27ae83e5ee8d 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,9 @@ class OneShotAnalysisState : public AnalysisState {
   /// Cache definitions of tensor values.
   DenseMap<Value, SetVector<Value>> cachedDefinitions;
 
+  /// Cache results of areNonConflictingSubsets checks.
+  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.

``````````

</details>


https://github.com/llvm/llvm-project/pull/189895