[mlir][linalg][bufferize] Split analysis-related code from BufferizationState/Options

[springerm]

Summary

This separates the analysis (and its helpers/data structures) more clearly from the rest of the bufferization. BufferizationAliasInfo is no longer passed to BufferizableOpInterface methods. Instead, the respective info can be queried from BufferizationState directly.

Depends On D117476

[springerm]

update

[nicolasvasilache]

Please mark as NFC in the commit message.

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

struct BufferizationOptions {
   using MemCpyFn =
       std::function<LogicalResult(OpBuilder &, Location, Value, Value)>;
 
   BufferizationOptions();
 
   // BufferizationOptions cannot be copied.
   BufferizationOptions(const BufferizationOptions &other) = delete;
 
-  /// Register a "post analysis" step. Such steps are executed after the
-  /// analysis, but before bufferization.
-  template <typename Step, typename... Args>
-  void addPostAnalysisStep(Args... args) {
-    postAnalysisSteps.emplace_back(
-        std::make_unique<Step>(std::forward<Args>(args)...));
-  }
-
   /// Return `true` if the op is allowed to be bufferized.
   bool isOpAllowed(Operation *op) const {
     if (!dialectFilter.hasValue())
       return true;
     return dialectFilter->contains(op->getDialect()->getNamespace());
   }
 
   /// Allow-list the given dialects in the dialect filter. Only ops from
struct BufferizationOptions {
   /// If set to `true`, does not modify the IR apart from adding attributes (for
   /// checking the results of the analysis) and post analysis steps.
   bool testAnalysisOnly = false;
 
   /// If set to `true`, the IR is annotated with details about RaW conflicts.
   /// For debugging only. Should be used together with `testAnalysisOnly`.
   bool printConflicts = false;
 
-  /// Registered post analysis steps.
-  PostAnalysisStepList postAnalysisSteps;
-
   /// Only bufferize ops from dialects that are allowed-listed by the filter.
   /// All other ops are ignored. This option controls the scope of partial
   /// bufferization.
   ///
   /// Note: If no filter is specified, all ops are bufferized (as long as they
   /// implement BufferizableOpInterface). If a filter is specified,
   /// `allowUnknownOps` should be enabled. Otherwise, bufferization would fail
   /// when encountering an op that is forbidden by the filter.
   Optional<DenseSet<StringRef>> dialectFilter;
 
 private:
   /// Allow-list a dialect in the dialect filter.
   template <typename DialectT>
   void addToDialectFilterImpl() {
     if (!dialectFilter.hasValue())
       dialectFilter.emplace();
     dialectFilter->insert(DialectT::getDialectNamespace());
   }
 };
 
+/// Options for analysis-enabled bufferization.
+struct AnalysisBufferizationOptions : public BufferizationOptions {
+  AnalysisBufferizationOptions() = default;
+
+  // AnalysisBufferizationOptions cannot be copied.
+  AnalysisBufferizationOptions(const AnalysisBufferizationOptions &) = delete;
+
+  /// Register a "post analysis" step. Such steps are executed after the
+  /// analysis, but before bufferization.
+  template <typename Step, typename... Args>
+  void addPostAnalysisStep(Args... args) {
+    postAnalysisSteps.emplace_back(
+        std::make_unique<Step>(std::forward<Args>(args)...));
+  }
+
+  /// Registered post analysis steps.
+  PostAnalysisStepList postAnalysisSteps;
+};
+
 /// Specify fine-grain relationship between buffers to enable more analysis.
 enum class BufferRelation {
   None,
   // TODO: ResultContainsOperand,
   // TODO: OperandContainsResult,
   Equivalent
 };
 
public:
   /// 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);
   }
 
-  /// Return true if `v1` and `v2` bufferize to aliasing buffers.
-  bool areAliasingBufferizedValues(Value v1, Value v2) const {
-    return aliasInfo.isEquivalent(v1, v2);
-  }
-
   /// Union the alias sets of `v1` and `v2`.
   void unionAliasSets(Value v1, Value v2) { aliasInfo.unionSets(v1, v2); }
 
   /// Union the equivalence classes of `v1` and `v2`.
   void unionEquivalenceClasses(Value v1, Value v2) {
     equivalentInfo.unionSets(v1, v2);
   }
 
