[mlir][bufferize][NFC] Remove BufferizationState

[springerm]

Summary

With the recent refactorings, this class is no longer needed. We can use BufferizationOptions in all places were BufferizationState was used.

Depends On D127652

mlir/docs/Bufferization.md

 ## What is One-Shot Bufferize?
 
 One-Shot Bufferize is a new tensor bufferization pass designed for IR in
 [destination-passing style](https://www.microsoft.com/en-us/research/wp-content/uploads/2016/11/dps-fhpc17.pdf),
 and with aggressive in-place bufferization.
 
 One-Shot Bufferize is:
 
-* **Monolithic**: A single MLIR pass does the entire
-work, whereas the previous bufferization in MLIR was split across multiple
-passes residing in different dialects. In One-Shot Bufferize,
-`BufferizableOpInterface` implementations are spread across different dialects.
-
-* A **whole-function at a time analysis**. In-place bufferization decisions are
-made by analyzing SSA use-def chains on tensors. Op interface implementations
-not only provide the rewrite logic from tensor ops to memref ops, but also
-helper methods for One-Shot Bufferize's analysis to query information about an
-op's bufferization/memory semantics.
-
-* **Extensible** via an op interface: All
-ops that implement `BufferizableOpInterface` can be bufferized.
-
-* **2-Pass**:
-Bufferization is internally broken down into 2 steps: First, analyze the entire
-IR and make bufferization decisions. Then, bufferize (rewrite) the IR. The
-analysis has access to exact SSA use-def information. It incrementally builds
-alias and equivalence sets and does not rely on a posteriori-alias analysis from
-preallocated memory.
-
-* **Greedy**: Operations are analyzed one-by-one and it is
-decided on the spot whether a tensor OpOperand must be copied or not. Heuristics
-determine the order of analysis.
-
-* **Modular**: The current One-Shot Analysis
-can be replaced with a different analysis. The result of the analysis are
-queried by the bufferization via `BufferizationState`, in particular
-`BufferizationState::isInPlace`. Any derived class of `BufferizationState` that
-implements a small number virtual functions can serve as a custom analysis. It
-is even possible to run One-Shot Bufferize without any analysis
-(`AlwaysCopyBufferizationState`), in which case One-Shot Bufferize behaves
-exactly like the old dialect conversion-based bufferization (i.e., copy every
-buffer before writing to it).
+*   **Monolithic**: A single MLIR pass does the entire work, whereas the
+    previous bufferization in MLIR was split across multiple passes residing in
+    different dialects. In One-Shot Bufferize, `BufferizableOpInterface`
+    implementations are spread across different dialects.
+
+*   A **whole-function at a time analysis**. In-place bufferization decisions
+    are made by analyzing SSA use-def chains on tensors. Op interface
+    implementations not only provide the rewrite logic from tensor ops to memref
+    ops, but also helper methods for One-Shot Bufferize's analysis to query
+    information about an op's bufferization/memory semantics.
+
+*   **Extensible** via an op interface: All ops that implement
+    `BufferizableOpInterface` can be bufferized.
+
+*   **2-Pass**: Bufferization is internally broken down into 2 steps: First,
+    analyze the entire IR and make bufferization decisions. Then, bufferize
+    (rewrite) the IR. The analysis has access to exact SSA use-def information.
+    It incrementally builds alias and equivalence sets and does not rely on a
+    posteriori-alias analysis from preallocated memory.
+
+*   **Greedy**: Operations are analyzed one-by-one and it is decided on the spot
+    whether a tensor OpOperand must be copied or not. Heuristics determine the
+    order of analysis.
+
+*   **Modular**: The current One-Shot Analysis can be replaced with a different
+    analysis. The result of the analysis are queried by the bufferization via
+    `AnalysisState`, in particular `AnalysisState::isInPlace`. Any derived class
+    of `AnalysisState` that implements a small number virtual functions can
+    serve as a custom analysis. It is even possible to run One-Shot Bufferize
+    without any analysis (`AlwaysCopyAnalysisState`), in which case One-Shot
+    Bufferize behaves exactly like the old dialect conversion-based
+    bufferization (i.e., copy every buffer before writing to it).
 
 To reduce complexity, One-Shot Bufferize should be
 [run after other transformations](https://llvm.discourse.group/t/rfc-linalg-on-tensors-update-and-comprehensive-bufferization-rfc/3373),
 typically as one of the last steps right before lowering memref ops. Many
 transformations are easier in tensor land; e.g., tile/fuse/… on tensors first,
 then bufferize the remaining IR.
 
 From an architecture perspective, One-Shot Bufferize consists of

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

struct BufferizationOptions {
   bool bufferizeFunctionBoundaries = false;
 
   /// Certain ops have aliasing OpOperand/OpResult invariants (e.g., scf.for).
   /// If this flag is set to `false`, those in variants are no longer enforced
   /// with buffer copies.
   ///
   /// Note: Deactivating this flag can lead to incorrect bufferization results
   /// when used incorrectly. This flag is useful with
-  /// `AlwaysCopyBufferizationState` which bufferizes all writing tensor
+  /// `AlwaysCopyAnalysisState` which bufferizes all writing tensor
   /// OpOperands out-of-place.
   bool enforceAliasingInvariants = true;
 
   /// This flag controls buffer types on function signatures.
   ///
   /// * InferLayoutMap: All function parameter types have a fully dynamic layout
   ///   map, but function result types are inferred from the body of the
   ///   function.

 private:
   /// Dialect-specific analysis state.
   DenseMap<StringRef, std::unique_ptr<DialectAnalysisState>> dialectState;
 
   /// A reference to current bufferization options.
   const BufferizationOptions &options;
 };
 
-/// BufferizationState provides helper functions for performing bufferization
-/// rewrites and handling memref buffers.
-struct BufferizationState {
-  BufferizationState(const BufferizationOptions &options) : options(options) {}
+/// Create an AllocTensorOp for the given shaped value (memref or tensor).
+/// If `copy` is set, the shaped value is copied. Otherwise, a tensor with
+/// undefined contents is allocated.
+Value allocateTensorForShapedValue(OpBuilder &b, Location loc,
+                                   Value shapedValue, bool escape,
+                                   bool copy = true);
 
-  /// Lookup the buffer for the given value. If the value was not bufferized
-  /// yet, wrap it in a ToMemrefOp. Otherwise, it is the result of a ToTensorOp,
-  /// from which the memref operand is returned.
-  Value getBuffer(RewriterBase &rewriter, Value value);
+Value getBuffer(RewriterBase &rewriter, Value value,
+                const BufferizationOptions &options);
 
-  /// Return the buffer type for a given Value (tensor) after bufferization.
-  ///
-  /// Note: Op implementations should preferrably call `getBuffer()->getType()`.
-  /// This function should only be used if `getBuffer` cannot be used.
-  BaseMemRefType getBufferType(Value value) const;
-
-  /// Return a reference to the BufferizationOptions.
-  const BufferizationOptions &getOptions() const { return options; }
-
-protected:
-  // BufferizationState should be passed as a reference.
-  BufferizationState(const BufferizationState &) = delete;
-
-private:
-  const BufferizationOptions &options;
-};
-
-/// Create an AllocTensorOp for the given shaped value (memref or tensor).
-/// If `copy` is set, the shaped value is copied. Otherwise, a tensor with
-/// undefined contents is allocated.
-Value allocateTensorForShapedValue(OpBuilder &b, Location loc,
-                                   Value shapedValue, bool escape,
-                                   bool copy = true);
+BaseMemRefType getBufferType(Value value, const BufferizationOptions &options);
 
 /// Replace an op with replacement values. The op is deleted. Tensor OpResults
 /// must be replaced with memref values.
 void replaceOpWithBufferizedValues(RewriterBase &rewriter, Operation *op,
                                    ValueRange values);
 
 /// Replace an op with a new op. The new op must have the same number of
 /// results as the replaced op. The new op may not return any tensor values.

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

let methods = [
               cast<BufferizableOpInterface>($_op.getOperation());
           return bufferizableOp.resolveTensorOpOperandConflicts(
               rewriter, state);
         }]
       >,
       InterfaceMethod<
         /*desc=*/[{
           Bufferize this op, i.e., rewrite it into a memref-based equivalent.
-          Buffers of tensor SSA values can be retrieved via `state.getBuffer`.
+          Buffers of tensor SSA values can be retrieved via `getBuffer`.
           Uses of tensor results of the existing tensor op can be replaced with
           `replaceOpWithBufferizedValues` or `replaceOpWithNewBufferizedOp`.
           These two functions automatically handle the tensor-to-memref type
           conversion.
 
           The implementation of this method must be consistent with the
           remaining methods, in particular `getAliasingOpOperand`. I.e., a
           tensor result `r` may only be replaced with:
           a) A buffer that aliases one of buffers in getAliasingOpOperand(r).
           b) Or: A newly allocated buffer.
 
-          Regions of an op should be inlined into the new op instead of cloning
-          them. This is not only more efficient, but also necessary so that no
-          analysis results are lost. (Bufferization decisions are tracked via
-          OpOperand pointers and cloned ops have new OpOperands.) If regions are
-          cloned instead of inlined, additional buffer copies may be inserted.
-
           This method will never be called on ops that do not have at least one
           tensor operand/result.
 
           The return value of this method indicates whether there was an error
           while bufferizing this op (such as failing to create a new buffer
           allocation op). The bufferization driver immediately stops bufferizing
           the input IR and returns `failure` in that case. If this op is
           expected to survive bufferization, `success` should be returned
           (together with `allow-unknown-ops` enabled).
         }],
         /*retType=*/"LogicalResult",
         /*methodName=*/"bufferize",
         /*args=*/(ins "RewriterBase &":$rewriter,
-                      "BufferizationState &":$state),
+                      "const BufferizationOptions &":$options),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           llvm_unreachable("bufferize not implemented");
           return failure();
         }]
       >,
       InterfaceMethod<
         /*desc=*/[{

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td

def Bufferization_AllocTensorOp : Bufferization_Op<"alloc_tensor",
 
   let arguments = (ins Variadic<Index>:$dynamicSizes,
                        Optional<AnyTensor>:$copy,
                        OptionalAttr<BoolAttr>:$escape);
 
   let results = (outs AnyTensor:$result);
 
   let extraClassDeclaration = [{
-    LogicalResult bufferize(RewriterBase &rewriter, BufferizationState &state);
+    LogicalResult bufferize(RewriterBase &rewriter,
+                            const BufferizationOptions &options);
 
     bool isMemoryWrite(OpResult opResult, const AnalysisState &state);
 
     bool bufferizesToMemoryRead(OpOperand &opOperand,
                                 const AnalysisState &state);
 
     bool bufferizesToMemoryWrite(OpOperand &opOperand,
                                  const AnalysisState &state);
let extraClassDeclaration = [{
     // One-Shot analysis has no information about the memref that is loaded from
     // by ToTensorOp. We have to assume that the loaded tensor may after
     // bufferization potentially alias with any other bufferized tensor. Since
     // ToTensorOp and ToMemrefOp have no aliasing OpOperand/OpResult pairs, this
     // cannot be encoded directly in the analysis. However, declaring ToTensorOp
     // results as not writable enforces a buffer copy and has the same effect.
 
     LogicalResult bufferize(RewriterBase &rewriter,
-                            BufferizationState &state) const {
+                            const BufferizationOptions &options) const {
       // to_tensor cannot be bufferized. However, other ops that are using
       // to_tensor's result will eventually be bufferized. At that point, they
       // will start using to_tensor's memref operand. Once all users of
       // to_tensor are bufferized, the op will not have any users anymore and
       // DCE away. In case of partial bufferization, to_memref(to_tensor(x))
       // constructs may be left over. These are folded by the canonicalizer or
       // FinalizingBufferize.
       return success();
let extraClassDeclaration = [{
     }
 
     SmallVector<OpResult> getAliasingOpResult(
         OpOperand &opOperand, const AnalysisState &state) const {
       return {};
     }
 
     LogicalResult bufferize(RewriterBase &rewriter,
-                            BufferizationState &state);
+                            const BufferizationOptions &options);
   }];
 
   let assemblyFormat = "$tensor attr-dict `:` type($memref)";
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 #endif // BUFFERIZATION_OPS

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

 #define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_BUFFERIZE_H
 
 #include "mlir/Transforms/DialectConversion.h"
 
 namespace mlir {
 namespace bufferization {
 
 class AnalysisState;
-struct BufferizationState;
 struct BufferizationOptions;
 class OpFilter;
 
 /// A helper type converter class that automatically populates the relevant
 /// materializations and type conversions for bufferization.
 class BufferizeTypeConverter : public TypeConverter {
 public:
   BufferizeTypeConverter();

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

 #define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTMODULEBUFFERIZE_H
 
 namespace mlir {
 
 struct LogicalResult;
 class ModuleOp;
 
 namespace bufferization {
-struct BufferizationState;
 class OneShotAnalysisState;
 struct OneShotBufferizationOptions;
 
 /// Analyze `moduleOp` and its nested ops. Bufferization decisions are stored in
 /// `state`.
 LogicalResult analyzeModuleOp(ModuleOp moduleOp, OneShotAnalysisState &state);
 
 /// Bufferize `op` and its nested ops that implement `BufferizableOpInterface`.

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

 using namespace mlir::bufferization;
 
 namespace {
 /// Bufferization of arith.constant. Replace with memref.get_global.
 struct ConstantOpInterface
     : public BufferizableOpInterface::ExternalModel<ConstantOpInterface,
                                                     arith::ConstantOp> {
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto constantOp = cast<arith::ConstantOp>(op);
 
     // Only ranked tensors are supported.
     if (!constantOp.getType().isa<RankedTensorType>())
       return failure();
 
     // Only constants inside a module are supported.
     auto moduleOp = constantOp->getParentOfType<ModuleOp>();
     if (!moduleOp)
       return failure();
 
     // Create global memory segment and replace tensor with memref pointing to
     // that memory segment.
     FailureOr<memref::GlobalOp> globalOp =
-        getGlobalFor(constantOp, state.getOptions().bufferAlignment);
+        getGlobalFor(constantOp, options.bufferAlignment);
     if (failed(globalOp))
       return failure();
     memref::GlobalOp globalMemref = globalOp.getValue();
     replaceOpWithNewBufferizedOp<memref::GetGlobalOp>(
         rewriter, op, globalMemref.type(), globalMemref.getName());
 
     return success();
   }
struct IndexCastOpInterface
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto castOp = cast<arith::IndexCastOp>(op);
     auto resultTensorType = castOp.getType().cast<TensorType>();
 
-    Value source = state.getBuffer(rewriter, castOp.getIn());
+    Value source = getBuffer(rewriter, castOp.getIn(), options);
     auto sourceType = source.getType().cast<BaseMemRefType>();
 
     // Result type should have same layout and address space as the source type.
     BaseMemRefType resultType;
     if (auto rankedMemRefType = sourceType.dyn_cast<MemRefType>()) {
       resultType = MemRefType::get(
           rankedMemRefType.getShape(), resultTensorType.getElementType(),
           rankedMemRefType.getLayout(), rankedMemRefType.getMemorySpace());
struct SelectOpInterface
   SmallVector<OpOperand *>
   getAliasingOpOperand(Operation *op, OpResult opResult,
                        const AnalysisState &state) const {
     return {&op->getOpOperand(1) /*true_value*/,
             &op->getOpOperand(2) /*false_value*/};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto selectOp = cast<arith::SelectOp>(op);
     Location loc = selectOp.getLoc();
 
     // TODO: It would be more efficient to copy the result of the `select` op
     // instead of its OpOperands. In the worst case, 2 copies are inserted at
     // the moment (one for each tensor). When copying the op result, only one
     // copy would be needed.
-    Value trueBuffer = state.getBuffer(rewriter, selectOp.getTrueValue());
-    Value falseBuffer = state.getBuffer(rewriter, selectOp.getFalseValue());
+    Value trueBuffer = getBuffer(rewriter, selectOp.getTrueValue(), options);
+    Value falseBuffer = getBuffer(rewriter, selectOp.getFalseValue(), options);
 
     // The "true" and the "false" operands must have the same type. If the
     // buffers have different types, they differ only in their layout map. Cast
     // both of them to the most dynamic MemRef type.
     if (trueBuffer.getType() != falseBuffer.getType()) {
       auto trueType = trueBuffer.getType().cast<MemRefType>();
       int64_t dynamicOffset = ShapedType::kDynamicStrideOrOffset;
       SmallVector<int64_t> dynamicStrides(trueType.getRank(),

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

 #ifndef NDEBUG
   auto rankedTensorType = tensor.getType().dyn_cast<RankedTensorType>();
   assert((!rankedTensorType || memrefType.cast<MemRefType>().getRank() ==
                                    rankedTensorType.getRank()) &&
          "to_memref would be invalid: mismatching ranks");
 #endif
 }
 
-Value BufferizationState::getBuffer(RewriterBase &rewriter, Value value) {
+Value bufferization::getBuffer(RewriterBase &rewriter, Value value,
+                               const BufferizationOptions &options) {
   auto tensorType = value.getType().dyn_cast<TensorType>();
   assert(tensorType && "unexpected non-tensor type");
 
   // Replace "%t = to_tensor %m" with %m.
   if (auto toTensorOp = value.getDefiningOp<bufferization::ToTensorOp>())
     return toTensorOp.memref();
 
   // Insert to_memref op.
   OpBuilder::InsertionGuard g(rewriter);
   setInsertionPointAfter(rewriter, value);
-  Type memrefType = getMemRefType(tensorType, getOptions());
+  Type memrefType = getMemRefType(tensorType, options);
   ensureToMemrefOpIsValid(value, memrefType);
   return rewriter.create<bufferization::ToMemrefOp>(value.getLoc(), memrefType,
                                                     value);
 }
 
 /// Return the buffer type for a given Value (tensor) after bufferization.
-BaseMemRefType BufferizationState::getBufferType(Value value) const {
+BaseMemRefType
+bufferization::getBufferType(Value value, const BufferizationOptions &options) {
   auto tensorType = value.getType().dyn_cast<TensorType>();
   assert(tensorType && "unexpected non-tensor type");
 
   if (auto toTensorOp = value.getDefiningOp<bufferization::ToTensorOp>())
     return toTensorOp.memref().getType().cast<BaseMemRefType>();
 
-  return getMemRefType(tensorType, getOptions());
+  return getMemRefType(tensorType, options);
 }
 
 void bufferization::replaceOpWithBufferizedValues(RewriterBase &rewriter,
                                                   Operation *op,
                                                   ValueRange values) {
   assert(values.size() == op->getNumResults() &&
          "expected one value per OpResult");
   OpBuilder::InsertionGuard g(rewriter);

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

void mlir::bufferization::populateDynamicDimSizes(
   }
 }
 
 //===----------------------------------------------------------------------===//
 // AllocTensorOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult AllocTensorOp::bufferize(RewriterBase &rewriter,
-                                       BufferizationState &state) {
+                                       const BufferizationOptions &options) {
   OpBuilder::InsertionGuard g(rewriter);
   Location loc = getLoc();
 
   // Nothing to do for dead AllocTensorOps.
   if (getOperation()->getUses().empty()) {
     rewriter.eraseOp(getOperation());
     return success();
   }
 
   // Create buffer allocation.
   Value copyBuffer;
   if (copy())
-    copyBuffer = state.getBuffer(rewriter, copy());
+    copyBuffer = getBuffer(rewriter, copy(), options);
   auto allocType =
       MemRefType::get(getType().getShape(), getType().getElementType());
   SmallVector<Value> dynamicDims = dynamicSizes();
   if (copy()) {
     assert(dynamicDims.empty() && "expected either `copy` or `dynamicDims`");
     populateDynamicDimSizes(rewriter, loc, copyBuffer, dynamicDims);
   }
   FailureOr<Value> alloc =
-      state.getOptions().createAlloc(rewriter, loc, allocType, dynamicDims);
+      options.createAlloc(rewriter, loc, allocType, dynamicDims);
   if (failed(alloc))
     return failure();
 
   // Create memory copy (if any).
   if (copy()) {
-    if (failed(
-            state.getOptions().createMemCpy(rewriter, loc, copyBuffer, *alloc)))
+    if (failed(options.createMemCpy(rewriter, loc, copyBuffer, *alloc)))
       return failure();
   }
 
   // Should the buffer be deallocated?
-  AnalysisState analysisState(state.getOptions());
+  AnalysisState analysisState(options);
   bool dealloc;
   if (escape().hasValue()) {
     dealloc = !*escape();
   } else {
     // No "escape" annotation found.
-    if (state.getOptions().createDeallocs) {
+    if (options.createDeallocs) {
       // Perform an ad-hoc analysis.
       dealloc = !analysisState.isTensorYielded(getResult());
     } else {
       dealloc = false;
     }
   }
 
   // Replace op.
   replaceOpWithBufferizedValues(rewriter, getOperation(), *alloc);
 
   // Create buffer deallocation (if requested).
   if (!dealloc)
     return success();
 
   rewriter.setInsertionPoint(rewriter.getInsertionBlock()->getTerminator());
-  if (failed(state.getOptions().createDealloc(rewriter, loc, *alloc)))
+  if (failed(options.createDealloc(rewriter, loc, *alloc)))
     return failure();
   return success();
 }
 
 bool AllocTensorOp::isMemoryWrite(OpResult opResult,
                                   const AnalysisState &state) {
   // AllocTensorOps do not write unless they have a `copy` value.
   return static_cast<bool>(copy());

 
 void ToMemrefOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
   results.add<DimOfCastOp, LoadOfToMemref, ToMemrefOfCast, TensorLoadToMemref>(
       context);
 }
 
 LogicalResult ToMemrefOp::bufferize(RewriterBase &rewriter,
-                                    BufferizationState &state) {
+                                    const BufferizationOptions &options) {
   // Fold to_memref(to_tensor(x)) to x. Insert a cast if necessary.
   (void)foldToMemrefToTensorPair(rewriter, *this);
   // Note: The return value of `bufferize` indicates whether there was an error
   // or not. (And not whether the pattern matched or not.)
   return success();
 }
 
 Optional<Operation *> CloneOp::buildDealloc(OpBuilder &builder, Value alloc) {

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

op->walk<WalkOrder::PreOrder>([&](Operation *op) {
     if (hasTensorSemantics(op))
       worklist.push_back(op);
   });
 
   // Keep track of all erased ops.
   DenseSet<Operation *> erasedOps;
 
   // Bufferize all ops.
-  BufferizationState bufferizationState(options);
   BufferizationRewriter rewriter(op->getContext(), erasedOps, toMemrefOps,
                                  worklist, options, opFilter);
   for (unsigned i = 0; i < worklist.size(); ++i) {
     Operation *op = worklist[i];
     // Skip ops that were erased.
     if (erasedOps.contains(op))
       continue;
     // Skip ops that are not bufferizable or not allowed.
     auto bufferizableOp = options.dynCastBufferizableOp(op);
     if (!bufferizableOp)
       continue;
     if (opFilter && !opFilter->isOpAllowed(op))
       continue;
     // Skip ops that no longer have tensor semantics.
     if (!hasTensorSemantics(op))
       continue;
     // Bufferize the op.
     rewriter.setInsertionPoint(op);
-    if (failed(bufferizableOp.bufferize(rewriter, bufferizationState)))
+    if (failed(bufferizableOp.bufferize(rewriter, options)))
       return op->emitError("failed to bufferize op");
   }
 
   // Fold all to_memref(to_tensor(x)) pairs.
   for (Operation *op : toMemrefOps) {
     rewriter.setInsertionPoint(op);
     (void)bufferization::foldToMemrefToTensorPair(rewriter,
                                                   cast<ToMemrefOp>(op));
   }
 
   /// Check the result of bufferization. Return an error if an op was not
   /// bufferized, unless partial bufferization is allowed.
-  if (bufferizationState.getOptions().allowUnknownOps)
+  if (options.allowUnknownOps)
     return success();
 
   for (Operation *op : worklist) {
     // Skip ops that are entirely gone.
     if (erasedOps.contains(op))
       continue;
     // Ops that no longer have tensor semantics (because they were updated
     // in-place) are allowed.

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

#endif
       return BufferRelation::Equivalent;
     }
     return BufferRelation::None;
   }
 
   /// All function arguments are writable. It is the responsibility of the
   /// CallOp to insert buffer copies where necessary.
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     func::CallOp callOp = cast<func::CallOp>(op);
     unsigned numResults = callOp.getNumResults();
     unsigned numOperands = callOp->getNumOperands();
     FuncOp funcOp = getCalledFunction(callOp);
     assert(funcOp && "expected CallOp to a FuncOp");
     FunctionType funcType = funcOp.getFunctionType();
 
     // Result types of the bufferized CallOp.
for (OpOperand &opOperand : callOp->getOpOperands()) {
       if (!tensorOperand.getType().isa<TensorType>()) {
         newOperands[idx] = tensorOperand;
         continue;
       }
 
       // Retrieve buffers for tensor operands.
       Value buffer = newOperands[idx];
       if (!buffer)
-        buffer = state.getBuffer(rewriter, opOperand.get());
+        buffer = getBuffer(rewriter, opOperand.get(), options);
 
       // Caller / callee type mismatch is handled with a CastOp.
       auto memRefType = funcType.getInput(idx);
       // Since we don't yet have a clear layout story, to_memref may
       // conservatively turn tensors into more dynamic memref than necessary.
       // If the memref type of the callee fails, introduce an extra memref.cast
       // that will either canonicalize away or fail compilation until we can do
       // something better.
struct ReturnOpInterface
   }
 
   SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
 #ifndef NDEBUG
     auto returnOp = cast<func::ReturnOp>(op);
     assert(isa<FuncOp>(returnOp->getParentOp()) &&
            "only support FuncOp parent for ReturnOp");
 #endif // NDEBUG
 
     // ReturnOps are bufferized as part of FuncOps.
     return success();
   }
 };
 
 struct FuncOpInterface
     : public BufferizableOpInterface::ExternalModel<FuncOpInterface, FuncOp> {
   /// Rewrite function bbArgs and return values into buffer form. This function
   /// bufferizes the function signature and the ReturnOp. When the entire
   /// function body has been bufferized, function return types can be switched
   /// to more concise memref types as part of `foldMemRefCasts`.
   ///
   /// All function bbArgs are writable unless they are explicitly marked as
   /// read-only. Callers must insert copies when needed.
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto funcOp = cast<FuncOp>(op);
     FunctionType funcType = funcOp.getFunctionType();
-    const OneShotBufferizationOptions &options =
-        static_cast<const OneShotBufferizationOptions &>(state.getOptions());
 
     // Construct the bufferized function type.
     SmallVector<Type> argTypes;
     for (const auto &it : llvm::enumerate(funcType.getInputs())) {
       Type argType = it.value();
       if (auto tensorType = argType.dyn_cast<TensorType>()) {
         argTypes.push_back(
             getBufferizedFunctionArgType(funcOp, it.index(), options));

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

 
 LogicalResult mlir::bufferization::bufferizeModuleOp(
     ModuleOp moduleOp, const OneShotAnalysisState &analysisState) {
   auto const &options = static_cast<const OneShotBufferizationOptions &>(
       analysisState.getOptions());
   assert(options.bufferizeFunctionBoundaries &&
          "expected that function boundary bufferization is activated");
   IRRewriter rewriter(moduleOp.getContext());
-  BufferizationState bufferizationState(options);
 
   // A list of functions in the order in which they are analyzed + bufferized.
   SmallVector<func::FuncOp> orderedFuncOps;
 
   // A mapping of FuncOps to their callers.
   FuncCallerMap callerMap;
 
   if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))

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

 #include "mlir/IR/Operation.h"
 
 using namespace mlir;
 using namespace linalg;
 using namespace mlir::bufferization;
 
 namespace {
 
-// TODO: Ops in the linalg dialect can directly implement this interface.
-
 /// Generic conversion for any LinalgOp on tensors.
 static LogicalResult bufferizeLinalgOp(RewriterBase &rewriter, LinalgOp op,
-                                       BufferizationState &state) {
+                                       const BufferizationOptions &options) {
   // Take a guard before anything else.
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPoint(op);
 
   // Nothing to do. This op is already bufferized.
   if (op.hasBufferSemantics())
     return success();
 
   // Ensure op has only tensors. Allow mixed tensor-buffer mode on a per-need
   // basis.
   if (!op.hasTensorSemantics())
     return op->emitError() << "op does not have tensor semantics";
 
   // New input operands for the cloned op.
   SmallVector<Value> newInputBuffers;
   newInputBuffers.reserve(op.getNumInputs());
   for (OpOperand *opOperand : op.getInputOperands()) {
     if (op.isScalar(opOperand)) {
       newInputBuffers.push_back(opOperand->get());
       continue;
     }
-    newInputBuffers.push_back(state.getBuffer(rewriter, opOperand->get()));
+    newInputBuffers.push_back(getBuffer(rewriter, opOperand->get(), options));
   }
 
   // New output operands for the cloned op.
   SmallVector<Value> newOutputBuffers;
   for (OpResult opResult : op->getOpResults()) {
     OpOperand *opOperand = op.getOutputOperand(opResult.getResultNumber());
-    Value resultBuffer = state.getBuffer(rewriter, opOperand->get());
+    Value resultBuffer = getBuffer(rewriter, opOperand->get(), options);
     newOutputBuffers.push_back(resultBuffer);
   }
 
   // Merge input/output operands.
   SmallVector<Value> newOperands = newInputBuffers;
   newOperands.append(newOutputBuffers.begin(), newOutputBuffers.end());
 
   // Set insertion point now that potential alloc/dealloc are introduced.
struct LinalgOpInterface
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
-    return bufferizeLinalgOp(rewriter, cast<LinalgOp>(op), state);
+                          const BufferizationOptions &options) const {
+    return bufferizeLinalgOp(rewriter, cast<LinalgOp>(op), options);
   }
 };
 
 /// Helper structure that iterates over all LinalgOps in `OpTys` and registers
 /// the `BufferizableOpInterface` with each of them.
 template <typename... Ops>
 struct LinalgOpInterfaceHelper {
   static void registerOpInterface(MLIRContext *ctx) {

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

bool isMemoryWrite(Operation *op, OpResult opResult,
     // in the analysis. This is a useful pattern for all ops that have tensor
     // OpResults but no tensor OpOperands. By default, `isMemoryWrite` is
     // implemented in terms of `bufferizesToMemoryWrite`, which does not work on
     // ops without OpOperands.
     return true;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto executeRegionOp = cast<scf::ExecuteRegionOp>(op);
 
     // Compute new result types.
     SmallVector<Type> newResultTypes;
     for (Type type : executeRegionOp->getResultTypes()) {
       if (auto tensorType = type.dyn_cast<TensorType>()) {
         // TODO: Infer the result type instead of computing it.
-        newResultTypes.push_back(getMemRefType(tensorType, state.getOptions()));
+        newResultTypes.push_back(getMemRefType(tensorType, options));
       } else {
         newResultTypes.push_back(type);
       }
     }
 
     // Create new op and move over region.
     auto newOp =
         rewriter.create<scf::ExecuteRegionOp>(op->getLoc(), newResultTypes);
bool isMemoryWrite(Operation *op, OpResult opResult,
     // * lastWrite = %0  (Note: The last write of %r would be a set: {%0, %1}.)
     // * conflictingWrite = %1
     //
     // For more details, check the "scf.IfOp" section of the design document.
     return true;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto ifOp = cast<scf::IfOp>(op);
 
     // Compute new types of the bufferized scf.if op.
     SmallVector<Type> newTypes;
     for (Type returnType : ifOp->getResultTypes()) {
       if (auto tensorType = returnType.dyn_cast<TensorType>()) {
         // TODO: Infer the result type instead of computing it.
-        newTypes.push_back(getMemRefType(tensorType, state.getOptions()));
+        newTypes.push_back(getMemRefType(tensorType, options));
       } else {
         newTypes.push_back(returnType);
       }
     }
 
     // Create new op.
     auto newIfOp =
         rewriter.create<scf::IfOp>(ifOp.getLoc(), newTypes, ifOp.getCondition(),
assert(memref::CastOp::areCastCompatible(buffer.getType(), type) &&
          "scf.while op bufferization: cast incompatible");
   return b.create<memref::CastOp>(buffer.getLoc(), type, buffer).getResult();
 }
 
 /// Helper function for loop bufferization. Return the bufferized values of the
 /// given OpOperands. If an operand is not a tensor, return the original value.
 static SmallVector<Value> getBuffers(RewriterBase &rewriter,
                                      MutableArrayRef<OpOperand> operands,
-                                     BufferizationState &state) {
+                                     const BufferizationOptions &options) {
   SmallVector<Value> result;
   for (OpOperand &opOperand : operands) {
     if (opOperand.get().getType().isa<TensorType>()) {
-      Value resultBuffer = state.getBuffer(rewriter, opOperand.get());
+      Value resultBuffer = getBuffer(rewriter, opOperand.get(), options);
       result.push_back(resultBuffer);
     } else {
       result.push_back(opOperand.get());
     }
   }
   return result;
 }
 
 /// Helper function for loop bufferization. Compute the buffer that should be
 /// yielded from a loop block (loop body or loop condition).
 static Value getYieldedBuffer(RewriterBase &rewriter, Value tensor,
-                              BaseMemRefType type, BufferizationState &state) {
+                              BaseMemRefType type,
+                              const BufferizationOptions &options) {
   assert(tensor.getType().isa<TensorType>() && "expected tensor");
   ensureToMemrefOpIsValid(tensor, type);
-  Value yieldedVal = state.getBuffer(rewriter, tensor);
+  Value yieldedVal = getBuffer(rewriter, tensor, options);
   return castBuffer(rewriter, yieldedVal, type);
 }
 
 /// Helper function for loop bufferization. Given a range of values, apply
 /// `func` to those marked in `tensorIndices`. Otherwise, store the unmodified
 /// value in the result vector.
 static SmallVector<Value>
 convertTensorValues(ValueRange values, const DenseSet<int64_t> &tensorIndices,
                     llvm::function_ref<Value(Value, int64_t)> func) {
   SmallVector<Value> result;
   for (const auto &it : llvm::enumerate(values)) {
     size_t idx = it.index();
     Value val = it.value();
     result.push_back(tensorIndices.contains(idx) ? func(val, idx) : val);
   }
   return result;
 }
 
 /// Helper function for loop bufferization. Given a list of pre-bufferization
 /// yielded values, compute the list of bufferized yielded values.
 SmallVector<Value> getYieldedValues(RewriterBase &rewriter, ValueRange values,
                                     TypeRange bufferizedTypes,
                                     const DenseSet<int64_t> &tensorIndices,
-                                    BufferizationState &state) {
+                                    const BufferizationOptions &options) {
   return convertTensorValues(
       values, tensorIndices, [&](Value val, int64_t index) {
         return getYieldedBuffer(rewriter, val,
                                 bufferizedTypes[index].cast<BaseMemRefType>(),
-                                state);
+                                options);
       });
 }
 
 /// Helper function for loop bufferization. Given a list of bbArgs of the new
 /// (bufferized) loop op, wrap the bufferized tensor args (now memrefs) into
 /// ToTensorOps, so that the block body can be moved over to the new op.
 SmallVector<Value>
 getBbArgReplacements(RewriterBase &rewriter, Block::BlockArgListType bbArgs,
LogicalResult resolveConflicts(Operation *op, RewriterBase &rewriter,
     }
 
     rewriter.updateRootInPlace(
         yieldOp, [&]() { yieldOp.getResultsMutable().assign(yieldValues); });
     return success();
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto forOp = cast<scf::ForOp>(op);
     Block *oldLoopBody = &forOp.getLoopBody().front();
 
     // Indices of all iter_args that have tensor type. These are the ones that
     // are bufferized.
     DenseSet<int64_t> indices = getTensorIndices(forOp.getInitArgs());
 
     // The new memref init_args of the loop.
     SmallVector<Value> initArgs =
-        getBuffers(rewriter, forOp.getIterOpOperands(), state);
+        getBuffers(rewriter, forOp.getIterOpOperands(), options);
 
     // Construct a new scf.for op with memref instead of tensor values.
     auto newForOp = rewriter.create<scf::ForOp>(
         forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(),
         forOp.getStep(), initArgs);
     newForOp->setAttrs(forOp->getAttrs());
     ValueRange initArgsRange(initArgs);
     TypeRange initArgsTypes(initArgsRange);
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
 
     // Move loop body to new loop.
     rewriter.mergeBlocks(oldLoopBody, loopBody, iterArgs);
 
     // Update scf.yield of new loop.
     auto yieldOp = cast<scf::YieldOp>(loopBody->getTerminator());
     rewriter.setInsertionPoint(yieldOp);
     SmallVector<Value> yieldValues = getYieldedValues(
-        rewriter, yieldOp.getResults(), initArgsTypes, indices, state);
+        rewriter, yieldOp.getResults(), initArgsTypes, indices, options);
     yieldOp.getResultsMutable().assign(yieldValues);
 
     // Replace loop results.
     replaceOpWithBufferizedValues(rewriter, op, newForOp->getResults());
 
     return success();
   }
 
LogicalResult resolveConflicts(Operation *op, RewriterBase &rewriter,
     rewriter.updateRootInPlace(yieldOp, [&]() {
       yieldOp.getResultsMutable().assign(afterYieldValues);
     });
 
     return success();
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto whileOp = cast<scf::WhileOp>(op);
 
     assert(whileOp.getBefore().getBlocks().size() == 1 &&
            "regions with multiple blocks not supported");
     Block *beforeBody = &whileOp.getBefore().front();
     assert(whileOp.getAfter().getBlocks().size() == 1 &&
            "regions with multiple blocks not supported");
     Block *afterBody = &whileOp.getAfter().front();
 
     // Indices of all bbArgs that have tensor type. These are the ones that
     // are bufferized. The "before" and "after" regions may have different args.
     DenseSet<int64_t> indicesBefore = getTensorIndices(whileOp.getInits());
     DenseSet<int64_t> indicesAfter =
         getTensorIndices(whileOp.getAfterArguments());
 
     // The new memref init_args of the loop.
     SmallVector<Value> initArgs =
-        getBuffers(rewriter, whileOp->getOpOperands(), state);
+        getBuffers(rewriter, whileOp->getOpOperands(), options);
 
     // The result types of a WhileOp are the same as the "after" bbArg types.
     SmallVector<Type> argsTypesAfter = llvm::to_vector(
         llvm::map_range(whileOp.getAfterArguments(), [&](BlockArgument bbArg) {
-          return state.getBufferType(bbArg).cast<Type>();
+          return getBufferType(bbArg, options).cast<Type>();
         }));
 
     // Construct a new scf.while op with memref instead of tensor values.
     ValueRange argsRangeBefore(initArgs);
     TypeRange argsTypesBefore(argsRangeBefore);
     auto newWhileOp = rewriter.create<scf::WhileOp>(whileOp.getLoc(),
                                                     argsTypesAfter, initArgs);
 
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     // Update scf.condition of new loop.
     auto newConditionOp = newWhileOp.getConditionOp();
     rewriter.setInsertionPoint(newConditionOp);
     // Only equivalent buffers or new buffer allocations may be yielded to the
     // "after" region.
     // TODO: This could be relaxed for better bufferization results.
     SmallVector<Value> newConditionArgs =
         getYieldedValues(rewriter, newConditionOp.getArgs(), argsTypesAfter,
-                         indicesAfter, state);
+                         indicesAfter, options);
     newConditionOp.getArgsMutable().assign(newConditionArgs);
 
     // Set up new iter_args and move the loop body block to the new op.
     // The old block uses tensors, so wrap the (memref) bbArgs of the new block
     // in ToTensorOps.
     rewriter.setInsertionPointToStart(newAfterBody);
     SmallVector<Value> newAfterArgs = getBbArgReplacements(
         rewriter, newWhileOp.getAfterArguments(), indicesAfter);
     rewriter.mergeBlocks(afterBody, newAfterBody, newAfterArgs);
 
     // Update scf.yield of the new loop.
     auto newYieldOp = newWhileOp.getYieldOp();
     rewriter.setInsertionPoint(newYieldOp);
     // Only equivalent buffers or new buffer allocations may be yielded to the
     // "before" region.
     // TODO: This could be relaxed for better bufferization results.
     SmallVector<Value> newYieldValues =
         getYieldedValues(rewriter, newYieldOp.getResults(), argsTypesBefore,
-                         indicesBefore, state);
+                         indicesBefore, options);
     newYieldOp.getResultsMutable().assign(newYieldValues);
 
     // Replace loop results.
     replaceOpWithBufferizedValues(rewriter, op, newWhileOp->getResults());
 
     return success();
   }
 
bool mustBufferizeInPlace(Operation *op, OpOperand &opOperand,
                             const AnalysisState &state) const {
     // Yield operands always bufferize inplace. Otherwise, an alloc + copy
     // may be generated inside the block. We should not return/yield allocations
     // when possible.
     return true;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto yieldOp = cast<scf::YieldOp>(op);
     if (!isa<scf::ExecuteRegionOp, scf::IfOp, scf::ForOp, scf::WhileOp>(
             yieldOp->getParentOp()))
       return yieldOp->emitError("unsupported scf::YieldOp parent");
     return success();
   }
 };
 
for (OpResult opResult : foreachThreadOp->getOpResults()) {
       rewriter.updateRootInPlace(destOperands.front()->getOwner(),
                                  [&]() { destOperands.front()->set(alloc); });
     }
 
     return success();
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &b,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     OpBuilder::InsertionGuard g(b);
     auto foreachThreadOp = cast<ForeachThreadOp>(op);
 
     // Gather new results of the ForeachThreadOp.
     SmallVector<Value> newResults;
     for (OpResult opResult : foreachThreadOp->getOpResults()) {
       OpOperand *insertDest =
           getInsertionDest(foreachThreadOp)[opResult.getResultNumber()];
       // Insert copies right before the PerformConcurrentlyOp terminator. They
       // should not be inside terminator (which would be the default insertion
       // point).
-      Value buffer = state.getBuffer(b, insertDest->get());
+      Value buffer = getBuffer(b, insertDest->get(), options);
       newResults.push_back(buffer);
     }
 
     // Create new ForeachThreadOp without any results and drop the automatically
     // introduced terminator.
     TypeRange newResultTypes;
     auto newForeachThreadOp =
         b.create<ForeachThreadOp>(foreachThreadOp.getLoc(), newResultTypes,
                                   foreachThreadOp.getNumThreads());
     newForeachThreadOp.getBody()->getTerminator()->erase();
 
     // Move over block contents of the old op.
     b.mergeBlocks(foreachThreadOp.getBody(), newForeachThreadOp.getBody(),
                   {newForeachThreadOp.getBody()->getArguments()});
 
     // Bufferize terminator.
     auto performConcurrentlyOp = cast<PerformConcurrentlyOp>(
         newForeachThreadOp.getBody()->getTerminator());
     b.setInsertionPoint(performConcurrentlyOp);
     unsigned resultCounter = 0;
     WalkResult walkResult =
         performConcurrentlyOp.walk([&](ParallelInsertSliceOp insertOp) {
           Location loc = insertOp.getLoc();
           Type srcType = getMemRefType(
-              insertOp.getSource().getType().cast<RankedTensorType>(),
-              state.getOptions());
+              insertOp.getSource().getType().cast<RankedTensorType>(), options);
           // ParallelInsertSliceOp bufferizes to a copy.
           auto srcMemref = b.create<bufferization::ToMemrefOp>(
               loc, srcType, insertOp.getSource());
           Value destMemref = newResults[resultCounter++];
           Value subview = b.create<memref::SubViewOp>(
               loc, destMemref, insertOp.getMixedOffsets(),
               insertOp.getMixedSizes(), insertOp.getMixedStrides());
           // This memcpy will fold away if everything bufferizes in-place.
-          if (failed(state.getOptions().createMemCpy(b, insertOp.getLoc(),
-                                                     srcMemref, subview)))
+          if (failed(options.createMemCpy(b, insertOp.getLoc(), srcMemref,
+                                          subview)))
             return WalkResult::interrupt();
           b.eraseOp(insertOp);
           return WalkResult::advance();
         });
     if (walkResult.wasInterrupted())
       return failure();
 
     // Replace the op.
     replaceOpWithBufferizedValues(b, op, newResults);
 
     return success();
   }
 };
 
 /// Nothing to do for PerformConcurrentlyOp.
 struct PerformConcurrentlyOpInterface
     : public BufferizableOpInterface::ExternalModel<
           PerformConcurrentlyOpInterface, PerformConcurrentlyOp> {
   LogicalResult bufferize(Operation *op, RewriterBase &b,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     llvm_unreachable("op does not have any tensor OpOperands / OpResults");
     return failure();
   }
 };
 
 /// Return true if the (ExtractSliceOp, ParallelInsertSliceOp) pair match (i.e.
 /// equivalent operand / result and same offset/sizes/strides specification).
 static bool areEquivalentExtractSliceOps(const AnalysisState &state,
struct ParallelInsertSliceOpInterface
   }
 
   LogicalResult resolveConflicts(Operation *op, RewriterBase &rewriter,
                                  const AnalysisState &state) const {
     return success();
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &b,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     // Will be bufferized as part of ForeachThreadOp.
     return failure();
   }
 
   // TODO: This is copied from TensorInterfaceImpl.cpp. Find a way to share
   // the code.
   bool isNotConflicting(Operation *op, OpOperand *uRead,
                         OpOperand *uConflictingWrite,

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

bool isMemoryWrite(Operation *op, OpResult opResult,
     // in the analysis. This is a useful pattern for all ops that have tensor
     // OpResults but no tensor OpOperands. By default, `isMemoryWrite` is
     // implemented in terms of `bufferizesToMemoryWrite`, which does not work on
     // ops without OpOperands.
     return true;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto assumingOp = cast<shape::AssumingOp>(op);
 
     // Compute new result types.
     SmallVector<Type> newResultTypes;
     for (Type type : assumingOp->getResultTypes()) {
       if (auto tensorType = type.dyn_cast<TensorType>()) {
         // TODO: Infer the result type instead of computing it.
-        newResultTypes.push_back(getMemRefType(tensorType, state.getOptions()));
+        newResultTypes.push_back(getMemRefType(tensorType, options));
       } else {
         newResultTypes.push_back(type);
       }
     }
 
     // Create new op and move over region.
     auto newOp = rewriter.create<shape::AssumingOp>(
         op->getLoc(), newResultTypes, assumingOp.getWitness());
bool mustBufferizeInPlace(Operation *op, OpOperand &opOperand,
                             const AnalysisState &state) const {
     // Yield operands always bufferize inplace. Otherwise, an alloc + copy
     // may be generated inside the block. We should not return/yield allocations
     // when possible.
     return true;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     // Op is bufferized as part of AssumingOp.
     return failure();
   }
 };
 
 } // namespace
 } // namespace shape
 } // namespace mlir
 
 void mlir::shape::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, shape::ShapeDialect *dialect) {
     shape::AssumingOp::attachInterface<AssumingOpInterface>(*ctx);
     shape::AssumingYieldOp::attachInterface<AssumingYieldOpInterface>(*ctx);
   });
 }

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

struct CastOpInterface
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto castOp = cast<tensor::CastOp>(op);
 
     // The result buffer still has the old (pre-cast) type.
-    Value resultBuffer = state.getBuffer(rewriter, castOp.source());
+    Value resultBuffer = getBuffer(rewriter, castOp.source(), options);
     auto sourceMemRefType = resultBuffer.getType().cast<BaseMemRefType>();
     Attribute memorySpace = sourceMemRefType.getMemorySpace();
     TensorType resultTensorType =
         castOp.getResult().getType().cast<TensorType>();
     MemRefLayoutAttrInterface layout;
 
     if (auto rankedMemRefType = sourceMemRefType.dyn_cast<MemRefType>())
       if (resultTensorType.isa<RankedTensorType>())
         layout = rankedMemRefType.getLayout();
 
     // Compute the new memref type.
-    Type resultMemRefType = getMemRefType(resultTensorType, state.getOptions(),
-                                          layout, memorySpace);
+    Type resultMemRefType =
+        getMemRefType(resultTensorType, options, layout, memorySpace);
 
     // Replace the op with a memref.cast.
     assert(memref::CastOp::areCastCompatible(resultBuffer.getType(),
                                              resultMemRefType) &&
            "CallOp::bufferize: cast incompatible");
     replaceOpWithNewBufferizedOp<memref::CastOp>(rewriter, op, resultMemRefType,
                                                  resultBuffer);
 
struct CollapseShapeOpInterface
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto collapseShapeOp = cast<tensor::CollapseShapeOp>(op);
     RankedTensorType tensorResultType = collapseShapeOp.getResultType();
-    Value buffer = state.getBuffer(rewriter, collapseShapeOp.src());
+    Value buffer = getBuffer(rewriter, collapseShapeOp.src(), options);
     auto bufferType = buffer.getType().cast<MemRefType>();
 
     if (tensorResultType.getRank() == 0) {
       // 0-d collapses must go through a different op builder.
       MemRefType resultType;
 
       if (bufferType.getLayout().isIdentity()) {
         // Standard layout: result type has no offset.
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
 
     // If the dims are not collapsible (due to an incompatible source layout
     // map), force an out-of-place bufferization, i.e., a buffer copy. This
     // newly allocated buffer will have no layout map and thus be collapsible.
     bool canBeCollapsed = memref::CollapseShapeOp::isGuaranteedCollapsible(
         bufferType, collapseShapeOp.getReassociationIndices());
     if (!canBeCollapsed) {
       // TODO: Create alloc_tensor ops during TensorCopyInsertion.
-      AnalysisState analysisState(state.getOptions());
+      AnalysisState analysisState(options);
       Value tensorAlloc = allocateTensorForShapedValue(
           rewriter, op->getLoc(), collapseShapeOp.src(),
           analysisState.isTensorYielded(collapseShapeOp.result()));
       auto memrefType =
           MemRefType::get(collapseShapeOp.getSrcType().getShape(),
                           collapseShapeOp.getSrcType().getElementType(),
                           AffineMap(), bufferType.getMemorySpaceAsInt());
       buffer = rewriter.create<bufferization::ToMemrefOp>(
struct DimOpInterface
   }
 
   SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto dimOp = cast<tensor::DimOp>(op);
-    auto v = state.getBuffer(rewriter, dimOp.source());
+    auto v = getBuffer(rewriter, dimOp.source(), options);
     replaceOpWithNewBufferizedOp<memref::DimOp>(rewriter, op, v, dimOp.index());
     return success();
   }
 };
 
 /// Bufferization of tensor.expand_shape. Replace with memref.expand_shape.
 struct ExpandShapeOpInterface
     : public BufferizableOpInterface::ExternalModel<ExpandShapeOpInterface,
struct ExpandShapeOpInterface
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto expandShapeOp = cast<tensor::ExpandShapeOp>(op);
     auto tensorResultType = expandShapeOp.getResultType();
-    auto buffer = state.getBuffer(rewriter, expandShapeOp.src());
+    auto buffer = getBuffer(rewriter, expandShapeOp.src(), options);
 
     // Memref result type is inferred by the builder based on reassociation
     // indices and result shape.
     replaceOpWithNewBufferizedOp<memref::ExpandShapeOp>(
         rewriter, op, tensorResultType.getShape(), buffer,
         expandShapeOp.getReassociationIndices());
     return success();
   }
struct ExtractSliceOpInterface
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::None;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
     Location loc = extractSliceOp.getLoc();
 
     // Even if this op was decided to bufferize out-of-place, do not insert the
     // buffer copy yet. This is done later in this function.
-    auto srcMemref = state.getBuffer(rewriter, extractSliceOp.source());
+    auto srcMemref = getBuffer(rewriter, extractSliceOp.source(), options);
     auto srcMemrefType = srcMemref.getType().cast<MemRefType>();
     auto dstTensorType =
         extractSliceOp.result().getType().cast<RankedTensorType>();
 
     // Expand offsets, sizes and strides to the full rank to handle the
     // rank-reducing case.
     SmallVector<OpFoldResult> mixedOffsets = extractSliceOp.getMixedOffsets();
     SmallVector<OpFoldResult> mixedSizes = extractSliceOp.getMixedSizes();
struct ExtractOpInterface
   }
 
   SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto extractOp = cast<tensor::ExtractOp>(op);
-    Value srcMemref = state.getBuffer(rewriter, extractOp.tensor());
+    Value srcMemref = getBuffer(rewriter, extractOp.tensor(), options);
     replaceOpWithNewBufferizedOp<memref::LoadOp>(rewriter, op, srcMemref,
                                                  extractOp.indices());
     return success();
   }
 };
 
 // Implements backtracking to traverse indices of the output buffer while
 // iterating over op.elements().
static void createStores(RewriterBase &rewriter, Location loc, int dim,
   }
 }
 
 /// Bufferization of tensor.from_elements.
 struct FromElementsOpInterface
     : public BufferizableOpInterface::ExternalModel<FromElementsOpInterface,
                                                     tensor::FromElementsOp> {
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto fromElementsOp = cast<tensor::FromElementsOp>(op);
 
     // Allocate a buffer for the result.
     Location loc = op->getLoc();
     auto tensorType = fromElementsOp.getType().cast<RankedTensorType>();
     auto shape = tensorType.getShape();
     // TODO: Create alloc_tensor ops during TensorCopyInsertion.
-    AnalysisState analysisState(state.getOptions());
+    AnalysisState analysisState(options);
     Value tensorAlloc = allocateTensorForShapedValue(
         rewriter, loc, fromElementsOp.result(),
         analysisState.isTensorYielded(fromElementsOp.result()),
         /*copy=*/false);
     auto memrefType =
         MemRefType::get(tensorType.getShape(), tensorType.getElementType());
     Value buffer = rewriter.create<bufferization::ToMemrefOp>(
         op->getLoc(), memrefType, tensorAlloc);
struct FromElementsOpInterface
   }
 };
 
 /// Bufferization of tensor.generate.
 struct GenerateOpInterface
     : public BufferizableOpInterface::ExternalModel<GenerateOpInterface,
                                                     tensor::GenerateOp> {
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto generateOp = cast<tensor::GenerateOp>(op);
     auto tensorType = generateOp.getType().cast<RankedTensorType>();
     // Allocate memory.
     Location loc = op->getLoc();
     // TODO: Create alloc_tensor ops during TensorCopyInsertion.
-    AnalysisState analysisState(state.getOptions());
+    AnalysisState analysisState(options);
     Value tensorAlloc = allocateTensorForShapedValue(
         rewriter, loc, generateOp.result(),
         analysisState.isTensorYielded(generateOp.result()),
         /*copy=*/false);
     auto memrefType =
         MemRefType::get(tensorType.getShape(), tensorType.getElementType());
     Value buffer = rewriter.create<bufferization::ToMemrefOp>(
         op->getLoc(), memrefType, tensorAlloc);
struct InsertOpInterface
 
   SmallVector<OpOperand *>
   getAliasingOpOperand(Operation *op, OpResult opResult,
                        const AnalysisState &state) const {
     return {&op->getOpOperand(1) /*dest*/};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto insertOp = cast<tensor::InsertOp>(op);
-    Value destMemref = state.getBuffer(rewriter, insertOp.dest());
+    Value destMemref = getBuffer(rewriter, insertOp.dest(), options);
     rewriter.create<memref::StoreOp>(insertOp.getLoc(), insertOp.scalar(),
                                      destMemref, insertOp.indices());
     replaceOpWithBufferizedValues(rewriter, op, destMemref);
     return success();
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
if (auto insertSliceOp = dyn_cast<InsertSliceOp>(conflictingWritingOp))
           hasMatchingExtractSliceOp(state, insertSliceOp.source(),
                                     insertSliceOp))
         return true;
 
     return false;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     // insert_slice ops arise from tiling and bufferizing them out-of-place is
     // generally a deal breaker. When used with loops, this ends up cloning the
     // whole tensor on every single iteration and is a symptom of a
     // catastrophically bad scheduling decision.
     // TODO: be very loud about it or even consider failing the pass.
     auto insertSliceOp = cast<tensor::InsertSliceOp>(op);
     Location loc = insertSliceOp.getLoc();
-    Value dstMemref = state.getBuffer(rewriter, insertSliceOp.dest());
+    Value dstMemref = getBuffer(rewriter, insertSliceOp.dest(), options);
 
     // Expand offsets, sizes and strides to the full rank to handle the
     // rank-reducing case.
     SmallVector<OpFoldResult> mixedOffsets = insertSliceOp.getMixedOffsets();
     SmallVector<OpFoldResult> mixedSizes = insertSliceOp.getMixedSizes();
     SmallVector<OpFoldResult> mixedStrides = insertSliceOp.getMixedStrides();
     OffsetSizeAndStrideOpInterface::expandToRank(
         dstMemref, mixedOffsets, mixedSizes, mixedStrides,
auto subviewMemRefType =
             mixedOffsets, mixedSizes, mixedStrides)
             .cast<MemRefType>();
     Value subView = rewriter.create<memref::SubViewOp>(
         loc, subviewMemRefType, dstMemref, mixedOffsets, mixedSizes,
         mixedStrides);
 
     // Copy tensor. If this tensor.insert_slice has a matching
     // tensor.extract_slice, the copy operation will eventually fold away.
-    auto srcMemref = state.getBuffer(rewriter, insertSliceOp.source());
-    if (failed(
-            state.getOptions().createMemCpy(rewriter, loc, srcMemref, subView)))
+    auto srcMemref = getBuffer(rewriter, insertSliceOp.source(), options);
+    if (failed(options.createMemCpy(rewriter, loc, srcMemref, subView)))
       return failure();
 
     replaceOpWithBufferizedValues(rewriter, op, dstMemref);
     return success();
   }
 };
 
 /// Bufferization of tensor.rank. Replace with memref.rank.
struct RankOpInterface
   }
 
   SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto rankOp = cast<tensor::RankOp>(op);
