[Mlir-commits] [mlir] [mlir][Bufferization] Accelerate bufferization pass (PR #160655)

[llvmlistbot at llvm.org]

https://github.com/mingzheTerapines updated https://github.com/llvm/llvm-project/pull/160655

>From aae42457501c44b6c21303cc014b849b84dc0b5b Mon Sep 17 00:00:00 2001
From: mingzheTerapines <mingzhe.zhang at terapines.com>
Date: Thu, 25 Sep 2025 15:12:26 +0800
Subject: [PATCH 1/2] [mlir][Bufferization] Accelerate bufferization pass

Accelerate bufferization pass by caching the result of getAliasingOpOperands function.
---
 .../IR/BufferizableOpInterface.h              |  3 +++
 .../IR/BufferizableOpInterface.cpp            | 23 +++++++++++++++----
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index f3b34f9fded7f..12ea3c81615a1 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
@@ -577,6 +577,9 @@ class AnalysisState {
   /// regions.
   DenseMap<std::pair<Operation *, Operation *>, bool>
       insideMutuallyExclusiveRegionsCache;
+
+  /// Cache for getAliasingOpOperands results to avoid expensive recomputation.
+  mutable DenseMap<Value, AliasingOpOperandList> aliasingOpOperandsCache;
 };
 
 /// BufferizationState provides information about the state of the IR during the
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index f7b0b87085f3d..0d8f3c331410d 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -119,6 +119,7 @@ bool AnalysisState::insideMutuallyExclusiveRegions(Operation *op0,
 void AnalysisState::resetCache() {
   enclosingRepetitiveRegionCache.clear();
   insideMutuallyExclusiveRegionsCache.clear();
+  aliasingOpOperandsCache.clear();
 }
 
 SymbolTableCollection &BufferizationState::getSymbolTables() {
@@ -413,12 +414,26 @@ static void setInsertionPointAfter(OpBuilder &b, Value value) {
 /// Determine which OpOperand* will alias with `value` if the op is bufferized
 /// in place. Return all tensor OpOperand* if the op is not bufferizable.
 AliasingOpOperandList AnalysisState::getAliasingOpOperands(Value value) const {
+  // Check cache first
+  auto it = aliasingOpOperandsCache.find(value);
+  if (it != aliasingOpOperandsCache.end()) {
+    return it->second;
+  }
+
+  AliasingOpOperandList result;
   if (Operation *op = getOwnerOfValue(value))
     if (auto bufferizableOp = getOptions().dynCastBufferizableOp(op))
-      return bufferizableOp.getAliasingOpOperands(value, *this);
-
-  // The op is not bufferizable.
-  return detail::unknownGetAliasingOpOperands(value);
+      result = bufferizableOp.getAliasingOpOperands(value, *this);
+    else
+      // The op is not bufferizable.
+      result = detail::unknownGetAliasingOpOperands(value);
+  else
+    // The op is not bufferizable.
+    result = detail::unknownGetAliasingOpOperands(value);
+
+  // Cache the result
+  aliasingOpOperandsCache[value] = result;
+  return result;
 }
 
 /// Determine which Values will alias with `opOperand` if the op is bufferized

>From 7869225e0bc7fbc50d30db6415364c98599fbeee Mon Sep 17 00:00:00 2001
From: mingzheTerapines <mingzhe.zhang at terapines.com>
Date: Fri, 26 Sep 2025 13:19:19 +0800
Subject: [PATCH 2/2] put code to lambda and check in debug build

---
 .../IR/BufferizableOpInterface.cpp            | 31 ++++++++++++-------
 1 file changed, 20 insertions(+), 11 deletions(-)

diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index 0d8f3c331410d..2ca21ab8f3404 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -414,24 +414,33 @@ static void setInsertionPointAfter(OpBuilder &b, Value value) {
 /// Determine which OpOperand* will alias with `value` if the op is bufferized
 /// in place. Return all tensor OpOperand* if the op is not bufferizable.
 AliasingOpOperandList AnalysisState::getAliasingOpOperands(Value value) const {
+  // Lambda to compute aliasing operands
+  auto computeAliasingOpOperands = [&]() -> AliasingOpOperandList {
+    AliasingOpOperandList result;
+    if (Operation *op = getOwnerOfValue(value))
+      if (auto bufferizableOp = getOptions().dynCastBufferizableOp(op))
+        result = bufferizableOp.getAliasingOpOperands(value, *this);
+      else
+        // The op is not bufferizable.
+        result = detail::unknownGetAliasingOpOperands(value);
+    else
+      // The op is not bufferizable.
+      result = detail::unknownGetAliasingOpOperands(value);
+    return result;
+  };
+
   // Check cache first
   auto it = aliasingOpOperandsCache.find(value);
   if (it != aliasingOpOperandsCache.end()) {
+#ifndef NDEBUG
+    assert(it->second == computeAliasingOpOperands() &&
+           "inconsistent cache result");
+#endif // NDEBUG
     return it->second;
   }
 
-  AliasingOpOperandList result;
-  if (Operation *op = getOwnerOfValue(value))
-    if (auto bufferizableOp = getOptions().dynCastBufferizableOp(op))
-      result = bufferizableOp.getAliasingOpOperands(value, *this);
-    else
-      // The op is not bufferizable.
-      result = detail::unknownGetAliasingOpOperands(value);
-  else
-    // The op is not bufferizable.
-    result = detail::unknownGetAliasingOpOperands(value);
-
   // Cache the result
+  AliasingOpOperandList result = computeAliasingOpOperands();
   aliasingOpOperandsCache[value] = result;
   return result;
 }