[mlir][linalg][bufferize][NFC] Rename functions in BufferizationState

[springerm]

Summary

The old function names (e.g., replaceOp) could have been confusing to users because they sound similar to rewriter functions, but have slightly different semantics.

Depends On D116448

[springerm]

update

[springerm]

update

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

public:
 
   /// Creates a memref deallocation. The given memref buffer must have been
   /// allocated using `createAlloc`.
   void createDealloc(OpBuilder &b, Location loc, Value allocatedBuffer) const;
 
   /// Creates a memcpy between two given buffers.
   void createMemCpy(OpBuilder &b, Location loc, Value from, Value to) const;
 
-  /// Replace an op with replacement values. The op is deleted. Tensor OpResults
-  /// must be replaced with memref values.
-  void replaceOp(RewriterBase &rewriter, Operation *op,
-                 ValueRange values) const;
-
-  /// 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) const {
-    Operation *newOp =
-        rewriter.create<OpTy>(op->getLoc(), std::forward<Args>(args)...);
-    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(RewriterBase &rewriter, Value tensor) const;
 
   /// 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
private:
 
   /// Dialect-specific bufferization state.
   DenseMap<StringRef, std::unique_ptr<DialectBufferizationState>> dialectState;
 
   /// A reference to current bufferization options.
   const BufferizationOptions &options;
 };
 
+/// Replace an op with replacement values. The op is deleted. Tensor OpResults
+/// must be replaced with memref values.
+void replaceOpWithBufferizedValues(RewriterBase &rewriter, Operation *op,
+                                   ValueRange values);
+
+/// Replace an op with a new op. Tensor OpResults must be replaced with memref
+/// values.
+template <typename OpTy, typename... Args>
+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;
+}
+
 /// 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 = {});
 
 /// Return a contiguous MemRefType (i.e. with canonical/empty layout map)

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

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);
     auto globalMemref = globalCreator.getGlobalFor(constantOp);
-    state.replaceOpWithNewOp<memref::GetGlobalOp>(
+    replaceOpWithNewBufferizedOp<memref::GetGlobalOp>(
         rewriter, op, globalMemref.type(), globalMemref.getName());
     return success();
   }
 
   bool isWritable(Operation *op, Value value,
                   const BufferizationState &state) const {
     // Memory locations returned by memref::GetGlobalOp may not be written to.
     assert(value.isa<OpResult>());

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

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(
-    RewriterBase &rewriter, Operation *op, ValueRange values) const {
+void mlir::linalg::comprehensive_bufferize::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.
     if (opResult.getUses().empty())
       continue;
 

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

static LogicalResult bufferizeLinalgOp(RewriterBase &rewriter, LinalgOp op,
   SmallVector<Value> newOperands = newInputBuffers;
   newOperands.append(newOutputBuffers.begin(), newOutputBuffers.end());
 
   // Set insertion point now that potential alloc/dealloc are introduced.
   rewriter.setInsertionPoint(op);
   op.clone(rewriter, op.getLoc(), /*resultTypes=*/TypeRange{}, newOperands);
 
   // Replace the results of the old op with the new output buffers.
-  state.replaceOp(rewriter, op, newOutputBuffers);
+  replaceOpWithBufferizedValues(rewriter, op, newOutputBuffers);
 
   return success();
 }
 
 /// 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.
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(rewriter, op, alloc);
+    replaceOpWithBufferizedValues(rewriter, op, alloc);
     return success();
   }
 };
 
 struct TiledLoopOpInterface
     : public BufferizableOpInterface::ExternalModel<TiledLoopOpInterface,
                                                     linalg::TiledLoopOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
for (auto it : llvm::zip(oldTerminator.values(), newOutputs)) {
           newTerminator.getLoc(), output.getType(), std::get<0>(it));
       state.createMemCpy(rewriter, newTerminator.getLoc(), toMemrefOp, output);
     }
 
     // Erase old terminator.
     rewriter.eraseOp(oldTerminator);
 
     // Replace results and delete old op.
