[mlir][linalg][bufferize] Group helpers in BufferizationState

[springerm]

Summary

This simplifies the signature of bufferize.

Depends On D113387

[springerm]

update

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

 /// Note: Unknown ops are handled conservatively and assumed to be writes.
 /// 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.
 Value findLastPrecedingWrite(Value value);
 
+struct BufferizationState;
+
 /// 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, const SmallVector<Value> &)>;
   using DeallocationFn = std::function<void(OpBuilder &, Location, Value)>;
   using MemCpyFn = std::function<void(OpBuilder &, Location, Value, Value)>;
   using CreateAllocDeallocFn =
-      std::function<Value(OpBuilder &, Location, Value,
-                          BufferizationAliasInfo &, AllocationCallbacks &)>;
+      std::function<Value(OpBuilder &, Location, Value, BufferizationState &)>;
 
   AllocationCallbacks(AllocationFn allocFn, DeallocationFn deallocFn,
                       MemCpyFn copyFn, CreateAllocDeallocFn allocDeallocFn)
       : allocationFn(allocFn), deallocationFn(deallocFn), memCpyFn(copyFn),
         createAllocDeallocFn(allocDeallocFn) {}
 
   /// A function that allocates memory.
   AllocationFn allocationFn;
 
   /// A function that deallocated memory. Must be allocated by `allocationFn`.
   DeallocationFn deallocationFn;
 
   /// A function that copies memory between two allocations.
   MemCpyFn memCpyFn;
 
   /// A function that creates an alloc-dealloc pair. This function may perform
   /// additional optimizations such as buffer allocation hoisting. This function
   /// calls `allocationFn` and `deallocationFn` to create (de)allocations.
   CreateAllocDeallocFn createAllocDeallocFn;
 };
 
+/// BufferizationState keeps track of bufferization state and provides access to
+/// the results of the analysis.
+struct BufferizationState {
+  BufferizationState(BufferizationAliasInfo &aliasInfo,
+                     AllocationCallbacks &allocationFns,
+                     BlockAndValueMapping &tensorToBufferMap)
+      : aliasInfo(aliasInfo), allocationFns(allocationFns),
+        tensorToBufferMap(tensorToBufferMap) {}
+
+  /// 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) const;
+
+  /// `aliasInfo` keeps track of aliasing and equivalent values.
+  BufferizationAliasInfo &aliasInfo;
+
+  /// `allocationFns` contains helper functions for creating alloc ops, dealloc
+  /// ops and memcpy ops.
+  AllocationCallbacks &allocationFns;
+
+  /// The mapping of tensors to buffers.
+  BlockAndValueMapping &tensorToBufferMap;
+};
+
 /// 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(OpBuilder &b, OpResult result,
-                      const BlockAndValueMapping &bvm,
-                      BufferizationAliasInfo &aliasInfo,
-                      AllocationCallbacks allocationFns);
+Value getResultBuffer(OpBuilder &b, OpResult result, BufferizationState &state);
 
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h.inc"
 
 namespace mlir {
struct AllocationHoistingBarrierOnly
 
   BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {
     return BufferRelation::None;
   }
 
   bool isWritable(Operation *op, Value value) const { return false; }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
-                          BlockAndValueMapping &bvm,
-                          BufferizationAliasInfo &aliasInfo,
-                          AllocationCallbacks &allocationFn) const {
+                          BufferizationState &state) const {
     auto isaTensor = [](Type t) { return t.isa<TensorType>(); };
     if (any_of(op->getOperandTypes(), isaTensor) ||
         any_of(op->getResultTypes(), isaTensor))
       return op->emitError() << "unsupported op with tensors";
     return success();
   }
 
   bool isAllocationHoistingBarrier(Operation *op) const { return true; }
 };
 
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_BUFFERIZABLEOPINTERFACE_H_

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

let methods = [
         /*methodName=*/"bufferRelation",
         /*args=*/(ins "OpOperand &":$opOperand),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           // Does not have to be implemented for ops without tensor OpOperands.
           llvm_unreachable("bufferRelation not implemented");
         }]
       >,
