[Mlir-commits] [mlir] [mlir][bufferization] Cache SymbolTableCollection for CallOp types (PR #176909)

[Prathamesh Tagore]

https://github.com/meshtag created https://github.com/llvm/llvm-project/pull/176909

Use the BufferizationState symbol table cache when resolving CallOp callee types in getBufferType(), avoiding repeated SymbolTableCollection creation. Add a const accessor (backed by a mutable cache) so const state can reuse the same tables. Completes a marked TODO.

>From 60a52c337fd3cf1613c572176f9be09e6f160999 Mon Sep 17 00:00:00 2001
From: Prathamesh Tagore <prathameshtagore at gmail.com>
Date: Wed, 14 Jan 2026 20:51:54 +0530
Subject: [PATCH] [mlir][bufferization] Cache SymbolTableCollection for CallOp
 types

Use the BufferizationState symbol table cache when resolving CallOp callee
types in getBufferType(), avoiding repeated SymbolTableCollection creation.
Add a const accessor (backed by a mutable cache) so const state can reuse
the same tables. Completes a marked TODO.
---
 .../mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h | 4 +++-
 .../Dialect/Bufferization/IR/BufferizableOpInterface.cpp    | 4 ++++
 .../Transforms/FuncBufferizableOpInterfaceImpl.cpp          | 6 ++----
 3 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index dd693a25fd54f..3f8392e3b8970 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
@@ -586,12 +586,14 @@ class BufferizationState {
 public:
   /// Get a reference to the collection of cached symbol tables.
   SymbolTableCollection &getSymbolTables();
+  /// Const overload so callers can reuse the cache from a const state.
+  SymbolTableCollection &getSymbolTables() const;
 
 private:
   /// The cached symbol tables.
   /// The user is expected to update / invalidate the cached symbol tables if
   /// the bufferized operation has the Symbol or SymbolTable traits.
-  SymbolTableCollection symbolTables;
+  mutable SymbolTableCollection symbolTables;
 };
 
 /// Create an AllocTensorOp for the given shaped value (memref or tensor).
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index 64474cf1fee0c..674d9d2716b3c 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -125,6 +125,10 @@ SymbolTableCollection &BufferizationState::getSymbolTables() {
   return symbolTables;
 }
 
+SymbolTableCollection &BufferizationState::getSymbolTables() const {
+  return symbolTables;
+}
+
 Region *bufferization::getNextEnclosingRepetitiveRegion(
     Region *region, const BufferizationOptions &options) {
   assert(isRepetitiveRegion(region, options) && "expected repetitive region");
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
index 8655ed3005a93..e43ab54a048b9 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
@@ -229,10 +229,8 @@ struct CallOpInterface
                 SmallVector<Value> &invocationStack) const {
     auto callOp = cast<func::CallOp>(op);
 
-    // TODO Avoid recomputing the symbol tables every time.
-    SymbolTableCollection symbolTable;
-
-    FuncOp funcOp = getCalledFunction(callOp, symbolTable);
+    // Reuse the cached symbol tables from the bufferization state.
+    FuncOp funcOp = getCalledFunction(callOp, state.getSymbolTables());
     assert(funcOp && "expected CallOp to a FuncOp");
 
     // If the callee was already bufferized, we can directly take the type from