[mlir][linalg][bufferize][NFC] Merge AllocationCallbacks into BufferizationOptions

[springerm]

Summary

Also move createAlloc and related helper functions out of BufferizationState. The goal is to make BufferizationState as small as possible. (Code cleanup)

Depends On D116687

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

 static constexpr int64_t kBufferAlignments = 128;
 
 class BufferizationAliasInfo;
 class BufferizableOpInterface;
 struct BufferizationOptions;
 class BufferizationState;
 struct PostAnalysisStep;
 
-/// 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<FailureOr<Value>(
-      OpBuilder &, Location, MemRefType, ArrayRef<Value>)>;
-  using DeallocationFn = std::function<void(OpBuilder &, Location, Value)>;
-  using MemCpyFn = std::function<void(OpBuilder &, Location, Value, Value)>;
-
-  AllocationCallbacks(AllocationFn allocFn, DeallocationFn deallocFn,
-                      MemCpyFn copyFn)
-      : allocationFn(std::move(allocFn)), deallocationFn(std::move(deallocFn)),
-        memCpyFn(std::move(copyFn)) {}
-
-  /// 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;
-};
-
-/// Return default allocation callbacks.
-std::unique_ptr<AllocationCallbacks> defaultAllocationCallbacks();
-
 /// PostAnalysisSteps can be registered with `BufferizationOptions` and are
 /// executed after the analysis, but before bufferization. They can be used to
 /// implement custom dialect-specific optimizations.
 struct PostAnalysisStep {
   virtual ~PostAnalysisStep() = default;
 
   /// Run the post analysis step. This function may modify the IR, but must keep
   /// `aliasInfo` consistent. Newly created operations and operations that
   /// should be re-analyzed must be added to `newOps`.
   virtual LogicalResult run(Operation *op, BufferizationState &state,
                             BufferizationAliasInfo &aliasInfo,
                             SmallVector<Operation *> &newOps) = 0;
 };
 
 using PostAnalysisStepList = std::vector<std::unique_ptr<PostAnalysisStep>>;
 
 /// Options for ComprehensiveBufferize.
 struct BufferizationOptions {
+  using AllocationFn = std::function<FailureOr<Value>(
+      OpBuilder &, Location, MemRefType, ArrayRef<Value>)>;
+  using DeallocationFn =
+      std::function<LogicalResult(OpBuilder &, Location, Value)>;
+  using MemCpyFn =
+      std::function<LogicalResult(OpBuilder &, Location, Value, Value)>;
+
   BufferizationOptions();
 
   // BufferizationOptions cannot be copied.
   BufferizationOptions(const BufferizationOptions &other) = delete;
 
   /// Register a "post analysis" step. Such steps are executed after the
   /// analysis, but before bufferization.
   template <typename Step, typename... Args>
struct BufferizationOptions {
   /// listed.
   BufferizableOpInterface dynCastBufferizableOp(Operation *op) const;
 
   /// Try to cast the given value to BufferizableOpInterface if the op is allow
   /// listed.
   BufferizableOpInterface dynCastBufferizableOp(Value value) const;
 
   /// Helper functions for allocation, deallocation, memory copying.
-  std::unique_ptr<AllocationCallbacks> allocationFns;
+  Optional<AllocationFn> allocationFn;
+  Optional<DeallocationFn> deallocationFn;
+  Optional<MemCpyFn> memCpyFn;
 
   /// Specifies whether returning newly allocated memrefs should be allowed.
   /// Otherwise, a pass failure is triggered.
   bool allowReturnMemref = false;
 
   /// Specifies whether not bufferizable ops are allowed in the input. If so,
   /// bufferization.to_memref and bufferization.to_tensor ops are inserted at
   /// the boundaries.
public:
   /// 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.
   SetVector<Value> findLastPrecedingWrite(Value value) const;
 
-  /// Creates a memref allocation.
-  FailureOr<Value> createAlloc(OpBuilder &b, Location loc, MemRefType type,
-                               ArrayRef<Value> dynShape) const;
-
-  /// Creates a memref allocation for the given shaped value. This function may
-  /// perform additional optimizations such as buffer allocation hoisting. If
-  /// `createDealloc`, a deallocation op is inserted at the point where the
-  /// allocation goes out of scope.
-  FailureOr<Value> createAlloc(OpBuilder &b, Location loc, Value shapedValue,
-                               bool deallocMemref) const;
-
-  /// 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;
-
   /// Return `true` if the given OpResult has been decided to bufferize inplace.
   bool isInPlace(OpOperand &opOperand) const;
 
   /// 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 was decided.
   FailureOr<Value>
   getBuffer(RewriterBase &rewriter, OpOperand &opOperand,
UnrankedMemRefType getUnrankedMemRefType(Type elementType,
                                          Attribute memorySpace = {});
 
 /// Return a MemRefType to which the `tensorType` can be bufferized in a
 /// composable fashion. The layout must be the most dynamic possible and
 /// canonicalize away once bufferization is finished.
 MemRefType getDynamicMemRefType(RankedTensorType tensorType,
                                 unsigned addressSpace = 0);
 
+/// Creates a memref allocation.
+FailureOr<Value> createAlloc(OpBuilder &b, Location loc, MemRefType type,
+                             ArrayRef<Value> dynShape,
+                             const BufferizationOptions &options);
+
+/// Creates a memref allocation for the given shaped value. This function may
+/// perform additional optimizations such as buffer allocation hoisting. If
+/// `createDealloc`, a deallocation op is inserted at the point where the
+/// allocation goes out of scope.
+FailureOr<Value> createAlloc(OpBuilder &b, Location loc, Value shapedValue,
+                             bool deallocMemref,
+                             const BufferizationOptions &options);
+
+/// Creates a memref deallocation. The given memref buffer must have been
+/// allocated using `createAlloc`.
+LogicalResult createDealloc(OpBuilder &b, Location loc, Value allocatedBuffer,
+                            const BufferizationOptions &options);
+
+/// Creates a memcpy between two given buffers.
+LogicalResult createMemCpy(OpBuilder &b, Location loc, Value from, Value to,
+                           const BufferizationOptions &options);
+
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h.inc"
 
 namespace mlir {
 namespace linalg {

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

 
 using namespace mlir;
 using namespace linalg::comprehensive_bufferize;
 
 //===----------------------------------------------------------------------===//
 // BufferizationOptions
 //===----------------------------------------------------------------------===//
 
-/// Default allocation function that is used by the comprehensive bufferization
-/// pass. The default currently creates a ranked memref using `memref.alloc`.
-static FailureOr<Value> defaultAllocationFn(OpBuilder &b, Location loc,
-                                            MemRefType type,
-                                            ArrayRef<Value> dynShape) {
-  Value allocated = b.create<memref::AllocOp>(
-      loc, type, dynShape, b.getI64IntegerAttr(kBufferAlignments));
-  return allocated;
-}
-
-/// Default deallocation function that is used by the comprehensive
-/// bufferization pass. It expects to recieve back the value called from the
-/// `defaultAllocationFn`.
-static void defaultDeallocationFn(OpBuilder &b, Location loc,
-                                  Value allocatedBuffer) {
-  b.create<memref::DeallocOp>(loc, allocatedBuffer);
-}
-
-/// Default memory copy function that is used by the comprehensive bufferization
-/// pass. Creates a `memref.copy` op.
-static void defaultMemCpyFn(OpBuilder &b, Location loc, Value from, Value to) {
-  b.create<memref::CopyOp>(loc, from, to);
-}
-
-std::unique_ptr<AllocationCallbacks>
-mlir::linalg::comprehensive_bufferize::defaultAllocationCallbacks() {
-  return std::make_unique<AllocationCallbacks>(
-      defaultAllocationFn, defaultDeallocationFn, defaultMemCpyFn);
-}
-
-// Default constructor for BufferizationOptions that sets all allocation
-// callbacks to their default functions.
-BufferizationOptions::BufferizationOptions()
-    : allocationFns(defaultAllocationCallbacks()) {}
+// Default constructor for BufferizationOptions.
+BufferizationOptions::BufferizationOptions() {}
 
 BufferizableOpInterface mlir::linalg::comprehensive_bufferize::
     BufferizationOptions::dynCastBufferizableOp(Operation *op) const {
   if (isOpAllowed(op))
     return dyn_cast<BufferizableOpInterface>(op);
   return nullptr;
 }
 
mlir::linalg::comprehensive_bufferize::BufferizationState::getBuffer(
   if (forceInPlace || aliasInfo.isInPlace(opOperand))
     return operandBuffer;
 
   // Bufferizing out-of-place: Allocate a new buffer.
   // 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);
+  FailureOr<Value> resultBuffer = createAlloc(rewriter, loc, operandBuffer,
+                                              options.createDeallocs, options);
   if (failed(resultBuffer))
     return failure();
   // Do not copy if the last preceding writes of `operand` are ops that do
   // 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.
   SetVector<Value> lastWrites = findLastPrecedingWrite(operand);
if (bufferizesToMemoryWrite(opOperand) && !bufferizesToMemoryRead(opOperand))
     return resultBuffer;
 
   if (customCopyInsertionPoint) {
     rewriter.setInsertionPoint(*customCopyInsertionPoint);
   } else {
     // The copy happens right before the op that is bufferized.
     rewriter.setInsertionPoint(op);
   }
-  createMemCpy(rewriter, loc, operandBuffer, *resultBuffer);
+  if (failed(
+          createMemCpy(rewriter, loc, operandBuffer, *resultBuffer, options)))
+    return failure();
 
   return resultBuffer;
 }
 
 void mlir::linalg::comprehensive_bufferize::replaceOpWithBufferizedValues(
     RewriterBase &rewriter, Operation *op, ValueRange values) {
   OpBuilder::InsertionGuard g(rewriter);
 
if (dynShape.empty())
     moveInsertionPointToAllocationHoistingBarrier(b);
 
   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.
-FailureOr<Value>
-mlir::linalg::comprehensive_bufferize::BufferizationState::createAlloc(
-    OpBuilder &b, Location loc, Value shapedValue, bool deallocMemref) const {
+FailureOr<Value> mlir::linalg::comprehensive_bufferize::createAlloc(
+    OpBuilder &b, Location loc, Value shapedValue, bool deallocMemref,
+    const BufferizationOptions &options) {
   // 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);
-  FailureOr<Value> allocated = createAlloc(b, loc, allocMemRefType, dynShape);
+  FailureOr<Value> allocated =
+      createAlloc(b, loc, allocMemRefType, dynShape, options);
   if (failed(allocated))
     return failure();
   Value casted = allocated.getValue();
   if (memRefType && memRefType != allocMemRefType) {
     assert(memref::CastOp::areCastCompatible(allocated.getValue().getType(),
                                              memRefType) &&
            "createAlloc: cast incompatible");
     casted = b.create<memref::CastOp>(loc, memRefType, allocated.getValue());
   }
 
   if (deallocMemref) {
     // 2. Create memory deallocation.
     b.setInsertionPoint(allocated.getValue().getParentBlock()->getTerminator());
-    createDealloc(b, loc, allocated.getValue());
+    if (failed(createDealloc(b, loc, allocated.getValue(), options)))
+      return failure();
   }
 
   return casted;
 }
 
 /// Create a memref allocation.
-FailureOr<Value>
-mlir::linalg::comprehensive_bufferize::BufferizationState::createAlloc(
-    OpBuilder &b, Location loc, MemRefType type,
-    ArrayRef<Value> dynShape) const {
-  return options.allocationFns->allocationFn(b, loc, type, dynShape);
+FailureOr<Value> mlir::linalg::comprehensive_bufferize::createAlloc(
+    OpBuilder &b, Location loc, MemRefType type, ArrayRef<Value> dynShape,
+    const BufferizationOptions &options) {
+  if (options.allocationFn)
+    return (*options.allocationFn)(b, loc, type, dynShape);
+
+  // Default bufferallocation via AllocOp.
+  Value allocated = b.create<memref::AllocOp>(
+      loc, type, dynShape, b.getI64IntegerAttr(kBufferAlignments));
+  return allocated;
 }
 
 /// Create a memref deallocation.
-void mlir::linalg::comprehensive_bufferize::BufferizationState::createDealloc(
-    OpBuilder &b, Location loc, Value allocatedBuffer) const {
-  return options.allocationFns->deallocationFn(b, loc, allocatedBuffer);
+LogicalResult mlir::linalg::comprehensive_bufferize::createDealloc(
+    OpBuilder &b, Location loc, Value allocatedBuffer,
+    const BufferizationOptions &options) {
+  if (options.deallocationFn)
+    return (*options.deallocationFn)(b, loc, allocatedBuffer);
+
+  // Default buffer deallocation via DeallocOp.
+  b.create<memref::DeallocOp>(loc, allocatedBuffer);
+  return success();
 }
 
 /// Create a memory copy between two memref buffers.
-void mlir::linalg::comprehensive_bufferize::BufferizationState::createMemCpy(
-    OpBuilder &b, Location loc, Value from, Value to) const {
-  return options.allocationFns->memCpyFn(b, loc, from, to);
+LogicalResult mlir::linalg::comprehensive_bufferize::createMemCpy(
+    OpBuilder &b, Location loc, Value from, Value to,
+    const BufferizationOptions &options) {
+  if (options.memCpyFn)
+    return (*options.memCpyFn)(b, loc, from, to);
+
+  b.create<memref::CopyOp>(loc, from, to);
+  return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific BlockAndValueMapping support with debugging.
 //===----------------------------------------------------------------------===//
 
 bool mlir::linalg::comprehensive_bufferize::isFunctionArgument(Value value) {
   auto bbArg = value.dyn_cast<BlockArgument>();

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

struct InitTensorOpInterface
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationState &state) const {
     auto initTensorOp = cast<linalg::InitTensorOp>(op);
 
     // The InitTensorOp may have been eliminated.
     if (initTensorOp->getUses().empty())
       return success();
 
-    FailureOr<Value> alloc = state.createAlloc(
-        rewriter, initTensorOp->getLoc(), initTensorOp.result(),
-        state.getOptions().createDeallocs);
+    FailureOr<Value> alloc =
+        createAlloc(rewriter, initTensorOp->getLoc(), initTensorOp.result(),
+                    state.getOptions().createDeallocs, state.getOptions());
     if (failed(alloc))
       return failure();
     replaceOpWithBufferizedValues(rewriter, op, *alloc);
     return success();
   }
 };
 
 /// Bufferization of linalg.tiled_loop. Replace with a new linalg.tiled_loop
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
 
     // Copy buffer of yielded tensor to output buffer. If everything bufferized
     // inplace, this copy will fold away.
     rewriter.setInsertionPoint(newTerminator);
     for (auto it : llvm::zip(oldTerminator.values(), newOutputs)) {
       Value output = std::get<1>(it);
       Value toMemrefOp = rewriter.create<bufferization::ToMemrefOp>(
           newTerminator.getLoc(), output.getType(), std::get<0>(it));
-      state.createMemCpy(rewriter, newTerminator.getLoc(), toMemrefOp, output);
+      if (failed(createMemCpy(rewriter, newTerminator.getLoc(), toMemrefOp,
+                              output, state.getOptions())))
+        return failure();
     }
 
     // Erase old terminator.
     rewriter.eraseOp(oldTerminator);
 
     // Replace results and delete old op.
     replaceOpWithBufferizedValues(rewriter, op, newResults);
 

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

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     auto dstTensorType =
         extractSliceOp.result().getType().cast<RankedTensorType>();
 
     // If not inplaceable, alloc.
     bool inplace = state.isInPlace(extractSliceOp->getOpOperand(0));
     Value alloc;
     if (!inplace) {
       FailureOr<Value> allocOrFailure =
-          state.createAlloc(rewriter, loc, extractSliceOp.result(),
-                            state.getOptions().createDeallocs);
+          createAlloc(rewriter, loc, extractSliceOp.result(),
+                      state.getOptions().createDeallocs, state.getOptions());
       if (failed(allocOrFailure))
         return failure();
       alloc = *allocOrFailure;
     }
 
     // Expand offsets, sizes and strides to the full rank to handle the
     // rank-reducing case.
     SmallVector<OpFoldResult> mixedOffsets = extractSliceOp.getMixedOffsets();
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     Value subView = rewriter.create<memref::SubViewOp>(
         loc, subviewMemRefType, srcMemref, mixedOffsets, mixedSizes,
         mixedStrides);
 
     // 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);
+        if (failed(createMemCpy(rewriter, extractSliceOp.getLoc(), subView,
+                                alloc, state.getOptions())))
+          return failure();
       subView = alloc;
     }
 
     replaceOpWithBufferizedValues(rewriter, op, subView);
     return success();
   }
 };
 
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     Value subView = rewriter.create<memref::SubViewOp>(
         loc, subviewMemRefType, *dstMemref, mixedOffsets, mixedSizes,
         mixedStrides);
 
     // Copy tensor. If this tensor.insert_slice has a matching
     // tensor.extract_slice, the copy operation will eventually fold away.
     Value srcMemref =
         *state.getBuffer(rewriter, insertSliceOp->getOpOperand(0) /*source*/);
-    state.createMemCpy(rewriter, loc, srcMemref, subView);
+    if (failed(createMemCpy(rewriter, loc, srcMemref, subView,
+                            state.getOptions())))
+      return failure();
 
     replaceOpWithBufferizedValues(rewriter, op, *dstMemref);
     return success();
   }
 };
 
 } // namespace tensor_ext
 } // namespace comprehensive_bufferize

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferizePass.cpp

static FailureOr<Value> allocationFnUsingAlloca(OpBuilder &b, Location loc,
   Value allocated = b.create<memref::AllocaOp>(
       loc, type, dynShape, b.getI64IntegerAttr(kBufferAlignments));
   return allocated;
 }
 
 void LinalgComprehensiveModuleBufferize::runOnOperation() {
   auto options = std::make_unique<BufferizationOptions>();
   if (useAlloca) {
-    options->allocationFns->allocationFn = allocationFnUsingAlloca;
-    options->allocationFns->deallocationFn = [](OpBuilder &b, Location loc,
-                                                Value v) {};
+    options->allocationFn = allocationFnUsingAlloca;
+    options->deallocationFn = [](OpBuilder &b, Location loc, Value v) {
+      return success();
+    };
   }
 
   options->allowReturnMemref = allowReturnMemref;
   options->allowUnknownOps = allowUnknownOps;
   options->analysisFuzzerSeed = analysisFuzzerSeed;
   options->testAnalysisOnly = testAnalysisOnly;
   options->printConflicts = printConflicts;
   options->createDeallocs = createDeallocs;