-      // TODO: Simplify method signature: Pass an OpBuilder and a
-      // BufferizationState object.
       InterfaceMethod<
         /*desc=*/[{
           Bufferize this op, i.e., rewrite it into a memref-based equivalent.
           `bvm` maps tensor values to memref values and this method should map
           tensor results to memref results after creating/modifying ops.
         }],
         /*retType=*/"LogicalResult",
         /*methodName=*/"bufferize",
         /*args=*/(ins "OpBuilder &":$b,
-                      "BlockAndValueMapping &":$bvm,
-                      "BufferizationAliasInfo &":$aliasInfo,
-                      "AllocationCallbacks &":$allocationFn),
+                      "BufferizationState &":$state),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           llvm_unreachable("bufferize not implemented");
           return failure();
         }]
       >,
       InterfaceMethod<
         /*desc=*/[{

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

 namespace comprehensive_bufferize {
 
 struct AllocationCallbacks;
 class BufferizationAliasInfo;
 
 // TODO: from some HW description.
 static constexpr int64_t kBufferAlignments = 128;
 
+struct BufferizationState;
+
 /// Analyze the `ops` to determine which OpResults are inplaceable.
 LogicalResult inPlaceAnalysis(SmallVector<Operation *> &ops,
                               BufferizationAliasInfo &aliasInfo,
                               const DominanceInfo &domInfo,
                               unsigned analysisFuzzerSeed = 0);
 
 // TODO: Do not expose those functions in the header file.
 /// Default allocation function that is used by the comprehensive bufferization

 /// Default memory copy function that is used by the comprehensive bufferization
 /// pass. Creates a `linalg.copy` op.
 void defaultMemCpyFn(OpBuilder &b, Location loc, Value from, Value to);
 
 /// Bufferize one particular op.
 /// `bufferizedFunctionTypes` (resp. `globalCreator`) are expected to be
 /// non-null if `op` is a CallOpInterface (resp. GlobalCreator).
 LogicalResult
-bufferizeOp(Operation *op, BlockAndValueMapping &bvm,
-            BufferizationAliasInfo &aliasInfo,
-            AllocationCallbacks allocationFns,
+bufferizeOp(Operation *op, BufferizationState &state,
             DenseMap<FuncOp, FunctionType> *bufferizedFunctionTypes = nullptr);
 
 /// Register external models implemented for the `BufferizableOpInterface`.
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 
 /// Try to eliminate InitTensorOps inside `funcOp`.
 ///
 /// * `rewriteFunc` generates the replacement for the InitTensorOp.

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

 //===- BufferizableOpInterface.cpp - Comprehensive Bufferize --------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
+#include "mlir/IR/AsmState.h"
 #include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/IR/Value.h"
 #include "llvm/Support/Debug.h"
 
 namespace mlir {
 namespace linalg {
 namespace comprehensive_bufferize {
 
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.cpp.inc"
 
Value mlir::linalg::comprehensive_bufferize::findLastPrecedingWrite(
   assert(result.size() == 1 && "expected exactly one result");
   return result.front();
 }
 
 /// 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 mlir::linalg::comprehensive_bufferize::getResultBuffer(
-    OpBuilder &b, OpResult result, const BlockAndValueMapping &bvm,
-    BufferizationAliasInfo &aliasInfo, AllocationCallbacks allocationFns) {
+    OpBuilder &b, OpResult result, BufferizationState &state) {
   OpBuilder::InsertionGuard guard(b);
   Operation *op = result.getOwner();
   SmallVector<OpOperand *> aliasingOperands = getAliasingOpOperand(result);
   assert(!aliasingOperands.empty() && "could not get aliasing OpOperand");
   OpOperand *opOperand = aliasingOperands.front();
   Value operand = opOperand->get();
-  Value operandBuffer = bvm.lookupOrNull(operand);
-  assert(operandBuffer && "operand buffer not found");
+  Value operandBuffer = state.lookupBuffer(operand);
   // Make sure that all OpOperands are the same buffer. If this is not the case,
   // we would have to materialize a memref value.
   // TODO: Should be looking for checking for "equivalent buffers" instead of
   // operator== here, but equivalent buffers for scf.if yield values are not
   // set up yet.
   if (!llvm::all_of(aliasingOperands, [&](OpOperand *o) {
-        return bvm.lookup(o->get()) == operandBuffer;
+        return state.lookupBuffer(o->get()) == operandBuffer;
       })) {
     op->emitError("result buffer is ambiguous");
     return Value();
   }
 
   // If bufferizing out-of-place, allocate a new buffer.
-  if (!aliasInfo.isInPlace(result)) {
+  if (!state.aliasInfo.isInPlace(result)) {
     // Ops with multiple aliasing operands can currently not bufferize
     // out-of-place.
     assert(
         aliasingOperands.size() == 1 &&
         "ops with multiple aliasing OpOperands cannot bufferize out-of-place");
     Location loc = op->getLoc();
     // Allocate the result buffer.
-    Value resultBuffer = allocationFns.createAllocDeallocFn(
-        b, loc, operand, aliasInfo, allocationFns);
+    Value resultBuffer =
+        state.allocationFns.createAllocDeallocFn(b, loc, operand, state);
     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 =
             lastWrite.getDefiningOp<BufferizableOpInterface>())
       if (!bufferizableOp.isMemoryWrite(lastWrite.cast<OpResult>()))
         skipCopy = true;
     // Do not copy if the copied data is never read.
     if (!isValueRead(result))
       skipCopy = true;
     // Do not copy if this op does not read the data, but writes it.
     if (bufferizesToMemoryWrite(*opOperand) &&
         !bufferizesToMemoryRead(*opOperand))
       skipCopy = true;
     if (!skipCopy) {
       // Set insertion point now that potential alloc/dealloc are introduced.
       b.setInsertionPoint(op);
-      allocationFns.memCpyFn(b, loc, operandBuffer, resultBuffer);
+      state.allocationFns.memCpyFn(b, loc, operandBuffer, resultBuffer);
     }
     return resultBuffer;
   }
 
   // Bufferizing in-place. No need to allocate a new buffer.
   return operandBuffer;
 }
+
+//===----------------------------------------------------------------------===//
+// 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");
+  return tensorToBufferMap.map(tensors, buffers);
+}
+
+/// Wrapper for better debugging.
+void mlir::linalg::comprehensive_bufferize::BufferizationState::mapBuffer(
+    Value tensor, Value buffer) {
+  assert(tensor && "unexpected empty tensor");
+  assert(tensor.getType().isa<TensorType>() && "unexpected non-tensor type");
+  return tensorToBufferMap.map(tensor, buffer);
+}
+
+/// Wrapper for better debugging.
+Value mlir::linalg::comprehensive_bufferize::BufferizationState::lookupBuffer(
+    Value tensor) const {
+  // TODO: if key comes from bbArg, forward.
+  assert(tensor.getType().isa<TensorType>() && "unexpected non-tensor type");
+  Value v = tensorToBufferMap.lookupOrNull(tensor);
+
+  if (!v) {
+    // Dump tensor for easier debugging.
+    tensor.dump();
+    llvm_unreachable("tensor is not mapped");
+    return Value();
+  }
+
+  return v;
+}

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

if (auto candidateOp = dyn_cast<ReturnOp>(b.getTerminator())) {
         return nullptr;
       returnOp = candidateOp;
     }
   }
   return returnOp;
 }
 
 //===----------------------------------------------------------------------===//
-// Bufferization-specific BlockAndValueMapping support with debugging.
-//===----------------------------------------------------------------------===//
-
-/// Wrapper for better debugging.
-static void map(BlockAndValueMapping &bvm, ValueRange keys, ValueRange values) {
-  assert(!keys.empty() && "Unexpected empty keys");
-  LDBG("\n\tMap: " << printValueInfo(keys.front())
-                   << "\n\tto: " << printValueInfo(values.front()) << '\n');
-  return bvm.map(keys, values);
-}
-
-/// Wrapper for better debugging.
-static void map(BlockAndValueMapping &bvm, Value key, Value value) {
-  LDBG("\n\tMap: " << printValueInfo(key) << "\n\tto: " << printValueInfo(value)
-                   << '\n');
-  return bvm.map(key, value);
-}
-
-/// Wrapper for better debugging.
-static Value lookup(const BlockAndValueMapping &bvm, Value key) {
-  // TODO: if key comes from bbArg, forward.
-  assert(key.getType().isa<TensorType>());
-  Value v = bvm.lookupOrNull(key);
-  if (v)
-    return v;
-
-  Operation *parentOp;
-  if (auto bbArg = key.dyn_cast<BlockArgument>()) {
-    if (isa<FuncOp>(key.getParentBlock()->getParentOp()))
-      parentOp = key.getParentBlock()->getParentOp();
-    else
-      parentOp = key.getParentBlock()->getParentOp()->getParentOfType<FuncOp>();
-  } else {
-    parentOp = key.getDefiningOp()->getParentOfType<FuncOp>();
-  }
-  LDBG("In func:\n" << *parentOp << "\nNO VALUE FOR KEY: " << key << '\n');
-  (void)parentOp;
-  return Value();
-}
-
-//===----------------------------------------------------------------------===//
 // Bufferization-specific attribute manipulation.
 // These are for testing and debugging only. Bufferization information is
 // stored in BufferizationAliasInfo. When run with `testAnalysisOnly`, the IR
 // is annotated with the results of the analysis (copied from
 // BufferizationAliasInfo), so that they can be checked in tests.
 //===----------------------------------------------------------------------===//
 
 /// Attribute marker to specify op results that can be bufferized inPlace.
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.
 static Value createNewAllocDeallocPairForShapedValue(
-    OpBuilder &b, Location loc, Value shapedValue,
-    BufferizationAliasInfo &aliasInfo, AllocationCallbacks &allocationFns) {
+    OpBuilder &b, Location loc, Value shapedValue, BufferizationState &state) {
   // 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 =
-      allocationFns.allocationFn(b, loc, allocMemRefType, dynShape);
+      state.allocationFns.allocationFn(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());
-    aliasInfo.insertNewBufferEquivalence(casted, allocated.getValue());
+    state.aliasInfo.insertNewBufferEquivalence(casted, allocated.getValue());
   }
 
   // 2. Create memory deallocation.
   b.setInsertionPoint(allocated.getValue().getParentBlock()->getTerminator());
-  allocationFns.deallocationFn(b, loc, allocated.getValue());
+  state.allocationFns.deallocationFn(b, loc, allocated.getValue());
   return casted;
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization as simple BlockAndValueMapping rewrites.
 //===----------------------------------------------------------------------===//
 
 /// 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.
 static LogicalResult
-bufferize(OpBuilder &b, CallOpInterface callOp, BlockAndValueMapping &bvm,
-          BufferizationAliasInfo &aliasInfo, AllocationCallbacks &allocationFns,
+bufferize(OpBuilder &b, CallOpInterface callOp, BufferizationState &state,
           DenseMap<FuncOp, FunctionType> &bufferizedFunctionTypes) {
   FuncOp funcOp = getCalledFunction(callOp);
   assert(isa<CallOp>(callOp.getOperation()) && funcOp &&
          "expected Callop to a FuncOp");
 
   // If nothing to do then we are done.
   if (!llvm::any_of(funcOp.getType().getInputs(), isaTensor) &&
       !llvm::any_of(funcOp.getType().getResults(), isaTensor))
for (OpOperand &returnOperand : returnOp->getOpOperands()) {
       if (!isaTensor(returnType)) {
         resultTypes.push_back(returnType);
         continue;
       }
 
       // If return operand is equivalent to some bbArg, no need to return it.
       Value returnVal = returnOperand.get();
       if (BlockArgument bbArg =
-              getEquivalentEnclosingFuncBBArg(returnVal, aliasInfo)) {
+              getEquivalentEnclosingFuncBBArg(returnVal, state.aliasInfo)) {
         Value oldRes = callOp->getResult(returnOperand.getOperandNumber());
         int64_t idx = bbArg.getArgNumber();
-        Value buffer = lookup(bvm, callOp->getOperand(idx));
-        assert(buffer && "expected bufferized value");
+        Value buffer = state.lookupBuffer(callOp->getOperand(idx));
         // Add CallOp operand/result equivalence: this is interprocedural info.
-        aliasInfo.insertNewBufferEquivalence(oldRes, buffer);
-        map(bvm, oldRes, buffer);
+        state.aliasInfo.insertNewBufferEquivalence(oldRes, buffer);
+        state.mapBuffer(oldRes, buffer);
         // Add a TensorLoadOp to kill all uses of the CallOp return.
         // Replace all uses of the CallOp results so we can erase the CallOp.
         // This TensorLoadOp must fold/DCE away or bufferization should be
         // considered failed.
         Value tensorLoad =
             b.create<memref::TensorLoadOp>(callOp.getLoc(), buffer);
         oldRes.replaceAllUsesWith(tensorLoad);
         // Add new op equivalence info.
-        aliasInfo.insertNewBufferEquivalence(tensorLoad, buffer);
-        map(bvm, tensorLoad, buffer);
+        state.aliasInfo.insertNewBufferEquivalence(tensorLoad, buffer);
+        state.mapBuffer(tensorLoad, buffer);
         continue;
       }
 
       // TODO: Need to hoist above function boundary.
-      if (Operation *allocOp = getEquivalentAlloc(returnVal, aliasInfo)) {
+      if (Operation *allocOp = getEquivalentAlloc(returnVal, state.aliasInfo)) {
         hoistedArguments.push_back(allocOp->getResult(0));
         continue;
       }
 
       // Other cases legitimately need to return a tensor, this is currently not
       // supported. For instance, if hoisting across function boundary has
       // failed, it may be due to e.g. data-dependent sizes. In such a case, we
       // would we need a better type than memref.
for (OpOperand &opOperand : callOp->getOpOperands()) {
     // Non-tensor operands are just copied.
     if (!tensorOperand.getType().isa<TensorType>()) {
       newOperands.push_back(tensorOperand);
       continue;
     }
 
     // Tensor operands are guaranteed to have been buferized.
     int64_t idx = opOperand.getOperandNumber();
-    Value buffer = lookup(bvm, tensorOperand);
-    assert(buffer && "expected bufferized value");
+    Value buffer = state.lookupBuffer(tensorOperand);
 
     // Caller / callee type mistmatch is handled with a CastOp.
     auto memRefType = bufferizedFuncType.getInput(idx);
     // Since we don't yet have a clear layout story, buffer_cast may
     // conservatively turn tensors into more dynamic memref than necessary.
     // If the memref type of the callee fails, introduce an extra memref.cast
     // that will either canonicalize away or fail compilation until we can do
     // something better.
     if (buffer.getType() != memRefType) {
       Value castBuffer =
           b.create<memref::CastOp>(callOp.getLoc(), memRefType, buffer);
       // Add new op equivalence info.
-      aliasInfo.insertNewBufferEquivalence(castBuffer, buffer);
-      map(bvm, tensorOperand, castBuffer);
+      state.aliasInfo.insertNewBufferEquivalence(castBuffer, buffer);
+      state.mapBuffer(tensorOperand, castBuffer);
       buffer = castBuffer;
     }
     newOperands.push_back(buffer);
   }
 
   // 4. Create the new CallOp.
   Operation *newCallOp = b.create<CallOp>(callOp.getLoc(), funcOp.sym_name(),
                                           resultTypes, newOperands);
   newCallOp->setAttrs(callOp->getAttrs());
   // Delete the op at the end of bufferization.
   return success();
 }
 
 /// FuncOp always creates TensorToMemRef ops.
 static LogicalResult bufferize(OpBuilder &b, FuncOp funcOp,
-                               BlockAndValueMapping &bvm,
-                               BufferizationAliasInfo &aliasInfo,
-                               AllocationCallbacks &allocationFn) {
+                               BufferizationState &state) {
   // Take a guard before anything else.
   OpBuilder::InsertionGuard g(b);
   b.setInsertionPointToStart(&funcOp.body().front());
   for (auto bbArg : funcOp.getArguments()) {
     auto tensorType = bbArg.getType().dyn_cast<TensorType>();
     if (!tensorType)
       continue;
     auto rankedTensorType = tensorType.dyn_cast<RankedTensorType>();
     // Cast the tensor to the most dynamic buffer possible. Further
     // canonicalizations will clean up.
     Type memRefType = rankedTensorType
                           ? getDynamicMemRefType(rankedTensorType)
                           : getContiguousOrUnrankedMemRefType(tensorType);
     Value bufferCast =
         b.create<memref::BufferCastOp>(funcOp.getLoc(), memRefType, bbArg);
-    aliasInfo.insertNewBufferEquivalence(bufferCast, bbArg);
-    map(bvm, bbArg, bufferCast);
+    state.aliasInfo.insertNewBufferEquivalence(bufferCast, bbArg);
+    state.mapBuffer(bbArg, bufferCast);
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization analyses.
 //===----------------------------------------------------------------------===//
 

 void mlir::linalg::comprehensive_bufferize::defaultMemCpyFn(OpBuilder &b,
                                                             Location loc,
                                                             Value from,
                                                             Value to) {
   b.create<CopyOp>(loc, from, to);
 }
 
 LogicalResult mlir::linalg::comprehensive_bufferize::bufferizeOp(
-    Operation *op, BlockAndValueMapping &bvm, BufferizationAliasInfo &aliasInfo,
-    AllocationCallbacks allocationFns,
+    Operation *op, BufferizationState &state,
     DenseMap<FuncOp, FunctionType> *bufferizedFunctionTypes) {
   OpBuilder b(op->getContext());
 
   // CallOps are handled separately.
   if (auto callOp = dyn_cast<CallOpInterface>(op)) {
     LDBG("Begin bufferize:\n" << callOp << '\n');
     if (!bufferizedFunctionTypes)
       llvm_unreachable(
           "null bufferizedFunctionTypes when bufferizing CallOpInterface");
-    return bufferize(b, callOp, bvm, aliasInfo, allocationFns,
-                     *bufferizedFunctionTypes);
+    return bufferize(b, callOp, state, *bufferizedFunctionTypes);
   }
 
   // Skip BufferCast and TensorLoad ops.
   if (isa<memref::BufferCastOp, memref::TensorLoadOp>(op))
     return success();
 
   // Bufferize using `BufferizableOpInterface`.
   if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
-    return bufferizableOp.bufferize(b, bvm, aliasInfo, allocationFns);
+    return bufferizableOp.bufferize(b, state);
 
   // Other op with tensors. No bufferization method specified.
   auto isaTensor = [](Type t) { return t.isa<TensorType>(); };
   if (any_of(op->getOperandTypes(), isaTensor) ||
       any_of(op->getResultTypes(), isaTensor))
     return op->emitError() << "unsupported op with tensors";
   return success();
 }
 
 static LogicalResult bufferizeFuncOpInternals(
-    FuncOp funcOp, BlockAndValueMapping &bvm, BufferizationAliasInfo &aliasInfo,
-    AllocationCallbacks &allocationFns,
+    FuncOp funcOp, BufferizationState &state,
     DenseMap<FuncOp, FunctionType> &bufferizedFunctionTypes) {
   LLVM_DEBUG(llvm::dbgs() << "\n\n");
   LDBG("Begin BufferizeFuncOpInternals:\n" << funcOp << '\n');
   OpBuilder b(funcOp->getContext());
 
   // Start by bufferizing `funcOp` arguments.
-  if (failed(bufferize(b, funcOp, bvm, aliasInfo, allocationFns)))
+  if (failed(bufferize(b, funcOp, state)))
     return failure();
 
   // Cannot erase ops during the traversal. Do that afterwards.
   SmallVector<Operation *> toErase;
 
   auto walkFunc = [&](Operation *op) -> WalkResult {
-    if (failed(bufferizeOp(op, bvm, aliasInfo, allocationFns,
-                           &bufferizedFunctionTypes)))
+    if (failed(bufferizeOp(op, state, &bufferizedFunctionTypes)))
       return failure();
 
     // Register post-walk erasure, if necessary.
     if (isa<CallOpInterface>(op))
       if (llvm::any_of(op->getOperandTypes(), isaTensor) ||
           llvm::any_of(op->getResultTypes(), isaTensor))
         toErase.push_back(op);
 
for (FuncOp funcOp : orderedFuncOps) {
     // bufferization phase.
     if (failed(eliminateInsertSliceAnchoredInitTensorOps(funcOp, aliasInfo,
                                                          domInfo)))
       return failure();
 
     // Bufferization phase.
     if (!options.testAnalysisOnly) {
       BlockAndValueMapping tensorToBufferMap;
-      if (failed(bufferizeFuncOpInternals(funcOp, tensorToBufferMap, aliasInfo,
-                                          *options.allocationFns,
-                                          bufferizedFunctionTypes)))
+      BufferizationState state(aliasInfo, *options.allocationFns,
+                               tensorToBufferMap);
+      if (failed(
+              bufferizeFuncOpInternals(funcOp, state, bufferizedFunctionTypes)))
         return failure();
     }
   }
   // Annotate operations if we only want to report the analysis.
   if (options.testAnalysisOnly) {
     annotateOpsWithBufferizationMarkers(moduleOp, aliasInfo);
     return success();
   }
struct ConstantOpInterface
     : public BufferizableOpInterface::ExternalModel<ConstantOpInterface,
                                                     arith::ConstantOp> {
   SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
                                                 OpResult opResult) const {
     return {};
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
-                          BlockAndValueMapping &bvm,
-                          BufferizationAliasInfo &aliasInfo,
-                          AllocationCallbacks &allocationFn) const {
+                          BufferizationState &state) const {
     auto constantOp = cast<arith::ConstantOp>(op);
     if (!isaTensor(constantOp.getResult().getType()))
       return success();
     assert(constantOp.getType().dyn_cast<RankedTensorType>() &&
            "not a constant ranked tensor");
     auto moduleOp = constantOp->getParentOfType<ModuleOp>();
     if (!moduleOp) {
       return constantOp.emitError(
           "cannot bufferize constants not within builtin.module op");
     }
     GlobalCreator globalCreator(moduleOp);
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPoint(constantOp);
 
     auto globalMemref = globalCreator.getGlobalFor(constantOp);
     Value memref = b.create<memref::GetGlobalOp>(
         constantOp.getLoc(), globalMemref.type(), globalMemref.getName());
-    aliasInfo.insertNewBufferEquivalence(memref, constantOp.getResult());
-    map(bvm, constantOp, memref);
+    state.aliasInfo.insertNewBufferEquivalence(memref, constantOp.getResult());
+    state.mapBuffer(constantOp, memref);
 
     return success();
   }
 
   bool isWritable(Operation *op, Value value) const {
     // Memory locations returned by memref::GetGlobalOp may not be written to.
     assert(value.isa<OpResult>());
     return false;
   }
 };
 
 } // namespace arith_ext
 
 // TODO: Ops in the linalg dialect can directly implement this interface.
 namespace linalg_ext {
 
 /// Helper function for LinalgOp bufferization.
 /// When allocating a new buffer, analyze whether `op` wants to read form that
 /// buffer. Only in that case, a copy of the result buffer may be needed.
-static LogicalResult allocateBuffersForResults(
-    OpBuilder &b, Location loc, LinalgOp op,
-    SmallVectorImpl<Value> &resultBuffers, BlockAndValueMapping &bvm,
-    BufferizationAliasInfo &aliasInfo, AllocationCallbacks &allocationFns) {
+static LogicalResult
+allocateBuffersForResults(OpBuilder &b, Location loc, LinalgOp op,
+                          SmallVectorImpl<Value> &resultBuffers,
+                          BufferizationState &state) {
   // Take a guard before anything else.
   OpBuilder::InsertionGuard g(b);
   b.setInsertionPoint(op);
 
   // TODO: provide the proper interface to iterate on OpResults and get the
   // matching OpOperands.
   for (OpOperand *opOperand : op.getOutputOperands()) {
     OpResult opResult = cast<BufferizableOpInterface>(op.getOperation())
                             .getAliasingOpResult(*opOperand);
     assert(opResult && "could not find correspond OpResult");
-    Value resultBuffer =
-        getResultBuffer(b, opResult, bvm, aliasInfo, allocationFns);
+    Value resultBuffer = getResultBuffer(b, opResult, state);
     if (!resultBuffer)
       return failure();
     resultBuffers.push_back(resultBuffer);
   }
 
   if (op->getNumResults())
-    map(bvm, op->getResults(), resultBuffers);
+    state.mapBuffer(op->getResults(), resultBuffers);
 
   return success();
 }
 
 /// Generic conversion for any LinalgOp on tensors.
 static LogicalResult bufferizeLinalgOp(OpBuilder &b, LinalgOp op,
-                                       BlockAndValueMapping &bvm,
-                                       BufferizationAliasInfo &aliasInfo,
-                                       AllocationCallbacks &allocationFns) {
+                                       BufferizationState &state) {
   // Take a guard before anything else.
   OpBuilder::InsertionGuard g(b);
 
   // Ensure op has only tensors. Allow mixed tensor-buffer mode on a per-need
   // basis.
   if (!op.hasTensorSemantics())
     return op->emitError() << "op does not have tensor semantics";
 
   Location loc = op.getLoc();
   SmallVector<Value> newInputBuffers;
   newInputBuffers.reserve(op.getNumInputs());
   for (OpOperand *opOperand : op.getInputOperands()) {
     if (op.isScalar(opOperand)) {
       newInputBuffers.push_back(opOperand->get());
       continue;
     }
-    newInputBuffers.push_back(lookup(bvm, opOperand->get()));
-    assert(newInputBuffers.back() && "missing buffer");
+    newInputBuffers.push_back(state.lookupBuffer(opOperand->get()));
   }
   SmallVector<Value> newOutputBuffers;
   // Try to allocate new buffers depending on op's inplace semantics.
-  if (failed(allocateBuffersForResults(b, loc, op, newOutputBuffers, bvm,
-                                       aliasInfo, allocationFns)))
+  if (failed(allocateBuffersForResults(b, loc, op, newOutputBuffers, state)))
     return failure();
 
   // Clone the newly bufferized op.
   SmallVector<Value> newOperands = newInputBuffers;
   newOperands.append(newOutputBuffers.begin(), newOutputBuffers.end());
 
   // Set insertion point now that potential alloc/dealloc are introduced.
   b.setInsertionPoint(op);
   op.clone(b, loc, /*resultTypes=*/TypeRange{}, newOperands);
 
   // Replace the results of the old op with the new output buffers.
   if (op->getNumResults())
-    map(bvm, op->getResults(), newOutputBuffers);
+    state.mapBuffer(op->getResults(), newOutputBuffers);
 
   // The original op will be DCE'd away later.
 
   return success();
 }
 
 template <typename OpTy>
 struct LinalgOpInterface
OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     return genericOp->getResult(outputOperandIndex - numOutputBuffers);
   }
 
   BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
-                          BlockAndValueMapping &bvm,
-                          BufferizationAliasInfo &aliasInfo,
-                          AllocationCallbacks &allocationFn) const {
-    return bufferizeLinalgOp(b, cast<LinalgOp>(op), bvm, aliasInfo,
-                             allocationFn);
+                          BufferizationState &state) const {
+    return bufferizeLinalgOp(b, cast<LinalgOp>(op), state);
   }
 };
 
 struct InitTensorOpInterface
     : public BufferizableOpInterface::ExternalModel<InitTensorOpInterface,
                                                     linalg::InitTensorOp> {
   SmallVector<OpOperand *> getAliasingOpOperand(Operation *op,
                                                 OpResult opResult) const {
     return {};
   }
 
   bool isMemoryWrite(Operation *op, OpResult opResult) const {
     // InitTensorOps allocate but do not write.
     return false;
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
-                          BlockAndValueMapping &bvm,
-                          BufferizationAliasInfo &aliasInfo,
-                          AllocationCallbacks &allocationFn) const {
+                          BufferizationState &state) const {
     auto initTensorOp = cast<linalg::InitTensorOp>(op);
 
     // The InitTensorOp may have been eliminated.
     if (initTensorOp->getUses().empty())
       return success();
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPoint(initTensorOp);
 
     Value alloc = createNewAllocDeallocPairForShapedValue(
-        b, initTensorOp->getLoc(), initTensorOp.result(), aliasInfo,
-        allocationFn);
-    map(bvm, initTensorOp.result(), alloc);
+        b, initTensorOp->getLoc(), initTensorOp.result(), state);
+    state.mapBuffer(initTensorOp.result(), alloc);
     return success();
   }
 };
 
 struct TiledLoopOpInterface
     : public BufferizableOpInterface::ExternalModel<TiledLoopOpInterface,
                                                     linalg::TiledLoopOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand) const {
bool isWritable(Operation *op, Value value) const {
     //   2. Or the matching iter operand is bufferized inplace and bbArg just
     //      bufferizes to that too.
     return true;
   }
 
   bool isAllocationHoistingBarrier(Operation *op) const { return true; }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
-                          BlockAndValueMapping &bvm,
-                          BufferizationAliasInfo &aliasInfo,
-                          AllocationCallbacks &allocationFn) const {
+                          BufferizationState &state) const {
     auto tiledLoopOp = cast<linalg::TiledLoopOp>(op);
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);
 
     // Allocate output buffers if needed, forward output tensor args to the
     // terminator.
     Operation *yieldOp = tiledLoopOp.getBody()->getTerminator();
for (Value oldOutputTensor : oldOutputs) {
         continue;
       }
 
       assert(oldOutputTensor.getType().isa<RankedTensorType>() &&
              "bufferizable output must be a ranked tensor");
 
       const OpResult &opResult = tiledLoopOp->getResult(resultIndex);
       OpOperand &yieldOperand = yieldOp->getOpOperand(resultIndex);
-      Value resultBuffer =
-          getResultBuffer(b, opResult, bvm, aliasInfo, allocationFn);
+      Value resultBuffer = getResultBuffer(b, opResult, state);
       if (!resultBuffer)
         return failure();
 
       // Insert mapping and aliasing info.
-      aliasInfo.createAliasInfoEntry(resultBuffer);
-      aliasInfo.insertNewBufferEquivalence(opResult, resultBuffer);
-      map(bvm, opResult, resultBuffer);
+      state.aliasInfo.createAliasInfoEntry(resultBuffer);
+      state.aliasInfo.insertNewBufferEquivalence(opResult, resultBuffer);
+      state.mapBuffer(opResult, resultBuffer);
 
       // Insert new operand and bbArg.
       tiledLoopOp->insertOperands(nextOutputOperandIndex, resultBuffer);
       BlockArgument newBufferBBArg =
           body->insertArgument(nextOutputBBArgIndex, resultBuffer.getType());
       BlockArgument oldTensorBBArg = body->getArgument(oldOutputBBArgIndex);
       // Insert mapping and aliasing info.
-      aliasInfo.createAliasInfoEntry(newBufferBBArg);
-      aliasInfo.insertNewBufferEquivalence(oldTensorBBArg, newBufferBBArg);
-      map(bvm, oldTensorBBArg, newBufferBBArg);
+      state.aliasInfo.createAliasInfoEntry(newBufferBBArg);
+      state.aliasInfo.insertNewBufferEquivalence(oldTensorBBArg,
+                                                 newBufferBBArg);
+      state.mapBuffer(oldTensorBBArg, newBufferBBArg);
 
       // Set operand of `linalg.yield` to the bbArg so it just canonicalizes
       // away later.
       yieldOperand.set(oldTensorBBArg);
 
       // Increment indices.
       ++numNewOutputBuffers;
       ++resultIndex;
LogicalResult bufferize(Operation *op, OpBuilder &b,
     int nextInputOperandIndex = numControlOperands + oldInputs.size();
     for (Value oldInputTensor : oldInputs) {
       if (!oldInputTensor.getType().isa<TensorType>()) {
         // Skip and increment the old bbarg index only.
         ++oldInputBBArgIndex;
         continue;
       }
 
-      Value inputBuffer = lookup(bvm, oldInputTensor);
-      assert(inputBuffer && " missing buffer for operand");
+      Value inputBuffer = state.lookupBuffer(oldInputTensor);
 
       // Insert new operand and bbArg.
       tiledLoopOp->insertOperands(nextInputOperandIndex, inputBuffer);
       BlockArgument newBufferBBArg =
           body->insertArgument(nextInputBBArgIndex, inputBuffer.getType());
       BlockArgument oldTensorBBArg = body->getArgument(oldInputBBArgIndex);
 
       // Insert mapping and aliasing info.
-      aliasInfo.createAliasInfoEntry(newBufferBBArg);
-      aliasInfo.insertNewBufferEquivalence(oldTensorBBArg, newBufferBBArg);
-      map(bvm, oldTensorBBArg, newBufferBBArg);
+      state.aliasInfo.createAliasInfoEntry(newBufferBBArg);
+      state.aliasInfo.insertNewBufferEquivalence(oldTensorBBArg,
+                                                 newBufferBBArg);
+      state.mapBuffer(oldTensorBBArg, newBufferBBArg);
 
       // Increment indices.
       ++numNewInputBuffers;
       ++oldInputBBArgIndex;
       ++nextInputBBArgIndex;
       ++nextInputOperandIndex;
     }
 
bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand) const {
     return false;
   }
 
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     return OpResult();
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
-                          BlockAndValueMapping &bvm,
-                          BufferizationAliasInfo &aliasInfo,
-                          AllocationCallbacks &allocationFn) const {
+                          BufferizationState &state) const {
     auto yieldOp = cast<linalg::YieldOp>(op);
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);
     // Cannot create IR past a yieldOp.
     b.setInsertionPoint(yieldOp);
 
     // No tensors -> success.
struct IfOpInterface
 
   bool mustBufferizeInPlace(Operation *op, OpResult opResult) const {
     // IfOp results always bufferize in-place. Since they have no OpOperands,
     // they are mostly ignored by the analysis once alias sets are set up.
     return true;
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
-                          BlockAndValueMapping &bvm,
-                          BufferizationAliasInfo &aliasInfo,
-                          AllocationCallbacks &allocationFn) const {
+                          BufferizationState &state) const {
     // scf::IfOp is bufferized after scf::YieldOp in the else branch.
     return success();
   }
 };
 
 /// Bufferize the scf::IfOp. This function is called after the YieldOp was
 /// bufferized.
 static LogicalResult bufferizeIfOp(scf::IfOp ifOp, OpBuilder &b,
-                                   BlockAndValueMapping &bvm,
-                                   BufferizationAliasInfo &aliasInfo,
-                                   AllocationCallbacks &allocationFn) {
+                                   BufferizationState &state) {
   // Take a guard before anything else.
   OpBuilder::InsertionGuard g(b);
   b.setInsertionPoint(ifOp);
 
   for (OpResult opResult : ifOp->getResults()) {
     if (!opResult.getType().isa<TensorType>())
       continue;
     // TODO: Atm we bail on unranked TensorType because we don't know how to
     // alloc an UnrankedMemRefType + its underlying ranked MemRefType.
     assert(opResult.getType().isa<RankedTensorType>() &&
            "unsupported unranked tensor");
 
-    Value resultBuffer =
-        getResultBuffer(b, opResult, bvm, aliasInfo, allocationFn);
+    Value resultBuffer = getResultBuffer(b, opResult, state);
     if (!resultBuffer)
       return failure();
 
-    aliasInfo.createAliasInfoEntry(resultBuffer);
-    map(bvm, opResult, resultBuffer);
+    state.aliasInfo.createAliasInfoEntry(resultBuffer);
+    state.mapBuffer(opResult, resultBuffer);
   }
 
   return success();
 }
 
 struct ForOpInterface
     : public BufferizableOpInterface::ExternalModel<ForOpInterface,
                                                     scf::ForOp> {
bool isWritable(Operation *op, Value value) const {
     //   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, OpBuilder &b,
-                          BlockAndValueMapping &bvm,
-                          BufferizationAliasInfo &aliasInfo,
-                          AllocationCallbacks &allocationFn) const {
+                          BufferizationState &state) const {
     // Note: This method is just setting up the mappings for the block arguments
     // and the result buffer. The op is bufferized after the scf::YieldOp.
 
     auto forOp = cast<scf::ForOp>(op);
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);
 
     for (OpResult opResult : forOp->getResults()) {
       if (!opResult.getType().isa<TensorType>())
         continue;
       // TODO: Atm we bail on unranked TensorType because we don't know how to
       // alloc an UnrankedMemRefType + its underlying ranked MemRefType.
       assert(opResult.getType().isa<RankedTensorType>() &&
              "unsupported unranked tensor");
 
       // TODO: More general: Matching bbArg does not bufferize to a read.
-      Value resultBuffer =
-          getResultBuffer(b, opResult, bvm, aliasInfo, allocationFn);
+      Value resultBuffer = getResultBuffer(b, opResult, state);
       if (!resultBuffer)
         return failure();
 
       OpOperand &opOperand = forOp.getOpOperandForResult(opResult);
       BlockArgument bbArg = forOp.getRegionIterArgForOpOperand(opOperand);
-      aliasInfo.createAliasInfoEntry(resultBuffer);
-      aliasInfo.insertNewBufferEquivalence(bbArg, resultBuffer);
-      map(bvm, bbArg, resultBuffer);
-      map(bvm, opResult, resultBuffer);
+      state.aliasInfo.createAliasInfoEntry(resultBuffer);
+      state.aliasInfo.insertNewBufferEquivalence(bbArg, resultBuffer);
+      state.mapBuffer(bbArg, resultBuffer);
+      state.mapBuffer(opResult, resultBuffer);
     }
 
     return success();
   }
 };
 
 /// Bufferize the scf::ForOp. This function is called after the YieldOp was
 /// bufferized.
 static LogicalResult bufferizeForOp(scf::ForOp forOp, OpBuilder &b,
-                                    BlockAndValueMapping &bvm,
-                                    BufferizationAliasInfo &aliasInfo,
-                                    AllocationCallbacks &allocationFn) {
+                                    BufferizationState &state) {
   auto yieldOp = cast<scf::YieldOp>(&forOp.region().front().back());
   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);
-    Value yieldedBuffer = lookup(bvm, operand.get());
-    Value bbArgBuffer = lookup(bvm, bbArg);
-    if (!aliasInfo.areEquivalentBufferizedValues(yieldedBuffer, bbArgBuffer)) {
+    Value yieldedBuffer = state.lookupBuffer(operand.get());
+    Value bbArgBuffer = state.lookupBuffer(bbArg);
+    if (!state.aliasInfo.areEquivalentBufferizedValues(yieldedBuffer,
+                                                       bbArgBuffer)) {
       // TODO: this could get resolved with copies but it can also turn into
       // swaps so we need to be careful about order of copies.
       return yieldOp->emitError()
              << "Yield operand #" << operand.getOperandNumber()
              << " does not bufferize to an equivalent buffer to the matching"
              << " enclosing scf::for operand";
     }
 
OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     return OpResult();
   }
 
   BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
-                          BlockAndValueMapping &bvm,
-                          BufferizationAliasInfo &aliasInfo,
-                          AllocationCallbacks &allocationFn) const {
+                          BufferizationState &state) const {
     auto yieldOp = cast<scf::YieldOp>(op);
 
     if (auto execOp = dyn_cast<scf::ExecuteRegionOp>(yieldOp->getParentOp())) {
       if (execOp->getNumResults() != 0)
         return execOp->emitError(
             "expected result-less scf.execute_region containing op");
       return success();
     }
 
     // Bufferize scf::IfOp after bufferizing the scf::YieldOp in the else
     // branch.
     if (auto ifOp = dyn_cast<scf::IfOp>(yieldOp->getParentOp())) {
       if (ifOp.elseYield() != yieldOp)
         return success();
-      return bufferizeIfOp(ifOp, b, bvm, aliasInfo, allocationFn);
+      return bufferizeIfOp(ifOp, b, state);
     }
 
     // Bufferize scf::ForOp after bufferizing the scf::YieldOp.
     if (auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp()))
-      return bufferizeForOp(forOp, b, bvm, aliasInfo, allocationFn);
+      return bufferizeForOp(forOp, b, state);
 
     return yieldOp->emitError("expected scf::ForOp parent for scf::YieldOp");
   }
 };
 
 } // namespace scf_ext
 
 namespace std_ext {
OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     return OpResult();
   }
 
   BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
