[mlir][linalg][bufferize][NFC] Remove RewriterBase from BufferizationState

[springerm]

Summary

This change simplifies BufferizationState. Having rewriter in BufferizationState could be confusing to users because a rewriter is also passed to each bufferize function and it is not obvious (by looking at the API) that these two rewriters are the same.

Depends On D116102

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

 ///   that allocate and/or deallocate memref buffers.
 /// * `lookupBuffer` returns the 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.
 /// * `replaceOp` replaces an op with new values.
 class BufferizationState {
 public:
-  BufferizationState(Operation *op, const BufferizationOptions &options,
-                     RewriterBase &rewriter);
+  BufferizationState(Operation *op, const BufferizationOptions &options);
 
   // BufferizationState should be passed as a reference.
   BufferizationState(const BufferizationState &) = delete;
 
   /// Determine which OpOperand* will alias with `result` if the op is
   /// bufferized in place. Return an empty vector if the op is not bufferizable.
   SmallVector<OpOperand *> getAliasingOpOperand(OpResult result);
 
public:
   /// 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);
 
   /// Replace an op with replacement values. The op is deleted. Tensor OpResults
   /// must be replaced with memref values.
-  void replaceOp(Operation *op, ValueRange values);
+  void replaceOp(RewriterBase &rewriter, Operation *op, ValueRange values);
 
   /// Replace an op with a new op. Tensor OpResults must be replaced with memref
   /// values.
   template <typename OpTy, typename... Args>
   OpTy replaceOpWithNewOp(RewriterBase &rewriter, Operation *op,
                           Args &&...args) {
     Operation *newOp =
         rewriter.create<OpTy>(op->getLoc(), std::forward<Args>(args)...);
-    replaceOp(op, newOp->getResults());
+    replaceOp(rewriter, op, newOp->getResults());
     return cast<OpTy>(newOp);
   }
 
   /// Lookup the memref buffer that is associated to the given tensor value.
   /// Asserts if no buffer is associated.
-  Value lookupBuffer(Value tensor);
+  Value lookupBuffer(RewriterBase &rewriter, Value tensor);
 
   /// Return `true` if the given OpResult has been decided to bufferize inplace.
   bool isInPlace(OpResult opResult) 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);
+  Value getResultBuffer(RewriterBase &rewriter, OpResult result);
 
   /// 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]);
   }
 
   /// Return a reference to the BufferizationOptions.
   const BufferizationOptions &getOptions() const { return options; }
 
-  /// Return a reference to the rewriter.
-  RewriterBase &getRewriter() { return rewriter; }
-
 private:
   friend LogicalResult
   runComprehensiveBufferize(Operation *op, const BufferizationOptions &options,
                             BufferizationState &state);
 
   friend LogicalResult
   runComprehensiveBufferize(ModuleOp moduleOp,
                             std::unique_ptr<BufferizationOptions> options);
 
   /// `aliasInfo` keeps track of aliasing and equivalent values. Only internal
   /// functions and `runComprehensiveBufferize` may access this object.
   BufferizationAliasInfo aliasInfo;
 
   /// Dialect-specific bufferization state.
   DenseMap<StringRef, std::unique_ptr<DialectBufferizationState>> dialectState;
 
   /// A reference to current bufferization options.
   const BufferizationOptions &options;
-
-  /// The OpBuilder used during bufferization.
-  RewriterBase &rewriter;
 };
 
 /// Bufferize all ops in the given region.
-LogicalResult bufferize(Region *region, BufferizationState &state);
+LogicalResult bufferize(RewriterBase &rewriter, Region *region,
+                        BufferizationState &state);
 
 /// Bufferize all ops in the given block.
-LogicalResult bufferize(Block *block, BufferizationState &state);
+LogicalResult bufferize(RewriterBase &rewriter, Block *block,
+                        BufferizationState &state);
 
 /// Bufferize the given op. If the op has no tensor OpOperands/OpResults, this
 /// function returns immediately. Otherwise, it calls the `bufferize` interface
 /// method of `BufferizableOpInterface`.
-LogicalResult bufferize(Operation *op, BufferizationState &state);
+LogicalResult bufferize(RewriterBase &rewriter, Operation *op,
+                        BufferizationState &state);
 
 /// Return a contiguous MemRefType (i.e. with canonical/empty layout map)
 /// with the same shape as `shapedType` and specified `layout` and
 /// `addressSpace`.
 MemRefType getContiguousMemRefType(ShapedType shapedType,
                                    MemRefLayoutAttrInterface layout = {},
                                    Attribute memorySpace = {});
 
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto isaTensor = [](Type t) { return t.isa<TensorType>(); };
     if (any_of(op->getOperandTypes(), isaTensor) ||
         any_of(op->getResultTypes(), isaTensor))
       if (!state.getOptions().allowUnknownOps)
         return op->emitError() << "unsupported op with tensors";
 
     for (Region &region : op->getRegions())