struct DialectBufferizationState {
   // Copying state is forbidden. Always pass as reference.
   DialectBufferizationState(const DialectBufferizationState &) = delete;
 };
 
 /// BufferizationState provides a variety of helper functions for dealing with
 /// tensor values and memref buffers.
 class BufferizationState {
 public:
-  BufferizationState(Operation *op, const BufferizationOptions &options);
-
-  // BufferizationState should be passed as a reference.
-  BufferizationState(const BufferizationState &) = delete;
-
   /// 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 OpResult if the op is not
   /// bufferizable.
   OpResult getAliasingOpResult(OpOperand &opOperand) const;
public:
   /// Furthermore, BlockArguments are also assumed to be writes. There is no
   /// analysis across block boundaries.
   ///
   /// Note: When reaching an end of the reverse SSA use-def chain, that value
   /// is returned regardless of whether it is a memory write or not.
   SetVector<Value> findLastPrecedingWrite(Value value) const;
 
   /// Return `true` if the given OpResult has been decided to bufferize inplace.
-  bool isInPlace(OpOperand &opOperand) const;
+  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;
template <typename StateT> StateT &getOrCreateDialectState(StringRef name) {
     if (!dialectState.count(name))
       dialectState[name] = std::make_unique<StateT>();
     return static_cast<StateT &>(*dialectState[name]);
   }
 
   /// Return a reference to the BufferizationOptions.
   const BufferizationOptions &getOptions() const { return options; }
 
-  /// Return a reference to the BufferizationAliasInfo.
-  BufferizationAliasInfo &getAliasInfo() { return aliasInfo; }
+protected:
+  BufferizationState(const BufferizationOptions &options);
 
-private:
-  /// `aliasInfo` keeps track of aliasing and equivalent values. Only internal
-  /// functions and `runComprehensiveBufferize` may access this object.
-  BufferizationAliasInfo aliasInfo;
+  // BufferizationState should be passed as a reference.
+  BufferizationState(const BufferizationState &) = delete;
+
+  ~BufferizationState() = default;
 
+private:
   /// Dialect-specific bufferization state.
   DenseMap<StringRef, std::unique_ptr<DialectBufferizationState>> dialectState;
 
   /// A reference to current bufferization options.
   const BufferizationOptions &options;
 };
 
+/// 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 {
+public:
+  AnalysisBufferizationState(Operation *op,
+                             const AnalysisBufferizationOptions &options);
+
+  AnalysisBufferizationState(const AnalysisBufferizationState &) = delete;
+
+  virtual ~AnalysisBufferizationState() = 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 `runComprehensiveBufferize` may access this object.
+  BufferizationAliasInfo aliasInfo;
+};
+
 /// Replace an op with replacement values. The op is deleted. Tensor OpResults
 /// must be replaced with memref values.
 void replaceOpWithBufferizedValues(RewriterBase &rewriter, Operation *op,
                                    ValueRange values);
 
 /// Replace an op with a new op. The new op must have the same number of
 /// results as the replaced op. The new op may not return any tensor values.
 template <typename OpTy, typename... Args>
