[mlir][linalg][bufferize][NFC] Clean up BufferizationState

[springerm]

Summary

Make fields private and clean up the interface. In particular, BufferizableOpInterface::bufferize no longer has access to aliasInfo. This was potentially dangerous because some of the ops registered in BufferizationAliasInfo may have been deleted.

Depends On D114930

[springerm]

rebase

[DavidSpickett]

@springerm This has produced a new warning on one of our bots https://lab.llvm.org/buildbot/#/builders/179/builds/2050/steps/11/logs/warnings__8_.

ComprehensiveBufferize.h:20:1: warning: struct 'BufferizationState' was previously declared as a class; this is valid, but may result in linker errors under the Microsoft C++ ABI [-Wmismatched-tags]

I think you just need to change the forward declaration in that file to "class".

[springerm]

In D114931#3175889, @DavidSpickett wrote:

@springerm This has produced a new warning on one of our bots https://lab.llvm.org/buildbot/#/builders/179/builds/2050/steps/11/logs/warnings__8_.

ComprehensiveBufferize.h:20:1: warning: struct 'BufferizationState' was previously declared as a class; this is valid, but may result in linker errors under the Microsoft C++ ABI [-Wmismatched-tags]

I think you just need to change the forward declaration in that file to "class".

Fixed. Thx for pointing that one.

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

 class FuncOp;
 
 namespace linalg {
 namespace comprehensive_bufferize {
 
 // TODO: from some HW description.
 static constexpr int64_t kBufferAlignments = 128;
 
-struct BufferizationState;
+class BufferizationAliasInfo;
+struct BufferizationOptions;
+class BufferizationState;
+struct PostAnalysisStep;
+
+// Pre-declaration of functions, so that it can be declared a friend below.
+LogicalResult runComprehensiveBufferize(
+    Operation *op, const BufferizationOptions &options,
+    BufferizationState &state,
+    const std::vector<std::unique_ptr<PostAnalysisStep>> &extraSteps);
 
 /// Callback functions that are used to allocate/deallocate/copy memory buffers.
 /// Comprehensive Bufferize provides default implementations of these functions.
 // TODO: Could be replaced with a "bufferization strategy" object with virtual
 // functions in the future.
 struct AllocationCallbacks {
   using AllocationFn = std::function<Optional<Value>(
       OpBuilder &, Location, MemRefType, ArrayRef<Value>)>;

 /// implement custom dialect-specific optimizations.
 struct PostAnalysisStep {
   virtual ~PostAnalysisStep() {}
 
   /// Run the post analysis step. This function may modify the IR, but must keep
   /// `aliasInfo` (inside `state`) consistent. Newly created operations and
   /// operations that should be re-analyzed must be stored in `newOps`.
   virtual LogicalResult run(Operation *op, BufferizationState &state,
+                            BufferizationAliasInfo &aliasInfo,
                             SmallVector<Operation *> &newOps) = 0;
 };
 
 using PostAnalysisStepList = std::vector<std::unique_ptr<PostAnalysisStep>>;
 
 /// Options for ComprehensiveBufferize.
 struct BufferizationOptions {
   BufferizationOptions();

 
 /// Dialect-specific bufferization state. Analysis/bufferization information
 /// that is specific to ops from a certain dialect can be stored in derived
 /// variants of this struct.
 struct DialectBufferizationState {
   virtual ~DialectBufferizationState() = default;
 };
 
-/// BufferizationState keeps track of bufferization state and provides access to
-/// the results of the analysis.
-struct BufferizationState {
+/// BufferizationState keeps track of memory buffers and provides a variety of
+/// helper functions for dealing with them. In particular,
+/// `BufferizableOpInterface::bufferize` implementation should utilize the
+/// following helper functions.
+///
+/// * `createAlloc` / `createDealloc` / `createAllocDeallocPair` creates ops
+///   that allocate and/or deallocate memref buffers.
+/// * `mapBuffer` maps a tensor value to a memref buffer during bufferization.
+/// * `lookupBuffer` returns the mapped memref buffer of a given tensor value.
+/// * `getResultBuffer` returns the memref buffer for a given tensor OpResult.
+///   Based on inplace bufferization decisions of the analysis, it may either
+///   directly return a mapped buffer or allocate a new brand new buffer.
+class BufferizationState {
+public:
   BufferizationState(Operation *op, const BufferizationOptions &options)
       : aliasInfo(op), options(options), builder(op->getContext()) {}
 
   // BufferizationState should be passed as a reference.
   BufferizationState(const BufferizationState &) = delete;
 
-  /// A function that creates an alloc-dealloc pair. This function may perform
-  /// additional optimizations such as buffer allocation hoisting. This function
-  /// calls `allocationFn` and `deallocationFn` to create (de)allocations.
-  Value createAllocDeallocFn(OpBuilder &builder, Location loc,
+  /// Creates a memref allocation.
+  Optional<Value> createAlloc(OpBuilder &b, Location loc, MemRefType type,
+                              ArrayRef<Value> dynShape);
+
+  /// Creates an alloc-dealloc pair. This function may perform additional
+  /// optimizations such as buffer allocation hoisting.
+  Value createAllocDeallocPair(OpBuilder &builder, Location loc,
                                Value shapedValue);
 
+  /// Creates a memref deallocation. The given memref buffer must have been
+  /// allocated using `createAlloc`.
+  void createDealloc(OpBuilder &b, Location loc, Value allocatedBuffer);
+
+  /// Creates a memcpy between two given buffers.
+  void createMemCpy(OpBuilder &b, Location loc, Value from, Value to);
+
   /// Map tensor values to memref buffers.
   void mapBuffer(ValueRange tensors, ValueRange buffers);
 
   /// Map a tensor value to a memref buffer.
   void mapBuffer(Value tensor, Value buffer);
 
   /// Lookup the memref buffer that is associated to the given tensor value.
   /// Asserts if no buffer is associated.
   Value lookupBuffer(Value tensor);
 
+  /// Return `true` if the given OpResult has been decided to bufferize inplace.
+  bool isInPlace(OpResult opResult) const;
+
   /// Return `true` if the given value is mapped.
   bool isMapped(Value value) const;
 
   /// 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.
   Value getResultBuffer(OpResult result);
 
   /// Mark `op` as obsolete, so that it is deleted after bufferization.
   void markOpObsolete(Operation *op);
 
   /// Erase all ops that were marked obsolete.
   void eraseObsoleteOps();
 
   /// Return dialect-specific bufferization state.
   template <typename StateT> StateT &getDialectState(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]);
   }
 
-  /// `aliasInfo` keeps track of aliasing and equivalent values.
+  /// Return a reference to the BufferizationOptions.
+  const BufferizationOptions &getOptions() const { return options; }
+
+  /// Return a reference to the OpBuilder.
+  OpBuilder &getBuilder() { return builder; }
+
+private:
+  friend LogicalResult runComprehensiveBufferize(
+      Operation *op, const BufferizationOptions &options,
+      BufferizationState &state,
+      const std::vector<std::unique_ptr<PostAnalysisStep>> &extraSteps);
+
+  /// `aliasInfo` keeps track of aliasing and equivalent values. Only internal
+  /// functions and `runComprehensiveBufferize` may access this object.
   BufferizationAliasInfo aliasInfo;
 
   /// The mapping of tensors to buffers.
   BlockAndValueMapping mapping;
 
   /// Obsolete ops that should be deleted after bufferization.
   SmallVector<Operation *> obsoleteOps;
 
bool isWritable(Operation *op, Value value, BufferizationState &state) const {
     return false;
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {
     auto isaTensor = [](Type t) { return t.isa<TensorType>(); };
     if (any_of(op->getOperandTypes(), isaTensor) ||
         any_of(op->getResultTypes(), isaTensor))
-      if (!state.options.allowUnknownOps)
+      if (!state.getOptions().allowUnknownOps)
         return op->emitError() << "unsupported op with tensors";
 
     for (Region &region : op->getRegions())
       if (failed(comprehensive_bufferize::bufferize(&region, state)))
         return failure();
 
     return success();
   }

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

struct InitTensorEliminationStep : public PostAnalysisStep {
   ///   `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, BufferizationState &state,
+      BufferizationAliasInfo &aliasInfo,
       std::function<bool(OpOperand &)> anchorMatchFunc,
       std::function<Value(OpBuilder &, Location, OpOperand &)> 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).
 struct InsertSliceAnchoredInitTensorEliminationStep
     : public InitTensorEliminationStep {
   LogicalResult run(Operation *op, BufferizationState &state,
+                    BufferizationAliasInfo &aliasInfo,
                     SmallVector<Operation *> &newOps) override;
 };
 
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 
 } // namespace linalg_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_LINALG_INTERFACE_IMPL_H

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

 namespace linalg {
 namespace comprehensive_bufferize {
 namespace scf_ext {
 
 /// Equivalence analysis for scf.for. Raise an error if iter_args are not
 /// equivalent to their corresponding loop yield values.
 struct AssertDestinationPassingStyle : public PostAnalysisStep {
   LogicalResult run(Operation *op, BufferizationState &state,
+                    BufferizationAliasInfo &aliasInfo,
                     SmallVector<Operation *> &newOps) override;
 };
 
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 
 } // namespace scf_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_SCF_INTERFACE_IMPL_H

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

 class DialectRegistry;
 
 namespace linalg {
 namespace comprehensive_bufferize {
 namespace tensor_ext {
 
 struct InplaceInsertSliceOpAnalysis : public PostAnalysisStep {
   LogicalResult run(Operation *op, BufferizationState &state,
+                    BufferizationAliasInfo &aliasInfo,
                     SmallVector<Operation *> &newOps) override;
 };
 
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 
 } // namespace tensor_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_TENSOR_INTERFACE_IMPL_H

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

if (!aliasInfo.isInPlace(result)) {
     assert(
         aliasingOperands.size() == 1 &&
         "ops with multiple aliasing OpOperands cannot bufferize out-of-place");
     Location loc = op->getLoc();
     // Move insertion point right after `operandBuffer`. That is where the
     // allocation should be inserted (in the absence of allocation hoisting).
     setInsertionPointAfter(builder, operandBuffer);
     // Allocate the result buffer.
-    Value resultBuffer = createAllocDeallocFn(builder, loc, operandBuffer);
+    Value resultBuffer = createAllocDeallocPair(builder, loc, operandBuffer);
     bool skipCopy = false;
     // Do not copy if the last preceding write of `operand` is an op that does
     // 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.
     Value lastWrite = findLastPrecedingWrite(operand);
     if (auto bufferizableOp =
for (Operation *op : ops)
     if (failed(bufferize(op, state)))
       return failure();
   return success();
 }
 
 LogicalResult
 mlir::linalg::comprehensive_bufferize::bufferize(Operation *op,
                                                  BufferizationState &state) {
-  OpBuilder &b = state.builder;
+  OpBuilder &b = state.getBuilder();
 
   // Check if op has tensor results or operands.
   auto isaTensor = [](Type t) { return t.isa<TensorType>(); };
   bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
   bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
   bool isFuncOp = isa<FuncOp>(op);
 
   // No tensor results or operands: Simply bufferize all nested ops.
   // Always pass through FuncOps.
   if (!hasTensorResult && !hasTensorOperand && !isFuncOp) {
     for (Region &region : op->getRegions())
       if (failed(bufferize(&region, state)))
         return failure();
     return success();
   }
 
   // Bufferize using `BufferizableOpInterface`. Interface implementations are
   // responsible for bufferizing nested ops.
   if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op)) {
     b.setInsertionPoint(op);
     return bufferizableOp.bufferize(b, state);
   }
 
   // `op` is an unbufferizable tensor op.
-  if (!state.options.allowUnknownOps)
+  if (!state.getOptions().allowUnknownOps)
     return op->emitError() << "unsupported op with tensors";
 
   // Replace all OpOperands with "to-tensor casted" bufferized values.
   for (OpOperand &operand : op->getOpOperands()) {
     if (operand.get().getType().isa<TensorType>() &&
         state.isMapped(operand.get())) {
       b.setInsertionPoint(op);
       Value toTensorOp = b.create<bufferization::ToTensorOp>(
static MemRefType getAllocationTypeAndShape(OpBuilder &b, Location loc,
 
   return allocMemRefType;
 }
 
 /// Create an Allocop/DeAllocOp pair, where the AllocOp is after
 /// `shapedValue.getDefiningOp` (or at the top of the block in case of a
 /// bbArg) and the DeallocOp is at the end of the block.
 Value mlir::linalg::comprehensive_bufferize::BufferizationState::
-    createAllocDeallocFn(OpBuilder &b, Location loc, Value shapedValue) {
+    createAllocDeallocPair(OpBuilder &b, Location loc, Value shapedValue) {
   // Take a guard before anything else.
   OpBuilder::InsertionGuard g(b);
 
   // 1. Create memory allocation.
   assert(shapedValue.getType().isa<ShapedType>());
   MemRefType memRefType = shapedValue.getType().dyn_cast<MemRefType>();
   SmallVector<Value> dynShape;
   // Note: getAllocationTypeAndShape also sets the insertion point.
   MemRefType allocMemRefType =
       getAllocationTypeAndShape(b, loc, shapedValue, dynShape);
-  Optional<Value> allocated =
-      options.allocationFns->allocationFn(b, loc, allocMemRefType, dynShape);
+  Optional<Value> allocated = createAlloc(b, loc, allocMemRefType, dynShape);
   // TODO: For now just assert the value is returned. Eventually need to
   // error-propagate.
   assert(allocated && "allocation failed");
   Value casted = allocated.getValue();
   if (memRefType && memRefType != allocMemRefType) {
     casted = b.create<memref::CastOp>(loc, memRefType, allocated.getValue());
   }
 
   // 2. Create memory deallocation.
   b.setInsertionPoint(allocated.getValue().getParentBlock()->getTerminator());
-  options.allocationFns->deallocationFn(b, loc, allocated.getValue());
+  createDealloc(b, loc, allocated.getValue());
   return casted;
 }
 
+/// Create a memref allocation.
+Optional<Value>
+mlir::linalg::comprehensive_bufferize::BufferizationState::createAlloc(
+    OpBuilder &b, Location loc, MemRefType type, ArrayRef<Value> dynShape) {
+  return options.allocationFns->allocationFn(b, loc, type, dynShape);
+}
+
+/// Create a memref deallocation.
+void mlir::linalg::comprehensive_bufferize::BufferizationState::createDealloc(
+    OpBuilder &b, Location loc, Value allocatedBuffer) {
+  return options.allocationFns->deallocationFn(b, loc, allocatedBuffer);
+}
+
+/// Create a memory copy between two memref buffers.
+void mlir::linalg::comprehensive_bufferize::BufferizationState::createMemCpy(
+    OpBuilder &b, Location loc, Value from, Value to) {
+  return options.allocationFns->memCpyFn(b, loc, from, to);
+}
+
 //===----------------------------------------------------------------------===//
 // Bufferization-specific BlockAndValueMapping support with debugging.
 //===----------------------------------------------------------------------===//
 
 /// Wrapper for better debugging.
 void mlir::linalg::comprehensive_bufferize::BufferizationState::mapBuffer(
     ValueRange tensors, ValueRange buffers) {
   assert(!tensors.empty() && "unexpected empty tensors");
Value mlir::linalg::comprehensive_bufferize::BufferizationState::lookupBuffer(
   assert((buffer.getType().isa<MemRefType>() ||
           buffer.getType().isa<UnrankedMemRefType>()) &&
          "expected that tensor is mapped to memref");
   return buffer;
 }
 
 bool mlir::linalg::comprehensive_bufferize::BufferizationState::isMapped(
     Value value) const {
+  assert(value.getType().isa<TensorType>() && "unexpected non-tensor type");
   return mapping.contains(value);
 }
 
+bool mlir::linalg::comprehensive_bufferize::BufferizationState::isInPlace(
+    OpResult opResult) const {
+  return aliasInfo.isInPlace(opResult);
+}
+
 void mlir::linalg::comprehensive_bufferize::BufferizationState::markOpObsolete(
     Operation *op) {
   obsoleteOps.push_back(op);
 }
 
 void mlir::linalg::comprehensive_bufferize::BufferizationState::
     eraseObsoleteOps() {
   for (Operation *op : obsoleteOps)

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

LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
   if (failed(inPlaceAnalysis(op, aliasInfo, state, domInfo,
                              options.analysisFuzzerSeed)))
     return failure();
   equivalenceAnalysis(op, aliasInfo);
 
   auto runPostAnalysisSteps = [&](const PostAnalysisStepList &steps) {
     for (const std::unique_ptr<PostAnalysisStep> &step : steps) {
       SmallVector<Operation *> newOps;
-      if (failed(step->run(op, state, newOps)))
+      if (failed(step->run(op, state, aliasInfo, newOps)))
         return failure();
       // Analyze ops that were created by the PostAnalysisStep.
       if (failed(inPlaceAnalysis(newOps, aliasInfo, state, domInfo)))
         return failure();
       equivalenceAnalysis(newOps, aliasInfo);
     }
     return success();
   };

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

struct InitTensorOpInterface
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {
     auto initTensorOp = cast<linalg::InitTensorOp>(op);
 
     // The InitTensorOp may have been eliminated.
     if (initTensorOp->getUses().empty())
       return success();
 
-    Value alloc = state.createAllocDeallocFn(b, initTensorOp->getLoc(),
+    Value alloc = state.createAllocDeallocPair(b, initTensorOp->getLoc(),
                                                initTensorOp.result());
     state.mapBuffer(initTensorOp.result(), alloc);
     return success();
   }
 };
 
 struct TiledLoopOpInterface
     : public BufferizableOpInterface::ExternalModel<TiledLoopOpInterface,
                                                     linalg::TiledLoopOp> {

 /// Try to eliminate InitTensorOps inside funcOp. 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::comprehensive_bufferize::linalg_ext::
     InitTensorEliminationStep::eliminateInitTensors(
         Operation *op, BufferizationState &state,
+        BufferizationAliasInfo &aliasInfo,
         std::function<bool(OpOperand &)> anchorMatchFunc,
         std::function<Value(OpBuilder &, Location, OpOperand &)> rewriteFunc,
         SmallVector<Operation *> &newOps) {
   OpBuilder b(op->getContext());
-  BufferizationAliasInfo &aliasInfo = state.aliasInfo;
 
   WalkResult status = op->walk([&](Operation *op) {
     for (OpOperand &operand : op->getOpOperands()) {
       // Is this a matching OpOperand?
       if (!anchorMatchFunc(operand))
         continue;
 
       SetVector<Value> maybeInitTensor =

 ///   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::comprehensive_bufferize::linalg_ext::
     InsertSliceAnchoredInitTensorEliminationStep::run(
         Operation *op, BufferizationState &state,
-        SmallVector<Operation *> &newOps) {
+        BufferizationAliasInfo &aliasInfo, SmallVector<Operation *> &newOps) {
   return eliminateInitTensors(
-      op, state,
+      op, state, aliasInfo,
       [&](OpOperand &operand) {
         auto insertSliceOp =
             dyn_cast<tensor::InsertSliceOp>(operand.getOwner());
         if (!insertSliceOp)
           return false;
         // Only inplace bufferized InsertSliceOps are eligible.
-        if (!state.aliasInfo.isInPlace(insertSliceOp->getOpResult(0)))
+        if (!aliasInfo.isInPlace(insertSliceOp->getOpResult(0)))
           return false;
         return &operand == &insertSliceOp->getOpOperand(0) /*source*/;
       },
       [](OpBuilder &b, Location loc, OpOperand &operand) {
         auto insertSliceOp = cast<tensor::InsertSliceOp>(operand.getOwner());
         auto extractOp = b.create<tensor::ExtractSliceOp>(
             loc, insertSliceOp.dest(), insertSliceOp.getMixedOffsets(),
             insertSliceOp.getMixedSizes(), insertSliceOp.getMixedStrides());

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

static ReturnOp getAssumedUniqueReturnOp(FuncOp funcOp) {
   return returnOp;
 }
 
 namespace {
 /// Store function BlockArguments that are equivalent to a returned value in
 /// ModuleBufferizationState.
 struct EquivalentFuncOpBBArgsAnalysis : public PostAnalysisStep {
   LogicalResult run(Operation *op, BufferizationState &state,
+                    BufferizationAliasInfo &aliasInfo,
                     SmallVector<Operation *> &newOps) override {
     ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
     auto funcOp = cast<FuncOp>(op);
 
     // Support only single return-terminated block in the function.
     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 (state.aliasInfo.areEquivalentBufferizedValues(returnVal.get(),
-                                                              bbArg))
+            if (aliasInfo.areEquivalentBufferizedValues(returnVal.get(), bbArg))
               moduleState.equivalentFuncArgs[returnVal.getOperandNumber()] =
                   bbArg.getArgNumber();
 
     return success();
   }
 };
 } // namespace
 

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

if (failed(comprehensive_bufferize::bufferize(loopBody, state)))
       return failure();
 
     return success();
   }
 };
 
 LogicalResult mlir::linalg::comprehensive_bufferize::scf_ext::
     AssertDestinationPassingStyle::run(Operation *op, BufferizationState &state,
+                                       BufferizationAliasInfo &aliasInfo,
                                        SmallVector<Operation *> &newOps) {
   LogicalResult status = success();
   op->walk([&](scf::YieldOp yieldOp) {
     auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp());
     if (!forOp)
       return WalkResult::advance();
 
     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 (!state.aliasInfo.areEquivalentBufferizedValues(operand.get(),
-                                                         bbArg)) {
+      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 an equivalent buffer to the matching"
             << " enclosing scf::for operand";
         return WalkResult::interrupt();

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

LogicalResult bufferize(Operation *op, OpBuilder &b,
     auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
     Location loc = extractSliceOp.getLoc();
     Value srcMemref = state.lookupBuffer(extractSliceOp.source());
     auto srcMemrefType = srcMemref.getType().cast<MemRefType>();
     auto dstTensorType =
         extractSliceOp.result().getType().cast<RankedTensorType>();
 
     // If not inplaceable, alloc.
-    bool inplace = state.aliasInfo.isInPlace(extractSliceOp->getResult(0));
+    bool inplace = state.isInPlace(extractSliceOp->getResult(0));
     Value alloc;
     if (!inplace)
-      alloc = state.createAllocDeallocFn(b, loc, extractSliceOp.result());
+      alloc = state.createAllocDeallocPair(b, loc, extractSliceOp.result());
 
     // Bufferize to subview.
     auto subviewMemRefType =
         memref::SubViewOp::inferRankReducedResultType(
             dstTensorType.getRank(), srcMemrefType,
             extractSliceOp.getMixedOffsets(), extractSliceOp.getMixedSizes(),
             extractSliceOp.getMixedStrides())
             .cast<MemRefType>();
     Value subView = b.create<memref::SubViewOp>(
         loc, subviewMemRefType, srcMemref, extractSliceOp.getMixedOffsets(),
         extractSliceOp.getMixedSizes(), extractSliceOp.getMixedStrides());
-    // Insert new alias.
-    state.aliasInfo.insertNewBufferAlias(subView, srcMemref);
 
     /// If not inplaceable, copy.
     if (!inplace) {
       // Do not copy if the copied data is never read.
       if (isValueRead(extractSliceOp.result()))
-        state.options.allocationFns->memCpyFn(b, extractSliceOp.getLoc(),
-                                              subView, alloc);
+        state.createMemCpy(b, extractSliceOp.getLoc(), subView, alloc);
       subView = alloc;
     }
 
     state.mapBuffer(extractSliceOp.result(), subView);
     return success();
   }
 };
 
if (needCopy) {
               insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes(),
               insertSliceOp.getMixedStrides())
               .cast<MemRefType>();
       Value subView = b.create<memref::SubViewOp>(
           loc, subviewMemRefType, dstMemref, insertSliceOp.getMixedOffsets(),
           insertSliceOp.getMixedSizes(), insertSliceOp.getMixedStrides());
       // Copy tensor.
       Value srcMemref = state.lookupBuffer(insertSliceOp.source());
-      state.options.allocationFns->memCpyFn(b, insertSliceOp.getLoc(),
-                                            srcMemref, subView);
+      state.createMemCpy(b, insertSliceOp.getLoc(), srcMemref, subView);
     }
 
     state.mapBuffer(insertSliceOp.result(), dstMemref);
     return success();
   }
 };
 
 } // namespace tensor_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 LogicalResult mlir::linalg::comprehensive_bufferize::tensor_ext::
     InplaceInsertSliceOpAnalysis::run(Operation *op, BufferizationState &state,
+                                      BufferizationAliasInfo &aliasInfo,
                                       SmallVector<Operation *> &newOps) {
   auto &tensorState = getTensorBufferizationState(state);
   op->walk([&](InsertSliceOp insertSliceOp) {
     // A copy of the source buffer is needed if either:
     //   - The producer of `source` is not inplace. This is the case where a
     //     slice is computed out of place into the inplace full tensor.
     //   - The result is not inplace. This is the case where the whole tensor is
     //     cloned and the clone needs to be updated.
-    if (isSourceEquivalentToAMatchingInplaceExtractSliceOp(state.aliasInfo,
+    if (isSourceEquivalentToAMatchingInplaceExtractSliceOp(aliasInfo,
                                                            insertSliceOp) &&
-        state.aliasInfo.isInPlace(insertSliceOp->getResult(0)))
+        state.isInPlace(insertSliceOp->getResult(0)))
       tensorState.insertSliceOpsWithoutCopy.insert(insertSliceOp);
   });
   return success();
 }
 
 void mlir::linalg::comprehensive_bufferize::tensor_ext::
     registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry) {
   registry.addOpInterface<tensor::CastOp, tensor_ext::CastOpInterface>();
   registry.addOpInterface<tensor::DimOp, tensor_ext::DimOpInterface>();
   registry.addOpInterface<tensor::ExtractSliceOp,
                           tensor_ext::ExtractSliceOpInterface>();
   registry.addOpInterface<tensor::ExtractOp, tensor_ext::ExtractOpInterface>();
   registry.addOpInterface<tensor::InsertOp, tensor_ext::InsertOpInterface>();
   registry.addOpInterface<tensor::InsertSliceOp,
                           tensor_ext::InsertSliceOpInterface>();
 }