-      if (failed(comprehensive_bufferize::bufferize(&region, state)))
+      if (failed(comprehensive_bufferize::bufferize(rewriter, &region, state)))
         return failure();
 
     return success();
   }
 
   bool isAllocationHoistingBarrier(Operation *op) const { return true; }
 };
 
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_BUFFERIZABLEOPINTERFACE_H_

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

Value mlir::linalg::comprehensive_bufferize::BufferizationState::
   // are currently no other ops where one OpResult may alias with multiple
   // OpOperands. Therefore, this function should return exactly one result at
   // the moment.
   assert(result.size() == 1 && "expected exactly one result");
   return result.front();
 }
 
 mlir::linalg::comprehensive_bufferize::BufferizationState::BufferizationState(
-    Operation *op, const BufferizationOptions &options, RewriterBase &rewriter)
-    : aliasInfo(op), options(options), rewriter(rewriter) {
+    Operation *op, const BufferizationOptions &options)
+    : aliasInfo(op), options(options) {
   // Set up alias sets for OpResults that must bufferize in-place. This should
   // be done before making any other bufferization decisions.
   op->walk([&](BufferizableOpInterface bufferizableOp) {
     if (!options.isOpAllowed(bufferizableOp))
       return WalkResult::skip();
     for (OpResult opResult : bufferizableOp->getOpResults()) {
       if (opResult.getType().isa<TensorType>())
         if (bufferizableOp.mustBufferizeInPlace(opResult, *this)) {
op->walk([&](BufferizableOpInterface bufferizableOp) {
     return WalkResult::advance();
   });
 }
 
 /// 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::BufferizationState::
-    getResultBuffer(OpResult result) {
+    getResultBuffer(RewriterBase &rewriter, OpResult result) {
   OpBuilder::InsertionGuard guard(rewriter);
   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 = lookupBuffer(operand);
+  Value operandBuffer = lookupBuffer(rewriter, 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 (aliasingOperands.size() > 1 &&
       !llvm::all_of(aliasingOperands, [&](OpOperand *o) {
-        return lookupBuffer(o->get()) == operandBuffer;
+        return lookupBuffer(rewriter, o->get()) == operandBuffer;
       })) {
     op->emitError("result buffer is ambiguous");
     return Value();
   }
 
   // If bufferizing out-of-place, allocate a new buffer.
   if (!aliasInfo.isInPlace(result)) {
     // Ops with multiple aliasing operands can currently not bufferize
if (!aliasInfo.isInPlace(result)) {
     return resultBuffer;
   }
 
   // Bufferizing in-place. No need to allocate a new buffer.
   return operandBuffer;
 }
 
 void mlir::linalg::comprehensive_bufferize::BufferizationState::replaceOp(
-    Operation *op, ValueRange values) {
+    RewriterBase &rewriter, Operation *op, ValueRange values) {
   OpBuilder::InsertionGuard g(rewriter);
 
   // Replace all OpResults with the given values.
   for (OpResult opResult : op->getOpResults()) {
     // Skip OpResult if it has no uses.
     if (opResult.getUses().empty())
       continue;
 
if (opResult.getType().isa<TensorType>()) {
           replacement.getLoc(), replacement);
     }
     opResult.replaceAllUsesWith(replacement);
   }
 
   rewriter.eraseOp(op);
 }
 
-LogicalResult
-mlir::linalg::comprehensive_bufferize::bufferize(Region *region,
-                                                 BufferizationState &state) {
+LogicalResult mlir::linalg::comprehensive_bufferize::bufferize(
+    RewriterBase &rewriter, Region *region, BufferizationState &state) {
   for (Block &block : *region)
-    if (failed(bufferize(&block, state)))
+    if (failed(bufferize(rewriter, &block, state)))
       return failure();
   return success();
 }
 
-LogicalResult
-mlir::linalg::comprehensive_bufferize::bufferize(Block *block,
-                                                 BufferizationState &state) {
+LogicalResult mlir::linalg::comprehensive_bufferize::bufferize(
+    RewriterBase &rewriter, Block *block, BufferizationState &state) {
   // Ops may get deleted during the traversal, so do not iterate over `block`
   // directly.
   SmallVector<Operation *> ops;
   ops.reserve(block->getOperations().size());
   for (Operation &op : *block)
     ops.push_back(&op);
   for (Operation *op : ops)
-    if (failed(bufferize(op, state)))
+    if (failed(bufferize(rewriter, op, state)))
       return failure();
   return success();
 }
 
-LogicalResult
-mlir::linalg::comprehensive_bufferize::bufferize(Operation *op,
-                                                 BufferizationState &state) {
-  RewriterBase &rewriter = state.getRewriter();
-
+LogicalResult mlir::linalg::comprehensive_bufferize::bufferize(
+    RewriterBase &rewriter, Operation *op, BufferizationState &state) {
   // 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 hasRegions = !op->getRegions().empty();
 
   // No tensor results/operands or regions. We are done.
   if (!hasTensorResult && !hasTensorOperand && !hasRegions)
     return success();
 
   // Bufferize using `BufferizableOpInterface`. Interface implementations are
   // responsible for bufferizing nested ops.
   if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(op)) {
     rewriter.setInsertionPoint(op);
     return bufferizableOp.bufferize(rewriter, state);
   }
 
   // `op` is an unbufferizable tensor op.
   if (!state.getOptions().allowUnknownOps)
     return op->emitError() << "unsupported op with tensors";
 
   // Bufferize all regions.
   for (Region &region : op->getRegions())
-    if (failed(bufferize(&region, state)))
+    if (failed(bufferize(rewriter, &region, state)))
       return failure();
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific scoped alloc/dealloc insertion support.
 //===----------------------------------------------------------------------===//

 bool mlir::linalg::comprehensive_bufferize::isFunctionArgument(Value value) {
   auto bbArg = value.dyn_cast<BlockArgument>();
   if (!bbArg)
     return false;
   return isa<FuncOp>(bbArg.getOwner()->getParentOp());
 }
 
 Value mlir::linalg::comprehensive_bufferize::BufferizationState::lookupBuffer(
-    Value tensor) {
+    RewriterBase &rewriter, Value tensor) {
   assert(tensor.getType().isa<TensorType>() && "unexpected non-tensor type");
 
   // Replace "%t = to_tensor %m" with %m.
   if (auto toTensorOp = tensor.getDefiningOp<bufferization::ToTensorOp>())
     return toTensorOp.memref();
 
   if (!isFunctionArgument(tensor)) {
     if (static_cast<bool>(options.dynCastBufferizableOp(tensor))) {

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

op->walk([&](Operation *op) {
     for (OpResult opResult : op->getResults())
       if (opResult.getType().isa<TensorType>())
         setInPlaceOpResult(opResult, aliasInfo.isInPlace(opResult));
   });
 }
 
 LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
     Operation *op, std::unique_ptr<BufferizationOptions> options) {
-  IRRewriter rewriter(op->getContext());
-  BufferizationState state(op, *options, rewriter);
+  BufferizationState state(op, *options);
   return runComprehensiveBufferize(op, *options, state);
 }
 
 LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
     Operation *op, const BufferizationOptions &options,
     BufferizationState &state) {
 
+  IRRewriter rewriter(op->getContext());
   DominanceInfo domInfo(op);
   BufferizationAliasInfo &aliasInfo = state.aliasInfo;
 
   if (failed(checkAliasInfoConsistency(op, domInfo, state, aliasInfo)))
     return failure();
 
   // If the analysis fails, just return.
   if (failed(inPlaceAnalysis(op, aliasInfo, state, domInfo,
LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
 
   // Annotate operations if we only want to report the analysis.
   if (options.testAnalysisOnly) {
     annotateOpsWithBufferizationMarkers(op, aliasInfo);
     return success();
   }
 
   // Bufferize the op and its nested ops.
-  if (failed(bufferize(op, state)))
+  if (failed(bufferize(rewriter, op, state)))
     return failure();
 
   return success();
 }

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

static LogicalResult bufferizeLinalgOp(RewriterBase &rewriter, LinalgOp op,
 
   SmallVector<Value> newInputBuffers;
   newInputBuffers.reserve(op.getNumInputs());
   for (OpOperand *opOperand : op.getInputOperands()) {
     if (op.isScalar(opOperand)) {
       newInputBuffers.push_back(opOperand->get());
       continue;
     }
-    newInputBuffers.push_back(state.lookupBuffer(opOperand->get()));
+    newInputBuffers.push_back(state.lookupBuffer(rewriter, opOperand->get()));
   }
 
   SmallVector<Value> newOutputBuffers;
   for (OpOperand *opOperand : op.getOutputOperands()) {
     OpResult opResult = op.getTiedOpResult(opOperand);
     assert(opResult && "could not find correspond OpResult");
-    Value resultBuffer = state.getResultBuffer(opResult);
+    Value resultBuffer = state.getResultBuffer(rewriter, opResult);
     if (!resultBuffer)
       return failure();
     newOutputBuffers.push_back(resultBuffer);
   }
 
   // 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.
   rewriter.setInsertionPoint(op);
   auto bufferizedOp = cast<LinalgOp>(op.clone(
       rewriter, op.getLoc(), /*resultTypes=*/TypeRange{}, newOperands));
 
   // Replace the results of the old op with the new output buffers.
-  state.replaceOp(op, newOutputBuffers);
+  state.replaceOp(rewriter, op, newOutputBuffers);
 
-  return comprehensive_bufferize::bufferize(bufferizedOp.getBlock(), state);
+  return comprehensive_bufferize::bufferize(rewriter, bufferizedOp.getBlock(),
+                                            state);
 }
 
 /// Linalg OpResults usually bufferize inplace with their tied (output
 /// OpOperands. However, if an output OpOperand is not used in the computation,
 /// it is better to bufferize inplace with an actually used input OpOperand;
 /// less memory will be touched that way.
 ///
 /// Example:
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     auto initTensorOp = cast<linalg::InitTensorOp>(op);
 
     // The InitTensorOp may have been eliminated.
     if (initTensorOp->getUses().empty())
       return success();
 
     Value alloc = state.createAllocDeallocPair(rewriter, initTensorOp->getLoc(),
                                                initTensorOp.result());
-    state.replaceOp(op, alloc);
+    state.replaceOp(rewriter, op, alloc);
     return success();
   }
 };
 
 struct TiledLoopOpInterface
     : public BufferizableOpInterface::ExternalModel<TiledLoopOpInterface,
                                                     linalg::TiledLoopOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
struct TiledLoopOpInterface
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto tiledLoopOp = cast<linalg::TiledLoopOp>(op);
 
     // Compute new inputs, outputs and results.
     SmallVector<Value> newInputs, newOutputs, newResults;
     for (Value value : tiledLoopOp.inputs()) {
       if (value.getType().isa<TensorType>()) {
-        newInputs.push_back(state.lookupBuffer(value));
+        newInputs.push_back(state.lookupBuffer(rewriter, value));
       } else {
         newInputs.push_back(value);
       }
     }
     int nextResultNum = 0;
     for (Value value : tiledLoopOp.outputs()) {
       if (value.getType().isa<TensorType>()) {
-        Value buffer =
-            state.getResultBuffer(tiledLoopOp->getResult(nextResultNum++));
+        Value buffer = state.getResultBuffer(
+            rewriter, tiledLoopOp->getResult(nextResultNum++));
         newOutputs.push_back(buffer);
         newResults.push_back(buffer);
       } else {
         newOutputs.push_back(value);
       }
     }
 
     // Create new TiledLoopOp.
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
 
     // Replace previous terminator with a new one that does not yield anything.
     Operation *oldTerminator = newTiledLoopOp.getBody()->getTerminator();
     rewriter.setInsertionPointToEnd(newTiledLoopOp.getBody());
     rewriter.create<linalg::YieldOp>(oldTerminator->getLoc());
     rewriter.eraseOp(oldTerminator);
 
     // Replace results and delete old op.
-    state.replaceOp(op, newResults);
+    state.replaceOp(rewriter, op, newResults);
 
     // Bufferize loop body.
-    return comprehensive_bufferize::bufferize(newTiledLoopOp.getBody(), state);
+    return comprehensive_bufferize::bufferize(rewriter,
+                                              newTiledLoopOp.getBody(), state);
   }
 };
 
 struct YieldOpInterface
     : public BufferizableOpInterface::ExternalModel<YieldOpInterface,
                                                     linalg::YieldOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               BufferizationState &state) const {

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

 /// assumes all CallOp perform the necessary work to clone operands so as to
 /// make them inplaceable. Reliance on this logic will need to be relaxed in the
 /// future.
 ///
 /// Note: Returning a memref currently fails bufferization. If such memrefs
 /// originate from an op with an Alloc effect, they could be hoisted in the
 /// future.
 static LogicalResult bufferizeFuncOpBoundary(FuncOp funcOp,
+                                             RewriterBase &rewriter,
                                              BufferizationState &state) {
   ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
 
   // If nothing to do then we are done.
   if (!llvm::any_of(funcOp.getType().getInputs(), isaTensor) &&
       !llvm::any_of(funcOp.getType().getResults(), isaTensor))
     return success();
 
for (OpOperand &returnOperand : returnOp->getOpOperands()) {
     }
 
     // If return operand is equivalent to some bbArg, no need to return it.
     if (moduleState.equivalentFuncArgs[funcOp].count(
             returnOperand.getOperandNumber()))
       continue;
 
     // Cast values at the call site if necessary.
-    returnValues.push_back(getNonCastedValue(state.lookupBuffer(returnVal)));
+    returnValues.push_back(
+        getNonCastedValue(state.lookupBuffer(rewriter, returnVal)));
   }
 
   // 2. Rewrite the terminator without the inPlace bufferizable values.
   ValueRange retValues{returnValues};
   FunctionType bufferizedFuncType = getOrCreateBufferizedFunctionType(
       funcOp, funcOp.getType().getInputs(), retValues.getTypes(),
       moduleState.bufferizedFunctionTypes);
   OpBuilder b(returnOp);
OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
     // bufferized and needs to inspect the BufferAliasInfo object. It can't
     // make a proper determination by itself and needs to be conservative.
     return OpResult();
   }
 
   /// In a first approximation, all the function arguments of a FuncOp are
   /// marked inplaceable. For now, it is the responsibility of the `callOp`
   /// bufferization to allow FuncOp that are inplaceable to write inPlace.
-  LogicalResult bufferize(Operation *op, OpBuilder &b,
+  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     CallOp callOp = cast<CallOp>(op);
     FuncOp funcOp = getCalledFunction(callOp);
     assert(isa<CallOp>(callOp.getOperation()) && funcOp &&
            "expected Callop to a FuncOp");
     ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
 
     // 1. Filter return types:
if (funcOp.body().empty()) {
 
         // If return operand is equivalent to some bbArg, no need to return it.
         if (moduleState.equivalentFuncArgs[funcOp].count(
                 returnOperand.getOperandNumber())) {
           int64_t idx =
               moduleState
                   .equivalentFuncArgs[funcOp][returnOperand.getOperandNumber()];
           Value oldRes = callOp->getResult(returnOperand.getOperandNumber());
-          Value buffer = state.lookupBuffer(callOp->getOperand(idx));
+          Value buffer = state.lookupBuffer(rewriter, callOp->getOperand(idx));
           // Add a ToTensorOp to kill all uses of the CallOp return.
           // Replace all uses of the CallOp results so we can erase the CallOp.
           // This ToTensorOp must fold/DCE away or bufferization should be
           // considered failed.
-          Value toTensorOp =
-              b.create<bufferization::ToTensorOp>(callOp.getLoc(), buffer);
+          Value toTensorOp = rewriter.create<bufferization::ToTensorOp>(
+              callOp.getLoc(), buffer);
           oldRes.replaceAllUsesWith(toTensorOp);
           continue;
         }
 
         resultTypes.push_back(returnType);
       }
     }
 
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 = state.lookupBuffer(tensorOperand);
+      Value buffer = state.lookupBuffer(rewriter, 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);
+        Value castBuffer = rewriter.create<memref::CastOp>(callOp.getLoc(),
+                                                           memRefType, buffer);
         buffer = castBuffer;
       }
       newOperands.push_back(buffer);
     }
 
     // 4. Create the new CallOp.
-    Operation *newCallOp = b.create<CallOp>(callOp.getLoc(), funcOp.sym_name(),
-                                            resultTypes, newOperands);
+    Operation *newCallOp = rewriter.create<CallOp>(
+        callOp.getLoc(), funcOp.sym_name(), resultTypes, newOperands);
     newCallOp->setAttrs(callOp->getAttrs());
 
     // 5. Delete the op at the end of bufferization.
     callOp->erase();
 
     return success();
   }
 };
bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
     return false;
   }
 
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                BufferizationState &state) const {
     return OpResult();
   }
 
-  LogicalResult bufferize(Operation *op, OpBuilder &b,
+  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto returnOp = cast<ReturnOp>(op);
     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 = state.lookupBuffer(operand.get());
-      Value returnTensor = b.create<bufferization::ToTensorOp>(
-          returnOp.getLoc(), v);
+      Value v = state.lookupBuffer(rewriter, operand.get());
+      Value returnTensor =
+          rewriter.create<bufferization::ToTensorOp>(returnOp.getLoc(), v);
       operand.set(returnTensor);
     }
     return success();
   }
 };
 
 struct FuncOpInterface
     : public BufferizableOpInterface::ExternalModel<FuncOpInterface, FuncOp> {
-  LogicalResult bufferize(Operation *op, OpBuilder &b,
+  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto funcOp = cast<FuncOp>(op);
 
     // Bufferize function body.
-    return comprehensive_bufferize::bufferize(&funcOp.body(), state);
+    return comprehensive_bufferize::bufferize(rewriter, &funcOp.body(), state);
   }
 
   /// Return `true` if the given function argument is writable.
   bool isWritable(Operation *op, Value value, BufferizationState &state) const {
     auto funcOp = cast<FuncOp>(op);
     BlockArgument bbArg = value.dyn_cast<BlockArgument>();
     assert(bbArg && "expected BlockArgument");
     ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
static void annotateOpsWithBufferizationMarkers(FuncOp funcOp,
   for (BlockArgument bbArg : funcOp.getArguments())
     if (bbArg.getType().isa<TensorType>())
       setInPlaceFuncArgument(bbArg, bufferizableOp.isWritable(bbArg, state));
 }
 
 LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
     ModuleOp moduleOp, std::unique_ptr<BufferizationOptions> options) {
   IRRewriter rewriter(moduleOp.getContext());
-  BufferizationState state(moduleOp, *options, rewriter);
+  BufferizationState state(moduleOp, *options);
   ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
   BufferizationAliasInfo &aliasInfo = state.aliasInfo;
 
   if (failed(getFuncOpsOrderedByCalls(moduleOp, moduleState.orderedFuncOps,
                                       moduleState.callerMap)))
     return failure();
 
   // Interestingly, all function args that are not visible outside of a module
LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
   }
 
   if (options->testAnalysisOnly)
     return success();
 
   for (FuncOp funcOp : moduleState.orderedFuncOps) {
     // Note: It would be good to apply cleanups here but we cannot as aliasInfo
     // would be invalidated.
-    if (failed(bufferizeFuncOpBoundary(funcOp, state)))
+    if (failed(bufferizeFuncOpBoundary(funcOp, rewriter, state)))
       return failure();
 
     if (!options->allowReturnMemref &&
         llvm::any_of(funcOp.getType().getResults(), [](Type t) {
           return t.isa<MemRefType, UnrankedMemRefType>();
         })) {
       funcOp->emitError("memref return type is unsupported");
       return failure();

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

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     // TODO: Add bufferization support when needed. scf.execute_region should be
     // bufferized similar to scf.if.
     auto executeRegionOp = cast<scf::ExecuteRegionOp>(op);
     bool hasTensorReturnType = any_of(
         op->getResultTypes(), [](Type t) { return t.isa<TensorType>(); });
     if (hasTensorReturnType)
       return op->emitError(
           "scf.execute_region with tensor result not supported");
-    return comprehensive_bufferize::bufferize(&executeRegionOp.getRegion(),
-                                              state);
+    return comprehensive_bufferize::bufferize(
+        rewriter, &executeRegionOp.getRegion(), state);
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const BufferizationAliasInfo &aliasInfo,
                                 BufferizationState &state) const {
     return BufferRelation::Equivalent;
   }
 };
for (OpOperand &operand : elseYieldOp->getOpOperands()) {
         Value toMemrefOp = rewriter.create<bufferization::ToMemrefOp>(
             operand.get().getLoc(), newTypes[operand.getOperandNumber()],
             operand.get());
         operand.set(toMemrefOp);
       }
     }
 
     // Replace op results.
