[mlir][linalg][bufferize][NFC] Simplify buffer API of BufferizationState

[springerm]

Summary

Instead of lookupBuffer and getResultBuffer, there is now a single getBuffer function. This simplifies the BufferizableOpInterface API and is less confusing to users. They could previously have called the wrong function.

Furthermore, since getBuffer now takes an OpOperand & instead of a Value, users can no longer accidentally use one of the previous two functions incorrectly, which would have resulted in missing buffer copies.

Depends On D115706

[springerm]

update

[springerm]

update

[springerm]

rebase

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

public:
   OpTy replaceOpWithNewBufferizedOp(RewriterBase &rewriter, Operation *op,
                                     Args &&...args) {
     auto newOp =
         rewriter.create<OpTy>(op->getLoc(), std::forward<Args>(args)...);
     replaceOpWithBufferizedValues(rewriter, op, newOp->getResults());
     return newOp;
   }
 
-  /// Lookup the memref buffer that is associated to the given tensor value.
-  /// Asserts if no buffer is associated.
-  Value lookupBuffer(RewriterBase &rewriter, Value tensor);
-
   /// Return `true` if the given OpResult has been decided to bufferize inplace.
   bool isInPlace(OpOperand &opOperand) const;
 
-  /// Return the result buffer (memref) for a given OpResult (tensor). Allocate
+  /// Return the buffer (memref) for a given OpOperand (tensor). Allocate
   /// a new buffer and copy over data from the existing buffer if out-of-place
-  /// bufferization is necessary.
-  FailureOr<Value> getResultBuffer(RewriterBase &rewriter, OpResult result);
+  /// bufferization was decided.
+  FailureOr<Value> getBuffer(RewriterBase &rewriter, OpOperand &opOperand,
+                             bool forceInPlace = false);
 
   /// 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]);
   }

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

for (OpOperand &opOperand : bufferizableOp->getOpOperands()) {
             aliasInfo.unionAliasSets(opOperand.get(), opResult);
           aliasInfo.markInPlace(opOperand);
         }
     }
     return WalkResult::advance();
   });
 }
 
+static Value lookupBuffer(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();
+
+  // Insert to_memref op.
+  OpBuilder::InsertionGuard g(rewriter);
+  setInsertionPointAfter(rewriter, tensor);
+  Type memrefType;
+  if (auto rankedTensorType = tensor.getType().dyn_cast<RankedTensorType>()) {
+    memrefType = getDynamicMemRefType(rankedTensorType);
+  } else {
+    memrefType = getUnrankedMemRefType(
+        tensor.getType().cast<TensorType>().getElementType());
+  }
+  return rewriter.create<bufferization::ToMemrefOp>(tensor.getLoc(), memrefType,
+                                                    tensor);
+}
+
 /// Return the result buffer (memref) for a given OpResult (tensor). Allocate
 /// a new buffer and copy over data from the existing buffer if out-of-place
 /// bufferization is necessary.
 FailureOr<Value>
-mlir::linalg::comprehensive_bufferize::BufferizationState::getResultBuffer(
-    RewriterBase &rewriter, OpResult result) {
+mlir::linalg::comprehensive_bufferize::BufferizationState::getBuffer(
+    RewriterBase &rewriter, OpOperand &opOperand, bool forceInPlace) {
   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(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(rewriter, o->get()) == operandBuffer;
-      }))
-    return FailureOr<Value>(op->emitError("result buffer is ambiguous"));
-
-  // If bufferizing out-of-place, allocate a new buffer.
-  if (!aliasInfo.isInPlace(*opOperand)) {
-    // 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();
-    // Move insertion point right after `operandBuffer`. That is where the
-    // allocation should be inserted (in the absence of allocation hoisting).
-    setInsertionPointAfter(rewriter, operandBuffer);
-    // Allocate the result buffer.
-    FailureOr<Value> resultBuffer =
-        createAlloc(rewriter, loc, operandBuffer, options.createDeallocs);
-    if (failed(resultBuffer))
-      return failure();
-    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 = options.dynCastBufferizableOp(lastWrite))
-      if (!bufferizableOp.isMemoryWrite(lastWrite.cast<OpResult>(), *this))
-        skipCopy = true;
-    // Do not copy if the copied data is never read. (Neither by this ops nor by
-    // any following op.)
-    if (!bufferizesToMemoryRead(*opOperand) && !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) {
-      // The copy happens right before the op that is bufferized.
-      rewriter.setInsertionPoint(op);
-      createMemCpy(rewriter, loc, operandBuffer, *resultBuffer);
-    }
-    return resultBuffer;
-  }
+  Operation *op = opOperand.getOwner();
+  Value operand = opOperand.get();
+  Value operandBuffer = lookupBuffer(rewriter, operand);
+
+  if (forceInPlace || aliasInfo.isInPlace(opOperand))
+    return operandBuffer;
+
+  // Bufferizing out-of-place: Allocate a new buffer.
+      return resultBuffer;
+  // Do not copy if the copied data is never read.
+  OpResult aliasingOpResult = getAliasingOpResult(opOperand);
+  if (aliasingOpResult && !bufferizesToMemoryRead(opOperand) &&
+      !isValueRead(aliasingOpResult))
+    return resultBuffer;
+  if (bufferizesToMemoryWrite(opOperand) && !bufferizesToMemoryRead(opOperand))
+    return resultBuffer;
+
 