-                          BlockAndValueMapping &bvm,
-                          BufferizationAliasInfo &aliasInfo,
-                          AllocationCallbacks &allocationFn) const {
+                          BufferizationState &state) const {
     llvm_unreachable("CallOps are handled separately");
     return failure();
   }
 };
 
 struct ReturnOpInterface
     : public BufferizableOpInterface::ExternalModel<ReturnOpInterface,
                                                     ReturnOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand) const {
     return false;
   }
 
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     return OpResult();
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
-                          BlockAndValueMapping &bvm,
-                          BufferizationAliasInfo &aliasInfo,
-                          AllocationCallbacks &allocationFn) const {
+                          BufferizationState &state) const {
     auto returnOp = cast<ReturnOp>(op);
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);
     // Cannot insert after returnOp.
     b.setInsertionPoint(returnOp);
 
     assert(isa<FuncOp>(returnOp->getParentOp()) &&
            "only support FuncOp parent for ReturnOp");
     for (OpOperand &operand : returnOp->getOpOperands()) {
       auto tensorType = operand.get().getType().dyn_cast<TensorType>();
       if (!tensorType)
         continue;
-      Value v = lookup(bvm, operand.get());
-      assert(v && "missing buffer for result");
+      Value v = state.lookupBuffer(operand.get());
       Value returnTensor = b.create<memref::TensorLoadOp>(returnOp.getLoc(), v);
       operand.set(returnTensor);
-      aliasInfo.insertNewBufferEquivalence(returnTensor, v);
-      map(bvm, returnTensor, v);
+      state.aliasInfo.insertNewBufferEquivalence(returnTensor, v);
+      state.mapBuffer(returnTensor, v);
     }
     return success();
   }
 };
 
 } // namespace std_ext
 
 namespace tensor_ext {
OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     return op->getResult(0);
   }
 
   BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
