[mlir][bufferize][NFC] Split BufferizationState into AnalysisState/BufferizationState

[springerm]

Summary

Depends On D121518

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

 #include "llvm/ADT/SetVector.h"
 
 namespace mlir {
 class BlockAndValueMapping;
 class DominanceInfo;
 
 namespace bufferization {
 
+class AnalysisState;
 class BufferizableOpInterface;
-class BufferizationState;
-struct DialectBufferizationState;
+struct DialectAnalysisState;
 
-/// Options for ComprehensiveBufferize.
+/// Options for BufferizableOpInterface-based bufferization.
 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.
+  /// Initializer function for analysis state.
+  using AnalysisStateInitFn = std::function<void(AnalysisState &)>;
+  /// Initializer function for dialect-specific analysis state.
   using DialectStateInitFn =
-      std::function<std::unique_ptr<DialectBufferizationState>()>;
+      std::function<std::unique_ptr<DialectAnalysisState>()>;
 
   /// 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;
+  /// Initializer functions for analysis state. These can be used to
+  /// initialize dialect-specific analysis state.
+  SmallVector<AnalysisStateInitFn> stateInitializers;
 
-  /// Add a bufferization state initializer that initializes the specified
-  /// dialect-specific bufferization state.
+  /// Add a analysis state initializer that initializes the specified
+  /// dialect-specific analysis state.
   void addDialectStateInitializer(StringRef name, const DialectStateInitFn &fn);
 
 private:
   /// Allow a dialect.
   template <typename DialectT>
   void allowDialectInFilterImpl() {
     allowDialectInFilter(DialectT::getDialectNamespace());
   }
enum class BufferRelation {
   // TODO: ResultContainsOperand,
   // TODO: OperandContainsResult,
   Equivalent
 };
 
 /// Return `true` if the given value is a BlockArgument of a FuncOp.
 bool isFunctionArgument(Value value);
 
-/// Dialect-specific bufferization state. Analysis/bufferization information
+/// Dialect-specific analysis state. Analysis/bufferization information
 /// that is specific to ops from a certain dialect can be stored in derived
 /// variants of this struct.
-struct DialectBufferizationState {
-  DialectBufferizationState() = default;
+struct DialectAnalysisState {
+  DialectAnalysisState() = default;
 
-  virtual ~DialectBufferizationState() = default;
+  virtual ~DialectAnalysisState() = default;
 
   // Copying state is forbidden. Always pass as reference.
-  DialectBufferizationState(const DialectBufferizationState &) = delete;
+  DialectAnalysisState(const DialectAnalysisState &) = delete;
 };
 
-/// BufferizationState provides a variety of helper functions for dealing with
-/// tensor values and memref buffers.
-class BufferizationState {
+/// AnalysisState provides a variety of helper functions for dealing with
+/// tensor values.
+class AnalysisState {
 public:
   /// Determine which OpOperand* will alias with `result` if the op is
   /// bufferized in place. Return an empty vector if the op is not bufferizable.
   SmallVector<OpOperand *> getAliasingOpOperand(OpResult result) const;
 
   /// Determine which OpResult will alias with `opOperand` if the op is
   /// bufferized in place. Return an empty vector if the op is not bufferizable.
   SmallVector<OpResult> getAliasingOpResult(OpOperand &opOperand) const;
public:
   SetVector<Value> findLastPrecedingWrite(Value value) const;
 
   /// Return `true` if the given OpResult has been decided to bufferize inplace.
   virtual bool isInPlace(OpOperand &opOperand) const = 0;
 
   /// Return true if `v1` and `v2` bufferize to equivalent buffers.
   virtual bool areEquivalentBufferizedValues(Value v1, Value v2) const = 0;
 
-  /// Return the buffer (memref) for a given OpOperand (tensor). Allocate
-  /// a new buffer and copy over data from the existing buffer if out-of-place
-  /// bufferization was decided.
-  FailureOr<Value>
-  getBuffer(RewriterBase &rewriter, OpOperand &opOperand,
-            bool forceInPlace = false,
-            Optional<Operation *> customCopyInsertionPoint = None) const;
-
-  /// Return dialect-specific bufferization state.
+  /// Return dialect-specific analysis state.
   template <typename StateT>
   Optional<const StateT *> getDialectState(StringRef name) const {
     auto it = dialectState.find(name);
     if (it == dialectState.end())
       return None;
     return static_cast<const StateT *>(it->getSecond().get());
   }
 
-  /// Return dialect-specific bufferization state or create one if none exists.
+  /// Return dialect-specific analysis state or create one if none exists.
   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) {
+                          std::unique_ptr<DialectAnalysisState> 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);
+  explicit AnalysisState(const BufferizationOptions &options);
 
-  // BufferizationState should be passed as a reference.
-  BufferizationState(const BufferizationState &) = delete;
+  // AnalysisState should be passed as a reference.
+  AnalysisState(const AnalysisState &) = delete;
 
-  ~BufferizationState() = default;
+  ~AnalysisState() = default;
 
 private:
-  /// Dialect-specific bufferization state.
-  DenseMap<StringRef, std::unique_ptr<DialectBufferizationState>> dialectState;
+  /// Dialect-specific analysis state.
+  DenseMap<StringRef, std::unique_ptr<DialectAnalysisState>> dialectState;
 
   /// A reference to current bufferization options.
   const BufferizationOptions &options;
 };
 
-/// This a "no analysis, always copy" BufferizationState. In the absence of an
+/// This a "no analysis, always copy" AnalysisState. In the absence of an
 /// analysis, a buffer must be copied each time it is written to. Therefore, all
 /// OpOperands that bufferize to a memory write must bufferize out-of-place.
-class AlwaysCopyBufferizationState : public BufferizationState {
+class AlwaysCopyAnalysisState : public AnalysisState {
 public:
-  explicit AlwaysCopyBufferizationState(const BufferizationOptions &options);
+  explicit AlwaysCopyAnalysisState(const BufferizationOptions &options);
 
-  AlwaysCopyBufferizationState(const AlwaysCopyBufferizationState &) = delete;
+  AlwaysCopyAnalysisState(const AlwaysCopyAnalysisState &) = delete;
 
-  virtual ~AlwaysCopyBufferizationState() = default;
+  virtual ~AlwaysCopyAnalysisState() = default;
 
   /// Return `true` if the given OpResult has been decided to bufferize inplace.
   bool isInPlace(OpOperand &opOperand) const override;
 
   /// Return true if `v1` and `v2` bufferize to equivalent buffers.
   bool areEquivalentBufferizedValues(Value v1, Value v2) const override;
 };
 
+/// BufferizationState provides helper functions for performing bufferization
+/// rewrites and handling memref buffers.
+struct BufferizationState {
+  BufferizationState(const AnalysisState &analysisState)
+      : analysisState(analysisState) {}
+
+  /// Return the buffer (memref) for a given OpOperand (tensor). Allocate
+  /// a new buffer and copy over data from the existing buffer if out-of-place
+  /// bufferization was decided.
+  FailureOr<Value>
+  getBuffer(RewriterBase &rewriter, OpOperand &opOperand,
+            bool forceInPlace = false,
+            Optional<Operation *> customCopyInsertionPoint = None) const;
+
+  /// Return a reference to the BufferizationOptions.
+  const BufferizationOptions &getOptions() const {
+    return analysisState.getOptions();
+  }
+
+  const AnalysisState &getAnalysisState() const { return analysisState; }
+
+protected:
+  // BufferizationState should be passed as a reference.
+  BufferizationState(const BufferizationState &) = delete;
+
+private:
+  const AnalysisState &analysisState;
+};
+
 /// 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);
 
 /// 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.

 /// BufferizableOpInterface for ops that are (not yet) bufferizable, but are
 /// known to be allocation hoisting barriers. All interface methods (except for
 /// `isAllocationHoistingBarrier`) are implemented conservatively.
 template <typename OpTy>
 struct AllocationHoistingBarrierOnly
     : public BufferizableOpInterface::ExternalModel<
           AllocationHoistingBarrierOnly<OpTy>, OpTy> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const BufferizationState &state) const {
+                              const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const BufferizationState &state) const {
+                               const AnalysisState &state) const {
     return true;
   }
 
   SmallVector<OpOperand *>
   getAliasingOpOperand(Operation *op, OpResult opResult,
-                       const BufferizationState &state) const {
+                       const AnalysisState &state) const {
     return {};
   }
 
-  SmallVector<OpResult>
-  getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                      const BufferizationState &state) const {
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
     return {};
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationState &state) const {
+                                const AnalysisState &state) const {
     return BufferRelation::None;
   }
 
   bool isWritable(Operation *op, Value value,
-                  const BufferizationState &state) const {
+                  const AnalysisState &state) const {
     return false;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) const {
     return failure();
   }
 
   bool isAllocationHoistingBarrier(Operation *op) const { return true; }
 };
 
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_IR_BUFFERIZABLEOPINTERFACE_H_

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

let methods = [
           even if it does actually not read; however, this can introduce
           unnecessary out-of-place bufferization decisions. The analysis of
           Comprehensive Bufferize considers OpOperands of unknown ops (that do
           not implement this interface) as reading OpOperands.
         }],
         /*retType=*/"bool",
         /*methodName=*/"bufferizesToMemoryRead",
         /*args=*/(ins "OpOperand &":$opOperand,
-                      "const BufferizationState &":$state),
+                      "const AnalysisState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           // Does not have to be implemented for ops without tensor OpOperands.
           llvm_unreachable("bufferizesToMemoryRead not implemented");
          }]
       >,
       InterfaceMethod<
         /*desc=*/[{
let methods = [
           even if it does actually not write; however, this can introduce
           unnecessary out-of-place bufferization decisions. The analysis of
           Comprehensive Bufferize considers OpOperands of unknown ops (that do
           not implement this interface) as writing OpOperands.
         }],
         /*retType=*/"bool",
         /*methodName=*/"bufferizesToMemoryWrite",
         /*args=*/(ins "OpOperand &":$opOperand,
-                      "const BufferizationState &":$state),
+                      "const AnalysisState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           // Does not have to be implemented for ops without tensor OpOperands.
           // Does not have to be implemented for OpOperands that do not have an
           // aliasing OpResult.
           llvm_unreachable("bufferizesToMemoryWrite not implemented");
          }]
       >,
       InterfaceMethod<
           /*desc=*/[{
             Return `true` if the given OpResult is a memory write. This is the
             case if in the following cases:
 
             * The corresponding aliasing OpOperand bufferizes to a memory write.
             * Or: There is no corresponding aliasing OpOperand.
 
             If the OpResult has multiple aliasing OpOperands, this method
             returns `true` if at least one of them bufferizes to a memory write.
           }],
           /*retType=*/"bool",
           /*methodName=*/"isMemoryWrite",
           /*args=*/(ins "OpResult":$opResult,
-                        "const BufferizationState &":$state),
+                        "const AnalysisState &":$state),
           /*methodBody=*/"",
           /*defaultImplementation=*/[{
             auto bufferizableOp =
                 cast<BufferizableOpInterface>($_op.getOperation());
             SmallVector<OpOperand*> opOperands =
               bufferizableOp.getAliasingOpOperand(opResult, state);
             if (opOperands.empty())
               return true;
let methods = [
           Return `true` if the given OpOperand must bufferize in-place. Alias
           sets and inplace attributes will be set up accordingly before making
           any other bufferization decisions. This method will never be called on
           OpOperands that do not have a tensor type.
         }],
         /*retType=*/"bool",
         /*methodName=*/"mustBufferizeInPlace",
         /*args=*/(ins "OpOperand &":$opOperand,
-                      "const BufferizationState &":$state),
+                      "const AnalysisState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           return false;
         }]
       >,
       InterfaceMethod<
         /*desc=*/[{
           Return the OpResult that aliases with a given OpOperand when
           bufferized in-place. This method will never be called on OpOperands
           that do not have a tensor type.
         }],
         /*retType=*/"SmallVector<OpResult>",
         /*methodName=*/"getAliasingOpResult",
         /*args=*/(ins "OpOperand &":$opOperand,
-                      "const BufferizationState &":$state),
+                      "const AnalysisState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           // Does not have to be implemented for ops without tensor OpOperands.
           llvm_unreachable("getAliasingOpResult not implemented");
         }]
       >,
       InterfaceMethod<
         /*desc=*/[{
           Return the OpOperands that alias with a given OpResult when
           bufferized in-place. This method will never be called on OpResults
           that do not have a tensor type.
 
           By default, this method is the inverse of `getAliasingOpResult`. Ops
           with a region that yield values may want to override this method to
           return the OpOperands that are yielded by the terminator.
 
           Note: This method can return multiple OpOperands, indicating that the
           given OpResult may at runtime alias with any of the OpOperands. This
           is useful for branches and for ops such as `arith.select`.
         }],
         /*retType=*/"SmallVector<OpOperand *>",
         /*methodName=*/"getAliasingOpOperand",
         /*args=*/(ins "OpResult":$opResult,
-                      "const BufferizationState &":$state),
+                      "const AnalysisState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           assert(opResult.getType().isa<TensorType>() &&
                  "expected OpResult with tensor type");
           SmallVector<OpOperand *> result;
           auto bufferizableOp =
               cast<BufferizableOpInterface>($_op.getOperation());
           for (OpOperand &opOperand : $_op.getOperation()->getOpOperands()) {
let methods = [
           OpResults that do not have at least one aliasing OpOperand.
 
           TODO: Support other relations such as "OpOperand is included in
           OpResult".
         }],
         /*retType=*/"BufferRelation",
         /*methodName=*/"bufferRelation",
         /*args=*/(ins "OpResult":$opResult,
-                      "const BufferizationState &":$state),
+                      "const AnalysisState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           // Does not have to be implemented for ops without tensor OpResults
           // that have an aliasing OpOperand.
           llvm_unreachable("bufferRelation not implemented");
         }]
       >,
       InterfaceMethod<
let methods = [
           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.
         }],
         /*retType=*/"LogicalResult",
         /*methodName=*/"bufferize",
         /*args=*/(ins "RewriterBase &":$rewriter,
-                      "const BufferizationState &":$state),
+                      "BufferizationState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           llvm_unreachable("bufferize not implemented");
           return failure();
         }]
       >,
       InterfaceMethod<
         /*desc=*/[{
let methods = [
 
           Whether a BlockArgument can be written to or not depends on the
           operation. This method conservatively returns `false`. This method
           will never be called on BlockArguments that do not have a tensor type.
         }],
         /*retType=*/"bool",
         /*methodName=*/"isWritable",
         /*args=*/(ins "Value":$value,
-                      "const BufferizationState &":$state),
+                      "const AnalysisState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           return value.isa<OpResult>();
         }]
       >,
       InterfaceMethod<
         /*desc=*/[{
           Return `true` if the op is an allocation hoisting barrier. Buffer
let methods = [
           If this method returns `true` the given OpOperands are not considered
           to be conflicting and do not force out-of-place bufferization. (There
           may still be other conflicts that do.)
         }],
         /*retType=*/"bool",
         /*methodName=*/"isNotConflicting",
         /*args=*/(ins "OpOperand *":$uRead,
                       "OpOperand *":$uWrite,
-                      "const BufferizationState &":$state),
+                      "const AnalysisState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           return false;
         }]
       >,
       InterfaceMethod<
         /*desc=*/[{
           Return `failure` if this op does not pass the analysis. This method
           is run during One-Shot Bufferize (after all post-analysis steps). If
           the op does not pass the analysis, bufferization is aborted.
 
           This method can be used to check expected invariants and limitations
           of the current bufferization implementation.
         }],
         /*retType=*/"LogicalResult",
         /*methodName=*/"verifyAnalysis",
