[mlir][bufferize][NFC] Move helper functions to BufferizationOptions

[springerm]

Summary

Move helper functions for creating allocs/deallocs/memcpys to BufferizationOptions.

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

struct BufferizationOptions {
   /// listed.
   BufferizableOpInterface dynCastBufferizableOp(Value value) const;
 
   /// Helper functions for allocation, deallocation, memory copying.
   Optional<AllocationFn> allocationFn;
   Optional<DeallocationFn> deallocationFn;
   Optional<MemCpyFn> memCpyFn;
 
+  /// Create a memref allocation with the given type and dynamic extents.
+  FailureOr<Value> createAlloc(OpBuilder &b, Location loc, MemRefType type,
+                               ValueRange dynShape) const;
+
+  /// Creates a memref deallocation. The given memref buffer must have been
+  /// allocated using `createAlloc`.
+  LogicalResult createDealloc(OpBuilder &b, Location loc,
+                              Value allocatedBuffer) const;
+
+  /// Creates a memcpy between two given buffers.
+  LogicalResult createMemCpy(OpBuilder &b, Location loc, Value from,
+                             Value to) const;
+
   /// 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.
   bool allowUnknownOps = false;
 
   /// Specifies whether function boundaries (ops in the func dialect) should be
   /// bufferized or not.
   bool bufferizeFunctionBoundaries = false;

 /// 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.
 BaseMemRefType getMemRefType(TensorType tensorType,
                              const BufferizationOptions &options,
                              MemRefLayoutAttrInterface layout = {},
                              Attribute memorySpace = {});
 
-/// 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);
-
 /// Try to hoist all new buffer allocations until the next hoisting barrier.
 LogicalResult hoistBufferAllocations(Operation *op,
                                      const BufferizationOptions &options);
 
 /// Create alloc/dealloc ops as specified in the bufferization options. If
 /// `onlyLeakingAlloc`, only those buffer allocations are processed for which no
 /// buffer deallocation can be created. `changed` is set to `true` if the IR was
 /// modified.

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

if (analysisState.bufferizesToMemoryWrite(opOperand) &&
     return resultBuffer;
 
   if (customCopyInsertionPoint) {
     rewriter.setInsertionPoint(*customCopyInsertionPoint);
   } else {
     // The copy happens right before the op that is bufferized.
     rewriter.setInsertionPoint(op);
   }
-  if (failed(
-          createMemCpy(rewriter, loc, operandBuffer, *resultBuffer, options)))
+  if (failed(options.createMemCpy(rewriter, loc, operandBuffer, *resultBuffer)))
     return failure();
 
   return resultBuffer;
 }
 