-                          BlockAndValueMapping &bvm,
-                          BufferizationAliasInfo &aliasInfo,
-                          AllocationCallbacks &allocationFn) const {
+                          BufferizationState &state) const {
     auto castOp = cast<tensor::CastOp>(op);
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPoint(castOp);
 
-    Value resultBuffer =
-        getResultBuffer(b, castOp->getResult(0), bvm, aliasInfo, allocationFn);
+    Value resultBuffer = getResultBuffer(b, castOp->getResult(0), state);
     if (!resultBuffer)
       return failure();
     Type sourceType = resultBuffer.getType();
     auto rankedMemRefType = sourceType.dyn_cast<MemRefType>();
     auto unrankedMemRefType = sourceType.dyn_cast<UnrankedMemRefType>();
     assert(rankedMemRefType || unrankedMemRefType);
     Attribute memorySpace = rankedMemRefType
                                 ? rankedMemRefType.getMemorySpace()
                                 : unrankedMemRefType.getMemorySpace();
     TensorType tensorType = castOp.getResult().getType().cast<TensorType>();
     MemRefLayoutAttrInterface layout =
         rankedMemRefType && tensorType.isa<RankedTensorType>()
             ? rankedMemRefType.getLayout()
             : MemRefLayoutAttrInterface();
     Type memRefType = getContiguousOrUnrankedMemRefType(
         castOp.getResult().getType(), layout, memorySpace);
     Value res =
         b.create<memref::CastOp>(castOp.getLoc(), memRefType, resultBuffer);
-    aliasInfo.insertNewBufferEquivalence(res, castOp.getResult());
-    map(bvm, castOp.getResult(), res);
+    state.aliasInfo.insertNewBufferEquivalence(res, castOp.getResult());
+    state.mapBuffer(castOp.getResult(), res);
     return success();
   }
 };
 
 struct DimOpInterface
     : public BufferizableOpInterface::ExternalModel<DimOpInterface,
                                                     tensor::DimOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand) const {
     return false;
   }
 
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     return OpResult();
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
-                          BlockAndValueMapping &bvm,
-                          BufferizationAliasInfo &aliasInfo,
-                          AllocationCallbacks &allocationFn) const {
+                          BufferizationState &state) const {
     auto dimOp = cast<tensor::DimOp>(op);
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPoint(dimOp);
 
     if (dimOp.source().getType().isa<RankedTensorType>()) {
-      Value v = lookup(bvm, dimOp.source());
-      assert(v && "missing buffer");
+      Value v = state.lookupBuffer(dimOp.source());
       dimOp.result().replaceAllUsesWith(
           b.create<memref::DimOp>(dimOp.getLoc(), v, dimOp.index()));
     }
     return success();
   }
 };
 
 struct ExtractSliceOpInterface