-  // Bufferizing in-place. No need to allocate a new buffer.
-  return operandBuffer;
-}
-
 void mlir::linalg::comprehensive_bufferize::BufferizationState::
     replaceOpWithBufferizedValues(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.

 
 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(
-    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();
-
-  // Insert to_memref op.
-  OpBuilder::InsertionGuard g(rewriter);
-  setInsertionPointAfter(rewriter, tensor);
-  Type memrefType;
-  if (auto rankedTensorType = tensor.getType().dyn_cast<RankedTensorType>()) {
-    memrefType = getDynamicMemRefType(rankedTensorType);
-  } else {
-    memrefType = getUnrankedMemRefType(
-        tensor.getType().cast<TensorType>().getElementType());
-  }
-  return rewriter.create<bufferization::ToMemrefOp>(tensor.getLoc(), memrefType,
-                                                    tensor);
-}
-
 bool mlir::linalg::comprehensive_bufferize::BufferizationState::isInPlace(
     OpOperand &opOperand) const {
   return aliasInfo.isInPlace(opOperand);
 }
 
 MemRefType mlir::linalg::comprehensive_bufferize::getContiguousMemRefType(
     ShapedType shapedType, MemRefLayoutAttrInterface layout,
     Attribute memorySpace) {

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

static LogicalResult bufferizeLinalgOp(RewriterBase &rewriter, LinalgOp op,
   // New input operands for the cloned 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(rewriter, opOperand->get()));
+    // Input operands are never written to.
+    newInputBuffers.push_back(
+        *state.getBuffer(rewriter, *opOperand, /*forceInPlace=*/true));
   }
 
   // New output operands for the cloned op.
   SmallVector<Value> newOutputBuffers;
-  for (OpOperand *opOperand : op.getOutputOperands()) {
-    OpResult opResult = op.getTiedOpResult(opOperand);
-    assert(opResult && "could not find correspond OpResult");
-    FailureOr<Value> resultBuffer = state.getResultBuffer(rewriter, opResult);
+  for (OpResult opResult : op->getOpResults()) {
+    SmallVector<OpOperand *> aliasingOpOperands =
+        state.getAliasingOpOperand(opResult);
+    assert(aliasingOpOperands.size() == 1 && "expected 1 OpOperand");
+    FailureOr<Value> resultBuffer =
+        state.getBuffer(rewriter, *aliasingOpOperands.front());
     if (failed(resultBuffer))
       return failure();
     newOutputBuffers.push_back(*resultBuffer);
   }
 
   // Merge input/output operands.
   SmallVector<Value> newOperands = newInputBuffers;
   newOperands.append(newOutputBuffers.begin(), newOutputBuffers.end());
struct TiledLoopOpInterface
   bool isAllocationHoistingBarrier(Operation *op) const { return true; }
 
   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(rewriter, value));
-      } else {
-        newInputs.push_back(value);
-      }
-    }
-    int nextResultNum = 0;
-    for (Value value : tiledLoopOp.outputs()) {
-      if (value.getType().isa<TensorType>()) {
-        FailureOr<Value> buffer = state.getResultBuffer(
-            rewriter, tiledLoopOp->getResult(nextResultNum++));
-        if (failed(buffer))
+    for (int i = tiledLoopOp.getNumControlOperands();
+         i < tiledLoopOp->getNumOperands(); ++i) {
+      OpOperand &operand = tiledLoopOp->getOpOperand(i);
+      Value rewrittenValue = operand.get();
+      if (rewrittenValue.getType().isa<TensorType>()) {
+        FailureOr<Value> bufferOrFailure = state.getBuffer(rewriter, operand);
+        if (failed(bufferOrFailure))
           return failure();
-        newOutputs.push_back(*buffer);
-        newResults.push_back(*buffer);
+        rewrittenValue = *bufferOrFailure;
+      }
+      if (i <
+          tiledLoopOp.getNumControlOperands() + tiledLoopOp.getNumInputs()) {
+        newInputs.push_back(rewrittenValue);
       } else {
-        newOutputs.push_back(value);
+        newOutputs.push_back(rewrittenValue);
+        if (operand.get().getType().isa<TensorType>())
+          newResults.push_back(rewrittenValue);
       }
     }
 
     // Create new TiledLoopOp.
     auto newTiledLoopOp = rewriter.create<TiledLoopOp>(
         tiledLoopOp.getLoc(), tiledLoopOp.lowerBound(),
         tiledLoopOp.upperBound(), tiledLoopOp.step(), newInputs, newOutputs,
         tiledLoopOp.iterator_types(), tiledLoopOp.distribution_types());

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

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(rewriter, returnVal)));
+        getNonCastedValue(*state.getBuffer(rewriter, returnOperand)));
   }
 
   // 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);
