[mlir][bufferize] Add BufferizationState initializers

[springerm]

Summary

Such initializer functions can be enqueued in BufferizationOptions. They can be used to set up dialect-specific bufferization state and/or options.

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

 namespace mlir {
 class BlockAndValueMapping;
 class DominanceInfo;
 class FuncOp;
 
 namespace bufferization {
 
 class BufferizableOpInterface;
-struct BufferizationOptions;
 class BufferizationState;
+struct DialectBufferizationState;
 
 /// Options for ComprehensiveBufferize.
 struct BufferizationOptions {
   /// Allocator function: Generate a memref allocation with the given type,
   /// dynamic extents and alignment.
   using AllocationFn = std::function<FailureOr<Value>(
       OpBuilder &, Location, MemRefType, ValueRange, unsigned int)>;
   /// Deallocator function: Deallocate a buffer that was allocated with
   /// AllocatorFn.
   using DeallocationFn =
       std::function<LogicalResult(OpBuilder &, Location, Value)>;
   /// Memcpy function: Generate a memcpy between two buffers.
   using MemCpyFn =
       std::function<LogicalResult(OpBuilder &, Location, Value, Value)>;
+  /// Initializer function for bufferization state.
+  using BufferizationStateInitFn = std::function<void(BufferizationState &)>;
+  /// Initializer function for dialect-specific bufferization state.
+  using DialectStateInitFn =
+      std::function<std::unique_ptr<DialectBufferizationState>()>;
 
   /// An op filter entry. Filters can be used to specify which ops should be
   /// processed by the bufferization.
   struct OpFilterEntry {
     /// If the filter function evaluates to `true`, the filter matches.
     using FilterFn = std::function<bool(Operation *)>;
 
     /// Filter type: A filter can either be a DENY filter or an ALLOW filter.
struct BufferizationOptions {
   /// BufferizableOpInterface). Otherwise, only filtered ops are bufferized.
   bool hasFilter = false;
 
   /// A list of op filters that determine whether an op should be processed or
   /// ignored by the bufferization. If `hasFilter`, only ops that are not
   /// DENY-filtered and have at least one matching ALLOW filter are processed.
   SmallVector<OpFilterEntry> opFilter;
 
+  /// Initializer functions for bufferization state. These can be used to
+  /// initialize dialect-specific bufferization state.
+  SmallVector<BufferizationStateInitFn> stateInitializers;
+
+  /// Add a bufferization state initializer that initializes the specified
+  /// dialect-specific bufferization state.
+  void addDialectStateInitializer(StringRef name, DialectStateInitFn fn);
+
 private:
   /// Allow a dialect.
   template <typename DialectT>
   void allowDialectInFilterImpl() {
     allowDialectInFilter(DialectT::getDialectNamespace());
   }
 
   /// Allow an op.
public:
   template <typename StateT>
   StateT &getOrCreateDialectState(StringRef name) {
     // Create state if it does not exist yet.
     if (!dialectState.count(name))
       dialectState[name] = std::make_unique<StateT>();
     return static_cast<StateT &>(*dialectState[name]);
   }
 
+  void insertDialectState(StringRef name,
+                          std::unique_ptr<DialectBufferizationState> state) {
+    assert(!dialectState.count(name) && "dialect state already initialized");
+    dialectState[name] = std::move(state);
+  }
+
   /// Return a reference to the BufferizationOptions.
   const BufferizationOptions &getOptions() const { return options; }
 
 protected:
   explicit BufferizationState(const BufferizationOptions &options);
 
   // BufferizationState should be passed as a reference.
   BufferizationState(const BufferizationState &) = delete;

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

 BufferizableOpInterface
 BufferizationOptions::dynCastBufferizableOp(Value value) const {
   if (auto bufferizableOp = value.getDefiningOp<BufferizableOpInterface>())
     if (isOpAllowed(bufferizableOp.getOperation()))
       return bufferizableOp;
   return nullptr;
 }
 
+void BufferizationOptions::addDialectStateInitializer(StringRef name,
+                                                      DialectStateInitFn fn) {
+  stateInitializers.push_back(
+      [=](BufferizationState &state) { state.insertDialectState(name, fn()); });
+}
+
 //===----------------------------------------------------------------------===//
 // Helper functions for BufferizableOpInterface
 //===----------------------------------------------------------------------===//
 
 static void setInsertionPointAfter(OpBuilder &b, Value value) {
   if (auto bbArg = value.dyn_cast<BlockArgument>()) {
     b.setInsertionPointToStart(bbArg.getOwner());
   } else {
return findValueInReverseUseDefChain(value, [&](Value value) {
     auto bufferizableOp = options.dynCastBufferizableOp(op);
     if (!bufferizableOp)
       return true;
     return bufferizableOp.isMemoryWrite(value.cast<OpResult>(), *this);
   });
 }
 
 BufferizationState::BufferizationState(const BufferizationOptions &options)
-    : options(options) {}
+    : options(options) {
+  for (const BufferizationOptions::BufferizationStateInitFn &fn :
+       options.stateInitializers)
+    fn(*this);
+}
 
 // bufferization.to_memref is not allowed to change the rank.
 static void ensureToMemrefOpIsValid(Value tensor, Type memrefType) {
 #ifndef NDEBUG
   auto rankedTensorType = tensor.getType().dyn_cast<RankedTensorType>();
   assert((!rankedTensorType || memrefType.cast<MemRefType>().getRank() ==
                                    rankedTensorType.getRank()) &&
          "to_memref would be invalid: mismatching ranks");