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

[llvmlistbot at llvm.org]

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

Accelerate bufferization pass by caching the result of getAliasingOpOperands function.

>From e3c7bdfae2cdedc54caec59f707bcf426921e959 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              | 30 ++++++++-----------
 .../IR/BufferizableOpInterface.cpp            | 23 +++++++++++---
 2 files changed, 31 insertions(+), 22 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index f3b34f9fded7f..b79be1a06cef4 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
@@ -61,8 +61,7 @@ struct AliasingValue {
   bool isDefinite;
 };
 
-template <typename T>
-class AliasList {
+template <typename T> class AliasList {
 public:
   /// Create an empty list of aliases.
   AliasList() = default;
@@ -124,8 +123,7 @@ class OpFilter {
   /// Allow the given dialects.
   ///
   /// This function adds one or multiple ALLOW entries.
-  template <typename... DialectTs>
-  void allowDialect() {
+  template <typename... DialectTs> void allowDialect() {
     // The following expands a call to allowDialectImpl for each dialect
     // in 'DialectTs'.
     (allowDialectImpl<DialectTs>(), ...);
@@ -134,8 +132,7 @@ class OpFilter {
   /// Deny the given dialects.
   ///
   /// This function adds one or multiple DENY entries.
-  template <typename... DialectTs>
-  void denyDialect() {
+  template <typename... DialectTs> void denyDialect() {
     (denyDialectImpl<DialectTs>(), ...);
   }
 
@@ -162,16 +159,14 @@ class OpFilter {
   /// Allow the given ops.
   ///
   /// This function adds one or multiple ALLOW entries.
-  template <typename... OpTys>
-  void allowOperation() {
+  template <typename... OpTys> void allowOperation() {
     (allowOperationImpl<OpTys>(), ...);
   }
 
   /// Deny the given ops.
   ///
   /// This function adds one or multiple DENY entries.
-  template <typename... OpTys>
-  void denyOperation() {
+  template <typename... OpTys> void denyOperation() {
     (denyOperationImpl<OpTys>(), ...);
   }
 
@@ -219,26 +214,22 @@ class OpFilter {
   }
 
   /// Allow a dialect.
-  template <typename DialectT>
-  void allowDialectImpl() {
+  template <typename DialectT> void allowDialectImpl() {
     allowDialect(DialectT::getDialectNamespace());
   }
 
   /// Deny a dialect.
-  template <typename DialectT>
-  void denyDialectImpl() {
+  template <typename DialectT> void denyDialectImpl() {
     denyDialect(DialectT::getDialectNamespace());
   }
 
   /// Allow an op.
-  template <typename OpTy>
-  void allowOperationImpl() {
+  template <typename OpTy> void allowOperationImpl() {
     allowOperation(OpTy::getOperationName());
   }
 
   /// Deny an op.
-  template <typename OpTy>
-  void denyOperationImpl() {
+  template <typename OpTy> void denyOperationImpl() {
     denyOperation(OpTy::getOperationName());
   }
 
@@ -577,6 +568,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