-        /*args=*/(ins "const BufferizationState &":$state),
+        /*args=*/(ins "const AnalysisState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           return success();
         }]
       >
   ];
 
   let extraClassDeclaration = [{
     /// Return `true` if the given OpOperand creates an alias but does neither
     /// read nor write. This implies that `bufferizesToMemoryRead` and
     /// `bufferizesToMemoryWrite` must return `false`. This method will never
     /// be called on OpOperands that do not have a tensor type.
     ///
     /// Examples of such ops are `tensor.extract_slice` and `tensor.cast`.
     bool bufferizesToAliasOnly(OpOperand &opOperand,
-                               const BufferizationState &state) {
+                               const AnalysisState &state) {
       auto bufferizableOp =
           cast<BufferizableOpInterface>(getOperation());
       return !bufferizableOp.bufferizesToMemoryRead(opOperand, state)
           && !bufferizableOp.bufferizesToMemoryWrite(opOperand, state)
           && !bufferizableOp.getAliasingOpResult(opOperand, state).empty();
     }
 
     // TODO: The following two attributes should belong to the tensor dialect.

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

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,
-                            const BufferizationState &state) const {
+                            BufferizationState &state) 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 failure();
     }
 
-    bool isWritable(Value value, const BufferizationState &state) const {
+    bool isWritable(Value value, const AnalysisState &state) const {
       // It is unknown whether the memref operand is writable or not.
       return false;
     }
   }];
 
   let assemblyFormat = "$memref attr-dict `:` type($memref)";
 
   let hasCanonicalizer = 1;
let extraClassDeclaration = [{
 
     // Note: ToMemrefOp / ToTensorOp are temporary ops that are inserted at the
     // bufferization boundary. When One-Shot bufferization is complete, there
     // should be no such ops left over. If `allowUnknownOps` (or after running a
     // partial bufferization pass), such ops may be part of the resulting IR,
     // but such IR may no longer be analyzable by One-Shot analysis.
 
     bool bufferizesToMemoryRead(OpOperand &opOperand,
-                                const BufferizationState &state) const {
+                                const AnalysisState &state) const {
       // It is unknown whether the resulting memref will be read or not.
       return true;
     }
 
     bool bufferizesToMemoryWrite(OpOperand &opOperand,
-                                 const BufferizationState &state) const {
+                                 const AnalysisState &state) const {
       // It is unknown whether the resulting MemRef will be written or not.
       return true;
     }
 
     bool mustBufferizeInPlace(OpOperand &opOperand,
-                              const BufferizationState &state) const {
+                              const AnalysisState &state) const {
       // ToMemrefOps always bufferize inplace.
       return true;
     }
 
     SmallVector<OpResult> getAliasingOpResult(
-        OpOperand &opOperand, const BufferizationState &state) const {
+        OpOperand &opOperand, const AnalysisState &state) const {
       return {};
     }
 
     LogicalResult bufferize(RewriterBase &rewriter,
-                            const BufferizationState &state);
+                            BufferizationState &state);
   }];
 
   let assemblyFormat = "$tensor attr-dict `:` type($memref)";
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 #endif // BUFFERIZATION_OPS

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

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

 /// unknown op (that does not implement `BufferizableOpInterface`) is found. No
 /// to_tensor/to_memref ops are inserted in that case.
 ///
 /// Note: The layout map chosen to bufferize is the most dynamic canonical
 /// strided layout of the proper rank. This ensures compatibility with expected
 /// layouts after transformations. Combinations of memref.cast +
 /// canonicalization are responsible for clean ups.
 // TODO: Extract `options` from `state` and pass as separate argument.
-LogicalResult bufferizeOp(Operation *op, const BufferizationState &state);
+LogicalResult bufferizeOp(Operation *op, const AnalysisState &analysisState);
+
+/// Bufferize `op` and its nested ops that implement `BufferizableOpInterface`.
+/// Reuse an existing `BufferizationState`.
+///
+/// Note: This function overload is useful for extending the bufferization.
+LogicalResult bufferizeOp(Operation *op,
+                          BufferizationState &bufferizationState);
 
 /// Bufferize `op` and its nested ops that implement `BufferizableOpInterface`.
 /// Buffers are duplicated and copied before any tensor use that bufferizes to
 /// a memory write.
 ///
 /// Note: This function bufferizes ops without utilizing analysis results. It
 /// can be used to implement partial bufferization passes.
 LogicalResult bufferizeOp(Operation *op, const BufferizationOptions &options);
 
-/// Populate the pattern set with a pattern that bufferizes ops that implement
-/// `BufferizableOpInterface`.
-void populateBufferizationPattern(const BufferizationState &state,
-                                  RewritePatternSet &patterns);
-
 BufferizationOptions getPartialBufferizationOptions();
 
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_BUFFERIZE_H

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

 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "llvm/ADT/EquivalenceClasses.h"
 
 namespace mlir {
 namespace bufferization {
 
-class AnalysisBufferizationState;
+struct OneShotBufferizationOptions;
 class BufferizationAliasInfo;
-struct AnalysisBufferizationOptions;
+class OneShotAnalysisState;
 
 /// PostAnalysisStepFns can be registered with `BufferizationOptions` and are
 /// executed after the analysis, but before bufferization. They can be used to
 /// implement custom dialect-specific optimizations. They may modify the IR, but
 /// must keep `aliasInfo` consistent. Newly created operations and operations
 /// that should be re-analyzed must be added to `newOps`.
 using PostAnalysisStepFn = std::function<LogicalResult(
-    Operation *, BufferizationState &, BufferizationAliasInfo &,
+    Operation *, AnalysisState &, BufferizationAliasInfo &,
     SmallVector<Operation *> &)>;
 
 using PostAnalysisStepList = SmallVector<PostAnalysisStepFn>;
 
 /// Options for analysis-enabled bufferization.
-struct AnalysisBufferizationOptions : public BufferizationOptions {
-  AnalysisBufferizationOptions() = default;
+struct OneShotBufferizationOptions : public BufferizationOptions {
+  OneShotBufferizationOptions() = default;
 
   /// Register a "post analysis" step. Such steps are executed after the
   /// analysis, but before bufferization.
   void addPostAnalysisStep(PostAnalysisStepFn fn) {
     postAnalysisSteps.push_back(fn);
   }
 
   /// Registered post analysis steps.
public:
   void insertNewBufferAlias(Value newValue, Value alias);
 
   /// Insert an info entry for `newValue` and merge its alias set with that of
   /// `alias`. Additionally, merge their equivalence classes.
   void insertNewBufferEquivalence(Value newValue, Value alias);
 
   /// Set the inPlace bufferization spec to true.
   /// Merge result's and operand's aliasing sets and iterate to a fixed point.
-  void bufferizeInPlace(OpOperand &operand, BufferizationState &state);
+  void bufferizeInPlace(OpOperand &operand, AnalysisState &state);
 
   /// Set the inPlace bufferization spec to false.
   void bufferizeOutOfPlace(OpOperand &operand);
 
   /// Return true if `v1` and `v2` bufferize to equivalent buffers.
   bool areEquivalentBufferizedValues(Value v1, Value v2) const {
     return equivalentInfo.isEquivalent(v1, v2);
   }
private:
   /// statically if two values are equivalent. In that case, the values are
   /// considered to be not equivalent.
   llvm::EquivalenceClasses<Value, ValueComparator> equivalentInfo;
 };
 
 /// State for analysis-enabled bufferization. This class keeps track of alias
 /// (via BufferizationAliasInfo) to decide if tensor OpOperands should bufferize
 /// in-place.
-class AnalysisBufferizationState : public BufferizationState {
+class OneShotAnalysisState : public AnalysisState {
 public:
-  AnalysisBufferizationState(Operation *op,
-                             const AnalysisBufferizationOptions &options);
+  OneShotAnalysisState(Operation *op,
+                       const OneShotBufferizationOptions &options);
 
-  AnalysisBufferizationState(const AnalysisBufferizationState &) = delete;
+  OneShotAnalysisState(const OneShotAnalysisState &) = delete;
 
-  virtual ~AnalysisBufferizationState() = default;
+  virtual ~OneShotAnalysisState() = default;
 
   /// Return a reference to the BufferizationAliasInfo.
   BufferizationAliasInfo &getAliasInfo() { return aliasInfo; }
 
   /// Return `true` if the given OpResult has been decided to bufferize inplace.
   bool isInPlace(OpOperand &opOperand) const override;
 
   /// Return true if `v1` and `v2` bufferize to equivalent buffers.
   bool areEquivalentBufferizedValues(Value v1, Value v2) const override;
 
 private:
   /// `aliasInfo` keeps track of aliasing and equivalent values. Only internal
   /// functions and `runOneShotBufferize` may access this object.
   BufferizationAliasInfo aliasInfo;
 };
 
 /// Analyze `op` and its nested ops. Bufferization decisions are stored in
 /// `state`.
-LogicalResult analyzeOp(Operation *op, AnalysisBufferizationState &state);
+LogicalResult analyzeOp(Operation *op, OneShotAnalysisState &state);
 
 /// Run One-Shot Bufferize on the given op: Analysis + Bufferization
 LogicalResult runOneShotBufferize(Operation *op,
-                                  const AnalysisBufferizationOptions &options);
+                                  const OneShotBufferizationOptions &options);
 
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTANALYSIS_H

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

 #ifndef MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_PASSES_H
 #define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_PASSES_H
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 namespace bufferization {
-struct AnalysisBufferizationOptions;
+struct OneShotBufferizationOptions;
 
 //===----------------------------------------------------------------------===//
 // Passes
 //===----------------------------------------------------------------------===//
 
 /// Creates an instance of the BufferDeallocation pass to free all allocated
 /// buffers.
 std::unique_ptr<Pass> createBufferDeallocationPass();

 
 /// Create a pass that bufferizes all ops that implement BufferizableOpInterface
 /// with One-Shot Bufferize.
 std::unique_ptr<Pass> createOneShotBufferizePass();
 
 /// Create a pass that bufferizes all ops that implement BufferizableOpInterface
 /// with One-Shot Bufferize and the specified bufferization options.
 std::unique_ptr<Pass>
-createOneShotBufferizePass(const AnalysisBufferizationOptions &options);
+createOneShotBufferizePass(const OneShotBufferizationOptions &options);
 
 /// Creates a pass that promotes heap-based allocations to stack-based ones.
 /// Only buffers smaller than the provided size are promoted.
 /// Dynamic shaped buffers are promoted up to the given rank.
 std::unique_ptr<Pass>
 createPromoteBuffersToStackPass(unsigned maxAllocSizeInBytes = 1024,
                                 unsigned maxRankOfAllocatedMemRef = 1);
 

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

 
 namespace mlir {
 
 class DialectRegistry;
 struct LogicalResult;
 class ModuleOp;
 
 namespace bufferization {
-struct AnalysisBufferizationOptions;
+struct OneShotBufferizationOptions;
 } // namespace bufferization
 
 namespace linalg {
 namespace comprehensive_bufferize {
 
 /// Run Module Bufferization on the given module. Performs a simple function
 /// call analysis to determine which function arguments are inplaceable. Then
 /// analyzes and bufferizes FuncOps one-by-one with One-Shot Bufferize.
 LogicalResult
 runModuleBufferize(ModuleOp moduleOp,
-                   bufferization::AnalysisBufferizationOptions options);
+                   bufferization::OneShotBufferizationOptions options);
 
 namespace std_ext {
 
 void registerModuleBufferizationExternalModels(DialectRegistry &registry);
 
 } // namespace std_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_MODULEBUFFERIZATION_H

mlir/include/mlir/Dialect/Linalg/Passes.h

 #define MLIR_DIALECT_LINALG_PASSES_H_
 
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 namespace bufferization {
-struct AnalysisBufferizationOptions;
+struct OneShotBufferizationOptions;
 } // namespace bufferization
 
 std::unique_ptr<Pass> createConvertElementwiseToLinalgPass();
 
 std::unique_ptr<Pass> createLinalgFoldUnitExtentDimsPass();
 
 std::unique_ptr<Pass> createLinalgElementwiseOpFusionPass();
 std::unique_ptr<Pass> createFoldReshapeOpsByLinearizationPass();

 /// analysis to determine which op operands and results may be bufferized in the
 /// same buffers. The analysis is performed on topologically sorted CallOp and
 /// FuncOp within a module. It provides analyses and bufferization across
 /// function boundaries. Within a function boundary, the analysis is performed
 /// on SSA use-def chains starting from function operands that are annotated
 /// with the 'inplaceable' attribute.
 std::unique_ptr<Pass> createLinalgComprehensiveModuleBufferizePass();
 std::unique_ptr<Pass> createLinalgComprehensiveModuleBufferizePass(
-    const bufferization::AnalysisBufferizationOptions &options);
+    const bufferization::OneShotBufferizationOptions &options);
 
 /// Create a pass to convert Linalg operations which work on tensors to use
 /// buffers instead.
 std::unique_ptr<OperationPass<FuncOp>> createLinalgBufferizePass();
 
 /// Create a pass to convert named Linalg operations to Linalg generic
 /// operations.
 std::unique_ptr<OperationPass<FuncOp>> createLinalgGeneralizationPass();

mlir/include/mlir/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.h

 /// * Only InitTensorOps that are anchored on a matching OpOperand as per
 ///   `anchorMatchFunc` are considered. "Anchored" means that there is a path
 ///   on the reverse SSA use-def chain, starting from the OpOperand and always
 ///   following the aliasing  OpOperand, that eventually ends at a single
 ///   InitTensorOp.
 /// * The result of `rewriteFunc` must usually be analyzed for inplacability.
 ///   This analysis can be skipped with `skipAnalysis`.
 LogicalResult
-eliminateInitTensors(Operation *op, bufferization::BufferizationState &state,
+eliminateInitTensors(Operation *op, bufferization::AnalysisState &state,
                      bufferization::BufferizationAliasInfo &aliasInfo,
                      AnchorMatchFn anchorMatchFunc, RewriteFn rewriteFunc,
                      SmallVector<Operation *> &newOps);
 
 /// Try to eliminate InitTensorOps inside `op` that are anchored on an
 /// InsertSliceOp, i.e., if it is eventually inserted into another tensor
 /// (and some other conditions are met).
 LogicalResult insertSliceAnchoredInitTensorEliminationStep(
-    Operation *op, bufferization::BufferizationState &state,
+    Operation *op, bufferization::AnalysisState &state,
     bufferization::BufferizationAliasInfo &aliasInfo,
     SmallVector<Operation *> &newOps);
 
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_BUFFERIZABLEOPINTERFACEIMPL_H

mlir/include/mlir/Dialect/SCF/BufferizableOpInterfaceImpl.h

 //===- BufferizableOpInterfaceImpl.h - Impl. of BufferizableOpInterface ---===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SCF_BUFFERIZABLEOPINTERFACEIMPL_H
 #define MLIR_DIALECT_SCF_BUFFERIZABLEOPINTERFACEIMPL_H
 
-#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
-
 namespace mlir {
 class DialectRegistry;
 
-namespace bufferization {
-class BufferizationState;
-class BufferizationAliasInfo;
-} // namespace bufferization
-
 namespace scf {
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 } // namespace scf
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SCF_BUFFERIZABLEOPINTERFACEIMPL_H

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,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) 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>();
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     memref::GlobalOp globalMemref = globalOp.getValue();
     replaceOpWithNewBufferizedOp<memref::GetGlobalOp>(
         rewriter, op, globalMemref.type(), globalMemref.getName());
 
     return success();
   }
 
   bool isWritable(Operation *op, Value value,
-                  const BufferizationState &state) const {
+                  const AnalysisState &state) const {
     // Memory locations returned by memref::GetGlobalOp may not be written to.
     assert(value.isa<OpResult>());
     return false;
   }
 };
 
 struct IndexCastOpInterface
     : public BufferizableOpInterface::ExternalModel<IndexCastOpInterface,
                                                     arith::IndexCastOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const BufferizationState &state) const {
+                              const AnalysisState &state) const {
     return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const BufferizationState &state) const {
+                               const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult>
-  getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                      const BufferizationState &state) const {
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
     return {op->getResult(0)};
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationState &state) const {
+                                const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) const {
     auto castOp = cast<arith::IndexCastOp>(op);
 
     Value source = *state.getBuffer(rewriter, op->getOpOperand(0) /*in*/);
     auto sourceType = source.getType().cast<BaseMemRefType>();
 
     // Result type should have same layout and address space as the source type.
     MemRefLayoutAttrInterface layout = {};
     if (auto rankedMemRefType = sourceType.dyn_cast<MemRefType>())
       layout = rankedMemRefType.getLayout();
     Type resultType =
         getMemRefType(castOp.getType().cast<TensorType>(), state.getOptions(),
                       layout, sourceType.getMemorySpace());
 
     replaceOpWithNewBufferizedOp<arith::IndexCastOp>(rewriter, op, resultType,
                                                      source);
     return success();
   }
 };
 
 /// Bufferization of arith.select. Just replace the operands.
 struct SelectOpInterface
     : public BufferizableOpInterface::ExternalModel<SelectOpInterface,
                                                     arith::SelectOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const BufferizationState &state) const {
+                              const AnalysisState &state) const {
     return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const BufferizationState &state) const {
+                               const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult>
-  getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                      const BufferizationState &state) const {
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
     return {op->getOpResult(0) /*result*/};
   }
 
   SmallVector<OpOperand *>
   getAliasingOpOperand(Operation *op, OpResult opResult,
-                       const BufferizationState &state) const {
+                       const AnalysisState &state) const {
     return {&op->getOpOperand(1) /*true_value*/,
             &op->getOpOperand(2) /*false_value*/};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) const {
     auto selectOp = cast<arith::SelectOp>(op);
 
     // `getBuffer` introduces copies if an OpOperand bufferizes out-of-place.
     // 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->getOpOperand(1) /*true_value*/);
     Value falseBuffer =
         *state.getBuffer(rewriter, selectOp->getOpOperand(2) /*false_value*/);
     replaceOpWithNewBufferizedOp<arith::SelectOp>(
         rewriter, op, selectOp.getCondition(), trueBuffer, falseBuffer);
     return success();
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationState &state) const {
+                                const AnalysisState &state) const {
     return BufferRelation::None;
   }
 };
 
 } // namespace
 
 void mlir::arith::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addOpInterface<ConstantOp, ConstantOpInterface>();
   registry.addOpInterface<IndexCastOp, IndexCastOpInterface>();
   registry.addOpInterface<SelectOp, SelectOpInterface>();
 }

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