-    state.replaceOp(op, newIfOp->getResults());
+    state.replaceOp(rewriter, op, newIfOp->getResults());
 
     // Bufferize then/else blocks.
-    if (failed(comprehensive_bufferize::bufferize(newIfOp.thenBlock(), state)))
+    if (failed(comprehensive_bufferize::bufferize(rewriter, newIfOp.thenBlock(),
+                                                  state)))
       return failure();
-    if (failed(comprehensive_bufferize::bufferize(newIfOp.elseBlock(), state)))
+    if (failed(comprehensive_bufferize::bufferize(rewriter, newIfOp.elseBlock(),
+                                                  state)))
       return failure();
 
     return success();
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const BufferizationAliasInfo &aliasInfo,
                                 BufferizationState &state) const {
auto convert = [&](ValueRange values,
         result.push_back(indices.contains(idx) ? func(val, idx) : val);
       }
       return result;
     };
 
     // Construct a new scf.for op with memref instead of tensor values.
     SmallVector<Value> initArgs =
         convert(forOp.getInitArgs(), [&](Value val, int64_t index) {
-          return state.getResultBuffer(forOp->getOpResult(index));
+          return state.getResultBuffer(rewriter, forOp->getOpResult(index));
         });
     auto newForOp = rewriter.create<scf::ForOp>(
         forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(),
         forOp.getStep(), initArgs);
     Block *loopBody = &newForOp.getLoopBody().front();
 
     // Set up new iter_args. The loop body uses tensors, so wrap the (memref)
     // iter_args of the new loop in ToTensorOps.
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     SmallVector<Value> yieldValues =
         convert(yieldOp.getResults(), [&](Value val, int64_t index) {
           return rewriter.create<bufferization::ToMemrefOp>(
               val.getLoc(), initArgs[index].getType(), val);
         });
     yieldOp.getResultsMutable().assign(yieldValues);
 
     // Replace loop results.