for (auto it : llvm::enumerate(callOp.getResultTypes())) {
         continue;
       }
 
       if (Optional<int64_t> bbArgIdx =
               getEquivalentFuncArgIdx(funcOp, moduleState, returnValIdx)) {
         // Return operands that are equivalent to some bbArg, are not
         // returned.
         Value buffer =
-            state.lookupBuffer(rewriter, callOp->getOperand(*bbArgIdx));
+            *state.getBuffer(rewriter, callOp->getOpOperand(*bbArgIdx),
+                             /*forceInPlace=*/true);
         replacementValues[returnValIdx] = buffer;
         newOperands[*bbArgIdx] = buffer;
         continue;
       }
 
       return callOp->emitError(
           "call to FuncOp that returns non-equivalent tensors not supported");
     }
for (OpOperand &opOperand : callOp->getOpOperands()) {
       if (!tensorOperand.getType().isa<TensorType>()) {
         newOperands[idx] = tensorOperand;
         continue;
       }
 
       // Retrieve buffers for tensor operands. Tensor operand buffers, who's
       // corresponding FuncOp bbArgs are equivalent to a returned tensor, were
       // already stored in `newOperands` during Step 1.
-      Value buffer = newOperands[idx]
-                         ? newOperands[idx]
-                         : state.lookupBuffer(rewriter, tensorOperand);
+      Value buffer = newOperands[idx] ? newOperands[idx]
+                                      : *state.getBuffer(rewriter, opOperand,
+                                                         /*forceInPlace=*/true);
 
       // Caller / callee type mistmatch is handled with a CastOp.
       auto memRefType = bufferizedFuncType.getInput(idx);
       // Since we don't yet have a clear layout story, to_memref 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.

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

auto convert = [&](ValueRange values,
         size_t idx = it.index();
         Value val = it.value();
         result.push_back(indices.contains(idx) ? func(val, idx) : val);
       }
       return result;
     };
 
     // Construct a new scf.for op with memref instead of tensor values.
-    bool resultBufferFailure = false;
-    SmallVector<Value> initArgs =
-        convert(forOp.getInitArgs(), [&](Value val, int64_t index) {
-          FailureOr<Value> resultBuffer =
-              state.getResultBuffer(rewriter, forOp->getOpResult(index));
-          if (failed(resultBuffer)) {
-            resultBufferFailure = true;
-            return Value();
-          }
-          return *resultBuffer;
-        });
-    if (resultBufferFailure)
+    SmallVector<Value> initArgs;
+    for (OpOperand &opOperand : forOp.getIterOpOperands()) {
+      if (opOperand.get().getType().isa<TensorType>()) {
+        FailureOr<Value> resultBuffer = state.getBuffer(rewriter, opOperand);
+        if (failed(resultBuffer))
           return failure();
+        initArgs.push_back(*resultBuffer);
+      } else {
+        initArgs.push_back(opOperand.get());
+      }
+    }
     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.
     rewriter.setInsertionPointToStart(loopBody);

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

struct CastOpInterface
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto castOp = cast<tensor::CastOp>(op);
 
     // The result buffer still has the old (pre-cast) type.
     FailureOr<Value> resultBuffer =
-        state.getResultBuffer(rewriter, castOp->getResult(0));
+        state.getBuffer(rewriter, castOp->getOpOperand(0) /*source*/);
     if (failed(resultBuffer))
       return failure();
     auto sourceMemRefType = resultBuffer->getType().cast<BaseMemRefType>();
     Attribute memorySpace = sourceMemRefType.getMemorySpace();
     TensorType resultTensorType =
         castOp.getResult().getType().cast<TensorType>();
     MemRefLayoutAttrInterface layout;
 