if (auto bufferizableOp = value.getDefiningOp<BufferizableOpInterface>())
     if (isOpAllowed(bufferizableOp.getOperation()))
       return bufferizableOp;
   return nullptr;
 }
 
 void BufferizationOptions::addDialectStateInitializer(
     StringRef name, const DialectStateInitFn &fn) {
   stateInitializers.push_back(
-      [=](BufferizationState &state) { state.insertDialectState(name, fn()); });
+      [=](AnalysisState &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 {
     b.setInsertionPointAfter(value.getDefiningOp());
   }
 }
 
 /// Determine which OpOperand* will alias with `result` if the op is bufferized
 /// in place. Return an empty vector if the op is not bufferizable.
 SmallVector<OpOperand *>
-BufferizationState::getAliasingOpOperand(OpResult result) const {
+AnalysisState::getAliasingOpOperand(OpResult result) const {
   if (Operation *op = result.getDefiningOp())
     if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
       return bufferizableOp.getAliasingOpOperand(result, *this);
   return {};
 }
 
 /// Determine which OpResult will alias with `opOperand` if the op is bufferized
 /// in place. Return an empty vector if the op is not bufferizable.
 SmallVector<OpResult>
-BufferizationState::getAliasingOpResult(OpOperand &opOperand) const {
+AnalysisState::getAliasingOpResult(OpOperand &opOperand) const {
   if (auto bufferizableOp =
           dyn_cast<BufferizableOpInterface>(opOperand.getOwner()))
     return bufferizableOp.getAliasingOpResult(opOperand, *this);
   return {};
 }
 
 /// Return true if `opOperand` bufferizes to a memory read. Return `true` if the
 /// op is not bufferizable.
-bool BufferizationState::bufferizesToMemoryRead(OpOperand &opOperand) const {
+bool AnalysisState::bufferizesToMemoryRead(OpOperand &opOperand) const {
   if (auto bufferizableOp =
           dyn_cast<BufferizableOpInterface>(opOperand.getOwner()))
     return bufferizableOp.bufferizesToMemoryRead(opOperand, *this);
 
   // Unknown op that returns a tensor. The inplace analysis does not support it.
   // Conservatively return true.
   return true;
 }
 
 /// Return true if `opOperand` bufferizes to a memory write. Return
 /// `true` if the op is not bufferizable.
-bool BufferizationState::bufferizesToMemoryWrite(OpOperand &opOperand) const {
+bool AnalysisState::bufferizesToMemoryWrite(OpOperand &opOperand) const {
   if (auto bufferizableOp =
           dyn_cast<BufferizableOpInterface>(opOperand.getOwner()))
     return bufferizableOp.bufferizesToMemoryWrite(opOperand, *this);
 
   // Unknown op that returns a tensor. The inplace analysis does not support it.
   // Conservatively return true.
   return true;
 }
 
 /// Return true if `opOperand` does neither read nor write but bufferizes to an
 /// alias. Return false if the op is not bufferizable.
-bool BufferizationState::bufferizesToAliasOnly(OpOperand &opOperand) const {
+bool AnalysisState::bufferizesToAliasOnly(OpOperand &opOperand) const {
   if (auto bufferizableOp =
           dyn_cast<BufferizableOpInterface>(opOperand.getOwner()))
     return bufferizableOp.bufferizesToAliasOnly(opOperand, *this);
 
   // Unknown op that returns a tensor. The inplace analysis does not support it.
   // Conservatively return false.
   return false;
 }
 
 /// Return true if the given value is read by an op that bufferizes to a memory
 /// read. Also takes into account ops that create an alias but do not read by
 /// themselves (e.g., ExtractSliceOp).
-bool BufferizationState::isValueRead(Value value) const {
+bool AnalysisState::isValueRead(Value value) const {
   assert(value.getType().isa<TensorType>() && "expected TensorType");
   SmallVector<OpOperand *> workingSet;
   for (OpOperand &use : value.getUses())
     workingSet.push_back(&use);
 
   while (!workingSet.empty()) {
     OpOperand *uMaybeReading = workingSet.pop_back_val();
     // Skip over all ops that neither read nor write (but create an alias).
     if (bufferizesToAliasOnly(*uMaybeReading))
       for (OpResult opResult : getAliasingOpResult(*uMaybeReading))
         for (OpOperand &use : opResult.getUses())
           workingSet.push_back(&use);
     if (bufferizesToMemoryRead(*uMaybeReading))
       return true;
   }
 
   return false;
 }
 
 // Starting from `value`, follow the use-def chain in reverse, always selecting
 // the aliasing OpOperands. Find and return Values for which `condition`
 // evaluates to true. OpOperands of such matching Values are not traversed any
 // further.
-llvm::SetVector<Value> BufferizationState::findValueInReverseUseDefChain(
+llvm::SetVector<Value> AnalysisState::findValueInReverseUseDefChain(
     Value value, llvm::function_ref<bool(Value)> condition) const {
   llvm::SetVector<Value> result, workingSet;
   workingSet.insert(value);
 
   while (!workingSet.empty()) {
     Value value = workingSet.pop_back_val();
     if (condition(value) || value.isa<BlockArgument>()) {
       result.insert(value);
for (OpOperand *o : opOperands)
       workingSet.insert(o->get());
   }
 
   return result;
 }
 
 // Find the Values of the last preceding write of a given Value.
 llvm::SetVector<Value>
-BufferizationState::findLastPrecedingWrite(Value value) const {
+AnalysisState::findLastPrecedingWrite(Value value) const {
   return findValueInReverseUseDefChain(value, [&](Value value) {
     Operation *op = value.getDefiningOp();
     if (!op)
       return true;
     auto bufferizableOp = options.dynCastBufferizableOp(op);
     if (!bufferizableOp)
       return true;
     return bufferizableOp.isMemoryWrite(value.cast<OpResult>(), *this);
   });
 }
 
-BufferizationState::BufferizationState(const BufferizationOptions &options)
+AnalysisState::AnalysisState(const BufferizationOptions &options)
     : options(options) {
-  for (const BufferizationOptions::BufferizationStateInitFn &fn :
+  for (const BufferizationOptions::AnalysisStateInitFn &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>();

 }
 
 /// Return the result buffer (memref) for a given OpResult (tensor). Allocate
 /// a new buffer and copy over data from the existing buffer if out-of-place
 /// bufferization is necessary.
 FailureOr<Value> BufferizationState::getBuffer(
     RewriterBase &rewriter, OpOperand &opOperand, bool forceInPlace,
     Optional<Operation *> customCopyInsertionPoint) const {
+  const BufferizationOptions &options = analysisState.getOptions();
   OpBuilder::InsertionGuard guard(rewriter);
   Operation *op = opOperand.getOwner();
   Location loc = op->getLoc();
   Value operand = opOperand.get();
   Value operandBuffer = lookupBuffer(rewriter, operand, options);
 
-  if (forceInPlace || isInPlace(opOperand))
+  if (forceInPlace || analysisState.isInPlace(opOperand))
     return operandBuffer;
 
   // Bufferizing out-of-place: Allocate a new buffer.
   // Move insertion point right after `operandBuffer`. That is where the
   // allocation should be inserted (in the absence of allocation hoisting).
   setInsertionPointAfter(rewriter, operandBuffer);
   // Allocate the result buffer.
   FailureOr<Value> resultBuffer = createAlloc(rewriter, loc, operandBuffer,
                                               options.createDeallocs, options);
   if (failed(resultBuffer))
     return failure();
   // Do not copy if the last preceding writes of `operand` are ops that do
   // not write (skipping ops that merely create aliases). E.g., InitTensorOp.
   // Note: If `findLastPrecedingWrite` reaches the end of the reverse SSA
   // use-def chain, it returns that value, regardless of whether it is a
   // memory write or not.
-  SetVector<Value> lastWrites = findLastPrecedingWrite(operand);
+  SetVector<Value> lastWrites = analysisState.findLastPrecedingWrite(operand);
   if (llvm::none_of(lastWrites, [&](Value lastWrite) {
         if (auto bufferizableOp = options.dynCastBufferizableOp(lastWrite))
           return bufferizableOp.isMemoryWrite(lastWrite.cast<OpResult>(),
-                                              *this);
+                                              analysisState);
         return true;
       }))
     return resultBuffer;
   // Do not copy if the copied data is never read.
-  SmallVector<OpResult> aliasingOpResults = getAliasingOpResult(opOperand);
-  if (!aliasingOpResults.empty() && !bufferizesToMemoryRead(opOperand) &&
-      llvm::none_of(aliasingOpResults,
-                    [&](OpResult opResult) { return isValueRead(opResult); }))
+  SmallVector<OpResult> aliasingOpResults =
+      analysisState.getAliasingOpResult(opOperand);
+  if (!aliasingOpResults.empty() &&
+      !analysisState.bufferizesToMemoryRead(opOperand) &&
+      llvm::none_of(aliasingOpResults, [&](OpResult opResult) {
+        return analysisState.isValueRead(opResult);
+      }))
     return resultBuffer;
   // Do not copy if this op does not read the data, but writes it.
-  if (bufferizesToMemoryWrite(opOperand) && !bufferizesToMemoryRead(opOperand))
+  if (analysisState.bufferizesToMemoryWrite(opOperand) &&
+      !analysisState.bufferizesToMemoryRead(opOperand))
     return resultBuffer;
 
   if (customCopyInsertionPoint) {
     rewriter.setInsertionPoint(*customCopyInsertionPoint);
   } else {
     // The copy happens right before the op that is bufferized.
     rewriter.setInsertionPoint(op);
   }
if (opResult.getType().isa<TensorType>()) {
           replacement.getLoc(), replacement);
     }
     replacements.push_back(replacement);
   }
 
   rewriter.replaceOp(op, replacements);
 }
 
-AlwaysCopyBufferizationState::AlwaysCopyBufferizationState(
+AlwaysCopyAnalysisState::AlwaysCopyAnalysisState(
     const BufferizationOptions &options)
-    : BufferizationState(options) {}
+    : AnalysisState(options) {}
 
 /// Return `true` if the given OpResult has been decided to bufferize inplace.
-bool AlwaysCopyBufferizationState::isInPlace(OpOperand &opOperand) const {
+bool AlwaysCopyAnalysisState::isInPlace(OpOperand &opOperand) const {
   // OpOperands that bufferize to a memory write are out-of-place, i.e., an
   // alloc and copy is inserted.
   return !bufferizesToMemoryWrite(opOperand);
 }
 
 /// Return true if `v1` and `v2` bufferize to equivalent buffers.
-bool AlwaysCopyBufferizationState::areEquivalentBufferizedValues(
-    Value v1, Value v2) const {
+bool AlwaysCopyAnalysisState::areEquivalentBufferizedValues(Value v1,
+                                                            Value v2) const {
   // There is no analysis, so we do not know if the values are equivalent. The
   // conservative answer is "false".
   return false;
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific scoped alloc/dealloc insertion support.
 //===----------------------------------------------------------------------===//

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

 
 void ToMemrefOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
   results.add<DimOfCastOp, LoadOfToMemref, ToMemrefOfCast, TensorLoadToMemref>(
       context);
 }
 
 LogicalResult ToMemrefOp::bufferize(RewriterBase &rewriter,
-                                    const BufferizationState &state) {
+                                    BufferizationState &state) {
   // Fold to_memref(to_tensor(x)) to x. Insert a cast if necessary.
   return foldToMemrefToTensorPair(rewriter, *this);
 }
 
 Optional<Operation *> CloneOp::buildDealloc(OpBuilder &builder, Value alloc) {
   return builder.create<memref::DeallocOp>(alloc.getLoc(), alloc)
       .getOperation();
 }

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

if (failed(applyFullConversion(func, target, std::move(patterns))))
       signalPassFailure();
   }
 };
 
 struct OneShotBufferizePass
     : public OneShotBufferizeBase<OneShotBufferizePass> {
   OneShotBufferizePass() : OneShotBufferizeBase<OneShotBufferizePass>() {}
 
-  explicit OneShotBufferizePass(const AnalysisBufferizationOptions &options)
+  explicit OneShotBufferizePass(const OneShotBufferizationOptions &options)
       : options(options) {}
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<bufferization::BufferizationDialect>();
   }
 
   void runOnOperation() override {
-    AnalysisBufferizationOptions opt;
+    OneShotBufferizationOptions opt;
     if (!options) {
       // Make new bufferization options if none were provided when creating the
       // pass.
       opt.allowReturnMemref = allowReturnMemref;
       opt.allowUnknownOps = allowUnknownOps;
       opt.analysisFuzzerSeed = analysisFuzzerSeed;
       opt.createDeallocs = createDeallocs;
       opt.fullyDynamicLayoutMaps = fullyDynamicLayoutMaps;
void runOnOperation() override {
     OpPassManager cleanupPipeline("builtin.module");
     cleanupPipeline.addPass(createCanonicalizerPass());
     cleanupPipeline.addPass(createCSEPass());
     cleanupPipeline.addPass(createLoopInvariantCodeMotionPass());
     (void)runPipeline(cleanupPipeline, moduleOp);
   }
 
 private:
-  llvm::Optional<AnalysisBufferizationOptions> options;
+  llvm::Optional<OneShotBufferizationOptions> options;
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::bufferization::createOneShotBufferizePass() {
   return std::make_unique<OneShotBufferizePass>();
 }
 
 std::unique_ptr<Pass> mlir::bufferization::createOneShotBufferizePass(
-    const AnalysisBufferizationOptions &options) {
+    const OneShotBufferizationOptions &options) {
   return std::make_unique<OneShotBufferizePass>(options);
 }
 
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::bufferization::createFinalizingBufferizePass() {
   return std::make_unique<FinalizingBufferizePass>();
 }
 
 //===----------------------------------------------------------------------===//
 // BufferizableOpInterface-based Bufferization
 //===----------------------------------------------------------------------===//
 
 static bool isaTensor(Type t) { return t.isa<TensorType>(); }
 
 /// Return true if the given op has a tensor result or a tensor operand.
 static bool hasTensorSemantics(Operation *op) {
   bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
   bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
   return hasTensorResult || hasTensorOperand;
 }
 
 /// Rewrite pattern that bufferizes bufferizable ops.
 struct BufferizationPattern
     : public OpInterfaceRewritePattern<BufferizableOpInterface> {
-  BufferizationPattern(MLIRContext *context, const BufferizationState &state,
+  BufferizationPattern(MLIRContext *context, BufferizationState &state,
                        PatternBenefit benefit = 1)
       : OpInterfaceRewritePattern<BufferizableOpInterface>(context, benefit),
-        state(state) {}
+        state(&state) {}
 
   LogicalResult matchAndRewrite(BufferizableOpInterface bufferizableOp,
                                 PatternRewriter &rewriter) const override {
+    const BufferizationOptions &options = state->getOptions();
+
     // No tensors => no buffers.
     if (!hasTensorSemantics(bufferizableOp.getOperation()))
       return failure();
-    if (!state.getOptions().isOpAllowed(bufferizableOp.getOperation()))
+    if (!options.isOpAllowed(bufferizableOp.getOperation()))
       return failure();
-    return bufferizableOp.bufferize(rewriter, state);
+    return bufferizableOp.bufferize(rewriter, *state);
   }
 
 private:
-  const BufferizationState &state;
+  BufferizationState *const state;
 };
 
 /// Check the result of bufferization. Return an error if an op was not
 /// bufferized, unless partial bufferization is allowed.
 static LogicalResult
 checkBufferizationResult(Operation *op, const BufferizationOptions &options) {
   if (!options.allowUnknownOps) {
     // Check if all ops were bufferized.
if (!options.allowUnknownOps) {
     if (failed(status))
       return status;
   }
 
   return success();
 }
 
 LogicalResult bufferization::bufferizeOp(Operation *op,
-                                         const BufferizationState &state) {
+                                         const AnalysisState &analysisState) {
+  BufferizationState bufferizationState(analysisState);
+  return bufferizeOp(op, bufferizationState);
+}
+
+LogicalResult
+bufferization::bufferizeOp(Operation *op,
+                           BufferizationState &bufferizationState) {
   // Bufferize the op and its nested ops.
   RewritePatternSet patterns(op->getContext());
-  populateBufferizationPattern(state, patterns);
+  patterns.add<BufferizationPattern>(patterns.getContext(), bufferizationState);
 
   // Bufferize ops top-to-bottom. When creating a new op, we should ideally
   // know the exact memref type of all operands. Otherwise, we have to use a
   // memref type with a fully dynamic layout map, which has to canonicalize
   // away. This is less efficient.
   //
   // Note: If "fullyDynamicLayoutMaps = false", we may have to insert buffer
   // copies to fold ("finalize") to_memref(to_tensor(x)) ops with non-cast-
   // compatible layout maps when doing a traversal other than top-to-bottom.
   // There are currently no canonicalization patterns to fold these away.
   GreedyRewriteConfig config;
   config.useTopDownTraversal = true;
 
   // TODO: Perform a preorder walk instead of the greedy pattern rewriter. This
   // would be more efficient because every bufferization pattern is guaranteed
   // to apply only a single time (otherwise, an assertion would be triggered).
   // However, there are restrictions wrt. erasing ops during a preorder walk,
   // which would likely require a larger refactoring.
   if (failed(applyPatternsAndFoldGreedily(op, std::move(patterns), config)))
     return failure();
 
-  return checkBufferizationResult(op, state.getOptions());
+  return checkBufferizationResult(op, bufferizationState.getOptions());
 }
 
 namespace {
-/// This a "no analysis, always copy" BufferizationState. In the absence of an
+/// This a "no analysis, always copy" AnalysisState. In the absence of an
 /// analysis, a buffer must be copied each time it is written to. Therefore, all
 /// OpOperands that bufferize to a memory write must bufferize out-of-place.
-class AlwaysCopyBufferizationState : public BufferizationState {
+class AlwaysCopyAnalysisState : public AnalysisState {
 public:
-  AlwaysCopyBufferizationState(const BufferizationOptions &options)
-      : BufferizationState(options) {}
+  AlwaysCopyAnalysisState(const BufferizationOptions &options)
+      : AnalysisState(options) {}
 
-  AlwaysCopyBufferizationState(const AlwaysCopyBufferizationState &) = delete;
+  AlwaysCopyAnalysisState(const AlwaysCopyAnalysisState &) = delete;
 
-  virtual ~AlwaysCopyBufferizationState() = default;
+  virtual ~AlwaysCopyAnalysisState() = default;
 
   /// Return `true` if the given OpResult has been decided to bufferize inplace.
   bool isInPlace(OpOperand &opOperand) const override {
     // OpOperands that bufferize to a memory write are out-of-place, i.e., an
     // alloc and copy is inserted.
     return !bufferizesToMemoryWrite(opOperand);
   }
 
   /// Return true if `v1` and `v2` bufferize to equivalent buffers.
   bool areEquivalentBufferizedValues(Value v1, Value v2) const override {
     // There is no analysis, so we do not know if the values are equivalent. The
     // conservative answer is "false".
     return false;
   }
 };
 } // namespace
 
 LogicalResult bufferization::bufferizeOp(Operation *op,
                                          const BufferizationOptions &options) {
-  AlwaysCopyBufferizationState state(options);
+  AlwaysCopyAnalysisState state(options);
   return bufferizeOp(op, state);
 }
 
-void bufferization::populateBufferizationPattern(
-    const BufferizationState &state, RewritePatternSet &patterns) {
-  patterns.add<BufferizationPattern>(patterns.getContext(), state);
-}
-
 BufferizationOptions bufferization::getPartialBufferizationOptions() {
   BufferizationOptions options;
   options.allowReturnMemref = true;
   options.allowUnknownOps = true;
   options.createDeallocs = false;
   options.fullyDynamicLayoutMaps = false;
   return options;
 }

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

 //    to bufferize inplace during the analysis phase.
 //
 // This file contains only the analysis. The actual bufferization is implemented
 // via `bufferizeOp` (Bufferize.h). For convenience, this file also contains a
 // helper function `runOneShotBufferize` that analyzes an op (and its nested
 // ops) and then bufferizes it.
 //
 // Inplace bufferization decisions are passed from the analysis to the
-// bufferization phase via `BufferizationState` and `BufferizationAliasInfo`.
+// bufferization phase via `AnalysisState` and `BufferizationAliasInfo`.
 // They can be printed for debugging purposes with `testAnalysisOnly`.
 //
 // Ops that do not implement `BufferizableOpInterface` can be analyzed but are
 // treated conservatively. E.g., the analysis has to assume that their tensor
 // OpOperands bufferize to memory writes. While such ops can be analyzed, they
 // are not bufferized and remain in the IR. to_tensor and to_memref ops are
 // inserted at the bufferization boundary.
 //

 
 /// Return `true` if a value was marked as in-place bufferized.
 bool BufferizationAliasInfo::isInPlace(OpOperand &operand) const {
   return inplaceBufferized.contains(&operand);
 }
 
 /// Set the inPlace bufferization spec to true.
 void BufferizationAliasInfo::bufferizeInPlace(OpOperand &operand,
-                                              BufferizationState &state) {
+                                              AnalysisState &state) {
   markInPlace(operand);
   for (OpResult result : state.getAliasingOpResult(operand))
     aliasInfo.unionSets(result, operand.get());
 }
 
 /// Set the inPlace bufferization spec to false.
 void BufferizationAliasInfo::bufferizeOutOfPlace(OpOperand &operand) {
   assert(!inplaceBufferized.contains(&operand) &&
for (auto mit = aliasInfo.member_begin(it), meit = aliasInfo.member_end();
        mit != meit; ++mit) {
     res.insert(static_cast<Value>(*mit));
   }
   return BufferizationAliasInfo::EquivalenceClassRangeType(
       aliasInfo.member_begin(it), aliasInfo.member_end());
 }
 
 //===----------------------------------------------------------------------===//
-// AnalysisBufferizationState
+// OneShotAnalysisState
 //===----------------------------------------------------------------------===//
 
-AnalysisBufferizationState::AnalysisBufferizationState(
-    Operation *op, const AnalysisBufferizationOptions &options)
-    : BufferizationState(options), aliasInfo(op) {
+OneShotAnalysisState::OneShotAnalysisState(
+    Operation *op, const OneShotBufferizationOptions &options)
+    : AnalysisState(options), aliasInfo(op) {
   // Set up alias sets for OpResults that must bufferize in-place. This should
   // be done before making any other bufferization decisions.
   op->walk([&](BufferizableOpInterface bufferizableOp) {
     if (!options.isOpAllowed(bufferizableOp))
       return WalkResult::skip();
     for (OpOperand &opOperand : bufferizableOp->getOpOperands()) {
       if (opOperand.get().getType().isa<TensorType>())
         if (bufferizableOp.mustBufferizeInPlace(opOperand, *this)) {
           for (OpResult opResult :
                bufferizableOp.getAliasingOpResult(opOperand, *this))
             aliasInfo.unionAliasSets(opOperand.get(), opResult);
           aliasInfo.markInPlace(opOperand);
         }
     }
     return WalkResult::advance();
   });
 }
 
-bool AnalysisBufferizationState::isInPlace(OpOperand &opOperand) const {
+bool OneShotAnalysisState::isInPlace(OpOperand &opOperand) const {
   return aliasInfo.isInPlace(opOperand);
 }
 
-bool AnalysisBufferizationState::areEquivalentBufferizedValues(Value v1,
-                                                               Value v2) const {
+bool OneShotAnalysisState::areEquivalentBufferizedValues(Value v1,
   return aliasInfo.areEquivalentBufferizedValues(v1, v2);
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific alias analysis.
 //===----------------------------------------------------------------------===//
 
 /// Return true if opOperand has been decided to bufferize in-place.
 static bool isInplaceMemoryWrite(OpOperand &opOperand,
                                  const BufferizationAliasInfo &aliasInfo,
-                                 BufferizationState &state) {
+                                 AnalysisState &state) {
   // OpOperands that do not bufferize to a memory write do not write in-place.
   if (!state.bufferizesToMemoryWrite(opOperand))
     return false;
   // Check current bufferization decisions.
   return aliasInfo.isInPlace(opOperand);
 }
 
 /// Return true if, under current bufferization decisions, the buffer of `value`
 /// is not writable.
 static bool aliasesNonWritableBuffer(Value value,
                                      const BufferizationAliasInfo &aliasInfo,
-                                     BufferizationState &state) {
+                                     AnalysisState &state) {
   bool foundNonWritableBuffer = false;
   aliasInfo.applyOnAliases(value, [&](Value v) {
     // Query BufferizableOpInterface to see if the value is writable.
     // TODO: Out-of-place bufferized value could be considered writable.
     if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(v))
       if (bufferizableOp && bufferizableOp.isWritable(v, state))
         return;
 
static bool aliasesNonWritableBuffer(Value value,
 
   return foundNonWritableBuffer;
 }
 
 /// Return true if the buffer to which `operand` would bufferize is equivalent
 /// to some buffer write.
 static bool aliasesInPlaceWrite(Value value,
                                 const BufferizationAliasInfo &aliasInfo,
-                                BufferizationState &state) {
+                                AnalysisState &state) {
   bool foundInplaceWrite = false;
   aliasInfo.applyOnAliases(value, [&](Value v) {
     for (auto &use : v.getUses()) {
       if (isInplaceMemoryWrite(use, aliasInfo, state)) {
         foundInplaceWrite = true;
         return;
       }
     }

 /// all given writes bufferize inplace.
 ///
 /// A conflict is: According to SSA use-def chains, a read R is supposed to read
 /// the result of a write W1. But because of bufferization decisions, R actually
 /// reads another write W2.
 static bool hasReadAfterWriteInterference(
     const DenseSet<OpOperand *> &usesRead,
     const DenseSet<OpOperand *> &usesWrite, const DominanceInfo &domInfo,
-    BufferizationState &state, const BufferizationAliasInfo &aliasInfo) {
+    AnalysisState &state, const BufferizationAliasInfo &aliasInfo) {
   const BufferizationOptions &options = state.getOptions();
 
   for (OpOperand *uRead : usesRead) {
     Operation *readingOp = uRead->getOwner();
 
     // Find most recent writes of uRead by following the SSA use-def chain.
     // E.g.:
     //

 /// If `checkConsistencyOnly` is true, this function checks if there is a
 /// read-after-write conflict without bufferizing `operand` inplace. This would
 /// indicate a problem with the current inplace bufferization decisions.
 ///
 /// Note: If `checkConsistencyOnly`, this function may be called with a null
 /// OpResult. In that case, only the consistency of bufferization decisions
 /// involving aliases of the given OpOperand are checked.
 static bool wouldCreateReadAfterWriteInterference(
-    OpOperand &operand, const DominanceInfo &domInfo, BufferizationState &state,
+    OpOperand &operand, const DominanceInfo &domInfo, AnalysisState &state,
     const BufferizationAliasInfo &aliasInfo,
     bool checkConsistencyOnly = false) {
   // Helper function to iterate on aliases of `root` and capture the reads.
   auto getAliasingReads = [&](DenseSet<OpOperand *> &res, Value root) {
     aliasInfo.applyOnAliases(root, [&](Value alias) {
       for (auto &use : alias.getUses())
         // Read to a value that aliases root.
         if (state.bufferizesToMemoryRead(use))
return hasReadAfterWriteInterference(usesRead, usesWrite, domInfo, state,
                                        aliasInfo);
 }
 
 /// Return true if bufferizing `opOperand` inplace would create a write to a
 /// non-writable buffer.
 static bool
 wouldCreateWriteToNonWritableBuffer(OpOperand &opOperand,
                                     const BufferizationAliasInfo &aliasInfo,
-                                    BufferizationState &state) {
+                                    AnalysisState &state) {
   // Certain buffers are not writeable:
   //   1. A function bbArg that is not inplaceable or
   //   2. A constant op.
   bool nonWritable =
       aliasesNonWritableBuffer(opOperand.get(), aliasInfo, state);
   if (!nonWritable)
     return false;
 
   // This is a problem only if the buffer is written to via some alias.
   bool hasWrite = aliasesInPlaceWrite(opOperand.get(), aliasInfo, state) ||
                   state.bufferizesToMemoryWrite(opOperand);
 
   for (OpResult opResult : state.getAliasingOpResult(opOperand))
     hasWrite |= aliasesInPlaceWrite(opResult, aliasInfo, state);
 
   return hasWrite;
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization analyses.
 //===----------------------------------------------------------------------===//
 
 /// Determine if `operand` can be bufferized in-place.
 static LogicalResult bufferizableInPlaceAnalysisImpl(
-    OpOperand &operand, BufferizationAliasInfo &aliasInfo,
-    BufferizationState &state, const DominanceInfo &domInfo) {
+    OpOperand &operand, BufferizationAliasInfo &aliasInfo, AnalysisState &state,
+    const DominanceInfo &domInfo) {
   bool foundInterference =
       wouldCreateWriteToNonWritableBuffer(operand, aliasInfo, state) ||
       wouldCreateReadAfterWriteInterference(operand, domInfo, state, aliasInfo);
 
   if (foundInterference)
     aliasInfo.bufferizeOutOfPlace(operand);
   else
     aliasInfo.bufferizeInPlace(operand, state);

 /// When bufferized inplace, an ExtractSliceOp does not by itself create any
 /// read or write from memory. Instead, it has the effect of merging the alias
 /// sets of the source and the result buffers.
 ///
 /// An analysis is required to ensure inplace bufferization would not result in
 /// RaW dependence violations.
 static LogicalResult inPlaceAnalysis(SmallVector<Operation *> &ops,
                                      BufferizationAliasInfo &aliasInfo,
-                                     BufferizationState &state,
+                                     AnalysisState &state,
                                      const DominanceInfo &domInfo,
                                      unsigned analysisFuzzerSeed = 0) {
   if (analysisFuzzerSeed) {
     // This is a fuzzer. For testing purposes only. Randomize the order in which
     // operations are analyzed. The bufferization quality is likely worse, but
     // we want to make sure that no assertions are triggered anywhere.
     std::mt19937 g(analysisFuzzerSeed);
     llvm::shuffle(ops.begin(), ops.end(), g);
static bool hasTensorSemantics(Operation *op) {
   bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
   bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
   return hasTensorResult || hasTensorOperand;
 }
 
 /// Analyze all ops that are contained in `op`.
 static LogicalResult inPlaceAnalysis(Operation *op,
                                      BufferizationAliasInfo &aliasInfo,
-                                     BufferizationState &state,
+                                     AnalysisState &state,
                                      const DominanceInfo &domInfo,
                                      unsigned analysisFuzzerSeed = 0) {
   // Collect ops so we can build our own reverse traversal.
   SmallVector<Operation *> ops;
   op->walk([&](Operation *op) {
     // No tensors => no buffers.
     if (!hasTensorSemantics(op))
       return;
     ops.push_back(op);
   });
 
   return inPlaceAnalysis(ops, aliasInfo, state, domInfo, analysisFuzzerSeed);
 }
 
 /// Analyze equivalence of tied OpResult/OpOperand pairs of the given ops.
 static void equivalenceAnalysis(SmallVector<Operation *> &ops,
                                 BufferizationAliasInfo &aliasInfo,
-                                BufferizationState &state) {
+                                AnalysisState &state) {
   for (Operation *op : ops)
     if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(op))
       for (OpResult opResult : op->getOpResults())
         if (opResult.getType().isa<TensorType>())
           for (OpOperand *opOperand :
                bufferizableOp.getAliasingOpOperand(opResult, state))
             if (state.isInPlace(*opOperand))
               if (bufferizableOp.bufferRelation(opResult, state) ==
                   BufferRelation::Equivalent)
                 aliasInfo.unionEquivalenceClasses(opResult, opOperand->get());
 }
 
 /// Analyze equivalence of tied OpResult/OpOperand pairs of all ops contained
 /// in `op`.
 static void equivalenceAnalysis(Operation *op,
                                 BufferizationAliasInfo &aliasInfo,
-                                BufferizationState &state) {
+                                AnalysisState &state) {
   // Traverse ops in PostOrder: Nested ops first, then enclosing ops.
   SmallVector<Operation *> ops;
   op->walk<WalkOrder::PostOrder>([&](Operation *op) {
     // No tensors => no buffers.
     if (none_of(op->getResultTypes(), isaTensor))
       return;
     ops.push_back(op);
   });
 
   equivalenceAnalysis(ops, aliasInfo, state);
 }
 
 /// Assert that the current bufferization decisions are consistent.
 static LogicalResult
 checkAliasInfoConsistency(Operation *op, const DominanceInfo &domInfo,
-                          BufferizationState &state,
+                          AnalysisState &state,
                           const BufferizationAliasInfo &aliasInfo) {
   const BufferizationOptions &options = state.getOptions();
   Operation *inconsistentOp = nullptr;
   WalkResult walkResult = op->walk([&](Operation *op) {
     if (auto bufferizableOp = options.dynCastBufferizableOp(op))
       for (OpOperand &opOperand : op->getOpOperands())
         if (opOperand.get().getType().isa<TensorType>()) {
           if (wouldCreateReadAfterWriteInterference(
if (walkResult.wasInterrupted())
     return inconsistentOp->emitError("input IR has RaW conflict");
   return success();
 }
 
 /// Annotate the IR with the result of the analysis. For testing/debugging only.
 static void
 annotateOpsWithBufferizationMarkers(Operation *op,
                                     const BufferizationAliasInfo &aliasInfo,
-                                    BufferizationState &state) {
+                                    AnalysisState &state) {
   op->walk([&](Operation *op) {
     if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(op))
       for (OpOperand &opOperand : op->getOpOperands())
         if (opOperand.get().getType().isa<TensorType>())
           setInPlaceOpOperand(opOperand, aliasInfo.isInPlace(opOperand));
   });
 }
 

 /// style because the yielded value %0 is defined in the parent block. The
 /// second scf.yield op does not satisfy destination-passing style because the
 /// yielded value %t is defined in the same block as the scf.yield op.
 // TODO: The current implementation checks for equivalent values instead of
 // aliasing values, which is stricter than needed. We can currently not check
 // for aliasing values because the analysis is a maybe-alias analysis and we
 // need a must-alias analysis here.
 static LogicalResult
-assertDestinationPassingStyle(Operation *op, BufferizationState &state,
+assertDestinationPassingStyle(Operation *op, AnalysisState &state,
                               BufferizationAliasInfo &aliasInfo,
                               SmallVector<Operation *> &newOps) {
   LogicalResult status = success();
   DominanceInfo domInfo(op);
   op->walk([&](Operation *returnOp) {
     if (!isRegionReturnLike(returnOp) ||
         !state.getOptions().isOpAllowed(returnOp))
       return WalkResult::advance();
op->walk([&](Operation *returnOp) {
 
     return WalkResult::advance();
   });
 
   return status;
 }
 
 LogicalResult bufferization::analyzeOp(Operation *op,
-                                       AnalysisBufferizationState &state) {
+                                       OneShotAnalysisState &state) {
   DominanceInfo domInfo(op);
   BufferizationAliasInfo &aliasInfo = state.getAliasInfo();
   const auto &options =
-      static_cast<const AnalysisBufferizationOptions &>(state.getOptions());
+      static_cast<const OneShotBufferizationOptions &>(state.getOptions());
 
   if (failed(checkAliasInfoConsistency(op, domInfo, state, aliasInfo)))
     return failure();
 
   // If the analysis fails, just return.
   if (failed(inPlaceAnalysis(op, aliasInfo, state, domInfo,
                              options.analysisFuzzerSeed)))
     return failure();
LogicalResult bufferization::analyzeOp(Operation *op,
 
   // Annotate operations if we only want to report the analysis.
   if (options.testAnalysisOnly)
     annotateOpsWithBufferizationMarkers(op, aliasInfo, state);
 
   return success(!failedAnalysis);
 }
 
-LogicalResult bufferization::runOneShotBufferize(
-    Operation *op, const AnalysisBufferizationOptions &options) {
-  AnalysisBufferizationState state(op, options);
+LogicalResult
+bufferization::runOneShotBufferize(Operation *op,
+                                   const OneShotBufferizationOptions &options) {
+  OneShotAnalysisState state(op, options);
   if (failed(analyzeOp(op, state)))
     return failure();
   if (options.testAnalysisOnly)
     return success();
   return bufferizeOp(op, state);
 }

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

 //
 // Module Bufferization is run via `runModuleBufferize(ModuleOp, ...)`. This
 // function analyzes the given module and determines the order of analysis and
 // bufferization: Functions that are called are processed before their
 // respective callers.
 //
 // After analyzing a FuncOp, additional information about its bbArgs is
 // gathered through PostAnalysisStepFns and stored in
-// `ModuleBufferizationState`.
+// `ModuleAnalysisState`.
 //
 // * `equivalentFuncOpBBArgsAnalysis` determines the equivalent bbArg for each
 //   tensor return value (if any).
 // * `funcOpBbArgReadWriteAnalysis` determines whether or not a tensor bbArg is
 //   read/written.
 //
 // Only tensors that are equivalent to some FuncOp bbArg may be returned.
 // Bufferization currently fails if other tensors (in particular tensors that

 using namespace tensor;
 using namespace comprehensive_bufferize;
 using namespace mlir::bufferization;
 
 namespace {
 /// The state of analysis of a FuncOp.
 enum class FuncOpAnalysisState { NotAnalyzed, InProgress, Analyzed };
 
-/// Extra bufferization state that is required for bufferization of function
+/// Extra analysis state that is required for bufferization of function
 /// boundaries.
-struct ModuleBufferizationState : public DialectBufferizationState {
+struct ModuleAnalysisState : public DialectAnalysisState {
   /// A mapping of ReturnOp OpOperand indices to equivalent FuncOp BBArg
   /// indices.
   DenseMap<FuncOp, DenseMap<int64_t, int64_t>> equivalentFuncArgs;
 
   /// A set of all read BlockArguments of FuncOps.
   // Note: BlockArgument knows about its owner, so we do not need to store
   // FuncOps here.
   DenseSet<BlockArgument> readBbArgs;
 
   /// A set of all written-to BlockArguments of FuncOps.
   DenseSet<BlockArgument> writtenBbArgs;
 
   /// Keep track of which FuncOps are fully analyzed or currently being
   /// analyzed.
   DenseMap<FuncOp, FuncOpAnalysisState> analyzedFuncOps;
 
   // A list of functions in the order in which they are analyzed + bufferized.
   SmallVector<FuncOp> orderedFuncOps;
 
   // A mapping of FuncOps to their callers.
   DenseMap<FuncOp, DenseSet<Operation *>> callerMap;
 };
 } // namespace
 
-/// Get ModuleBufferizationState.
-static const ModuleBufferizationState &
-getModuleBufferizationState(const BufferizationState &state) {
-  Optional<const ModuleBufferizationState *> maybeState =
-      state.getDialectState<ModuleBufferizationState>(
+/// Get ModuleAnalysisState.
+static const ModuleAnalysisState &
+getModuleAnalysisState(const AnalysisState &state) {
+  Optional<const ModuleAnalysisState *> maybeState =
+      state.getDialectState<ModuleAnalysisState>(
           func::FuncDialect::getDialectNamespace());
-  assert(maybeState.hasValue() && "ModuleBufferizationState does not exist");
+  assert(maybeState.hasValue() && "ModuleAnalysisState does not exist");
   return **maybeState;
 }
 
-/// Get or create ModuleBufferizationState.
-static ModuleBufferizationState &
-getModuleBufferizationState(BufferizationState &state) {
-  return state.getOrCreateDialectState<ModuleBufferizationState>(
+/// Get or create ModuleAnalysisState.
+static ModuleAnalysisState &getModuleAnalysisState(AnalysisState &state) {
+  return state.getOrCreateDialectState<ModuleAnalysisState>(
       func::FuncDialect::getDialectNamespace());
 }
 
 /// Return the state (phase) of analysis of the FuncOp.
-static FuncOpAnalysisState
-getFuncOpAnalysisState(const BufferizationState &state, FuncOp funcOp) {
-  const ModuleBufferizationState &moduleState =
-      getModuleBufferizationState(state);
+static FuncOpAnalysisState getFuncOpAnalysisState(const AnalysisState &state,
+                                                  FuncOp funcOp) {
+  const ModuleAnalysisState &moduleState = getModuleAnalysisState(state);
   auto it = moduleState.analyzedFuncOps.find(funcOp);
   if (it == moduleState.analyzedFuncOps.end())
     return FuncOpAnalysisState::NotAnalyzed;
   return it->second;
 }
 
 /// Return the unique ReturnOp that terminates `funcOp`.
 /// Return nullptr if there is no such unique ReturnOp.
static void annotateEquivalentReturnBbArg(OpOperand &returnVal,
   }
   equivBbArgs[returnVal.getOperandNumber()] = bbArg.getArgNumber();
 
   OpBuilder b(op->getContext());
   op->setAttr(kEquivalentArgsAttr, b.getI64ArrayAttr(equivBbArgs));
 }
 
 /// Store function BlockArguments that are equivalent to a returned value in
-/// ModuleBufferizationState.
+/// ModuleAnalysisState.
 static LogicalResult
-equivalentFuncOpBBArgsAnalysis(Operation *op, BufferizationState &state,
+equivalentFuncOpBBArgsAnalysis(Operation *op, AnalysisState &state,
                                BufferizationAliasInfo &aliasInfo,
                                SmallVector<Operation *> &newOps) {
-  ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
+  ModuleAnalysisState &moduleState = getModuleAnalysisState(state);
 
   // Support only single return-terminated block in the function.
   auto funcOp = cast<FuncOp>(op);
   func::ReturnOp returnOp = getAssumedUniqueReturnOp(funcOp);
   assert(returnOp && "expected func with single return op");
 
   for (OpOperand &returnVal : returnOp->getOpOperands())
     if (returnVal.get().getType().isa<RankedTensorType>())
       for (BlockArgument bbArg : funcOp.getArguments())
         if (bbArg.getType().isa<RankedTensorType>())
           if (aliasInfo.areEquivalentBufferizedValues(returnVal.get(), bbArg)) {
             moduleState
                 .equivalentFuncArgs[funcOp][returnVal.getOperandNumber()] =
                 bbArg.getArgNumber();
             if (state.getOptions().testAnalysisOnly)
               annotateEquivalentReturnBbArg(returnVal, bbArg);
           }
 
   return success();
 }
 
 /// Return true if the buffer of the given tensor value is written to. Must not
 /// be called for values inside not yet analyzed functions. (Post-analysis
 /// steps do not have to be run yet, i.e., "in progress" is also OK.)
-static bool isValueWritten(Value value, const BufferizationState &state,
+static bool isValueWritten(Value value, const AnalysisState &state,
                            const BufferizationAliasInfo &aliasInfo) {
 #ifndef NDEBUG
   assert(value.getType().isa<TensorType>() && "expected TensorType");
   FuncOp funcOp;
   if (auto bbArg = value.dyn_cast<BlockArgument>()) {
     Operation *owner = bbArg.getOwner()->getParentOp();
     funcOp = isa<FuncOp>(owner) ? cast<FuncOp>(owner)
                                 : owner->getParentOfType<FuncOp>();
static void annotateFuncArgAccess(FuncOp funcOp, BlockArgument bbArg,
   }
   funcOp.setArgAttr(bbArg.getArgNumber(), "bufferization.access", accessType);
 }
 
 /// Determine which FuncOp bbArgs are read and which are written. If this
 /// PostAnalysisStepFn is run on a function with unknown ops, it will
 /// conservatively assume that such ops bufferize to a read + write.
 static LogicalResult
-funcOpBbArgReadWriteAnalysis(Operation *op, BufferizationState &state,
+funcOpBbArgReadWriteAnalysis(Operation *op, AnalysisState &state,
                              BufferizationAliasInfo &aliasInfo,
                              SmallVector<Operation *> &newOps) {
-  ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
+  ModuleAnalysisState &moduleState = getModuleAnalysisState(state);
   auto funcOp = cast<FuncOp>(op);
 
   // If the function has no body, conservatively assume that all args are
   // read + written.
   if (funcOp.getBody().empty()) {
     for (BlockArgument bbArg : funcOp.getArguments()) {
       moduleState.readBbArgs.insert(bbArg);
       moduleState.writtenBbArgs.insert(bbArg);
getBufferizedFunctionType(MLIRContext *ctx, TypeRange argumentTypes,
   return FunctionType::get(ctx, argTypes, retTypes);
 }
 
 /// Gather equivalence info of CallOps.
 /// Note: This only adds new equivalence info if `funcOp` was already analyzed.
 // TODO: This does not handle cyclic function call graphs etc.
 static void equivalenceAnalysis(FuncOp funcOp,
                                 BufferizationAliasInfo &aliasInfo,
-                                ModuleBufferizationState &moduleState) {
+                                ModuleAnalysisState &moduleState) {
   funcOp->walk([&](func::CallOp callOp) {
     FuncOp calledFunction = getCalledFunction(callOp);
     assert(calledFunction && "could not retrieved called FuncOp");
 
     // No equivalence info available for the called function.
     if (!moduleState.equivalentFuncArgs.count(calledFunction))
       return WalkResult::skip();
 

 /// future.
 ///
 /// Note: Returning a memref currently fails bufferization. If such memrefs
 /// originate from an op with an Alloc effect, they could be hoisted in the
 /// future.
 static LogicalResult bufferizeFuncOpBoundary(FuncOp funcOp,
                                              RewriterBase &rewriter,
                                              BufferizationState &state) {
-  ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
+  const ModuleAnalysisState &moduleState =
+      getModuleAnalysisState(state.getAnalysisState());
 
   // If nothing to do then we are done.
   if (!llvm::any_of(funcOp.getType().getInputs(), isaTensor) &&
       !llvm::any_of(funcOp.getType().getResults(), isaTensor))
     return success();
 
   // Get the bufferized FunctionType for funcOp or construct it if not yet
   // available.
for (OpOperand &returnOperand : returnOp->getOpOperands()) {
 
     // If not a renturn tensor type just forward it.
     if (!returnVal.getType().isa<RankedTensorType>()) {
       returnValues.push_back(returnVal);
       continue;
     }
 
     // If return operand is equivalent to some bbArg, no need to return it.
-    if (moduleState.equivalentFuncArgs[funcOp].count(
-            returnOperand.getOperandNumber()))
+    auto funcOpIt = moduleState.equivalentFuncArgs.find(funcOp);
+    if (funcOpIt != moduleState.equivalentFuncArgs.end() &&
+        funcOpIt->second.count(returnOperand.getOperandNumber()))
       continue;
 
     // Cast values at the call site if necessary.
     returnValues.push_back(
         getNonCastedValue(*state.getBuffer(rewriter, returnOperand)));
   }
 
   // 2. Rewrite the terminator without the inPlace bufferizable values.

 namespace mlir {
 namespace linalg {
 namespace comprehensive_bufferize {
 namespace std_ext {
 
 /// Return the index of the bbArg in the given FuncOp that is equivalent to the
 /// specified return value (if any).
 static Optional<int64_t>
-getEquivalentFuncArgIdx(FuncOp funcOp, const ModuleBufferizationState &state,
+getEquivalentFuncArgIdx(FuncOp funcOp, const ModuleAnalysisState &state,
                         int64_t returnValIdx) {
   auto funcOpIt = state.equivalentFuncArgs.find(funcOp);
   if (funcOpIt == state.equivalentFuncArgs.end())
     // No equivalence info stores for funcOp.
     return None;
 
   auto retValIt = funcOpIt->getSecond().find(returnValIdx);
   if (retValIt == funcOpIt->getSecond().end())
     // Return value has no equivalent bbArg.
     return None;
 
   return retValIt->getSecond();
 }
 
 struct CallOpInterface
     : public BufferizableOpInterface::ExternalModel<CallOpInterface,
                                                     func::CallOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const BufferizationState &state) const {
+                              const AnalysisState &state) const {
     func::CallOp callOp = cast<func::CallOp>(op);
     FuncOp funcOp = getCalledFunction(callOp);
     assert(funcOp && "expected CallOp to a FuncOp");
 
-    const ModuleBufferizationState &moduleState =
-        getModuleBufferizationState(state);
+    const ModuleAnalysisState &moduleState = getModuleAnalysisState(state);
     if (getFuncOpAnalysisState(state, funcOp) != FuncOpAnalysisState::Analyzed)
       // FuncOp not analyzed yet. Assume that OpOperand is read.
       return true;
 
     return moduleState.readBbArgs.contains(
         funcOp.getArgument(opOperand.getOperandNumber()));
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const BufferizationState &state) const {
+                               const AnalysisState &state) const {
     func::CallOp callOp = cast<func::CallOp>(op);
     FuncOp funcOp = getCalledFunction(callOp);
     assert(funcOp && "expected CallOp to a FuncOp");
 
-    const ModuleBufferizationState &moduleState =
-        getModuleBufferizationState(state);
+    const ModuleAnalysisState &moduleState = getModuleAnalysisState(state);
     if (getFuncOpAnalysisState(state, funcOp) != FuncOpAnalysisState::Analyzed)
       // FuncOp not analyzed yet. Assume that OpOperand is written.
       return true;
 
     return moduleState.writtenBbArgs.contains(
         funcOp.getArgument(opOperand.getOperandNumber()));
   }
 
-  SmallVector<OpResult>
-  getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                      const BufferizationState &state) const {
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
     func::CallOp callOp = cast<func::CallOp>(op);
     FuncOp funcOp = getCalledFunction(callOp);
     assert(funcOp && "expected CallOp to a FuncOp");
-    const ModuleBufferizationState &moduleState =
-        getModuleBufferizationState(state);
+    const ModuleAnalysisState &moduleState = getModuleAnalysisState(state);
 
     SmallVector<OpResult> result;
     for (int64_t resultIdx = 0; resultIdx < callOp->getNumResults();
          ++resultIdx)
       if (Optional<int64_t> maybeArgNumber =
               getEquivalentFuncArgIdx(funcOp, moduleState, resultIdx))
         if (*maybeArgNumber == opOperand.getOperandNumber())
           result.push_back(callOp->getOpResult(resultIdx));
 
     return result;
   }
 
   SmallVector<OpOperand *>
   getAliasingOpOperand(Operation *op, OpResult opResult,
-                       const BufferizationState &state) const {
+                       const AnalysisState &state) const {
     func::CallOp callOp = cast<func::CallOp>(op);
     FuncOp funcOp = getCalledFunction(callOp);
     assert(funcOp && "expected CallOp to a FuncOp");
-    const ModuleBufferizationState &moduleState =
-        getModuleBufferizationState(state);
+    const ModuleAnalysisState &moduleState = getModuleAnalysisState(state);
 
     // TODO: We should be looking for aliasing block arguments here. The current
     // condition is actually stronger than neccesary. Once we check for aliasing
     // block arguments, we may be multiple.
     if (Optional<int64_t> maybeArgNumber = getEquivalentFuncArgIdx(
             funcOp, moduleState, opResult.getResultNumber()))
       return {&op->getOpOperand(*maybeArgNumber)};
 
     // Note: Returning a non-equivalent tensor from a FuncOp is currently not
     // supported an will fail bufferization.
     return {};
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationState &state) const {
+                                const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
   /// In a first approximation, all the function arguments of a FuncOp are
   /// marked inplaceable. For now, it is the responsibility of the `callOp`
   /// bufferization to allow FuncOp that are inplaceable to write inPlace.
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) 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");
-    const ModuleBufferizationState &moduleState =
-        getModuleBufferizationState(state);
+    const ModuleAnalysisState &moduleState =
+        getModuleAnalysisState(state.getAnalysisState());
 
     // Result types of the bufferized CallOp.
     SmallVector<Type> resultTypes;
     // Replacement values for the existing CallOp. These are usually the results
     // of the bufferized CallOp, unless a tensor result folds onto an operand.
     SmallVector<Value> replacementValues(numResults, Value());
     // For non-tensor results: A mapping from return val indices of the old
     // CallOp to return val indices of the bufferized CallOp.
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     return success();
   }
 };
 
 struct ReturnOpInterface
     : public BufferizableOpInterface::ExternalModel<ReturnOpInterface,
                                                     func::ReturnOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const BufferizationState &state) const {
+                              const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const BufferizationState &state) const {
+                               const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult>
-  getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                      const BufferizationState &state) const {
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
     return {};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) const {
 #ifndef NDEBUG
     auto returnOp = cast<func::ReturnOp>(op);
     assert(isa<FuncOp>(returnOp->getParentOp()) &&
            "only support FuncOp parent for ReturnOp");
 #endif // NDEBUG
     return failure();
   }
 };
 
 struct FuncOpInterface
     : public BufferizableOpInterface::ExternalModel<FuncOpInterface, FuncOp> {
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) const {
     return failure();
   }
 
   /// Return `true` if the given function argument is writable.
   bool isWritable(Operation *op, Value value,
-                  const BufferizationState &state) const {
+                  const AnalysisState &state) const {
     auto funcOp = cast<FuncOp>(op);
     BlockArgument bbArg = value.dyn_cast<BlockArgument>();
     assert(bbArg && "expected BlockArgument");
 
     // "linalg.inplaceable" overrides other writability decisions. This is
     // currently used for testing only.
     if (BoolAttr inplaceAttr = funcOp.getArgAttrOfType<BoolAttr>(
             bbArg.getArgNumber(),

 static void setInPlaceFuncArgument(BlockArgument bbArg, bool inPlace) {
   auto funcOp = cast<FuncOp>(bbArg.getOwner()->getParentOp());
   funcOp.setArgAttr(bbArg.getArgNumber(),
                     BufferizableOpInterface::kInplaceableAttrName,
                     BoolAttr::get(bbArg.getContext(), inPlace));
 }
 
 /// Annotate the IR with the result of the analysis. For testing/debugging only.
-static void
-annotateOpsWithBufferizationMarkers(FuncOp funcOp,
-                                    const BufferizationState &state) {
+static void annotateOpsWithBufferizationMarkers(FuncOp funcOp,
+                                                const AnalysisState &state) {
   auto bufferizableOp = cast<BufferizableOpInterface>(funcOp.getOperation());
   for (BlockArgument bbArg : funcOp.getArguments())
     if (bbArg.getType().isa<TensorType>())
       setInPlaceFuncArgument(bbArg, bufferizableOp.isWritable(bbArg, state));
 }
 
 LogicalResult mlir::linalg::comprehensive_bufferize::runModuleBufferize(
-    ModuleOp moduleOp, AnalysisBufferizationOptions options) {
+    ModuleOp moduleOp, OneShotBufferizationOptions options) {
   IRRewriter rewriter(moduleOp.getContext());
-  AnalysisBufferizationState state(moduleOp, options);
-  ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
-  BufferizationAliasInfo &aliasInfo = state.getAliasInfo();
+  OneShotAnalysisState analysisState(moduleOp, options);
+  BufferizationState bufferizationState(analysisState);
+  ModuleAnalysisState &moduleState = getModuleAnalysisState(analysisState);
+  BufferizationAliasInfo &aliasInfo = analysisState.getAliasInfo();
 
   if (failed(getFuncOpsOrderedByCalls(moduleOp, moduleState.orderedFuncOps,
                                       moduleState.callerMap)))
     return failure();
 
   // Collect bbArg/return value information after the analysis.
   options.addPostAnalysisStep(equivalentFuncOpBBArgsAnalysis);
   options.addPostAnalysisStep(funcOpBbArgReadWriteAnalysis);
 
   // Analyze ops.
   for (FuncOp funcOp : moduleState.orderedFuncOps) {
     // No body => no analysis.
     if (funcOp.getBody().empty())
       continue;
 
     // Now analyzing function.
     moduleState.analyzedFuncOps[funcOp] = FuncOpAnalysisState::InProgress;
 
     // Analyze funcOp.
-    if (failed(analyzeOp(funcOp, state)))
+    if (failed(analyzeOp(funcOp, analysisState)))
       return failure();
 
     // Gather equivalence info for CallOps.
     // TODO: Make this a post-analysis step.
     equivalenceAnalysis(funcOp, aliasInfo, moduleState);
 
     // Mark op as fully analyzed.
     moduleState.analyzedFuncOps[funcOp] = FuncOpAnalysisState::Analyzed;
 
     // Add annotations to function arguments.
     if (options.testAnalysisOnly)
-      annotateOpsWithBufferizationMarkers(funcOp, state);
+      annotateOpsWithBufferizationMarkers(funcOp, analysisState);
   }
 
   if (options.testAnalysisOnly)
     return success();
 
   // Bufferize function bodies.
   for (FuncOp funcOp : moduleState.orderedFuncOps) {
     // No body => no analysis.
     if (funcOp.getBody().empty())
       continue;
 
-    if (failed(bufferizeOp(funcOp, state)))
+    if (failed(bufferizeOp(funcOp, bufferizationState)))
       return failure();
   }
 
   // Bufferize function boundaries.
   for (FuncOp funcOp : moduleState.orderedFuncOps) {
     // Note: It would be good to apply cleanups here but we cannot as aliasInfo
     // would be invalidated.
-    if (failed(bufferizeFuncOpBoundary(funcOp, rewriter, state)))
+    if (failed(bufferizeFuncOpBoundary(funcOp, rewriter, bufferizationState)))
       return failure();
 
     if (!options.allowReturnMemref &&
         llvm::any_of(funcOp.getType().getResults(), [](Type t) {
           return t.isa<MemRefType, UnrankedMemRefType>();
         })) {
       funcOp->emitError("memref return type is unsupported");
       return failure();

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

 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,
-                                       const BufferizationState &state) {
+                                       BufferizationState &state) {
   // 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();
 
for (OpOperand *opOperand : op.getInputOperands()) {
     newInputBuffers.push_back(
         *state.getBuffer(rewriter, *opOperand, /*forceInPlace=*/true));
   }
 
   // New output operands for the cloned op.
   SmallVector<Value> newOutputBuffers;
   for (OpResult opResult : op->getOpResults()) {
     SmallVector<OpOperand *> aliasingOpOperands =
-        state.getAliasingOpOperand(opResult);
+        state.getAnalysisState().getAliasingOpOperand(opResult);
     assert(aliasingOpOperands.size() == 1 && "expected 1 OpOperand");
     FailureOr<Value> resultBuffer =
         state.getBuffer(rewriter, *aliasingOpOperands.front());
     if (failed(resultBuffer))
       return failure();
     newOutputBuffers.push_back(*resultBuffer);
   }
 

 
 /// Bufferization of linalg.generic. Replace with a new linalg.generic that
 /// operates entirely on memrefs.
 template <typename OpTy>
 struct LinalgOpInterface
     : public BufferizableOpInterface::ExternalModel<LinalgOpInterface<OpTy>,
                                                     OpTy> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const BufferizationState &state) const {
+                              const AnalysisState &state) const {
     // Operand is read if it is used in the computation.
     auto genericOp = cast<linalg::LinalgOp>(op);
     return genericOp.payloadUsesValueFromOperand(&opOperand);
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const BufferizationState &state) const {
+                               const AnalysisState &state) const {
     // Operand is written to if it has an aliasing OpResult.
     auto bufferizableOp = cast<BufferizableOpInterface>(op);
     return !bufferizableOp.getAliasingOpResult(opOperand, state).empty();
   }
 
   SmallVector<OpOperand *>
   getAliasingOpOperand(Operation *op, OpResult opResult,
-                       const BufferizationState &state) const {
+                       const AnalysisState &state) const {
     auto genericOp = cast<linalg::LinalgOp>(op);
 
     // By default, the i-th OpResult may alias with the i-th "out" tensor.
     if (state.getOptions().alwaysAliasingWithDest)
       return {genericOp.getOutputOperand(opResult.getResultNumber())};
 
     // We can try to be smart and alias in-place with an "in" tensor if the
     // corresponding "out" tensor is not used in the computation.
     // Aliasing OpOperand/OpResult pairs are computed by `computeAliasingPairs`.
     DenseMap<OpOperand *, OpResult> pairs = computeAliasingPairs(genericOp);
     for (OpOperand *opOperand : genericOp.getInputAndOutputOperands())
       if (pairs[opOperand] == opResult)
         return {opOperand};
     return {};
   }
 
-  SmallVector<OpResult>
-  getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                      const BufferizationState &state) const {
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
     auto genericOp = cast<linalg::LinalgOp>(op);
 
     // By default, the i-th "out" tensor may alias with the i-th OpResult.
     if (state.getOptions().alwaysAliasingWithDest) {
       if (genericOp.isOutputTensor(&opOperand))
         return {genericOp.getTiedOpResult(&opOperand)};
       return {};
     }
 
     // We can try to be smart. See comment in `getAliasingOpOperand`.
     // Aliasing OpOperand/OpResult pairs are computed by `computeAliasingPairs`.
     DenseMap<OpOperand *, OpResult> pairs = computeAliasingPairs(genericOp);
     if (!pairs.count(&opOperand))
       return {};
     return {pairs[&opOperand]};
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationState &state) const {
+                                const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) const {
     return bufferizeLinalgOp(rewriter, cast<LinalgOp>(op), state);
   }
 };
 
 struct InitTensorOpInterface
     : public BufferizableOpInterface::ExternalModel<InitTensorOpInterface,
                                                     linalg::InitTensorOp> {
   bool isMemoryWrite(Operation *op, OpResult opResult,
-                     const BufferizationState &state) const {
+                     const AnalysisState &state) const {
     // InitTensorOps allocate but do not write.
     return false;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) const {
     auto initTensorOp = cast<linalg::InitTensorOp>(op);
 
     // The InitTensorOp may have been eliminated.
     if (initTensorOp->getUses().empty())
       return success();
 
     FailureOr<Value> alloc =
         createAlloc(rewriter, initTensorOp->getLoc(), initTensorOp.result(),

 }
 
 /// Try to eliminate InitTensorOps inside `op`. An InitTensorOp is replaced
 /// with the the result of `rewriteFunc` if it is anchored on a matching
 /// OpOperand. "Anchored" means that there is a path on the reverse SSA use-def
 /// chain, starting from the OpOperand and always following the aliasing
 /// OpOperand, that eventually ends at a single InitTensorOp.
 LogicalResult mlir::linalg::eliminateInitTensors(
-    Operation *op, BufferizationState &state, BufferizationAliasInfo &aliasInfo,
+    Operation *op, AnalysisState &state, BufferizationAliasInfo &aliasInfo,
     AnchorMatchFn anchorMatchFunc, RewriteFn rewriteFunc,
     SmallVector<Operation *> &newOps) {
   OpBuilder b(op->getContext());
 
   WalkResult status = op->walk([&](Operation *op) {
     for (OpOperand &operand : op->getOpOperands()) {
       // Skip operands that do not bufferize inplace.
       if (!aliasInfo.isInPlace(operand))

 /// * On the reverse use-def chain path from the InsertSliceOp to the
 ///   InitTensorOp, all ops were decided to bufferize inplace and the buffer
 ///   relation is "equivalent" (TODO: can be relaxed if needed).
 /// * The reverse use-def chain has exactly one end, which is the InitTensorOp.
 ///
 /// Note that the newly inserted ExtractSliceOp may have to bufferize
 /// out-of-place due to RaW conflicts.
 LogicalResult mlir::linalg::insertSliceAnchoredInitTensorEliminationStep(
-    Operation *op, BufferizationState &state, BufferizationAliasInfo &aliasInfo,
+    Operation *op, AnalysisState &state, BufferizationAliasInfo &aliasInfo,
     SmallVector<Operation *> &newOps) {
   return eliminateInitTensors(
       op, state, aliasInfo,
       /*anchorMatchFunc=*/
       [&](OpOperand &operand, SmallVector<Value> &neededValues) {
         auto insertSliceOp =
             dyn_cast<tensor::InsertSliceOp>(operand.getOwner());
         if (!insertSliceOp)

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

struct LinalgComprehensiveModuleBufferize
     : public LinalgComprehensiveModuleBufferizeBase<
           LinalgComprehensiveModuleBufferize> {
   LinalgComprehensiveModuleBufferize() = default;
 
   LinalgComprehensiveModuleBufferize(
       const LinalgComprehensiveModuleBufferize &p) = default;
 
   explicit LinalgComprehensiveModuleBufferize(
-      const AnalysisBufferizationOptions &options)
+      const OneShotBufferizationOptions &options)
       : options(options) {}
 
   void runOnOperation() override;
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry
         .insert<bufferization::BufferizationDialect, linalg::LinalgDialect,
                 memref::MemRefDialect, tensor::TensorDialect,
                 vector::VectorDialect, scf::SCFDialect,
                 arith::ArithmeticDialect, func::FuncDialect, AffineDialect>();
     affine_ext::registerBufferizableOpInterfaceExternalModels(registry);
     arith::registerBufferizableOpInterfaceExternalModels(registry);
     linalg::registerBufferizableOpInterfaceExternalModels(registry);
     scf::registerBufferizableOpInterfaceExternalModels(registry);
     std_ext::registerModuleBufferizationExternalModels(registry);
     tensor::registerBufferizableOpInterfaceExternalModels(registry);
     vector::registerBufferizableOpInterfaceExternalModels(registry);
   }
 
 private:
-  llvm::Optional<AnalysisBufferizationOptions> options;
+  llvm::Optional<OneShotBufferizationOptions> options;
 };
 } // namespace
 
 static void applyEnablingTransformations(ModuleOp moduleOp) {
   RewritePatternSet patterns(moduleOp.getContext());
   patterns.add<GeneralizePadOpPattern>(moduleOp.getContext());
   (void)applyPatternsAndFoldGreedily(moduleOp, std::move(patterns));
 }
 
 static FailureOr<Value> allocationFnUsingAlloca(OpBuilder &b, Location loc,
                                                 MemRefType type,
                                                 ValueRange dynShape,
                                                 unsigned int bufferAlignment) {
   Value allocated = b.create<memref::AllocaOp>(
       loc, type, dynShape, b.getI64IntegerAttr(bufferAlignment));
   return allocated;
 }
 
 void LinalgComprehensiveModuleBufferize::runOnOperation() {
-  AnalysisBufferizationOptions opt;
+  OneShotBufferizationOptions opt;
   if (!options) {
     // Make new bufferization options if none were provided when creating the
     // pass.
     if (useAlloca) {
       opt.allocationFn = allocationFnUsingAlloca;
       opt.deallocationFn = [](OpBuilder &b, Location loc, Value v) {
         return success();
       };
void LinalgComprehensiveModuleBufferize::runOnOperation() {
   (void)runPipeline(cleanupPipeline, moduleOp);
 }
 
 std::unique_ptr<Pass> mlir::createLinalgComprehensiveModuleBufferizePass() {
   return std::make_unique<LinalgComprehensiveModuleBufferize>();
 }
 
 std::unique_ptr<Pass> mlir::createLinalgComprehensiveModuleBufferizePass(
-    const AnalysisBufferizationOptions &options) {
+    const OneShotBufferizationOptions &options) {
   return std::make_unique<LinalgComprehensiveModuleBufferize>(options);
 }

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

 
 /// Bufferization of scf.execute_region. Can be analyzed, but bufferization not
 /// fully implemented at the moment.
 struct ExecuteRegionOpInterface
     : public BufferizableOpInterface::ExternalModel<ExecuteRegionOpInterface,
                                                     scf::ExecuteRegionOp> {
   SmallVector<OpOperand *>
   getAliasingOpOperand(Operation *op, OpResult opResult,
-                       const BufferizationState &state) const {
+                       const AnalysisState &state) const {
     // ExecuteRegionOps do not have tensor OpOperands. The yielded value can be
     // any SSA value that is in scope. To allow for use-def chain traversal
     // through ExecuteRegionOps in the analysis, the corresponding yield value
     // is considered to be aliasing with the result.
     auto executeRegionOp = cast<scf::ExecuteRegionOp>(op);
     size_t resultNum = std::distance(op->getOpResults().begin(),
                                      llvm::find(op->getOpResults(), opResult));
     // TODO: Support multiple blocks.
     assert(executeRegionOp.getRegion().getBlocks().size() == 1 &&
            "expected exactly 1 block");
     auto yieldOp = dyn_cast<scf::YieldOp>(
         executeRegionOp.getRegion().front().getTerminator());
     assert(yieldOp && "expected scf.yield terminator in scf.execute_region");
     return {&yieldOp->getOpOperand(resultNum)};
   }
 
   // TODO: For better bufferization results, this could return `true` only if
   // there is a memory write in the region.
   bool isMemoryWrite(Operation *op, OpResult opResult,
-                     const BufferizationState &state) const {
+                     const AnalysisState &state) const {
     // Similar to scf.if, results of this op are always considered memory writes
     // 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,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) 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>()) {
         newResultTypes.push_back(getMemRefType(tensorType, state.getOptions()));
       } else {
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
 
     // Replace old op.
     rewriter.replaceOp(executeRegionOp, newResults);
 
     return success();
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationState &state) const {
+                                const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 };
 
 /// Bufferization of scf.if. Replace with a new scf.if that yields memrefs.
 struct IfOpInterface
     : public BufferizableOpInterface::ExternalModel<IfOpInterface, scf::IfOp> {
   SmallVector<OpOperand *>
   getAliasingOpOperand(Operation *op, OpResult opResult,
-                       const BufferizationState &state) const {
+                       const AnalysisState &state) const {
     // IfOps do not have tensor OpOperands. The yielded value can be any SSA
     // value that is in scope. To allow for use-def chain traversal through
     // IfOps in the analysis, both corresponding yield values from the then/else
     // branches are considered to be aliasing with the result.
     auto ifOp = cast<scf::IfOp>(op);
     size_t resultNum = std::distance(op->getOpResults().begin(),
                                      llvm::find(op->getOpResults(), opResult));
     return {&ifOp.thenYield()->getOpOperand(resultNum),
             &ifOp.elseYield()->getOpOperand(resultNum)};
   }
 
   // TODO: For better bufferization results, this could return `true` only if
   // there is a memory write in one (or both) of the branches. Since this is not
   // allowed at the moment, we should never encounter scf.ifs that yield
   // unmodified tensors. Such scf.yield ops could just fold away.
   bool isMemoryWrite(Operation *op, OpResult opResult,
-                     const BufferizationState &state) const {
+                     const AnalysisState &state) const {
     // IfOp results are always considered memory writes in the analysis. This
     // design decision simplifies the analysis considerably. E.g., consider the
     // following test case:
     //
     // %0 = "some_writing_op" : tensor<?xf32>
     // %r = scf.if %c -> (tensor<?xf32>) {
     //   scf.yield %0
     // } else {
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,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) 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>()) {
         newTypes.push_back(getMemRefType(tensorType, state.getOptions()));
       } else {
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
 
     // Replace op results.
     replaceOpWithBufferizedValues(rewriter, op, newIfOp->getResults());
 
     return success();
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationState &state) const {
+                                const AnalysisState &state) const {
     // IfOp results are equivalent to their corresponding yield values if both
     // yield values are equivalent to each other.
     auto bufferizableOp = cast<BufferizableOpInterface>(op);
     SmallVector<OpOperand *> yieldValues =
         bufferizableOp.getAliasingOpOperand(opResult, state);
     assert(yieldValues.size() == 2 && "expected 2 yield values");
     bool equivalentYields = state.areEquivalentBufferizedValues(
         yieldValues[0]->get(), yieldValues[1]->get());
     return equivalentYields ? BufferRelation::Equivalent : BufferRelation::None;
   }
 };
 
 /// Bufferization of scf.for. Replace with a new scf.for that operates on
 /// memrefs.
 struct ForOpInterface
     : public BufferizableOpInterface::ExternalModel<ForOpInterface,
                                                     scf::ForOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const BufferizationState &state) const {
+                              const AnalysisState &state) const {
     // scf::ForOp alone doesn't bufferize to a memory read, one of the uses of
     // its matching bbArg may.
     auto forOp = cast<scf::ForOp>(op);
     return state.isValueRead(forOp.getRegionIterArgForOpOperand(opOperand));
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const BufferizationState &state) const {
+                               const AnalysisState &state) const {
     // Tensor iter_args of scf::ForOps are always considered as a write. This is
     // to simplify the analysis.
     // TODO: Consider doing sth. like isValueWritten.
     return true;
   }
 
-  SmallVector<OpResult>
-  getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                      const BufferizationState &state) const {
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
     auto forOp = cast<scf::ForOp>(op);
     if (!opOperand.get().getType().isa<RankedTensorType>())
       return {};
     return {forOp.getResultForOpOperand(opOperand)};
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationState &state) const {
+                                const AnalysisState &state) const {
     // ForOp results are equivalent to their corresponding init_args if the
     // corresponding iter_args and yield values are equivalent.
     auto forOp = cast<scf::ForOp>(op);
     OpOperand &forOperand = forOp.getOpOperandForResult(opResult);
     auto bbArg = forOp.getRegionIterArgForOpOperand(forOperand);
     auto yieldOp = cast<scf::YieldOp>(&forOp.getLoopBody().front().back());
     bool equivalentYield = state.areEquivalentBufferizedValues(
         bbArg, yieldOp->getOperand(opResult.getResultNumber()));
     return equivalentYield ? BufferRelation::Equivalent : BufferRelation::None;
   }
 
   bool isWritable(Operation *op, Value value,
-                  const BufferizationState &state) const {
+                  const AnalysisState &state) const {
     // Interestingly, scf::ForOp's bbArg can **always** be viewed
     // inplace from the perspective of ops nested under:
     //   1. Either the matching iter operand is not bufferized inplace and an
     //      alloc + optional copy makes the bbArg itself inplaceable.
     //   2. Or the matching iter operand is bufferized inplace and bbArg just
     //      bufferizes to that too.
     return true;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) 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;
     for (const auto &it : llvm::enumerate(forOp.getInitArgs()))
       if (it.value().getType().isa<TensorType>())
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     return success();
   }
 
   /// Assert that yielded values of an scf.for op are aliasing with their
   /// corresponding bbArgs. This is required because the i-th OpResult of an
   /// scf.for op is currently assumed to alias with the i-th iter_arg (in the
   /// absence of conflicts).
   LogicalResult verifyAnalysis(Operation *op,
-                               const BufferizationState &state) const {
+                               const AnalysisState &state) const {
     auto forOp = cast<scf::ForOp>(op);
     auto yieldOp =
         cast<scf::YieldOp>(forOp.getLoopBody().front().getTerminator());
     for (OpOperand &operand : yieldOp->getOpOperands()) {
       auto tensorType = operand.get().getType().dyn_cast<TensorType>();
       if (!tensorType)
         continue;
 

 };
 
 /// Bufferization of scf.yield. Bufferized as part of their enclosing ops, so
 /// this is for analysis only.
 struct YieldOpInterface
     : public BufferizableOpInterface::ExternalModel<YieldOpInterface,
                                                     scf::YieldOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const BufferizationState &state) const {
+                              const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const BufferizationState &state) const {
+                               const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult>
-  getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                      const BufferizationState &state) const {
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
     if (isa<scf::IfOp>(op->getParentOp()))
       return {op->getParentOp()->getResult(opOperand.getOperandNumber())};
     if (isa<scf::ExecuteRegionOp>(op->getParentOp()))
       return {op->getParentOp()->getResult(opOperand.getOperandNumber())};
     return {};
   }
 
   bool mustBufferizeInPlace(Operation *op, OpOperand &opOperand,
-                            const BufferizationState &state) const {
+                            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,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) const {
     auto yieldOp = cast<scf::YieldOp>(op);
     if (!isa<scf::ExecuteRegionOp, scf::IfOp, scf::ForOp>(
             yieldOp->getParentOp()))
       return yieldOp->emitError("unsupported scf::YieldOp parent");
     return success();
   }
 };
 

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

 namespace {
 
 /// Bufferization of shape.assuming.
 struct AssumingOpInterface
     : public BufferizableOpInterface::ExternalModel<AssumingOpInterface,
                                                     shape::AssumingOp> {
   SmallVector<OpOperand *>
   getAliasingOpOperand(Operation *op, OpResult opResult,
-                       const BufferizationState &state) const {
+                       const AnalysisState &state) const {
     // AssumingOps do not have tensor OpOperands. The yielded value can be any
     // SSA value that is in scope. To allow for use-def chain traversal through
     // AssumingOps in the analysis, the corresponding yield value is considered
     // to be aliasing with the result.
     auto assumingOp = cast<shape::AssumingOp>(op);
     size_t resultNum = std::distance(op->getOpResults().begin(),
                                      llvm::find(op->getOpResults(), opResult));
     // TODO: Support multiple blocks.
     assert(assumingOp.getDoRegion().getBlocks().size() == 1 &&
            "expected exactly 1 block");
     auto yieldOp = dyn_cast<shape::AssumingYieldOp>(
         assumingOp.getDoRegion().front().getTerminator());
     assert(yieldOp && "expected shape.assuming_yield terminator");
     return {&yieldOp->getOpOperand(resultNum)};
   }
 
   // TODO: For better bufferization results, this could return `true` only if
   // there is a memory write in the region.
   bool isMemoryWrite(Operation *op, OpResult opResult,
-                     const BufferizationState &state) const {
+                     const AnalysisState &state) const {
     // Similar to scf.if, results of this op are always considered memory writes
     // 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,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) 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>()) {
         newResultTypes.push_back(getMemRefType(tensorType, state.getOptions()));
       } else {
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
 
     // Replace old op.
     rewriter.replaceOp(assumingOp, newResults);
 
     return success();
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationState &state) const {
+                                const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 };
 
 /// Bufferization of shape.assuming_yield. Bufferized as part of their enclosing
 /// ops, so this is for analysis only.
 struct AssumingYieldOpInterface
     : public BufferizableOpInterface::ExternalModel<AssumingYieldOpInterface,
                                                     shape::AssumingOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const BufferizationState &state) const {
+                              const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const BufferizationState &state) const {
+                               const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult>
-  getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                      const BufferizationState &state) const {
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
     assert(isa<shape::AssumingOp>(op->getParentOp()) &&
            "expected that parent is an AssumingOp");
     return {op->getParentOp()->getResult(opOperand.getOperandNumber())};
   }
 
   bool mustBufferizeInPlace(Operation *op, OpOperand &opOperand,
-                            const BufferizationState &state) const {
+                            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,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) const {
     // Op is bufferized as part of AssumingOp.
     return failure();
   }
 };
 
 } // namespace
 } // namespace shape
 } // namespace mlir
 
 void mlir::shape::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addOpInterface<shape::AssumingOp, AssumingOpInterface>();
   registry.addOpInterface<shape::AssumingYieldOp, AssumingYieldOpInterface>();
 }

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

 namespace mlir {
 namespace tensor {
 namespace {
 
 struct CastOpInterface
     : public BufferizableOpInterface::ExternalModel<CastOpInterface,
                                                     tensor::CastOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const BufferizationState &state) const {
+                              const AnalysisState &state) const {
     return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const BufferizationState &state) const {
+                               const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult>
-  getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                      const BufferizationState &state) const {
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
     return {op->getResult(0)};
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationState &state) const {
+                                const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) const {
     auto castOp = cast<tensor::CastOp>(op);
 
     // The result buffer still has the old (pre-cast) type.
     FailureOr<Value> resultBuffer =
         state.getBuffer(rewriter, castOp->getOpOperand(0) /*source*/);
     if (failed(resultBuffer))
       return failure();
     auto sourceMemRefType = resultBuffer->getType().cast<BaseMemRefType>();
struct CastOpInterface
   }
 };
 
 /// Bufferization of tensor.collapse_shape. Replace with memref.collapse_shape.
 struct CollapseShapeOpInterface
     : public BufferizableOpInterface::ExternalModel<CollapseShapeOpInterface,
                                                     tensor::CollapseShapeOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const BufferizationState &state) const {
+                              const AnalysisState &state) const {
     return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const BufferizationState &state) const {
+                               const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult>
-  getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                      const BufferizationState &state) const {
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
     if (&opOperand == &op->getOpOperand(0) /*src*/)
       return {op->getOpResult(0)};
     return {};
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationState &state) const {
+                                const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) const {
     auto collapseShapeOp = cast<tensor::CollapseShapeOp>(op);
     Value buffer =
         *state.getBuffer(rewriter, collapseShapeOp->getOpOperand(0) /*src*/);
     Type resultType =
         getMemRefType(collapseShapeOp.getResultType(), state.getOptions());
     replaceOpWithNewBufferizedOp<memref::CollapseShapeOp>(
         rewriter, op, resultType, buffer, collapseShapeOp.reassociation());
     return success();
   }
 };
 
 /// Bufferization of tensor.dim. Replace with memref.dim.
 struct DimOpInterface
     : public BufferizableOpInterface::ExternalModel<DimOpInterface,
                                                     tensor::DimOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const BufferizationState &state) const {
+                              const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const BufferizationState &state) const {
+                               const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult>
-  getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                      const BufferizationState &state) const {
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
     return {};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) const {
     auto dimOp = cast<tensor::DimOp>(op);
     Value v = *state.getBuffer(rewriter, dimOp->getOpOperand(0) /*source*/);
     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,
                                                     tensor::ExpandShapeOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const BufferizationState &state) const {
+                              const AnalysisState &state) const {
     return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const BufferizationState &state) const {
+                               const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult>
-  getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                      const BufferizationState &state) const {
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
     if (&opOperand == &op->getOpOperand(0) /*src*/)
       return {op->getOpResult(0)};
     return {};
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationState &state) const {
+                                const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) const {
     auto expandShapeOp = cast<tensor::ExpandShapeOp>(op);
     Value buffer =
         *state.getBuffer(rewriter, expandShapeOp->getOpOperand(0) /*src*/);
     Type resultType =
         getMemRefType(expandShapeOp.getResultType(), state.getOptions());
     replaceOpWithNewBufferizedOp<memref::ExpandShapeOp>(
         rewriter, op, resultType, buffer, expandShapeOp.reassociation());
     return success();
   }
 };
 
 /// Bufferization of tensor.extract_slice. Replace with memref.subview.
 struct ExtractSliceOpInterface
     : public BufferizableOpInterface::ExternalModel<ExtractSliceOpInterface,
                                                     tensor::ExtractSliceOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const BufferizationState &state) const {
+                              const AnalysisState &state) const {
     return false;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const BufferizationState &state) const {
+                               const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult>
-  getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                      const BufferizationState &state) const {
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
     if (&opOperand == &op->getOpOperand(0) /*source*/)
       return {op->getOpResult(0)};
     return {};
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationState &state) const {
+                                const AnalysisState &state) const {
     return BufferRelation::None;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) const {
     auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
     Location loc = extractSliceOp.getLoc();
     Value srcMemref =
         *state.getBuffer(rewriter, extractSliceOp->getOpOperand(0) /*source*/,
                          /*forceInPlace=*/true);
     auto srcMemrefType = srcMemref.getType().cast<MemRefType>();
     auto dstTensorType =
         extractSliceOp.result().getType().cast<RankedTensorType>();
 
     // If not inplaceable, alloc.
-    bool inplace = state.isInPlace(extractSliceOp->getOpOperand(0));
+    bool inplace =
+        state.getAnalysisState().isInPlace(extractSliceOp->getOpOperand(0));
     Value alloc;
     if (!inplace) {
       FailureOr<Value> allocOrFailure =
           createAlloc(rewriter, loc, extractSliceOp.result(),
                       state.getOptions().createDeallocs, state.getOptions());
       if (failed(allocOrFailure))
         return failure();
       alloc = *allocOrFailure;
auto subviewMemRefType = memref::SubViewOp::inferRankReducedResultType(
                                  .cast<MemRefType>();
     Value subView = rewriter.create<memref::SubViewOp>(
         loc, subviewMemRefType, srcMemref, mixedOffsets, mixedSizes,
         mixedStrides);
 
     // If not inplaceable, copy.
     if (!inplace) {
       // Do not copy if the copied data is never read.
-      if (state.isValueRead(extractSliceOp.result()))
+      if (state.getAnalysisState().isValueRead(extractSliceOp.result()))
         if (failed(createMemCpy(rewriter, extractSliceOp.getLoc(), subView,
                                 alloc, state.getOptions())))
           return failure();
       subView = alloc;
     }
 
     replaceOpWithBufferizedValues(rewriter, op, subView);
     return success();
   }
 };
 
 /// Bufferization of tensor.extract. Replace with memref.load.
 struct ExtractOpInterface
     : public BufferizableOpInterface::ExternalModel<ExtractOpInterface,
                                                     tensor::ExtractOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const BufferizationState &state) const {
+                              const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const BufferizationState &state) const {
+                               const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult>
-  getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                      const BufferizationState &state) const {
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
     return {};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) const {
     auto extractOp = cast<tensor::ExtractOp>(op);
     Value srcMemref =
         *state.getBuffer(rewriter, extractOp->getOpOperand(0) /*tensor*/);
     replaceOpWithNewBufferizedOp<memref::LoadOp>(rewriter, op, srcMemref,
                                                  extractOp.indices());
     return success();
   }
 };
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,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) 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();
     MemRefType resultType = getContiguousMemRefType(tensorType);
     FailureOr<Value> maybeBuffer =
struct FromElementsOpInterface
   }
 };
 
 /// Bufferization of tensor.generate.
 struct GenerateOpInterface
     : public BufferizableOpInterface::ExternalModel<GenerateOpInterface,
                                                     tensor::GenerateOp> {
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) const {
     auto generateOp = cast<tensor::GenerateOp>(op);
 
     // Allocate memory.
     Location loc = op->getLoc();
     MemRefType memrefType =
         getContiguousMemRefType(generateOp.getType().cast<RankedTensorType>());
     FailureOr<Value> maybeResult =
         createAlloc(rewriter, loc, memrefType, generateOp.dynamicExtents(),
struct GenerateOpInterface
   }
 };
 
 /// Bufferization of tensor.insert. Replace with memref.store.
 struct InsertOpInterface
     : public BufferizableOpInterface::ExternalModel<InsertOpInterface,
                                                     tensor::InsertOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const BufferizationState &state) const {
+                              const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const BufferizationState &state) const {
+                               const AnalysisState &state) const {
     return true;
   }
 
-  SmallVector<OpResult>
-  getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                      const BufferizationState &state) const {
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
     assert(&opOperand == &op->getOpOperand(1) /*dest*/ &&
            "expected dest OpOperand");
     return {op->getOpResult(0)};
   }
 
   SmallVector<OpOperand *>
   getAliasingOpOperand(Operation *op, OpResult opResult,
-                       const BufferizationState &state) const {
+                       const AnalysisState &state) const {
     return {&op->getOpOperand(1) /*dest*/};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) const {
     auto insertOp = cast<tensor::InsertOp>(op);
     FailureOr<Value> destMemref =
         state.getBuffer(rewriter, insertOp->getOpOperand(1) /*dest*/);
     if (failed(destMemref))
       return failure();
     rewriter.create<memref::StoreOp>(insertOp.getLoc(), insertOp.scalar(),
                                      *destMemref, insertOp.indices());
     replaceOpWithBufferizedValues(rewriter, op, *destMemref);
     return success();
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationState &state) const {
+                                const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 };
 
 /// Return true if the (ExtractSliceOp, InsertSliceOp) pair match (i.e.
 /// equivalent operand / result and same offset/sizes/strides specification).
 ///
 /// This is one particular type of relationship between ops on tensors that
 /// reduce to an equivalence on buffers. This should be generalized and
 /// exposed as interfaces on the proper types.
-static bool areEquivalentExtractSliceOps(const BufferizationState &state,
+static bool areEquivalentExtractSliceOps(const AnalysisState &state,
                                          ExtractSliceOp st, InsertSliceOp sti) {
   if (!st || !sti)
     return false;
   if (sti != sti &&
       !state.areEquivalentBufferizedValues(st.source(), sti.dest()))
     return false;
   if (!sameOffsetsSizesAndStrides(st, sti, isEqualConstantIntOrValue))
     return false;
   return true;
 }
 
 /// Return true if `value` is originating from an ExtractSliceOp that matches
 /// the given InsertSliceOp.
-static bool hasMatchingExtractSliceOp(const BufferizationState &state,
-                                      Value value, InsertSliceOp insertOp) {
+static bool hasMatchingExtractSliceOp(const AnalysisState &state, Value value,
+                                      InsertSliceOp insertOp) {
   auto condition = [&](Value val) {
     if (auto extractOp = val.getDefiningOp<ExtractSliceOp>())
       if (areEquivalentExtractSliceOps(state, extractOp, insertOp))
         return true;
     return false;
   };
 
   return llvm::all_of(state.findValueInReverseUseDefChain(value, condition),
                       condition);
 }
 
 /// Bufferization of tensor.insert_slice. Replace with a memory copy. Under
 /// certain circumstances, this op can also be a no-op.
 struct InsertSliceOpInterface
     : public BufferizableOpInterface::ExternalModel<InsertSliceOpInterface,
                                                     tensor::InsertSliceOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const BufferizationState &state) const {
+                              const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const BufferizationState &state) const {
+                               const AnalysisState &state) const {
     return &opOperand == &op->getOpOperand(1) /*dest*/;
   }
 
-  SmallVector<OpResult>
-  getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                      const BufferizationState &state) const {
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
     if (&opOperand == &op->getOpOperand(1) /*dest*/)
       return {op->getResult(0)};
     return {};
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationState &state) const {
+                                const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
   bool isNotConflicting(Operation *op, OpOperand *uRead,
                         OpOperand *uConflictingWrite,
-                        const BufferizationState &state) const {
+                        const AnalysisState &state) const {
     Operation *readingOp = uRead->getOwner();
     Operation *conflictingWritingOp = uConflictingWrite->getOwner();
 
     // Special rules for matching ExtractSliceOp/InsertSliceOp pairs. If
     // uRead is an InsertSliceOp...
     if (auto insertSliceOp = dyn_cast<InsertSliceOp>(readingOp)) {
       // As an example, consider the following IR.
       //
if (auto insertSliceOp = dyn_cast<InsertSliceOp>(conflictingWritingOp))
           hasMatchingExtractSliceOp(state, insertSliceOp.source(),
                                     insertSliceOp))
         return true;
 
     return false;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) 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();
 
struct InsertSliceOpInterface
   }
 };
 
 /// Bufferization of tensor.rank. Replace with memref.rank.
 struct RankOpInterface
     : public BufferizableOpInterface::ExternalModel<RankOpInterface,
                                                     tensor::RankOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const BufferizationState &state) const {
+                              const AnalysisState &state) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const BufferizationState &state) const {
+                               const AnalysisState &state) const {
     return false;
   }
 
-  SmallVector<OpResult>
-  getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                      const BufferizationState &state) const {
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
     return {};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) const {
     auto rankOp = cast<tensor::RankOp>(op);
     Value v = *state.getBuffer(rewriter, rankOp->getOpOperand(0) /*source*/);
     replaceOpWithNewBufferizedOp<memref::RankOp>(rewriter, op, rankOp.getType(),
                                                  v);
     return success();
   }
 };
 

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

 namespace {
 
 /// Bufferization of vector.transfer_read. Replaced with a new
 /// vector.transfer_read that operates on a memref.
 struct TransferReadOpInterface
     : public BufferizableOpInterface::ExternalModel<TransferReadOpInterface,
                                                     vector::TransferReadOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const BufferizationState &state) const {
+                              const AnalysisState &state) const {
     assert(opOperand.get().getType().isa<RankedTensorType>() &&
            "only tensor types expected");
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const BufferizationState &state) const {
+                               const AnalysisState &state) const {
     assert(opOperand.get().getType().isa<RankedTensorType>() &&
            "only tensor types expected");
     return false;
   }
 
-  SmallVector<OpResult>
-  getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                      const BufferizationState &state) const {
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
     return {};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) const {
     auto readOp = cast<vector::TransferReadOp>(op);
     assert(readOp.getShapedType().isa<TensorType>() &&
            "only tensor types expected");
 
     // TransferReadOp always reads from the bufferized op.source().
     Value buffer =
         *state.getBuffer(rewriter, readOp->getOpOperand(0) /*source*/);
     replaceOpWithNewBufferizedOp<vector::TransferReadOp>(
         rewriter, readOp, readOp.getVectorType(), buffer, readOp.indices(),
         readOp.permutation_map(), readOp.padding(), readOp.mask(),
         readOp.in_boundsAttr());
     return success();
   }
 };
 
 /// Bufferization of vector.transfer_write. Replace with a new
 /// vector.transfer_write that operates on a memref.
 struct TransferWriteOpInterface
     : public BufferizableOpInterface::ExternalModel<TransferWriteOpInterface,
                                                     vector::TransferWriteOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
-                              const BufferizationState &state) const {
+                              const AnalysisState &state) const {
     assert(opOperand.get().getType().isa<TensorType>() &&
            "only tensor types expected");
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
-                               const BufferizationState &state) const {
+                               const AnalysisState &state) const {
     assert(opOperand.get().getType().isa<TensorType>() &&
            "only tensor types expected");
     return true;
   }
 
-  SmallVector<OpResult>
-  getAliasingOpResult(Operation *op, OpOperand &opOperand,
-                      const BufferizationState &state) const {
+  SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
+                                            const AnalysisState &state) const {
     assert(opOperand.get().getType().isa<TensorType>() &&
            "only tensor types expected");
     return {op->getOpResult(0)};
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationState &state) const {
+                                const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
-                          const BufferizationState &state) const {
+                          BufferizationState &state) 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.
     // Leave the previous transfer_write to dead code as it still has uses at
     // this point.
     FailureOr<Value> resultBuffer =