return &opOperand == &op->getOpOperand(0) /*source*/
                : OpResult();
   }
 
   BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {
     return BufferRelation::None;
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
-                          BlockAndValueMapping &bvm,
-                          BufferizationAliasInfo &aliasInfo,
-                          AllocationCallbacks &allocationFn) const {
+                          BufferizationState &state) const {
     auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);
 
     LDBG("bufferize: " << *extractSliceOp << '\n');
 
     Location loc = extractSliceOp.getLoc();
     // Bail if source was not bufferized.
-    Value srcMemref = lookup(bvm, extractSliceOp.source());
-    if (!srcMemref)
-      return failure();
+    Value srcMemref = state.lookupBuffer(extractSliceOp.source());
     auto srcMemrefType = srcMemref.getType().cast<MemRefType>();
     auto dstTensorType =
         extractSliceOp.result().getType().cast<RankedTensorType>();
 
     // If not inplaceable, alloc.
     Value alloc;
-    if (!aliasInfo.isInPlace(extractSliceOp->getResult(0)))
+    if (!state.aliasInfo.isInPlace(extractSliceOp->getResult(0)))
       alloc = createNewAllocDeallocPairForShapedValue(
-          b, loc, extractSliceOp.result(), aliasInfo, allocationFn);
+          b, loc, extractSliceOp.result(), state);
 
     // Set insertion point now that potential alloc/dealloc are introduced.
     b.setInsertionPoint(extractSliceOp);
 
     // 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.