-    state.replaceOp(op, newForOp->getResults());
+    state.replaceOp(rewriter, op, newForOp->getResults());
 
     // Bufferize loop body.
-    if (failed(comprehensive_bufferize::bufferize(loopBody, state)))
+    if (failed(comprehensive_bufferize::bufferize(rewriter, loopBody, state)))
       return failure();
 
     return success();
   }
 };
 
 // TODO: Evolve toward matching ReturnLike ops. Check for aliasing values that
 // do not bufferize inplace. (Requires a few more changes for ConstantOp,

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

BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 BufferizationState &state) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto castOp = cast<tensor::CastOp>(op);
 
-    Value resultBuffer = state.getResultBuffer(castOp->getResult(0));
+    Value resultBuffer = state.getResultBuffer(rewriter, castOp->getResult(0));
     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()
OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
     return OpResult();
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto dimOp = cast<tensor::DimOp>(op);
     if (!dimOp.source().getType().isa<RankedTensorType>())
       return dimOp.emitError("unranked tensor not supported");
-    Value v = state.lookupBuffer(dimOp.source());
+    Value v = state.lookupBuffer(rewriter, dimOp.source());
     state.replaceOpWithNewOp<memref::DimOp>(rewriter, op, v, dimOp.index());
     return success();
   }
 };
 
 struct ExtractSliceOpInterface
     : public BufferizableOpInterface::ExternalModel<ExtractSliceOpInterface,
                                                     tensor::ExtractSliceOp> {
BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 BufferizationState &state) const {
     return BufferRelation::None;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
     Location loc = extractSliceOp.getLoc();
-    Value srcMemref = state.lookupBuffer(extractSliceOp.source());
+    Value srcMemref = state.lookupBuffer(rewriter, extractSliceOp.source());
     auto srcMemrefType = srcMemref.getType().cast<MemRefType>();
     auto dstTensorType =
         extractSliceOp.result().getType().cast<RankedTensorType>();
 
     // If not inplaceable, alloc.
     bool inplace = state.isInPlace(extractSliceOp->getResult(0));
     Value alloc;
     if (!inplace)
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     /// If not inplaceable, copy.
     if (!inplace) {
       // Do not copy if the copied data is never read.
       if (state.isValueRead(extractSliceOp.result()))
         state.createMemCpy(rewriter, extractSliceOp.getLoc(), subView, alloc);
       subView = alloc;
     }
 
-    state.replaceOp(op, subView);
+    state.replaceOp(rewriter, op, subView);
     return success();
   }
 };
 
 struct ExtractOpInterface
     : public BufferizableOpInterface::ExternalModel<ExtractOpInterface,
                                                     tensor::ExtractOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
