[mlir][linalg][bufferize] Add MemCpyFn to AllocationCallbacks struct

[springerm]

Summary

This in preparation of decoupling BufferizableOpInterface, Comprehensive Bufferize and dialects.

The goal of this commit is to make getResultBuffer (and other bufferize functions) independent of LinalgOps.

Depends On D112902

[springerm]

As an alternative to this revision, I could also just use memref::CopyOp to get rid of the Linalg dependency.

[springerm]

no change

[springerm]

rebase

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

 Optional<Value> defaultAllocationFn(OpBuilder &b, Location loc,
                                     Value shapedValue);
 
 /// Default deallocation function that is used by the comprehensive
 /// bufferization pass. It expects to recieve back the value called from the
 /// `defaultAllocationFn`.
 void defaultDeallocationFn(OpBuilder &b, Location loc, Value allocatedBuffer);
 
+/// Default memory copy function that is used by the comprehensive bufferization
+/// pass. Creates a `linalg.copy` op.
+void defaultMemCpyFn(OpBuilder &b, Location loc, Value from, Value to);
+
 /// Callback functions that are used by the comprehensive bufferization pass to
 /// allocate/deallocate memory. These default to use the
 /// `defaultAllocationFn`/`defaultDeallocationFn`, but can be overridden by the
 /// caller. The `deallocationFn` is gauranteed to recieve the `Value` returned
 /// by the `allocationFn`.
 struct AllocationCallbacks {
-  std::function<Optional<Value>(OpBuilder &b, Location loc, Value shapedValue)>
-      allocationFn = defaultAllocationFn;
-  std::function<void(OpBuilder &b, Location loc, Value v)> deallocationFn =
-      defaultDeallocationFn;
+  using AllocationFn =
+      std::function<Optional<Value>(OpBuilder &, Location, 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(allocFn), deallocationFn(deallocFn), memCpyFn(copyFn) {}
+
+  AllocationFn allocationFn;
+  DeallocationFn deallocationFn;
+  MemCpyFn memCpyFn;
 };
 
 /// Bufferize one particular op.
 /// `bufferizedFunctionTypes` (resp. `globalCreator`) are expected to be
 /// non-null if `op` is a CallOpInterface (resp. GlobalCreator).
 LogicalResult
 bufferizeOp(Operation *op, BlockAndValueMapping &bvm,
             BufferizationAliasInfo &aliasInfo,
             AllocationCallbacks allocationFns,
             DenseMap<FuncOp, FunctionType> *bufferizedFunctionTypes = nullptr);
 
 void registerBufferiableOpInterfaceExternalModels(DialectRegistry &registry);
 } // namespace linalg
 } // namespace mlir
 
 #endif // define MLIR_DIALECT_LINALG_TRANSFORMS_COMPREHENSIVE_BUFFERIZE_H

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

if (needCopy) {
     if (isInitTensorOp(operand))
       skipCopy = true;
     // Do not copy if the copied data is never read.
     if (!isValueRead(result))
       skipCopy = true;
     if (!skipCopy) {
       // Set insertion point now that potential alloc/dealloc are introduced.
       b.setInsertionPoint(op);
-      b.create<CopyOp>(loc, operandBuffer, resultBuffer);
+      allocationFns.memCpyFn(b, loc, operandBuffer, resultBuffer);
     }
     return resultBuffer;
   }
 
   // Bufferizing in-place. No need to allocate a new buffer.
   return operandBuffer;
 }
 

 
 void mlir::linalg::defaultDeallocationFn(OpBuilder &b, Location loc,
                                          Value allocatedBuffer) {
   OpBuilder::InsertionGuard g(b);
   b.setInsertionPoint(allocatedBuffer.getParentBlock()->getTerminator());
   b.create<memref::DeallocOp>(loc, allocatedBuffer);
 }
 
+void mlir::linalg::defaultMemCpyFn(OpBuilder &b, Location loc, Value from,
+                                   Value to) {
+  b.create<CopyOp>(loc, from, to);
+}
+
 LogicalResult mlir::linalg::bufferizeOp(
     Operation *op, BlockAndValueMapping &bvm, BufferizationAliasInfo &aliasInfo,
     AllocationCallbacks allocationFns,
     DenseMap<FuncOp, FunctionType> *bufferizedFunctionTypes) {
   OpBuilder b(op->getContext());
 
   // CallOps are handled separately.
   if (auto callOp = dyn_cast<CallOpInterface>(op)) {
if (!allocationFns) {
     // The allocation functions to use needs to be set here. The flag for the
     // pass and flag for the use of alloca map to LLVM command line
     // options. These being static global objects have no set order in which
     // they are defined. So ideally this should be in the constructor, but the
     // constructor might be called before the flag is initialized using the
     // command line option. So this is set up at the start of the pass.
     if (useAlloca) {
       AllocationCallbacks allocaAllocationFns = {
-          allocationFnUsingAlloca, [](OpBuilder &b, Location loc, Value v) {}};
+          allocationFnUsingAlloca, [](OpBuilder &b, Location loc, Value v) {},
+          defaultMemCpyFn};
       allocationFns =
           std::make_unique<AllocationCallbacks>(std::move(allocaAllocationFns));
     } else {
-      allocationFns = std::make_unique<AllocationCallbacks>();
+      allocationFns = std::make_unique<AllocationCallbacks>(
+          defaultAllocationFn, defaultDeallocationFn, defaultMemCpyFn);
     }
   }
   ModuleOp moduleOp = getOperation();
   applyEnablingTransformations(moduleOp);
 
   SmallVector<FuncOp> orderedFuncOps;
   DenseMap<FuncOp, DenseSet<Operation *>> callerMap;
   DenseMap<FuncOp, FunctionType> bufferizedFunctionTypes;
Value subView = b.create<memref::SubViewOp>(
         extractSliceOp.getMixedSizes(), extractSliceOp.getMixedStrides());
     // Insert new alias.
     aliasInfo.insertNewBufferAlias(subView, srcMemref);
 
     /// If not inplaceable, copy.
     if (alloc) {
       // Do not copy if the copied data is never read.
       if (isValueRead(extractSliceOp.result()))
-        b.create<CopyOp>(extractSliceOp.getLoc(), subView, alloc);
+        allocationFn.memCpyFn(b, extractSliceOp.getLoc(), subView, alloc);
       subView = alloc;
     }
 
     map(bvm, extractSliceOp.result(), subView);
     return success();
   }
 };
 
if (!aliasInfo.isSourceEquivalentToAMatchingInplaceExtractSliceOp(
       LDBG("insert_slice needs extra source copy: " << insertSliceOp.source()
                                                     << " -> copy\n");
       // Take a subview of the dst.
       Value subView = b.create<memref::SubViewOp>(
           loc, subviewMemRefType, dstMemref, insertSliceOp.getMixedOffsets(),
           insertSliceOp.getMixedSizes(), insertSliceOp.getMixedStrides());
       // Insert new alias.
       aliasInfo.insertNewBufferAlias(subView, dstMemref);
-      b.create<CopyOp>(insertSliceOp.getLoc(), srcMemref, subView);
+      allocationFn.memCpyFn(b, insertSliceOp.getLoc(), srcMemref, subView);
     }
 
     map(bvm, insertSliceOp.result(), dstMemref);
 
     return success();
   }
 };