struct DimOpInterface
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                BufferizationState &state) const {
     return OpResult();
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     auto dimOp = cast<tensor::DimOp>(op);
-    Value v = state.lookupBuffer(rewriter, dimOp.source());
+    Value v = *state.getBuffer(rewriter, dimOp->getOpOperand(0) /*source*/);
     state.replaceOpWithNewBufferizedOp<memref::DimOp>(rewriter, op, v,
                                                       dimOp.index());
     return success();
   }
 };
 
 /// Bufferization of tensor.extract_slice. Replace with memref.subview.
 struct ExtractSliceOpInterface
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(rewriter, extractSliceOp.source());
+    Value srcMemref =
+        *state.getBuffer(rewriter, extractSliceOp->getOpOperand(0) /*source*/,
+                         /*forceInPlace=*/true);
     auto srcMemrefType = srcMemref.getType().cast<MemRefType>();
     auto dstTensorType =
         extractSliceOp.result().getType().cast<RankedTensorType>();
 
     // If not inplaceable, alloc.
     bool inplace = state.isInPlace(extractSliceOp->getOpOperand(0));
     Value alloc;
     if (!inplace) {
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(rewriter, extractOp.tensor());
+    Value srcMemref =
+        *state.getBuffer(rewriter, extractOp->getOpOperand(0) /*tensor*/);
     state.replaceOpWithNewBufferizedOp<memref::LoadOp>(rewriter, op, srcMemref,
                                                        extractOp.indices());
     return success();
   }
 };
 
 /// Bufferization of tensor.insert. Replace with memref.store.
 struct 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);
     FailureOr<Value> destMemref =
-        state.getResultBuffer(rewriter, insertOp->getOpResult(0));
+        state.getBuffer(rewriter, insertOp->getOpOperand(1) /*dest*/);
     if (failed(destMemref))
       return failure();
     rewriter.create<memref::StoreOp>(insertOp.getLoc(), insertOp.scalar(),
                                      *destMemref, insertOp.indices());
     state.replaceOpWithBufferizedValues(rewriter, op, *destMemref);
     return success();
   }
 
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();
 
     // When bufferizing out-of-place, `getResultBuffer` allocates.
     FailureOr<Value> dstMemref =
-        state.getResultBuffer(rewriter, insertSliceOp->getResult(0));
+        state.getBuffer(rewriter, insertSliceOp->getOpOperand(1) /*dest*/);
     if (failed(dstMemref))
       return failure();
 
     // 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. If this tensor.insert_slice has a matching
     // tensor.extract_slice, the copy operation will eventually fold away.
-    Value srcMemref = state.lookupBuffer(rewriter, insertSliceOp.source());
+    Value srcMemref =
+        *state.getBuffer(rewriter, insertSliceOp->getOpOperand(0) /*source*/);
     state.createMemCpy(rewriter, loc, srcMemref, subView);
 
     state.replaceOpWithBufferizedValues(rewriter, op, *dstMemref);
     return success();
   }
 };
 
 } // namespace tensor_ext

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(rewriter, readOp.source());
+    Value buffer =
+        *state.getBuffer(rewriter, readOp->getOpOperand(0) /*source*/);
     state.replaceOpWithNewBufferizedOp<vector::TransferReadOp>(
         rewriter, readOp, readOp.getVectorType(), buffer, readOp.indices(),
         readOp.permutation_map(), readOp.padding(), readOp.mask(),
         readOp.in_boundsAttr());
     return success();
   }
 };
 
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     auto writeOp = cast<vector::TransferWriteOp>(op);
     assert(writeOp.getShapedType().isa<TensorType>() &&
            "only tensor types expected");
 
     // Create a new transfer_write on buffer that doesn't have a return value.
     // Leave the previous transfer_write to dead code as it still has uses at
     // this point.
     FailureOr<Value> resultBuffer =
-        state.getResultBuffer(rewriter, op->getResult(0));
+        state.getBuffer(rewriter, op->getOpOperand(1) /*source*/);
     if (failed(resultBuffer))
       return failure();
     rewriter.create<vector::TransferWriteOp>(
         writeOp.getLoc(), writeOp.vector(), *resultBuffer, writeOp.indices(),
         writeOp.permutation_mapAttr(), writeOp.in_boundsAttr());
     state.replaceOpWithBufferizedValues(rewriter, op, *resultBuffer);
 
     return success();