Remove caching option and separate extension

Author: mscuttari

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

@@ -639,16 +639,13 @@ class BufferizationState {
     return const_cast<BufferizationState *>(this)->getExtension<Ty>();
   }
 
+  /// Get a reference to the collection of cached symbol tables.
+  SymbolTableCollection &getSymbolTables();
+
+private:
   /// Extensions attached to the state, identified by the TypeID of their type.
   /// Only one extension of any given type is allowed.
   DenseMap<TypeID, std::unique_ptr<Extension>> extensions;
-};
-
-/// Extra bufferization state that is required for bufferization of operations
-/// declaring a symbol or a symbol table.
-struct SymbolBufferizationState : public BufferizationState::Extension {
-  SymbolBufferizationState(BufferizationState &state)
-      : BufferizationState::Extension(state) {}
 
   /// The cached symbol tables.
   /// The user is expected to update / invalidate the cached symbol tables if

mlir/include/mlir/Dialect/Bufferization/Transforms/BufferUtils.h

@@ -127,11 +127,6 @@ FailureOr<memref::GlobalOp> getGlobalFor(arith::ConstantOp constantOp,
                                          uint64_t alignment,
                                          Attribute memorySpace = {});
 
-FailureOr<memref::GlobalOp> getGlobalFor(arith::ConstantOp op,
-                                         BufferizationState &state,
-                                         uint64_t alignment,
-                                         Attribute memorySpace);
-
 void removeSymbol(Operation *op, BufferizationState &state);
 
 void insertSymbol(Operation *op, BufferizationState &state);

mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h

@@ -52,12 +52,6 @@ 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

mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp

@@ -47,7 +47,8 @@ struct ConstantOpInterface
     // Create global memory segment and replace tensor with memref pointing to
     // that memory segment.
     FailureOr<memref::GlobalOp> globalOp =
-        getGlobalFor(constantOp, state, options.bufferAlignment, memorySpace);
+        getGlobalFor(constantOp, state.getSymbolTables(),
+                     options.bufferAlignment, memorySpace);
     if (failed(globalOp))
       return failure();
     memref::GlobalOp globalMemref = *globalOp;

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

@@ -127,6 +127,10 @@ void AnalysisState::resetCache() {
 
 BufferizationState::Extension::~Extension() = default;
 
+SymbolTableCollection &BufferizationState::getSymbolTables() {
+  return symbolTables;
+}
+
 Region *bufferization::getNextEnclosingRepetitiveRegion(
     Region *region, const BufferizationOptions &options) {
   assert(isRepetitiveRegion(region, options) && "expected repetitive region");

mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp

@@ -85,10 +85,6 @@ 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";

mlir/lib/Dialect/Bufferization/Transforms/BufferUtils.cpp

@@ -161,40 +161,17 @@ bufferization::getGlobalFor(arith::ConstantOp constantOp,
 }
 
 namespace mlir::bufferization {
-FailureOr<memref::GlobalOp> getGlobalFor(arith::ConstantOp op,
-                                         BufferizationState &state,
-                                         uint64_t alignment,
-                                         Attribute memorySpace) {
-  if (auto *symbolBufferizationState =
-          state.getExtension<SymbolBufferizationState>()) {
-    // Use the cached symbol tables.
-    return getGlobalFor(op, symbolBufferizationState->symbolTables, alignment,
-                        memorySpace);
-  }
-
-  SymbolTableCollection symbolTables;
-  return getGlobalFor(op, symbolTables, alignment, memorySpace);
-}
-
 void removeSymbol(Operation *op, BufferizationState &state) {
-  if (auto *symbolBufferizationState =
-          state.getExtension<SymbolBufferizationState>()) {
-    SymbolTable &symbolTable =
-        symbolBufferizationState->symbolTables.getSymbolTable(
-            op->getParentWithTrait<OpTrait::SymbolTable>());
+  SymbolTable &symbolTable = state.getSymbolTables().getSymbolTable(
+      op->getParentWithTrait<OpTrait::SymbolTable>());
 
-    symbolTable.remove(op);
-  }
+  symbolTable.remove(op);
 }
 
 void insertSymbol(Operation *op, BufferizationState &state) {
-  if (auto *symbolBufferizationState =
-          state.getExtension<SymbolBufferizationState>()) {
-    SymbolTable &symbolTable =
-        symbolBufferizationState->symbolTables.getSymbolTable(
-            op->getParentWithTrait<OpTrait::SymbolTable>());
+  SymbolTable &symbolTable = state.getSymbolTables().getSymbolTable(
+      op->getParentWithTrait<OpTrait::SymbolTable>());
 
-    symbolTable.insert(op);
-  }
+  symbolTable.insert(op);
 }
 } // namespace mlir::bufferization

mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp

@@ -163,10 +163,6 @@ struct OneShotBufferizePass
 
     BufferizationState state;
 
-    if (opt.cacheSymbolTables) {
-      state.addExtension<SymbolBufferizationState>();
-    }
-
     BufferizationStatistics statistics;
     ModuleOp moduleOp = getOperation();
     if (opt.bufferizeFunctionBoundaries) {

mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp

@@ -116,11 +116,6 @@ class SparsificationAndBufferizationPass
 
     bufferization::BufferizationState bufferizationState;
 
-    if (updatedOptions.cacheSymbolTables) {
-      bufferizationState
-          .addExtension<bufferization::SymbolBufferizationState>();
-    }
-
     if (failed(bufferization::bufferizeModuleOp(cast<ModuleOp>(getOperation()),
                                                 updatedOptions,
                                                 bufferizationState)))