 /// Return the buffer type for a given OpOperand (tensor) after bufferization.
 BaseMemRefType BufferizationState::getBufferType(OpOperand &opOperand) const {
   Value tensor = opOperand.get();
bool AlwaysCopyAnalysisState::isTensorYielded(Value tensor) const {
   return true;
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific scoped alloc/dealloc insertion support.
 //===----------------------------------------------------------------------===//
 
 /// Create a memref allocation with the given type and dynamic extents.
-static FailureOr<Value> createAlloc(OpBuilder &b, Location loc, MemRefType type,
-                                    ValueRange dynShape,
-                                    const BufferizationOptions &options) {
-  if (options.allocationFn)
-    return (*options.allocationFn)(b, loc, type, dynShape,
-                                   options.bufferAlignment);
+FailureOr<Value> BufferizationOptions::createAlloc(OpBuilder &b, Location loc,
+                                                   MemRefType type,
+                                                   ValueRange dynShape) const {
+  if (allocationFn)
+    return (*allocationFn)(b, loc, type, dynShape, bufferAlignment);
 
   // Default bufferallocation via AllocOp.
   Value allocated = b.create<memref::AllocOp>(
-      loc, type, dynShape, b.getI64IntegerAttr(options.bufferAlignment));
+      loc, type, dynShape, b.getI64IntegerAttr(bufferAlignment));
   return allocated;
 }
 
 /// Creates a memref deallocation. The given memref buffer must have been
 /// allocated using `createAlloc`.
-LogicalResult
-bufferization::createDealloc(OpBuilder &b, Location loc, Value allocatedBuffer,
-                             const BufferizationOptions &options) {
-  if (options.deallocationFn)
-    return (*options.deallocationFn)(b, loc, allocatedBuffer);
+LogicalResult BufferizationOptions::createDealloc(OpBuilder &b, Location loc,
+                                                  Value allocatedBuffer) const {
+  if (deallocationFn)
+    return (*deallocationFn)(b, loc, allocatedBuffer);
 
   // Default buffer deallocation via DeallocOp.
   b.create<memref::DeallocOp>(loc, allocatedBuffer);
   return success();
 }
 
 /// Compute the type of the `memref` to use for allocating the buffer for
 /// `shapedValue`. Also returns (by reference in `dynShape`), the value for the
skipDealloc = !getOptions().createDeallocs ||
                   getAnalysisState().isTensorYielded(shapedValue);
   }
 
   // Create the buffer allocation.
   return createBufferAllocation(b, loc, allocMemRefType, dynShape, skipDealloc);
 }
 
 /// Create a memory copy between two memref buffers.
-LogicalResult bufferization::createMemCpy(OpBuilder &b, Location loc,
-                                          Value from, Value to,
-                                          const BufferizationOptions &options) {
-  if (options.memCpyFn)
-    return (*options.memCpyFn)(b, loc, from, to);
+LogicalResult BufferizationOptions::createMemCpy(OpBuilder &b, Location loc,
+                                                 Value from, Value to) const {
+  if (memCpyFn)
+    return (*memCpyFn)(b, loc, from, to);
 
   b.create<memref::CopyOp>(loc, from, to);
   return success();
 }
 
 LogicalResult
 bufferization::createAllocDeallocOps(Operation *op,
                                      const BufferizationOptions &options,
WalkResult status = op->walk([&](memref::AllocaOp allocaOp) {
     bool skipDealloc = allocaOp->hasAttr(kSkipDeallocAttr);
     if (onlyLeakingAllocs && !skipDealloc)
       return WalkResult::skip();
 
     // Create alloc.
     Block *block = allocaOp->getBlock();
     rewriter.setInsertionPoint(allocaOp);
     FailureOr<Value> alloc =
-        createAlloc(rewriter, allocaOp->getLoc(), allocaOp.getType(),
-                    allocaOp.dynamicSizes(), options);
+        options.createAlloc(rewriter, allocaOp->getLoc(), allocaOp.getType(),
+                            allocaOp.dynamicSizes());
     if (failed(alloc))
       return WalkResult::interrupt();
     rewriter.replaceOp(allocaOp, *alloc);
     if (changed)
       *changed = true;
 
     // Stop here if the buffer should not be deallocated.
     if (skipDealloc)
       return WalkResult::advance();
 
     // Create dealloc.
     rewriter.setInsertionPoint(block->getTerminator());
-    if (failed(createDealloc(rewriter, alloc->getLoc(), *alloc, options)))
+    if (failed(options.createDealloc(rewriter, alloc->getLoc(), *alloc)))
       return WalkResult::interrupt();
 
     return WalkResult::advance();
   });
 
   return success(!status.wasInterrupted());
 }
 

mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp

for (OpOperand &returnOperand : returnOp->getOpOperands()) {
           Value toMemrefOp = rewriter.create<bufferization::ToMemrefOp>(
               loc,
               getMemRefType(returnVal.getType().cast<TensorType>(), options),
               returnVal);
           BlockArgument equivBbArg = funcOp.getArgument(*equivBbArgIdx);
           // Note: This copy will fold away. It must be inserted here to ensure
           // that `returnVal` still has at least one use and does not fold away.
           if (failed(
-                  createMemCpy(rewriter, loc, toMemrefOp, equivBbArg, options)))
+                  options.createMemCpy(rewriter, loc, toMemrefOp, equivBbArg)))
             return funcOp->emitError("could not generate copy for bbArg");
           continue;
         }
       }
 
       returnValues.push_back(*state.getBuffer(rewriter, returnOperand));
     }
 

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

static Value getYieldedBuffer(RewriterBase &rewriter, Value tensor,
   // TODO: There are cases in which it is not neccessary to return a new
   // buffer allocation. E.g., when equivalent values are yielded in a
   // different order. This could be resolved with copies.
   Optional<Value> yieldedAlloc = state.createAlloc(
       rewriter, tensor.getLoc(), yieldedVal, /*deallocMemref=*/false);
   // TODO: We should rollback, but for now just assume that this always
   // succeeds.
   assert(yieldedAlloc.hasValue() && "could not create alloc");
-  LogicalResult copyStatus = bufferization::createMemCpy(
-      rewriter, tensor.getLoc(), yieldedVal, *yieldedAlloc, state.getOptions());
+  LogicalResult copyStatus = state.getOptions().createMemCpy(
+      rewriter, tensor.getLoc(), yieldedVal, *yieldedAlloc);
   (void)copyStatus;
   assert(succeeded(copyStatus) && "could not create memcpy");
 
   // The iter_arg memref type may have a layout map. Cast the new buffer
   // to the same type if needed.
   return castBuffer(rewriter, *yieldedAlloc, type);
 }
 

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

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.getAnalysisState().isValueRead(extractSliceOp.result()))
-        if (failed(createMemCpy(rewriter, extractSliceOp.getLoc(), subView,
-                                alloc, state.getOptions())))
+        if (failed(state.getOptions().createMemCpy(
+                rewriter, extractSliceOp.getLoc(), subView, alloc)))
           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.
     auto srcMemref =
         state.getBuffer(rewriter, insertSliceOp->getOpOperand(0) /*source*/);
-    if (failed(srcMemref) || failed(createMemCpy(rewriter, loc, *srcMemref,
-                                                 subView, state.getOptions())))
+    if (failed(srcMemref) || failed(state.getOptions().createMemCpy(
+                                 rewriter, loc, *srcMemref, subView)))
       return failure();
 
     replaceOpWithBufferizedValues(rewriter, op, *dstMemref);
     return success();
   }
 };
 
 /// Bufferization of tensor.rank. Replace with memref.rank.