struct AllocationHoistingBarrierOnly
   }
 
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                const BufferizationState &state) const {
     return OpResult();
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationAliasInfo &aliasInfo,
                                 const BufferizationState &state) const {
     return BufferRelation::None;
   }
 
   bool isWritable(Operation *op, Value value,
                   const BufferizationState &state) const {
     return false;
   }

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

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 BufferizationAliasInfo &":$aliasInfo,
                       "const BufferizationState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           // Does not have to be implemented for ops without tensor OpResults
           // that have an aliasing OpOperand.
           llvm_unreachable("bufferRelation not implemented");
         }]
       >,
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 BufferizationAliasInfo &":$aliasInfo),
+                      "const BufferizationState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           return false;
         }]
       >
   ];
 
   let extraClassDeclaration = [{

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

 
 #include "mlir/IR/BuiltinOps.h"
 
 namespace mlir {
 
 namespace linalg {
 namespace comprehensive_bufferize {
 
+class AnalysisBufferizationState;
 class BufferizationAliasInfo;
-struct BufferizationOptions;
+struct AnalysisBufferizationOptions;
 class BufferizationState;
 
 /// Analyze `op` and its nested ops. Bufferization decisions are stored in
 /// `state`.
-LogicalResult analyzeOp(Operation *op, BufferizationState &state);
+LogicalResult analyzeOp(Operation *op, AnalysisBufferizationState &state);
 
 /// Bufferize `op` and its nested ops. Bufferization decisions are stored in
 /// `state`.
 LogicalResult bufferizeOp(Operation *op, const BufferizationState &state);
 
 /// Run Comprehensive Bufferize on the given op: Analysis + Bufferization
-LogicalResult
-runComprehensiveBufferize(Operation *op,
-                          std::unique_ptr<BufferizationOptions> options);
+LogicalResult runComprehensiveBufferize(
+    Operation *op, std::unique_ptr<AnalysisBufferizationOptions> options);
 
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_COMPREHENSIVEBUFFERIZE_H

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

 
 class DialectRegistry;
 struct LogicalResult;
 class ModuleOp;
 
 namespace linalg {
 namespace comprehensive_bufferize {
 
-struct BufferizationOptions;
+struct AnalysisBufferizationOptions;
 
 /// 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 Comprehensive Bufferization.
-LogicalResult
-runComprehensiveBufferize(ModuleOp moduleOp,
-                          std::unique_ptr<BufferizationOptions> options);
+LogicalResult runComprehensiveBufferize(
+    ModuleOp moduleOp, std::unique_ptr<AnalysisBufferizationOptions> 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/lib/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.cpp

return findValueInReverseUseDefChain(value, [&](Value value) {
     auto bufferizableOp = options.dynCastBufferizableOp(op);
     if (!bufferizableOp)
       return true;
     return bufferizableOp.isMemoryWrite(value.cast<OpResult>(), *this);
   });
 }
 
 mlir::linalg::comprehensive_bufferize::BufferizationState::BufferizationState(
-    Operation *op, const BufferizationOptions &options)
-    : aliasInfo(op), options(options) {
+    const BufferizationOptions &options)
+    : options(options) {}
+
+mlir::linalg::comprehensive_bufferize::AnalysisBufferizationState::
+    AnalysisBufferizationState(Operation *op,
+                               const AnalysisBufferizationOptions &options)
+    : BufferizationState(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)) {
mlir::linalg::comprehensive_bufferize::BufferizationState::getBuffer(
     RewriterBase &rewriter, OpOperand &opOperand, bool forceInPlace,
     Optional<Operation *> customCopyInsertionPoint) const {
   OpBuilder::InsertionGuard guard(rewriter);
   Operation *op = opOperand.getOwner();
   Location loc = op->getLoc();
   Value operand = opOperand.get();
   Value operandBuffer = lookupBuffer(rewriter, operand);
 
-  if (forceInPlace || aliasInfo.isInPlace(opOperand))
+  if (forceInPlace || 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,

 
 bool mlir::linalg::comprehensive_bufferize::isFunctionArgument(Value value) {
   auto bbArg = value.dyn_cast<BlockArgument>();
   if (!bbArg)
     return false;
   return isa<FuncOp>(bbArg.getOwner()->getParentOp());
 }
 
-bool mlir::linalg::comprehensive_bufferize::BufferizationState::isInPlace(
-    OpOperand &opOperand) const {
+bool mlir::linalg::comprehensive_bufferize::AnalysisBufferizationState::
+    isInPlace(OpOperand &opOperand) const {
   return aliasInfo.isInPlace(opOperand);
 }
 
+bool mlir::linalg::comprehensive_bufferize::AnalysisBufferizationState::
+    areEquivalentBufferizedValues(Value v1, Value v2) const {
+  return aliasInfo.areEquivalentBufferizedValues(v1, v2);
+}
+
 MemRefType mlir::linalg::comprehensive_bufferize::getContiguousMemRefType(
     ShapedType shapedType, MemRefLayoutAttrInterface layout,
     Attribute memorySpace) {
   return MemRefType::get(shapedType.getShape(), shapedType.getElementType(),
                          layout, memorySpace);
 }
 
 UnrankedMemRefType mlir::linalg::comprehensive_bufferize::getUnrankedMemRefType(

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

for (OpOperand *uConflictingWrite : usesWrite) {
       // No conflict if the reading use equals the use of the conflicting write.
       // A use cannot conflict with itself. Note: Just being the same op is not
       // enough. It has to be the same use.
       if (uConflictingWrite == uRead)
         continue;
 
       // No conflict if the op interface says so.
       if (auto bufferizableOp = options.dynCastBufferizableOp(readingOp))
-        if (bufferizableOp.isNotConflicting(uRead, uConflictingWrite, state,
-                                            aliasInfo))
+        if (bufferizableOp.isNotConflicting(uRead, uConflictingWrite, state))
           continue;
 
       if (conflictingWritingOp != readingOp)
         if (auto bufferizableOp =
                 options.dynCastBufferizableOp(conflictingWritingOp))
-          if (bufferizableOp.isNotConflicting(uRead, uConflictingWrite, state,
-                                              aliasInfo))
+          if (bufferizableOp.isNotConflicting(uRead, uConflictingWrite, state))
             continue;
 
       // Ops are not conflicting if they are in mutually exclusive regions.
       if (insideMutuallyExclusiveRegions(readingOp, conflictingWritingOp))
         continue;
 
       // Check all possible last writes.
       for (Value lastWrite : lastWrites) {
static void equivalenceAnalysis(SmallVector<Operation *> &ops,
                                 BufferizationState &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, aliasInfo, state) ==
+              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,
if (!options.allowUnknownOps) {
 
     if (failed(status))
       return status;
   }
 
   return success();
 }
 
-LogicalResult
-mlir::linalg::comprehensive_bufferize::analyzeOp(Operation *op,
-                                                 BufferizationState &state) {
+LogicalResult mlir::linalg::comprehensive_bufferize::analyzeOp(
+    Operation *op, AnalysisBufferizationState &state) {
   DominanceInfo domInfo(op);
   BufferizationAliasInfo &aliasInfo = state.getAliasInfo();
-  const BufferizationOptions &options = state.getOptions();
+  const auto &options =
+      static_cast<const AnalysisBufferizationOptions &>(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 mlir::linalg::comprehensive_bufferize::bufferizeOp(
   patterns.add<BufferizationPattern>(op->getContext(), state);
   if (failed(applyPatternsAndFoldGreedily(op, std::move(patterns))))
     return failure();
 
   return checkBufferizationResult(op, state.getOptions());
 }
 
 LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
-    Operation *op, std::unique_ptr<BufferizationOptions> options) {
-  BufferizationState state(op, *options);
+    Operation *op, std::unique_ptr<AnalysisBufferizationOptions> options) {
+  AnalysisBufferizationState state(op, *options);
   if (failed(analyzeOp(op, state)))
     return failure();
   if (options->testAnalysisOnly)
     return success();
   return bufferizeOp(op, state);
 }

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

OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
     auto genericOp = cast<linalg::LinalgOp>(op);
 
     // Aliasing OpOperand/OpResult pairs are computed by `computeAliasingPairs`.
     DenseMap<OpOperand *, OpResult> pairs = computeAliasingPairs(genericOp);
     return pairs[&opOperand];
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationAliasInfo &aliasInfo,
                                 const BufferizationState &state) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationState &state) const {
     return bufferizeLinalgOp(rewriter, cast<LinalgOp>(op), state);
   }
OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                const BufferizationState &state) const {
     auto tiledLoopOp = cast<linalg::TiledLoopOp>(op);
 
     // Output operands are tied to their corresponding OpResults.
     return tiledLoopOp.getTiedOpResult(opOperand);
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationAliasInfo &aliasInfo,
                                 const BufferizationState &state) const {
     return BufferRelation::Equivalent;
   }
 
   bool isWritable(Operation *op, Value value,
                   const BufferizationState &state) const {
     // Interestingly, linalg::TiledLoopOp's bbArgs can **always** be viewed
     // inplace from the perspective of nested ops:

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

if (Optional<int64_t> maybeArgNumber = getEquivalentFuncArgIdx(
       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 BufferizationAliasInfo &aliasInfo,
                                 const BufferizationState &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,
annotateOpsWithBufferizationMarkers(FuncOp funcOp,
                                     const BufferizationState &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::runComprehensiveBufferize(
-    ModuleOp moduleOp, std::unique_ptr<BufferizationOptions> options) {
+    ModuleOp moduleOp, std::unique_ptr<AnalysisBufferizationOptions> options) {
   IRRewriter rewriter(moduleOp.getContext());
-  BufferizationState state(moduleOp, *options);
+  AnalysisBufferizationState state(moduleOp, *options);
   ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
   BufferizationAliasInfo &aliasInfo = state.getAliasInfo();
 
   if (failed(getFuncOpsOrderedByCalls(moduleOp, moduleState.orderedFuncOps,
                                       moduleState.callerMap)))
     return failure();
 
   // Collect bbArg/return value information after the analysis.

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

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
 
     // Replace old op.
     rewriter.replaceOp(executeRegionOp, newResults);
 
     return success();
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationAliasInfo &aliasInfo,
                                 const BufferizationState &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> {
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
 
     // Replace op results.
     replaceOpWithBufferizedValues(rewriter, op, newIfOp->getResults());
 
     return success();
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationAliasInfo &aliasInfo,
                                 const BufferizationState &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 = aliasInfo.areEquivalentBufferizedValues(
+    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
OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                const BufferizationState &state) const {
     auto forOp = cast<scf::ForOp>(op);
     if (!opOperand.get().getType().isa<RankedTensorType>())
       return OpResult();
     return forOp.getResultForOpOperand(opOperand);
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationAliasInfo &aliasInfo,
                                 const BufferizationState &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 = aliasInfo.areEquivalentBufferizedValues(
+    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 {
     // Interestingly, scf::ForOp's bbArg can **always** be viewed
     // inplace from the perspective of ops nested under:
op->walk([&](scf::ForOp forOp) {
     for (OpOperand &operand : yieldOp->getOpOperands()) {
       auto tensorType = operand.get().getType().dyn_cast<TensorType>();
       if (!tensorType)
         continue;
 
       OpOperand &forOperand = forOp.getOpOperandForResult(
           forOp->getResult(operand.getOperandNumber()));
       auto bbArg = forOp.getRegionIterArgForOpOperand(forOperand);
-      if (!aliasInfo.areAliasingBufferizedValues(operand.get(), bbArg)) {
+      // Note: This is overly strict. We should check for aliasing bufferized
+      // values. But we don't have a "must-alias" analysis yet.
+      if (!aliasInfo.areEquivalentBufferizedValues(operand.get(), bbArg)) {
         // TODO: this could get resolved with copies but it can also turn into
         // swaps so we need to be careful about order of copies.
         status =
             yieldOp->emitError()
             << "Yield operand #" << operand.getOperandNumber()
             << " does not bufferize to a buffer that is aliasing the matching"
             << " enclosing scf::for operand";
         return WalkResult::interrupt();

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

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     Value falseBuffer =
         *state.getBuffer(rewriter, selectOp->getOpOperand(2) /*false_value*/);
     replaceOpWithNewBufferizedOp<SelectOp>(
         rewriter, op, selectOp.getCondition(), trueBuffer, falseBuffer);
     return success();
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationAliasInfo &aliasInfo,
                                 const BufferizationState &state) const {
     return BufferRelation::None;
   }
 };
 
 } // namespace std_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 void mlir::linalg::comprehensive_bufferize::std_ext::
     registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry) {
   registry.addOpInterface<SelectOp, std_ext::SelectOpInterface>();
 }

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

struct CastOpInterface
   }
 
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                const BufferizationState &state) const {
     return op->getResult(0);
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationAliasInfo &aliasInfo,
                                 const BufferizationState &state) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationState &state) const {
     auto castOp = cast<tensor::CastOp>(op);
 
struct ExtractSliceOpInterface
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                const BufferizationState &state) const {
     return &opOperand == &op->getOpOperand(0) /*source*/
                ? op->getResult(0)
                : OpResult();
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationAliasInfo &aliasInfo,
                                 const BufferizationState &state) const {
     return BufferRelation::None;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationState &state) const {
     auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
     Location loc = extractSliceOp.getLoc();
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 BufferizationAliasInfo &aliasInfo,
                                 const BufferizationState &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 BufferizationAliasInfo &aliasInfo,
+static bool areEquivalentExtractSliceOps(const BufferizationState &state,
                                          ExtractSliceOp st, InsertSliceOp sti) {
   if (!st || !sti)
     return false;
-  if (!aliasInfo.areEquivalentBufferizedValues(st.source(), sti.dest()))
+  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 BufferizationAliasInfo &aliasInfo,
-                                      const BufferizationState &state,
+static bool hasMatchingExtractSliceOp(const BufferizationState &state,
                                       Value value, InsertSliceOp insertOp) {
   auto condition = [&](Value val) {
     if (auto extractOp = val.getDefiningOp<ExtractSliceOp>())
-      if (areEquivalentExtractSliceOps(aliasInfo, extractOp, insertOp))
+      if (areEquivalentExtractSliceOps(state, extractOp, insertOp))
         return true;
     return false;
   };
 
   return llvm::all_of(state.findValueInReverseUseDefChain(value, condition),
                       condition);
 }
 
struct InsertSliceOpInterface
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                const BufferizationState &state) const {
     return &opOperand == &op->getOpOperand(1) /*dest*/
                ? op->getResult(0)
                : OpResult();
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationAliasInfo &aliasInfo,
                                 const BufferizationState &state) const {
     return BufferRelation::Equivalent;
   }
 
   bool isNotConflicting(Operation *op, OpOperand *uRead,
                         OpOperand *uConflictingWrite,
-                        const BufferizationState &state,
-                        const BufferizationAliasInfo &aliasInfo) const {
+                        const BufferizationState &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.
       //
       // %0 = tensor.extract_slice %t[%a, %b][%c, %d][1, 1] {inplace = [true] }
       // %1 = linalg.fill %cst, %0 {inplace= [true] }
       // %2 = tensor.insert_slice %1 into %t[%a, %b][%c, %d][1, 1]
       //     {inplace= [true] }
 
       // TODO: Use insertSliceOp.getDestOpOperand etc. when available.
       if (uRead == &insertSliceOp->getOpOperand(1) /*dest*/ &&
-          hasMatchingExtractSliceOp(aliasInfo, state, uConflictingWrite->get(),
+          hasMatchingExtractSliceOp(state, uConflictingWrite->get(),
                                     insertSliceOp))
         // Case 1: The main insight is that InsertSliceOp reads only part of
         // the destination tensor. The overwritten area is not read. If
         // uConflictingWrite writes into exactly the memory location that is
         // being read by uRead, this is not a conflict.
         //
         // In the above example:
         // uRead             = OpOperand 1 (%t) of tensor.insert_slice
         // uConflictingWrite = OpOperand 1 (%0) of linalg.fill
         //
         // The read of %t does not conflict with the write of the FillOp
         // (same aliases!) because the area that the FillOp operates on is
         // exactly the one that is *not* read via %t.
         return true;
 
       if (uRead == &insertSliceOp->getOpOperand(0) /*source*/ &&
           uConflictingWrite == &insertSliceOp->getOpOperand(1) /*dest*/ &&
-          hasMatchingExtractSliceOp(aliasInfo, state, uRead->get(),
-                                    insertSliceOp))
+          hasMatchingExtractSliceOp(state, uRead->get(), insertSliceOp))
         // Case 2: The read of the source tensor and the write to the dest
         // tensor via an InsertSliceOp is not a conflict if the read is
         // reading exactly that part of an equivalent tensor that the
         // InsertSliceOp is writing.
         //
         // In the above example:
         // uRead             = OpOperand 0 (%1) of tensor.insert_slice
         // uConflictingWrite = OpOperand 1 (%t) of tensor.insert_slice
if (auto insertSliceOp = dyn_cast<InsertSliceOp>(conflictingWritingOp))
       // uRead             = OpOperand 0 (%1) of vector.transfer_read
       // uConflictingWrite = OpOperand 1 (%t) of tensor.insert_slice
       // lastWrite         = %1
       //
       // This is not a conflict because the InsertSliceOp overwrites the
       // memory segment of %1 with the exact same data. (Effectively, there
       // is no memory write here.)
       if (uConflictingWrite == &insertSliceOp->getOpOperand(1) /*dest*/ &&
-          aliasInfo.areEquivalentBufferizedValues(uRead->get(),
-                                                  insertSliceOp.source()) &&
-          hasMatchingExtractSliceOp(aliasInfo, state, insertSliceOp.source(),
+          state.areEquivalentBufferizedValues(uRead->get(),
+          hasMatchingExtractSliceOp(state, insertSliceOp.source(),
                                     insertSliceOp))
         return true;
 
     return false;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationState &state) const {

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

struct TransferWriteOpInterface
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                const BufferizationState &state) const {
     assert(opOperand.get().getType().isa<TensorType>() &&
            "only tensor types expected");
     return op->getOpResult(0);
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
-                                const BufferizationAliasInfo &aliasInfo,
                                 const BufferizationState &state) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationState &state) const {
     auto writeOp = cast<vector::TransferWriteOp>(op);
     assert(writeOp.getShapedType().isa<TensorType>() &&

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

static FailureOr<Value> allocationFnUsingAlloca(OpBuilder &b, Location loc,
                                                 MemRefType type,
                                                 ArrayRef<Value> dynShape) {
   Value allocated = b.create<memref::AllocaOp>(
       loc, type, dynShape, b.getI64IntegerAttr(kBufferAlignments));
   return allocated;
 }
 
 void LinalgComprehensiveModuleBufferize::runOnOperation() {
-  auto options = std::make_unique<BufferizationOptions>();
+  auto options = std::make_unique<AnalysisBufferizationOptions>();
   if (useAlloca) {
     options->allocationFn = allocationFnUsingAlloca;
     options->deallocationFn = [](OpBuilder &b, Location loc, Value v) {
       return success();
     };
   }
 
   options->allowReturnMemref = allowReturnMemref;

mlir/test/lib/Dialect/Linalg/TestComprehensiveBufferize.cpp

struct TestComprehensiveFunctionBufferize
   Option<bool> createDeallocs{
       *this, "create-deallocs",
       llvm::cl::desc("Specify if buffers should be deallocated"),
       llvm::cl::init(true)};
 };
 } // namespace
 
 void TestComprehensiveFunctionBufferize::runOnOperation() {
-  auto options = std::make_unique<BufferizationOptions>();
+  auto options = std::make_unique<AnalysisBufferizationOptions>();
 
   if (!allowReturnMemref)
     options->addPostAnalysisStep<scf_ext::AssertScfForAliasingProperties>();
 
   options->allowReturnMemref = allowReturnMemref;
   options->allowUnknownOps = allowUnknownOps;
   options->testAnalysisOnly = testAnalysisOnly;
   options->analysisFuzzerSeed = analysisFuzzerSeed;