Add extension mechanism to BufferizationState

Author: mscuttari

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

@@ -582,12 +582,77 @@ class AnalysisState {
 /// bufferization process.
 class BufferizationState {
 public:
-  /// Get the cached symbol tables.
-  /// The user is expected to update / invalidate the cached symbol tables if
-  /// the bufferized operation have the Symbol or SymbolTable traits.
-  SymbolTableCollection &getSymbolTables();
+  /// Base class for BufferizationState extensions that allow BufferizationState
+  /// to contain user-specified information in the state object. The extension
+  /// mechanism of BufferizationState mirrors the one of OneShotAnalysisState.
+  class Extension {
+  public:
+    /// Base virtual destructor.
+    // Out-of-line definition ensures symbols are emitted in a single object
+    // file.
+    virtual ~Extension();
+
+  protected:
+    /// Constructs an extension of the given state object.
+    Extension(BufferizationState &state) : state(state) {}
+
+    /// Provides read-only access to the parent OneShotAnalysisState object.
+    const BufferizationState &getBufferizationState() const { return state; }
+
+  private:
+    /// Back-reference to the state that is being extended.
+    BufferizationState &state;
+  };
 
-private:
+  /// Adds a new Extension of the type specified as template parameter,
+  /// constructing it with the arguments provided. The extension is owned by the
+  /// BufferizationState. It is expected that the state does not already have an
+  /// extension of the same type. Extension constructors are expected to take a
+  /// reference to BufferizationState as first argument, automatically supplied
+  /// by this call.
+  template <typename Ty, typename... Args>
+  Ty &addExtension(Args &&...args) {
+    static_assert(std::is_base_of<Extension, Ty>::value,
+                  "only a class derived from "
+                  "BufferizationState::Extension is allowed");
+    auto ptr = std::make_unique<Ty>(*this, std::forward<Args>(args)...);
+    auto result = extensions.try_emplace(TypeID::get<Ty>(), std::move(ptr));
+    assert(result.second && "extension already added");
+    return *static_cast<Ty *>(result.first->second.get());
+  }
+
+  /// Returns the extension of the specified type.
+  template <typename Ty>
+  Ty *getExtension() {
+    static_assert(std::is_base_of<Extension, Ty>::value,
+                  "only a class derived from "
+                  "BufferizationState::Extension is allowed");
+    auto iter = extensions.find(TypeID::get<Ty>());
+    if (iter == extensions.end())
+      return nullptr;
+    return static_cast<Ty *>(iter->second.get());
+  }
+
+  /// Returns the extension of the specified type.
+  template <typename Ty>
+  const Ty *getExtension() const {
+    return const_cast<BufferizationState *>(this)->getExtension<Ty>();
+  }
+
+  /// 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
+  /// the bufferized operation has the Symbol or SymbolTable traits.
   SymbolTableCollection symbolTables;
 };
 

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

@@ -29,6 +29,7 @@ class GlobalOp;
 } // namespace memref
 
 namespace bufferization {
+class BufferizationState;
 
 /// A simple analysis that detects allocation operations.
 class BufferPlacementAllocs {
@@ -126,6 +127,15 @@ 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);
+
 } // namespace bufferization
 } // namespace mlir
 

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

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

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

@@ -125,10 +125,6 @@ void AnalysisState::resetCache() {
   insideMutuallyExclusiveRegionsCache.clear();
 }
 
-SymbolTableCollection &BufferizationState::getSymbolTables() {
-  return symbolTables;
-}
-
 Region *bufferization::getNextEnclosingRepetitiveRegion(
     Region *region, const BufferizationOptions &options) {
   assert(isRepetitiveRegion(region, options) && "expected repetitive region");

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

@@ -159,3 +159,42 @@ bufferization::getGlobalFor(arith::ConstantOp constantOp,
   global->moveBefore(&moduleOp.front());
   return global;
 }
+
+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.remove(op);
+  }
+}
+
+void insertSymbol(Operation *op, BufferizationState &state) {
+  if (auto *symbolBufferizationState =
+          state.getExtension<SymbolBufferizationState>()) {
+    SymbolTable &symbolTable =
+        symbolBufferizationState->symbolTables.getSymbolTable(
+            op->getParentWithTrait<OpTrait::SymbolTable>());
+
+    symbolTable.insert(op);
+  }
+}
+} // namespace mlir::bufferization

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

@@ -9,6 +9,7 @@
 #include "mlir/Dialect/MLProgram/Transforms/BufferizableOpInterfaceImpl.h"
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/Transforms/BufferUtils.h"
 #include "mlir/Dialect/MLProgram/IR/MLProgram.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 
@@ -58,10 +59,7 @@ struct GlobalOpInterface
     if (!globalOp.getValue().has_value())
       return globalOp.emitError("global op must have a value");
 
-    SymbolTable &symbolTable = state.getSymbolTables().getSymbolTable(
-        globalOp->getParentWithTrait<OpTrait::SymbolTable>());
-
-    symbolTable.remove(globalOp);
+    bufferization::removeSymbol(globalOp, state);
 
     auto tensorType = cast<TensorType>(globalOp.getType());
     auto memrefType = getMemRefTypeWithStaticIdentityLayout(tensorType);
@@ -74,7 +72,7 @@ struct GlobalOpInterface
         /*constant=*/!globalOp.getIsMutable(),
         /*alignment=*/nullptr);
 
-    symbolTable.insert(replacement);
+    bufferization::insertSymbol(replacement, state);
     return success();
   }
 };