[mlir][linalg][bufferize] Add RTTI to BufferizationState

[springerm]

Summary

The main purpose of this change is to add an assertion for quickly detecting incorrect usage of ComprehensiveBufferize. (I.e., registering the CallOp external model, without using ModuleBufferization.)

Note: ExtensibleRTTI ("RTTI for open class hierarchies") was chosen on purpose here to avoid refering to "ModuleBufferization" in BufferizableOpInterface.

Depends On D114226

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h

 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/EquivalenceClasses.h"
 #include "llvm/ADT/SetVector.h"
+#include "llvm/Support/ExtensibleRTTI.h"
 
 namespace mlir {
 class BlockAndValueMapping;
 class DominanceInfo;
 class FuncOp;
 
 namespace linalg {
 namespace comprehensive_bufferize {
struct AllocationCallbacks {
   /// A function that creates an alloc-dealloc pair. This function may perform
   /// additional optimizations such as buffer allocation hoisting. This function
   /// calls `allocationFn` and `deallocationFn` to create (de)allocations.
   CreateAllocDeallocFn createAllocDeallocFn;
 };
 
 /// BufferizationState keeps track of bufferization state and provides access to
 /// the results of the analysis.
-struct BufferizationState {
-  BufferizationState(ModuleOp moduleOp, AllocationCallbacks &allocationFns)
-      : aliasInfo(moduleOp), allocationFns(allocationFns) {}
+///
+/// Note: BufferizationState uses RTTIExtends, so that extensions of the
+/// bufferization (such as ModuleBufferization) can define their own derived
+/// structs without having to mention their type IDs/enums here.
+struct BufferizationState
+    : public llvm::RTTIExtends<BufferizationState, llvm::RTTIRoot> {
+  BufferizationState(Operation *rootOp, AllocationCallbacks &allocationFns)
+      : aliasInfo(rootOp), allocationFns(allocationFns) {}
 
   // BufferizationState should be passed as a reference.
   BufferizationState(const BufferizationState &) = delete;
 
+  /// Type ID for RTTI.
+  static char ID;
+
   /// Map tensor values to memref buffers.
   void mapBuffer(ValueRange tensors, ValueRange buffers);
 
   /// Map a value to another value.
   void mapValue(Value from, Value to);
 
   /// Map a tensor value to a memref buffer.
   void mapBuffer(Value tensor, Value buffer);

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.cpp

MemRefType mlir::linalg::comprehensive_bufferize::getDynamicMemRefType(
   int64_t dynamicOffset = ShapedType::kDynamicStrideOrOffset;
   SmallVector<int64_t> dynamicStrides(tensorType.getRank(),
                                       ShapedType::kDynamicStrideOrOffset);
   AffineMap stridedLayout = makeStridedLinearLayoutMap(
       dynamicStrides, dynamicOffset, tensorType.getContext());
   return MemRefType::get(tensorType.getShape(), tensorType.getElementType(),
                          stridedLayout, addressSpace);
 }
+
+char mlir::linalg::comprehensive_bufferize::BufferizationState::ID = 0;

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp

 using namespace mlir;
 using namespace linalg;
 using namespace tensor;
 using namespace comprehensive_bufferize;
 
 namespace {
 /// A specialization of BufferizationState that keeps track of additional
 /// state required for bufferization of function boundaries.
-struct ModuleBufferizationState : public BufferizationState {
-  using BufferizationState::BufferizationState;
+struct ModuleBufferizationState
+    : public llvm::RTTIExtends<ModuleBufferizationState, BufferizationState> {
+  using llvm::RTTIExtends<ModuleBufferizationState,
+                          BufferizationState>::RTTIExtends;
+
+  /// Type ID for RTTI.
+  static char ID;
 
   /// A map for looking up bufferized function types.
   DenseMap<FuncOp, FunctionType> bufferizedFunctionTypes;
 };
+
+char ModuleBufferizationState::ID = 0;
 } // namespace
 
 static bool isaTensor(Type t) { return t.isa<TensorType>(); }
 
 /// Remove the attribute that triggers inplace bufferization on a FuncOp
 /// argument `bbArg`.
 static void removeBufferizationFuncArguments(BlockArgument bbArg) {
   auto funcOp = cast<FuncOp>(bbArg.getOwner()->getParentOp());
LogicalResult bufferize(Operation *op, OpBuilder &b,
     }
 
     // 2. Compute bufferized FunctionType.
     SmallVector<Type> argumentTypes{callOp->getOperandTypes()};
     ValueRange hoistedArgs{hoistedArguments};
     llvm::append_range(argumentTypes, hoistedArgs.getTypes());
     // Get the bufferized FunctionType for funcOp or construct it if not yet
     // available.
-    // TODO: Assert that `state` is a ModuleBufferizationState.
+    assert(isa<ModuleBufferizationState>(state) &&
+           "cannot use CallOp bufferization interface impl without "
+           "ModuleBufferization");
     FunctionType bufferizedFuncType = getOrCreateBufferizedFunctionType(
         funcOp, argumentTypes, resultTypes,
         static_cast<ModuleBufferizationState &>(state).bufferizedFunctionTypes);
 
     // 3. Rewrite tensor operands as memrefs based on `bufferizedFuncType`.
     SmallVector<Value> newOperands;
     newOperands.reserve(callOp->getNumOperands());
     for (OpOperand &opOperand : callOp->getOpOperands()) {