struct ExtractOpInterface
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                BufferizationState &state) const {
     return OpResult();
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto extractOp = cast<tensor::ExtractOp>(op);
-    Value srcMemref = state.lookupBuffer(extractOp.tensor());
+    Value srcMemref = state.lookupBuffer(rewriter, extractOp.tensor());
     state.replaceOpWithNewOp<memref::LoadOp>(rewriter, op, srcMemref,
                                              extractOp.indices());
     return success();
   }
 };
 
 struct InsertOpInterface
     : public BufferizableOpInterface::ExternalModel<InsertOpInterface,
getAliasingOpOperand(Operation *op, OpResult opResult,
                        BufferizationState &state) const {
     return {&op->getOpOperand(1) /*dest*/};
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto insertOp = cast<tensor::InsertOp>(op);
     Location loc = insertOp.getLoc();
-    Value destMemref = state.getResultBuffer(insertOp->getOpResult(0));
+    Value destMemref =
+        state.getResultBuffer(rewriter, insertOp->getOpResult(0));
     rewriter.create<memref::StoreOp>(loc, insertOp.scalar(), destMemref,
                                      insertOp.indices());
-    state.replaceOp(op, destMemref);
+    state.replaceOp(rewriter, op, destMemref);
     return success();
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const BufferizationAliasInfo &aliasInfo,
                                 BufferizationState &state) const {
     return BufferRelation::Equivalent;
   }
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     // whole tensor on every single iteration and is a symptom of a
     // catastrophically bad scheduling decision.
     // TODO: be very loud about it or even consider failing the pass.
     auto insertSliceOp = cast<tensor::InsertSliceOp>(op);
     Location loc = insertSliceOp.getLoc();
     TensorBufferizationState &tensorState = getTensorBufferizationState(state);
 
     // When bufferizing out-of-place, `getResultBuffer` allocates.
