[mlir][linalg] BufferizeToAllocation: Add custom memcpy op

[springerm]

Summary

Add a new option that allows users to specify a memcpy op: "memref.tensor_store", "memref.copy" or "linalg.copy".

Depends On: D154967

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

let description = [{
     - `tensor.pad` is lowered to an allocation, followed by a `linalg.fill` and
       and a buffer copy (all on memrefs).
     - `vector.mask` is bufferized together with its region. The allocation is
       placed in front of the `vector.mask` op.
 
     An optional memory space attribute can be specified for the materialized
     buffer allocation.
 
+    If a memory copy is needed, a "memref.tensor_store" is used when possible.
+    This is an op with tensor semantics that will bufferize to a memory copy
+    later. Which concrete op will be used for the memory copy is up to the
+    bufferization framework. Alternatively, a custom memcpy op can be specified
+    via `memcpy_op`. Currently supported are "memref.copy" and "linalg.copy".
+    In that case, the source of each memcpy must not have a custom memory space.
+    Furthermore, because the future buffer layout unknown for a given tensor,
+    a fully dynamic layout is assumed for best compatibility. Users should use
+    "memref.tensor_store" when possible.
+
     #### Return modes
 
     This operation consumes the `target` handle and produces the
     `allocated_buffer` and `new_ops` handles. It always succeeds.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target,
-                       OptionalAttr<AnyAttr>:$memory_space);
+                       OptionalAttr<AnyAttr>:$memory_space,
+                       DefaultValuedAttr<StrAttr, "\"memref.tensor_store\"">:
+                           $memcpy_op);
+  let hasVerifier = 1;
   let results = (outs Transform_AnyValue:$allocated_buffer,
                       Transform_AnyOpType:$new_ops);
   let assemblyFormat = "$target attr-dict `:` type($target)";
   
   let builders = [
     OpBuilder<(ins "Value":$target, "Attribute":$memorySpace)>,
     OpBuilder<(ins "Value":$target, "int64_t":$memorySpace)>
   ];

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

 
 /// Return vector::CombiningKind for the given op.
 std::optional<vector::CombiningKind> getCombinerOpKind(Operation *combinerOp);
 
 //===----------------------------------------------------------------------===//
 // Bufferization-related transforms.
 //===----------------------------------------------------------------------===//
 
+struct BufferizeToAllocationOptions {
+  enum class MemcpyOp { MemrefTensorStore = 0, MemrefCopy = 1, LinalgCopy = 2 };
+
+  MemcpyOp memcpyOp = MemcpyOp::MemrefTensorStore;
+};
+
 /// Materialize a buffer allocation for the given tensor.pad op and lower the
 /// op to linalg.fill/linalg.generic + memref.tensor_store. E.g.:
 ///
 /// %0 = tensor.pad low[%l] high[%h] %t ...
 ///
 /// is lowered to:
 ///
 /// %alloc = memref.alloc
 /// linalg.fill ... outs(%alloc)
 /// %subview = memref.subview %alloc [%l] [...] [1]
 /// memref.tensor_store %t, %subview
 /// %0 = bufferization.to_tensor %alloc restrict writable
 ///
 /// In addition to rewriting the IR as shown above, this function returns the
 /// newly allocated buffer. The `insertionPoint` parameter can be used to
 /// specify a custom insertion point for the buffer allocation.