-    auto v = state.getBuffer(rewriter, rankOp.tensor());
+    auto v = getBuffer(rewriter, rankOp.tensor(), options);
     replaceOpWithNewBufferizedOp<memref::RankOp>(rewriter, op, rankOp.getType(),
                                                  v);
     return success();
   }
 };
 
 /// Bufferization of tensor.reshape. Replace with memref.reshape.
 struct ReshapeOpInterface
struct ReshapeOpInterface
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto reshapeOp = cast<tensor::ReshapeOp>(op);
-    Value srcBuffer = state.getBuffer(rewriter, reshapeOp.source());
-    Value shapeBuffer = state.getBuffer(rewriter, reshapeOp.shape());
+    Value srcBuffer = getBuffer(rewriter, reshapeOp.source(), options);
+    Value shapeBuffer = getBuffer(rewriter, reshapeOp.shape(), options);
     auto resultTensorType = reshapeOp.getResult().getType().cast<TensorType>();
-    auto resultMemRefType = getMemRefType(resultTensorType, state.getOptions());
+    auto resultMemRefType = getMemRefType(resultTensorType, options);
     replaceOpWithNewBufferizedOp<memref::ReshapeOp>(
         rewriter, op, resultMemRefType, srcBuffer, shapeBuffer);
     return success();
   }
 };
 
 } // namespace
 } // namespace tensor

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