-    aliasInfo.insertNewBufferAlias(subView, srcMemref);
+    state.aliasInfo.insertNewBufferAlias(subView, srcMemref);
 
     /// If not inplaceable, copy.
     if (alloc) {
       // Do not copy if the copied data is never read.
       if (isValueRead(extractSliceOp.result()))
-        allocationFn.memCpyFn(b, extractSliceOp.getLoc(), subView, alloc);
+        state.allocationFns.memCpyFn(b, extractSliceOp.getLoc(), subView,
+                                     alloc);
       subView = alloc;
     }
 
-    map(bvm, extractSliceOp.result(), subView);
+    state.mapBuffer(extractSliceOp.result(), subView);
     return success();
   }
 };
 
 struct ExtractOpInterface
     : public BufferizableOpInterface::ExternalModel<ExtractOpInterface,
                                                     tensor::ExtractOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand) const {
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand) const {
     return false;
   }
 
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     return OpResult();
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
-                          BlockAndValueMapping &bvm,
-                          BufferizationAliasInfo &aliasInfo,
-                          AllocationCallbacks &allocationFn) const {
+                          BufferizationState &state) const {
     auto extractOp = cast<tensor::ExtractOp>(op);
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPoint(extractOp);
 
     Location loc = extractOp.getLoc();
-    Value srcMemref = lookup(bvm, extractOp.tensor());
+    Value srcMemref = state.lookupBuffer(extractOp.tensor());
     Value l = b.create<memref::LoadOp>(loc, srcMemref, extractOp.indices());
     extractOp.replaceAllUsesWith(l);
     return success();
   }
 };
 
 /// Return true if the source of a `insertSliceOp` bufferizes to an
 /// equivalent ExtractSliceOp that bufferizes inplace.
return &opOperand == &op->getOpOperand(1) /*dest*/
                : OpResult();
   }
 
   BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
-                          BlockAndValueMapping &bvm,
-                          BufferizationAliasInfo &aliasInfo,
-                          AllocationCallbacks &allocationFn) const {
+                          BufferizationState &state) const {
     auto insertSliceOp = cast<tensor::InsertSliceOp>(op);
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPoint(insertSliceOp);
 
     LDBG("bufferize: " << *insertSliceOp << '\n');
 
     Location loc = insertSliceOp.getLoc();
     // Since insert_slice arise from tiling and introducing loops, this
     // case 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.
     // Alloc a copy for `insertSliceOp.dest()`, it will become the result
     // buffer.
-    Value dstMemref = getResultBuffer(b, insertSliceOp->getResult(0), bvm,
-                                      aliasInfo, allocationFn);
+    Value dstMemref = getResultBuffer(b, insertSliceOp->getResult(0), state);
     if (!dstMemref)
       return failure();
     auto dstMemrefType = dstMemref.getType().cast<MemRefType>();
 
-    Value srcMemref = lookup(bvm, insertSliceOp.source());
-    if (!srcMemref)
-      return failure();
+    Value srcMemref = state.lookupBuffer(insertSliceOp.source());
     auto subviewMemRefType =
         memref::SubViewOp::inferRankReducedResultType(
             insertSliceOp.getSourceType().getRank(), dstMemrefType,
             insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes(),
             insertSliceOp.getMixedStrides())
             .cast<MemRefType>();
 
     // 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.
     // TODO: Is this necessary?
-    if (!isSourceEquivalentToAMatchingInplaceExtractSliceOp(aliasInfo,
+    if (!isSourceEquivalentToAMatchingInplaceExtractSliceOp(state.aliasInfo,
                                                             insertSliceOp) ||
-        !aliasInfo.isInPlace(insertSliceOp->getResult(0))) {
+        !state.aliasInfo.isInPlace(insertSliceOp->getResult(0))) {
       LDBG("insert_slice needs extra source copy: " << insertSliceOp.source()
                                                     << " -> copy\n");
       // Take a subview of the dst.
       Value subView = b.create<memref::SubViewOp>(
           loc, subviewMemRefType, dstMemref, insertSliceOp.getMixedOffsets(),
           insertSliceOp.getMixedSizes(), insertSliceOp.getMixedStrides());
       // Insert new alias.
-      aliasInfo.insertNewBufferAlias(subView, dstMemref);
-      allocationFn.memCpyFn(b, insertSliceOp.getLoc(), srcMemref, subView);
+      state.aliasInfo.insertNewBufferAlias(subView, dstMemref);
+      state.allocationFns.memCpyFn(b, insertSliceOp.getLoc(), srcMemref,
+                                   subView);
     }
 
-    map(bvm, insertSliceOp.result(), dstMemref);
+    state.mapBuffer(insertSliceOp.result(), dstMemref);
 
     return success();
   }
 };
 
 } // namespace tensor_ext
 
 namespace vector_ext {
bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand) const {
     return false;
   }
 
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     return OpResult();
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
-                          BlockAndValueMapping &bvm,
-                          BufferizationAliasInfo &aliasInfo,
-                          AllocationCallbacks &allocationFn) const {
+                          BufferizationState &state) const {
     auto transferReadOp = cast<vector::TransferReadOp>(op);
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPoint(op);
 
     if (transferReadOp.getShapedType().isa<MemRefType>())
       return failure();
 
     // TransferReadOp always reads from the bufferized op.source().
-    Value v = lookup(bvm, transferReadOp.source());
-    assert(v && "missing buffer");
+    Value v = state.lookupBuffer(transferReadOp.source());
     transferReadOp.sourceMutable().assign(v);
     return success();
   }
 };
 
 struct TransferWriteOpInterface
     : public BufferizableOpInterface::ExternalModel<TransferWriteOpInterface,
                                                     vector::TransferWriteOp> {
OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     return op->getOpResult(0);
   }
 
   BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
-                          BlockAndValueMapping &bvm,
-                          BufferizationAliasInfo &aliasInfo,
-                          AllocationCallbacks &allocationFn) const {
+                          BufferizationState &state) const {
     auto writeOp = cast<vector::TransferWriteOp>(op);
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPoint(op);
 
     if (writeOp.getShapedType().isa<MemRefType>())
       return failure();
 
     // 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.
-    Value resultBuffer =
-        getResultBuffer(b, op->getResult(0), bvm, aliasInfo, allocationFn);
+    Value resultBuffer = getResultBuffer(b, op->getResult(0), state);
     if (!resultBuffer)
       return failure();
     b.create<vector::TransferWriteOp>(
         writeOp.getLoc(), writeOp.vector(), resultBuffer, writeOp.indices(),
         writeOp.permutation_map(),
         writeOp.in_bounds() ? *writeOp.in_bounds() : ArrayAttr());
-    map(bvm, op->getResult(0), resultBuffer);
+    state.mapBuffer(op->getResult(0), resultBuffer);
 
     return success();
   }
 };
 
 } // namespace vector_ext
 
 namespace {