-Value bufferizeToAllocation(RewriterBase &rewriter, tensor::PadOp padOp,
-                            Attribute memorySpace = {},
+Value bufferizeToAllocation(RewriterBase &rewriter,
+                            const BufferizeToAllocationOptions &options,
+                            tensor::PadOp padOp, Attribute memorySpace = {},
                             Operation *insertionPoint = nullptr);
 
 /// Materialize a buffer allocation for the given vector.mask op and bufferize
 /// the op, including its region. E.g.:
 ///
 /// %0 = vector.mask {
 ///   vector.transfer_write %v, %t : vector<16xf32>, tensor<?xf32>
 /// } : vector<16xi1> -> tensor<?xf32>
 ///
 /// is lowered to:
 ///
 /// %alloc = memref.alloc
 /// memref.tensor_store %t, %subview
 /// vector.mask {
 ///   vector.transfer_write %arg0, %alloc : vector<16xf32>, memref<?xf32>
 /// } : vector<16xi1>
 /// %0 = bufferization.to_tensor %alloc restrict writable
 ///
 /// In addition to rewriting the IR as shown above, this function returns the
 /// newly allocated buffer. The `insertionPoint` parameter can be used to
 /// specify a custom insertion point for the buffer allocation.
-Value bufferizeToAllocation(RewriterBase &rewriter, vector::MaskOp maskOp,
-                            Attribute memorySpace = {},
+Value bufferizeToAllocation(RewriterBase &rewriter,
+                            const BufferizeToAllocationOptions &options,
+                            vector::MaskOp maskOp, Attribute memorySpace = {},
                             Operation *insertionPoint = nullptr);
 
 /// Bufferize the given op with tensor semantics and materialize the result in
 /// a newly allocated buffer.
 ///
 /// Only bufferizable ops that bufferize to a memory write or have an
 /// aliasing OpOperand (and do not themselves bufferize to an allocation) are
 /// supported. They are bufferized using their BufferizableOpInterface
 /// implementation.
 ///
 /// Selected ops that bufferize to an allocation (or need special handling) are
 /// also supported:
 /// - tensor.pad
 /// - vector.mask
 ///
 /// This function returns the newly allocated buffer. The `insertionPoint`
 /// parameter can be used to specify a custom insertion point for the buffer
 /// allocation.
-Value bufferizeToAllocation(RewriterBase &rewriter, Operation *op,
-                            Attribute memorySpace = {},
+Value bufferizeToAllocation(RewriterBase &rewriter,
+                            const BufferizeToAllocationOptions &options,
+                            Operation *op, Attribute memorySpace = {},
                             Operation *insertionPoint = nullptr);
 
 /// Try to eliminate tensor::EmptyOps inside `op` that are anchored on a
 /// LinalgOp. This transforms looks for LinalgOps that have an unused output
 /// operand and an input operand that is rooted in a tensor::EmptyOp. The
 /// tensor::EmptyOp uses are replaced with the output operand and the two
 /// operands of the LinalgOp are swapped.
 ///

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

DiagnosedSilenceableFailure transform::BufferizeToAllocationOp::apply(
     transform::TransformResults &results, transform::TransformState &state) {
   // Attach listener to keep track of newly created ops.
   OpBuilder::Listener *previousListener = rewriter.getListener();
   auto resetListener =
       llvm::make_scope_exit([&]() { rewriter.setListener(previousListener); });
   NewOpsListener newOpsListener(previousListener);
   rewriter.setListener(&newOpsListener);
 
+  linalg::BufferizeToAllocationOptions options;
+  if (getMemcpyOp() == "memref.tensor_store") {
+    options.memcpyOp =
+        linalg::BufferizeToAllocationOptions::MemcpyOp::MemrefTensorStore;
+  } else if (getMemcpyOp() == "memref.copy") {
+    options.memcpyOp =
+        linalg::BufferizeToAllocationOptions::MemcpyOp::MemrefCopy;
+  } else if (getMemcpyOp() == "linalg.copy") {
+    options.memcpyOp =
+        linalg::BufferizeToAllocationOptions::MemcpyOp::LinalgCopy;
+  } else {
+    llvm_unreachable("invalid memcpy op");
+  }
+
   // Bufferize ops.
   Attribute memorySpace =
       getMemorySpace().has_value() ? getMemorySpace().value() : Attribute();
   SmallVector<Value> allocatedBuffers;
   for (Operation *op : state.getPayloadOps(getTarget())) {
-    Value buffer = linalg::bufferizeToAllocation(rewriter, op, memorySpace);
+    Value buffer =
+        linalg::bufferizeToAllocation(rewriter, options, op, memorySpace);
     if (!buffer) {
       DiagnosedSilenceableFailure diag = emitSilenceableError()
                                          << "failed to bufferize operation";
       diag.attachNote(op->getLoc()) << "target payload op";
       return diag;
     }
     allocatedBuffers.push_back(buffer);
   }
 
   // Set results.
   results.setValues(cast<OpResult>(getAllocatedBuffer()), allocatedBuffers);
   results.set(cast<OpResult>(getNewOps()), newOpsListener.getNewOps());
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::BufferizeToAllocationOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   consumesHandle(getTarget(), effects);
   producesHandle(getAllocatedBuffer(), effects);
   producesHandle(getNewOps(), effects);
   modifiesPayload(effects);
 }
 
+LogicalResult transform::BufferizeToAllocationOp::verify() {
+  if (getMemcpyOp() != "memref.tensor_store" &&
+      getMemcpyOp() != "memref.copy" && getMemcpyOp() != "linalg.copy")
+    return emitOpError() << "unsupported memcpy op";
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // DecomposeOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::DecomposeOp::applyToOne(transform::TransformRewriter &rewriter,
                                    LinalgOp target,
                                    transform::ApplyToEachResultList &results,

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

static Value createInserts(RewriterBase &rewriter, Location loc, int dim,
   for (int i = 0; i < shape[dim]; ++i) {
     indices[dim] = constants[i];
     destination = createInserts(rewriter, loc, dim + 1, destination, shape,
                                 constants, elementIt, indices);
   }
   return destination;
 }
 
+/// Create a memcpy from the given source tensor to the given destination
+/// memref. The copy op type can be specified in the `options`.
+static void createMemcpy(OpBuilder &b, Location loc, Value tensorSource,
+                         Value memrefDest,
+                         const linalg::BufferizeToAllocationOptions &options) {
+  auto tensorType = dyn_cast<RankedTensorType>(tensorSource.getType());
+  assert(tensorType && "expected ranked tensor");
+  assert(memrefDest.getType().isa<MemRefType>() && "expected ranked memref");
+
+  switch (options.memcpyOp) {
+  case linalg::BufferizeToAllocationOptions::MemcpyOp::MemrefTensorStore:
+    // Note: This is the preferred way of memcpy'ing because no layout map
+    // and/or memory space must be specified for the source.
+    b.create<memref::TensorStoreOp>(loc, tensorSource, memrefDest);
+    break;
+  case linalg::BufferizeToAllocationOptions::MemcpyOp::MemrefCopy: {
+    // TODO: Support custom memory space on source.
+    // We do not know the layout map of the source yet, so use a fully dynamic
+    // layout for best compatibility.
+    Value toMemref = b.create<bufferization::ToMemrefOp>(
+        loc, bufferization::getMemRefTypeWithFullyDynamicLayout(tensorType),
+        tensorSource, /*readOnly=*/true);
+    b.create<memref::CopyOp>(loc, toMemref, memrefDest);
+  } break;
+  case linalg::BufferizeToAllocationOptions::MemcpyOp::LinalgCopy: {
+    // TODO: Support custom memory space on source.
+    // We do not know the layout map of the source yet, so use a fully dynamic
+    // layout for best compatibility.
+    Value toMemref = b.create<bufferization::ToMemrefOp>(
+        loc, bufferization::getMemRefTypeWithFullyDynamicLayout(tensorType),
+        tensorSource, /*readOnly=*/true);
+    b.create<linalg::CopyOp>(loc, toMemref, memrefDest);
+  } break;
+  };
+}
+
 static Operation *movePaddingToFillOrGenericOp(RewriterBase &rewriter,
                                                Location loc, PadOp padOp,
                                                Value dest) {
   OpBuilder::InsertionGuard g(rewriter);
   RankedTensorType resultType = padOp.getResultType();
 
   // Examine the yielded value to decide if a linalg.generic is neede or a
   // linalg.fill is sufficient.
static Value createAllocationForTensor(RewriterBase &rewriter, Location loc,
 
   // Place deallocation at the end of the block.
   rewriter.setInsertionPoint(rewriter.getInsertionBlock()->getTerminator());
   rewriter.create<memref::DeallocOp>(loc, alloc);
 
   return alloc;
 }
 
-Value linalg::bufferizeToAllocation(RewriterBase &rewriter, PadOp padOp,
-                                    Attribute memorySpace,
-                                    Operation *insertionPoint) {
+Value linalg::bufferizeToAllocation(
+    RewriterBase &rewriter, const linalg::BufferizeToAllocationOptions &options,
+    PadOp padOp, Attribute memorySpace, Operation *insertionPoint) {
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPoint(insertionPoint ? insertionPoint : padOp);
   Location loc = padOp.getLoc();
 
   // Create buffer allocation.
   Value alloc =
       createAllocationForTensor(rewriter, loc, padOp.getResult(), memorySpace);
   rewriter.setInsertionPoint(padOp);
 
   if (!padOp.hasZeroLowPad() || !padOp.hasZeroHighPad()) {
     // Create linalg.fill or linalg.generic. Not needed if there is no padding.
     Operation *fillOp =
         movePaddingToFillOrGenericOp(rewriter, loc, padOp, alloc);
     rewriter.setInsertionPointAfter(fillOp);
   }
 
   // Create memref.tensor_store.
   SmallVector<OpFoldResult> sizes =
       getMixedSizes(rewriter, loc, padOp.getSource());
   SmallVector<OpFoldResult> strides(padOp.getResultType().getRank(),
                                     rewriter.getIndexAttr(1));
   Value subview = rewriter.create<memref::SubViewOp>(
       loc, alloc, /*offsets=*/padOp.getMixedLowPad(), sizes, strides);
-  rewriter.create<memref::TensorStoreOp>(loc, padOp.getSource(), subview);
+  createMemcpy(rewriter, loc, padOp.getSource(), subview, options);
 
   // Create bufferization.to_tensor with "restrict" and "writable". The returned
   // tensor is a new buffer allocation, so it does not alias with any buffer.
   Value toTensorOp = rewriter.create<bufferization::ToTensorOp>(
       loc, alloc, /*restrict=*/true, /*writable=*/true);
   rewriter.replaceOp(padOp, toTensorOp);
   return alloc;
 }
 
-Value linalg::bufferizeToAllocation(RewriterBase &rewriter,
-                                    vector::MaskOp maskOp,
-                                    Attribute memorySpace,
-                                    Operation *insertionPoint) {
+Value linalg::bufferizeToAllocation(
+    RewriterBase &rewriter, const linalg::BufferizeToAllocationOptions &options,
+    vector::MaskOp maskOp, Attribute memorySpace, Operation *insertionPoint) {
   assert(llvm::range_size(maskOp.getMaskBlock()->without_terminator()) == 1 &&
          "expected single masked op");
   OpBuilder::InsertionGuard g(rewriter);
-  bufferization::BufferizationOptions options;
+  bufferization::BufferizationOptions bufferizationOptions;
   Operation *yieldOp = maskOp.getMaskRegion().front().getTerminator();
   assert(isa<vector::YieldOp>(yieldOp) && "expected yield op terminator");
 
   // Bufferize maskable op. By default, place the buffer allocation right before
   // the mask op.
   Value alloc = bufferizeToAllocation(
-      rewriter, maskOp.getMaskableOp(), memorySpace,
+      rewriter, options, maskOp.getMaskableOp(), memorySpace,
       /*insertionPoint=*/insertionPoint ? insertionPoint : maskOp);
 
   // Bufferize terminator.
   rewriter.setInsertionPoint(yieldOp);
   if (failed(cast<bufferization::BufferizableOpInterface>(yieldOp).bufferize(
-          rewriter, options)))
+          rewriter, bufferizationOptions)))
     return nullptr;
 
   // Erase dead to_tensor ops inside of the mask op. This is necessary because
   // there only be one op (apart from the terminator) inside the mask op.
   // TODO: Remove dead to_tensor ops more aggressively during bufferization.
   SmallVector<Operation *> toTensorOps;
   maskOp.walk([&](bufferization::ToTensorOp toTensorOp) {
     if (toTensorOp->getUses().empty())
       toTensorOps.push_back(toTensorOp.getOperation());
   });
   for (Operation *op : toTensorOps)
     rewriter.eraseOp(op);
 
   // Bufferize mask op.
   SmallVector<OpOperand *> resultUses;
   for (Value result : maskOp.getResults())
     if (isa<TensorType>(result.getType()))
       for (OpOperand &use : result.getUses())
         resultUses.push_back(&use);
   rewriter.setInsertionPoint(maskOp);
   if (failed(cast<bufferization::BufferizableOpInterface>(maskOp.getOperation())
-                 .bufferize(rewriter, options)))
+                 .bufferize(rewriter, bufferizationOptions)))
     return nullptr;
 
   // Set "restrict" attribute, indicating that no other tensor aliases with
   // this tensor. That is because we just allocated a new buffer for the tensor.
   for (OpOperand *resultUse : resultUses) {
     auto toTensorOp =
         resultUse->get().getDefiningOp<bufferization::ToTensorOp>();
     assert(toTensorOp && "expected to_tensor op");
mlir::linalg::rewriteInDestinationPassingStyle(RewriterBase &rewriter,
   SmallVector<OpFoldResult> sliceStrides(resultType.getRank(),
                                          rewriter.getIndexAttr(1));
   auto insertSliceOp = rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
       padOp, padOp.getSource(), fillOp->getResult(0),
       /*offsets=*/padOp.getMixedLowPad(), sliceSizes, sliceStrides);
   return insertSliceOp.getOperation();
 }
 
-Value linalg::bufferizeToAllocation(RewriterBase &rewriter, Operation *op,
-                                    Attribute memorySpace,
-                                    Operation *insertionPoint) {
+Value linalg::bufferizeToAllocation(
+    RewriterBase &rewriter, const linalg::BufferizeToAllocationOptions &options,
+    Operation *op, Attribute memorySpace, Operation *insertionPoint) {
   using namespace bufferization;
 
   // Call specialized overload for certain ops.
   if (auto padOp = dyn_cast<tensor::PadOp>(op))
-    return bufferizeToAllocation(rewriter, padOp, memorySpace);
+    return bufferizeToAllocation(rewriter, options, padOp, memorySpace);
   if (auto maskOp = dyn_cast<vector::MaskOp>(op))
-    return bufferizeToAllocation(rewriter, maskOp, memorySpace);
+    return bufferizeToAllocation(rewriter, options, maskOp, memorySpace);
 
   // Only bufferizable ops are supported.
   auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op);
   if (!bufferizableOp)
     return nullptr;
-  BufferizationOptions options;
-  AnalysisState state(options);
+  BufferizationOptions bufferizationOptions;
+  AnalysisState state(bufferizationOptions);
 
   // Gather tensor results.
   SmallVector<OpResult> tensorResults;
   for (OpResult result : op->getResults()) {
     if (!result.getType().isa<TensorType>())
       continue;
     // Unranked tensors are not supported
     if (!isa<RankedTensorType>(result.getType()))
Value linalg::bufferizeToAllocation(
   SmallVector<Value> allocs;
   for (OpOperand *operand : outOfPlaceOperands) {
     Value alloc = createAllocationForTensor(rewriter, op->getLoc(),
                                             operand->get(), memorySpace);
     allocs.push_back(alloc);
     if (!state.findDefinitions(operand->get()).empty()) {
       // Initialize buffer with a copy of the operand data. Not needed if the
       // tensor is uninitialized.
-      rewriter.create<memref::TensorStoreOp>(op->getLoc(), operand->get(),
-                                             alloc);
+      createMemcpy(rewriter, op->getLoc(), operand->get(), alloc, options);
     }
     rewriter.updateRootInPlace(op, [&]() {
       operand->set(rewriter.create<ToTensorOp>(op->getLoc(), alloc));
     });
   }
 
   // Bufferize the op.
   rewriter.setInsertionPoint(op);
-  if (failed(bufferizableOp.bufferize(rewriter, options)))
+  if (failed(bufferizableOp.bufferize(rewriter, bufferizationOptions)))
     return nullptr;
 
   // Set "restrict" attribute, indicating that no other tensor aliases with
   // this tensor. That is because we just allocated a new buffer for the tensor.
   for (OpOperand *resultUse : resultUses) {
     auto toTensorOp = resultUse->get().getDefiningOp<ToTensorOp>();
     assert(toTensorOp && "expected to_tensor op");
     rewriter.updateRootInPlace(toTensorOp, [&]() {

mlir/test/Dialect/Linalg/transform-op-bufferize-to-allocation.mlir

^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %2, %new = transform.structured.bufferize_to_allocation %0 : !transform.any_op
 
   // Ensure that one linalg.fill was generated.
   %fill_op = transform.select "linalg.fill" in %new : (!transform.any_op) -> !transform.any_op
   // expected-remark @below{{1}}
   test_print_number_of_associated_payload_ir_ops %fill_op : !transform.any_op
 
-  // Ensure that one memref.tensor_store was generated.
+  // Ensure that one linalg.copy was generated.
   %tensor_store = transform.select "memref.tensor_store" in %new : (!transform.any_op) -> !transform.any_op
   // expected-remark @below{{1}}
   test_print_number_of_associated_payload_ir_ops %tensor_store : !transform.any_op
 }
 
 // -----
 
+// CHECK-LABEL: func @tensor_pad_constant_with_custom_copy(
+//   CHECK-NOT:   memref.tensor_store
+//   CHECK-NOT:   memref.copy
+//       CHECK:   linalg.copy
+func.func @tensor_pad_constant_with_custom_copy(
+    %t: tensor<?x10xindex>, %l2: index, %h1: index, %h2: index)
+        -> tensor<?x?xindex>
+{
+  %0 = tensor.pad %t low[5, %l2] high[%h1, %h2] {
+  ^bb0(%arg0: index, %arg1: index):
+    %c = arith.constant 50 : index
+    tensor.yield %c : index
+  } : tensor<?x10xindex> to tensor<?x?xindex>
+  return %0 : tensor<?x?xindex>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 3, memcpy_op = "linalg.copy"}: !transform.any_op
+
+  // Ensure that one linalg.fill was generated.
+  %fill_op = transform.select "linalg.fill" in %new : (!transform.any_op) -> !transform.any_op
+  // expected-remark @below{{1}}
+  test_print_number_of_associated_payload_ir_ops %fill_op : !transform.any_op
+
+  // Ensure that one linalg.copy was generated.
+  %linalg_copy = transform.select "linalg.copy" in %new : (!transform.any_op) -> !transform.any_op
+  // expected-remark @below{{1}}
+  test_print_number_of_associated_payload_ir_ops %linalg_copy : !transform.any_op
+
+  // Make sure that One-Shot Bufferize can bufferize the rest.
+  %4 = transform.bufferization.one_shot_bufferize %arg1 : (!transform.any_op) -> !transform.any_op
+}
+
+// -----
+
 // CHECK-LABEL: func @tensor_pad_constant(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?x10xindex>
 //       CHECK:   %[[src:.*]] = bufferization.to_memref %[[t]]
 //       CHECK:   %[[alloc:.*]] = memref.alloc
 //       CHECK:   %[[subview:.*]] = memref.subview %[[alloc]]
 //       CHECK:   memref.copy %[[src]], %[[subview]]
 //       CHECK:   bufferization.to_tensor %[[alloc]] restrict writable
 func.func @tensor_pad_constant(%t: tensor<?x10xindex>, %l2: index, %h1: index,