struct TransferReadOpInterface
   }
 
   SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     return {};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto readOp = cast<vector::TransferReadOp>(op);
     assert(readOp.getShapedType().isa<TensorType>() &&
            "only tensor types expected");
-    Value buffer = state.getBuffer(rewriter, readOp.getSource());
+    Value buffer = getBuffer(rewriter, readOp.getSource(), options);
     replaceOpWithNewBufferizedOp<vector::TransferReadOp>(
         rewriter, readOp, readOp.getVectorType(), buffer, readOp.getIndices(),
         readOp.getPermutationMap(), readOp.getPadding(), readOp.getMask(),
         readOp.getInBoundsAttr());
     return success();
   }
 };
 
struct TransferWriteOpInterface
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          BufferizationState &state) const {
+                          const BufferizationOptions &options) const {
     auto writeOp = cast<vector::TransferWriteOp>(op);
     assert(writeOp.getShapedType().isa<TensorType>() &&
            "only tensor types expected");
 
     // Create a new transfer_write on buffer that doesn't have a return value.
-    Value resultBuffer = state.getBuffer(rewriter, writeOp.getSource());
+    Value resultBuffer = getBuffer(rewriter, writeOp.getSource(), options);
     rewriter.create<vector::TransferWriteOp>(
         writeOp.getLoc(), writeOp.getVector(), resultBuffer,
         writeOp.getIndices(), writeOp.getPermutationMapAttr(),
         writeOp.getInBoundsAttr());
     replaceOpWithBufferizedValues(rewriter, op, resultBuffer);
 
     return success();
   }