llvm
diff --git a/‎mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
+6 b/‎mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
+6
diff --git a/‎mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
+5 b/‎mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
+5
diff --git a/‎mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
+5 b/‎mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
+5
diff --git a/‎mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
+5 b/‎mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
+5
@@ -52,6 +52,12 @@ struct OneShotBufferizationOptions : public BufferizationOptions {
   /// `AnalysisHeuristic::Fuzzer`. The fuzzer should be used only with
   /// `testAnalysisOnly = true`.
   unsigned analysisFuzzerSeed = 0;
+
+  /// Enable caching of symbol tables. If enabled, the SymbolBufferizationState
+  /// class is attached to the bufferization state and the user is required to
+  /// keep the cached symbol tables consistent with respect to the performed
+  /// bufferizations.
+  bool cacheSymbolTables = false;
 };
 
 /// State for analysis-enabled bufferization. This class keeps track of alias
 
@@ -85,6 +85,10 @@ transform::OneShotBufferizeOp::apply(transform::TransformRewriter &rewriter,
   auto payloadOps = state.getPayloadOps(getTarget());
   BufferizationState bufferizationState;
 
+  if (options.cacheSymbolTables) {
+    bufferizationState.addExtension<SymbolBufferizationState>();
+  }
+
   for (Operation *target : payloadOps) {
     if (!isa<ModuleOp, FunctionOpInterface>(target))
       return emitSilenceableError() << "expected module or function target";
@@ -166,6 +170,7 @@ class BufferizationTransformDialectExtension
     registerTransformOps<
 #define GET_OP_LIST
 #include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp.inc"
+
         >();
   }
 };
 
@@ -162,6 +162,11 @@ struct OneShotBufferizePass
     }
 
     BufferizationState state;
+
+    if (opt.cacheSymbolTables) {
+      state.addExtension<SymbolBufferizationState>();
+    }
+
     BufferizationStatistics statistics;
     ModuleOp moduleOp = getOperation();
     if (opt.bufferizeFunctionBoundaries) {
 
@@ -116,6 +116,11 @@ class SparsificationAndBufferizationPass
 
     bufferization::BufferizationState bufferizationState;
 
+    if (updatedOptions.cacheSymbolTables) {
+      bufferizationState
+          .addExtension<bufferization::SymbolBufferizationState>();
+    }
+
     if (failed(bufferization::bufferizeModuleOp(cast<ModuleOp>(getOperation()),
                                                 updatedOptions,
                                                 bufferizationState)))