-    state.replaceOp(rewriter, op, newResults);
+    replaceOpWithBufferizedValues(rewriter, op, newResults);
 
     return success();
   }
 };
 
 struct YieldOpInterface
     : public BufferizableOpInterface::ExternalModel<YieldOpInterface,
                                                     linalg::YieldOp> {

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

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     // Get replacement values for non-tensor / non-equivalent results.
     for (unsigned i = 0; i < replacementValues.size(); ++i) {
       if (replacementValues[i])
         continue;
       replacementValues[i] = newCallOp->getResult(*retValMapping[i]);
     }
 
     // 5. Replace the old op with the new op.
-    state.replaceOp(rewriter, callOp, replacementValues);
+    replaceOpWithBufferizedValues(rewriter, callOp, replacementValues);
 
     return success();
   }
 };
 
 struct ReturnOpInterface
     : public BufferizableOpInterface::ExternalModel<ReturnOpInterface,
                                                     ReturnOp> {

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

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(rewriter, op, newIfOp->getResults());
+    replaceOpWithBufferizedValues(rewriter, op, newIfOp->getResults());
 
     return success();
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const BufferizationAliasInfo &aliasInfo,
                                 const BufferizationState &state) const {
     // IfOp results are equivalent to their corresponding yield values if both
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(rewriter, op, newForOp->getResults());
+    replaceOpWithBufferizedValues(rewriter, op, newForOp->getResults());
 
     return success();
   }
 };
 
 // TODO: Evolve toward matching ReturnLike ops. Check for aliasing values that
 // do not bufferize inplace. (Requires a few more changes for ConstantOp,
 // InitTensorOp, CallOp.)

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

Attribute memorySpace = rankedMemRefType
                                 : 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);
-    state.replaceOpWithNewOp<memref::CastOp>(rewriter, op, memRefType,
+    replaceOpWithNewBufferizedOp<memref::CastOp>(rewriter, op, memRefType,
                                                  resultBuffer);
     return success();
   }
 };
 
 struct DimOpInterface
     : public BufferizableOpInterface::ExternalModel<DimOpInterface,
                                                     tensor::DimOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
struct DimOpInterface
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const 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(rewriter, dimOp.source());
-    state.replaceOpWithNewOp<memref::DimOp>(rewriter, op, v, dimOp.index());
+    replaceOpWithNewBufferizedOp<memref::DimOp>(rewriter, op, v, dimOp.index());
     return success();
   }
 };
 
 struct ExtractSliceOpInterface
     : public BufferizableOpInterface::ExternalModel<ExtractSliceOpInterface,
                                                     tensor::ExtractSliceOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
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(rewriter, op, subView);
+    replaceOpWithBufferizedValues(rewriter, op, subView);
     return success();
   }
 };
 
 struct ExtractOpInterface
     : public BufferizableOpInterface::ExternalModel<ExtractOpInterface,
                                                     tensor::ExtractOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                const BufferizationState &state) const {
     return OpResult();
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationState &state) const {
     auto extractOp = cast<tensor::ExtractOp>(op);
     Value srcMemref = state.lookupBuffer(rewriter, extractOp.tensor());
-    state.replaceOpWithNewOp<memref::LoadOp>(rewriter, op, srcMemref,
+    replaceOpWithNewBufferizedOp<memref::LoadOp>(rewriter, op, srcMemref,
                                                  extractOp.indices());
     return success();
   }
 };
 
 struct InsertOpInterface
     : public BufferizableOpInterface::ExternalModel<InsertOpInterface,
                                                     tensor::InsertOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
struct InsertOpInterface
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationState &state) const {
     auto insertOp = cast<tensor::InsertOp>(op);
     Location loc = insertOp.getLoc();
     Value destMemref =
         state.getResultBuffer(rewriter, insertOp->getOpResult(0));
     rewriter.create<memref::StoreOp>(loc, insertOp.scalar(), destMemref,
                                      insertOp.indices());
-    state.replaceOp(rewriter, op, destMemref);
+    replaceOpWithBufferizedValues(rewriter, op, destMemref);
     return success();
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const BufferizationAliasInfo &aliasInfo,
                                 const BufferizationState &state) const {
     return BufferRelation::Equivalent;
   }
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());
     state.createMemCpy(rewriter, insertSliceOp.getLoc(), srcMemref, subView);
 
-    state.replaceOp(rewriter, op, dstMemref);
+    replaceOpWithBufferizedValues(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,
                           const 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 read = rewriter.create<vector::TransferReadOp>(
-        readOp.getLoc(), readOp.getVectorType(), buffer, readOp.indices(),
+    replaceOpWithNewBufferizedOp<vector::TransferReadOp>(
+        rewriter, readOp, readOp.getVectorType(), buffer, readOp.indices(),
         readOp.permutation_map(), readOp.padding(), readOp.mask(),
         readOp.in_boundsAttr());
-    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,
     // Leave the previous transfer_write to dead code as it still has uses at
     // this point.
     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(rewriter, op, resultBuffer);
+    replaceOpWithBufferizedValues(rewriter, op, resultBuffer);
 
     return success();
   }
 };
 
 } // namespace vector_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg