[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] [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