-    Value dstMemref = state.getResultBuffer(insertSliceOp->getResult(0));
+    Value dstMemref =
+        state.getResultBuffer(rewriter, insertSliceOp->getResult(0));
     if (!dstMemref)
       return failure();
 
     bool needCopy =
         !tensorState.insertSliceOpsWithoutCopy.contains(insertSliceOp);
     if (needCopy) {
       // Take a subview of the dst.
       auto dstMemrefType = dstMemref.getType().cast<MemRefType>();
       auto subviewMemRefType =
           memref::SubViewOp::inferRankReducedResultType(
               insertSliceOp.getSourceType().getRank(), dstMemrefType,
               insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes(),
               insertSliceOp.getMixedStrides())
               .cast<MemRefType>();
       Value subView = rewriter.create<memref::SubViewOp>(
           loc, subviewMemRefType, dstMemref, insertSliceOp.getMixedOffsets(),
           insertSliceOp.getMixedSizes(), insertSliceOp.getMixedStrides());
       // Copy tensor.
-      Value srcMemref = state.lookupBuffer(insertSliceOp.source());
+      Value srcMemref = state.lookupBuffer(rewriter, insertSliceOp.source());
       state.createMemCpy(rewriter, insertSliceOp.getLoc(), srcMemref, subView);
     }
 
-    state.replaceOp(op, dstMemref);
+    state.replaceOp(rewriter, op, dstMemref);
     return success();
   }
 };
 
 } // namespace tensor_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir

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

struct TransferReadOpInterface
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto readOp = cast<vector::TransferReadOp>(op);
     assert(readOp.getShapedType().isa<TensorType>() &&
            "only tensor types expected");
 
     // TransferReadOp always reads from the bufferized op.source().
-    Value buffer = state.lookupBuffer(readOp.source());
+    Value buffer = state.lookupBuffer(rewriter, readOp.source());
     Value read = rewriter.create<vector::TransferReadOp>(
         readOp.getLoc(), readOp.getVectorType(), buffer, readOp.indices(),
         readOp.permutation_map(), readOp.padding(), readOp.mask(),
         readOp.in_boundsAttr());
-    state.replaceOp(op, read);
+    state.replaceOp(rewriter, op, read);
     return success();
   }
 };
 
 struct TransferWriteOpInterface
     : public BufferizableOpInterface::ExternalModel<TransferWriteOpInterface,
                                                     vector::TransferWriteOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto writeOp = cast<vector::TransferWriteOp>(op);
     assert(writeOp.getShapedType().isa<TensorType>() &&
            "only tensor types expected");
 
     // Create a new transfer_write on buffer that doesn't have a return value.
     // Leave the previous transfer_write to dead code as it still has uses at
     // this point.
-    Value resultBuffer = state.getResultBuffer(op->getResult(0));
+    Value resultBuffer = state.getResultBuffer(rewriter, op->getResult(0));
     if (!resultBuffer)
       return failure();
     rewriter.create<vector::TransferWriteOp>(
         writeOp.getLoc(), writeOp.vector(), resultBuffer, writeOp.indices(),
         writeOp.permutation_mapAttr(), writeOp.in_boundsAttr());
-    state.replaceOp(op, resultBuffer);
+    state.replaceOp(rewriter, op, resultBuffer);
 
     return success();
   }
 };
 
 } // namespace vector_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg