[mlir][linalg] Padding transformation: Write back result to original destination

[springerm]

Summary

Copy back the padded result to the original destination of the computation. This is important for bufferization, to ensure that the result of the computation does not suddenly materialize in a different buffer due to padding.

A bufferization.copy_tensor is inserted for every (unpadded) result. Such ops bufferize to memcpys, but they fold away, should the padding fold away.

Depends On: D153552

[springerm]

update

[springerm]

rebase

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

def PadOp : Op<Transform_Dialect, "structured.pad",
   let arguments =
     (ins TransformHandleTypeInterface:$target,
          DefaultValuedAttr<ArrayAttr, "{}">:$padding_values,
          DefaultValuedAttr<I64ArrayAttr, "{}">:$padding_dimensions,
          OptionalAttr<I64ArrayAttr>:$pad_to_multiple_of,
          DefaultValuedAttr<I64ArrayAttr, "{}">:$pack_paddings,
          DefaultValuedAttr<
           TypedArrayAttrBase<I64ArrayAttr, "array of arrays of i64">,
-          "{}">:$transpose_paddings);
+          "{}">:$transpose_paddings,
+         DefaultValuedAttr<BoolAttr, "true">:$copy_back);
   let results = (outs TransformHandleTypeInterface:$transformed);
 
   let assemblyFormat =
     "$target attr-dict `:` "
     "custom<SemiFunctionType>(type($target), type($transformed))";
   let hasVerifier = 1;
 
   let extraClassDeclaration = [{

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

 
 /// Pad the iterator dimensions `paddingDimensions` of all `opToPad` operands
 /// to a static bounding box. `padToMultipleOf` indicates that each padding
 /// dimension should be padded to the specified multiple. If the derived padding
 /// sizes should not be rounded up to any multiple, use "1". Use `paddingValues`
 /// and `packPaddings` to set padding value and nofold attribute of the created
 /// tensor::PadOps, respectively. Update `paddedOp` to the cloned operation with
 /// statically shaped `paddingDimensions` and return the extracted dynamically
-/// shaped results. If padding fails, return failure.
+/// shaped results. If padding fails, return failure. If `copyBack` is set, the
+/// unpadded result is copied back into the original destination tensor.
 FailureOr<SmallVector<Value>>
 rewriteAsPaddedOp(RewriterBase &rewriter, LinalgOp opToPad,
-                  const LinalgPaddingOptions &options, LinalgOp &paddedOp);
+                  const LinalgPaddingOptions &options, LinalgOp &paddedOp,
+                  bool copyBack);
 
 namespace detail {
 
 /// Helper struct to hold the results of building a packing loop nest.
 struct PackingResult {
   SmallVector<OpFoldResult> offsets, sizes, strides;
   SmallVector<Value> clonedLoopIvs, leadingPackedTensorIndexings;
   GenericOp maybeTransposeOp;

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

transform::PadOp::applyToOne(transform::TransformRewriter &rewriter,
   options.paddingDimensions = extractFromI64ArrayAttr(getPaddingDimensions());
   SmallVector<int64_t> padToMultipleOf(options.paddingDimensions.size(), 1);
   if (getPadToMultipleOf().has_value())
     padToMultipleOf = extractFromI64ArrayAttr(*getPadToMultipleOf());
   options.padToMultipleOf = padToMultipleOf;
   options.paddingValues = paddingValues;
   options.packPaddings = packPaddings;
   FailureOr<SmallVector<Value>> result =
-      rewriteAsPaddedOp(rewriter, target, options, paddedOp);
+      rewriteAsPaddedOp(rewriter, target, options, paddedOp, getCopyBack());
   if (succeeded(result)) {
     // We need to perform our own replacement here because this API is still
     // used in patterns that "pad and hoist", for which the replacement values
     // need to be different.
     // TODO: clean this up and stop "pad and hoist" behavior more globally now
     // that we have more composable abstractions.
     rewriter.replaceOp(target, *result);
     results.push_back(paddedOp);

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

 //===- Padding.cpp - Padding of Linalg ops --------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
 
 #define DEBUG_TYPE "linalg-padding"
 
 using namespace mlir;
 using namespace mlir::linalg;
LLVM_DEBUG(DBGS() << "--SUCCESS, makeComposedPadHighOp with type: "
                     << paddedTensorType);
   return makeComposedPadHighOp(rewriter, opToPad->getLoc(), paddedTensorType,
                                opOperand->get(), paddingValue, nofold);
 }
 
 FailureOr<SmallVector<Value>>
 linalg::rewriteAsPaddedOp(RewriterBase &rewriter, LinalgOp opToPad,
                           const LinalgPaddingOptions &options,
-                          LinalgOp &paddedOp) {
+                          LinalgOp &paddedOp, bool copyBack) {
   LLVM_DEBUG(DBGS() << "Start rewriteAsPaddedOp : " << opToPad << "\n");
   Location loc = opToPad->getLoc();
 
   // TODO: there are cases where we may still want to pad to larger sizes.
   if (!opToPad.hasTensorSemantics())
     return rewriter.notifyMatchFailure(opToPad,
                                        "expected operation on tensors");
 
for (const auto &en : llvm::enumerate(paddedOp->getResults())) {
     int64_t resultNumber = en.index();
     int64_t rank = cast<RankedTensorType>(paddedResult.getType()).getRank();
     SmallVector<OpFoldResult> offsets(rank, rewriter.getIndexAttr(0));
     SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));
     paddedSubtensorResults.push_back(rewriter.create<tensor::ExtractSliceOp>(
         loc, paddedResult, offsets, reifiedResultShapes[resultNumber],
         strides));
   }
+
+  if (!copyBack)
     return paddedSubtensorResults;
+
+  // Copy back unpadded results to the original destination (i.e., inits of the
+  // linalg op), so that the destination buffer of the computation does not
+  // change. If the padding folds away, this will materizalize as a memcpy
+  // between two identical buffers, which will then also fold away.
+  SmallVector<Value> copiedBack;
+  for (auto it :
+       llvm::zip(paddedSubtensorResults, opToPad.getDpsInitOperands())) {
+    copiedBack.push_back(rewriter.create<bufferization::CopyTensorOp>(
+        loc, std::get<0>(it), std::get<1>(it)->get()));
+  }
+  return copiedBack;
 }
 
 FailureOr<LinalgOp>
 mlir::linalg::padAndHoistLinalgOp(RewriterBase &rewriter, LinalgOp linalgOp,
                                   const LinalgPaddingOptions &options) {
   if (!linalgOp.hasTensorSemantics())
     return rewriter.notifyMatchFailure(
         linalgOp, "only applies to Linalg ops with tensor semantics");
 
   // Pad the operation.
   LinalgOp paddedOp;
-  FailureOr<SmallVector<Value>> newResults =
-      rewriteAsPaddedOp(rewriter, linalgOp, options, paddedOp);
+  FailureOr<SmallVector<Value>> newResults = rewriteAsPaddedOp(
+      rewriter, linalgOp, options, paddedOp, /*copyBack=*/false);
   if (failed(newResults))
     return rewriter.notifyMatchFailure(linalgOp,
                                        "failed to rewrite as a padded op");
 
   // Hoist the padding.
   for (const auto &en : enumerate(options.hoistPaddings)) {
     if (static_cast<int64_t>(en.index()) >= paddedOp->getNumOperands())
       break;

mlir/test/Dialect/Linalg/transform-op-hoist-pad.mlir

^bb1(%arg1: !transform.any_op):
   %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1
     : (!transform.any_op) -> !transform.any_op
 
 
   %matmul_l1, %loops_l1 = transform.structured.tile_to_scf_for %matmul [5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 
   %matmul_padded = transform.structured.pad %matmul_l1 {
     padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
-    padding_dimensions=[0, 1, 2]
+    padding_dimensions=[0, 1, 2],
+    copy_back = false
   } : (!transform.any_op) -> !transform.any_op
 
   // In this case, the pad op is actually empty: we only tile the first dimension
   // and it does not have an impact on the RHS operand.
   // expected-error @below {{incompatible payload operation name}}
   %pad = transform.get_producer_of_operand %matmul_padded[1]
     : (!transform.any_op) -> !transform.op<"tensor.pad">
 
^bb1(%arg1: !transform.any_op):
   %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1
     : (!transform.any_op) -> !transform.any_op
 
 
   %matmul_l1, %loops_l1 = transform.structured.tile_to_scf_for %matmul [5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 
   %matmul_padded = transform.structured.pad %matmul_l1 {
     padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
-    padding_dimensions=[0, 1, 2]
+    padding_dimensions=[0, 1, 2],
+    copy_back = false
   } : (!transform.any_op) -> !transform.any_op
 
   %pad = transform.get_producer_of_operand %matmul_padded[2]
     : (!transform.any_op) -> !transform.op<"tensor.pad">
 
   // We do not know yet how to hoist the init.
   // expected-error @below {{transform.structured.hoist_pad failed to apply}}
   transform.structured.hoist_pad %pad by 1 loops
^bb1(%arg1: !transform.any_op):
   %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1
     : (!transform.any_op) -> !transform.any_op
 
 
   %matmul_l1, %loops_l1 = transform.structured.tile_to_scf_for %matmul [5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 
   %matmul_padded = transform.structured.pad %matmul_l1 {
     padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
-    padding_dimensions=[0, 1, 2]
+    padding_dimensions=[0, 1, 2],
+    copy_back = false
   } : (!transform.any_op) -> !transform.any_op
 
   %pad = transform.get_producer_of_operand %matmul_padded[0]
     : (!transform.any_op) -> !transform.any_op
 
   transform.structured.hoist_pad %pad by 1 loops
      : (!transform.any_op) -> !transform.any_op
 }
^bb1(%arg1: !transform.any_op):
   %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1
     : (!transform.any_op) -> !transform.any_op
 
 
   %matmul_l1, %loops_l1 = transform.structured.tile_to_scf_for %matmul [5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 
   %matmul_padded = transform.structured.pad %matmul_l1 {
     padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
-    padding_dimensions=[0, 1, 2]
+    padding_dimensions=[0, 1, 2],
+    copy_back = false
   } : (!transform.any_op) -> !transform.any_op
 
   %pad = transform.get_producer_of_operand %matmul_padded[0]
     : (!transform.any_op) -> !transform.any_op
 
   transform.structured.hoist_pad %pad by 1 loops, transpose by [1, 0]
      : (!transform.any_op) -> !transform.any_op
 }
^bb1(%arg1: !transform.any_op):
   %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1
     : (!transform.any_op) -> !transform.any_op
 
 
   %matmul_l1, %loops_l1:2 = transform.structured.tile_to_scf_for %matmul [5, 0, 7] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
 
   %matmul_padded = transform.structured.pad %matmul_l1 {
     padding_values=[0.0: f32, 0.0 : f32, 0.0 : f32],
-    padding_dimensions=[0, 1, 2]
+    padding_dimensions=[0, 1, 2],
+    copy_back = false
   } : (!transform.any_op) -> !transform.any_op
 
   %pad = transform.get_producer_of_operand %matmul_padded[2]
     : (!transform.any_op) -> !transform.op<"tensor.pad">
 
   transform.structured.hoist_pad %pad by 1 loops
      : (!transform.op<"tensor.pad">) -> !transform.any_op
 }

mlir/test/Dialect/Linalg/transform-op-pad.mlir

 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %1 = transform.structured.pad %0 {
     padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     pack_paddings=[1, 1, 1]
   } : (!transform.any_op) -> !transform.any_op
 }
+
+// -----
+
+#map = affine_map<()[s0] -> (-s0 + 12, 7)>
+
+// CHECK-LABEL: @pack_everything
+func.func @pack_everything(%arg0: tensor<24x12xf32>,
+                           %arg1: tensor<12x25xf32>,
+                           %arg2: tensor<24x25xf32>,
+                           %iv0 : index, %iv1 : index, %iv2 : index) -> tensor<24x25xf32> {
+  %0 = affine.min #map()[%iv2]
+
+  //      CHECK: %[[T0:.*]] = tensor.extract_slice %
+  //      CHECK: %[[T1:.*]] = tensor.extract_slice %
+  //      CHECK: %[[T2:.*]] = tensor.extract_slice %
+  %1 = tensor.extract_slice %arg0[%iv0, %iv2] [4, %0] [1, 1] : tensor<24x12xf32> to tensor<4x?xf32>
+  %2 = tensor.extract_slice %arg1[%iv2, %iv1] [%0, 5] [1, 1] : tensor<12x25xf32> to tensor<?x5xf32>
+  %3 = tensor.extract_slice %arg2[%iv0, %iv1] [4, 5] [1, 1] : tensor<24x25xf32> to tensor<4x5xf32>
+
+  //  CHECK-DAG: %[[CST:.*]] = arith.constant 0.
+  //  CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+
+  //      CHECK: %[[PAD0:.*]] = tensor.pad %[[T0]] nofold
+  //      CHECK: %[[PAD1:.*]] = tensor.pad %[[T1]] nofold
+  //      CHECK: %[[PAD2:.*]] = tensor.pad %[[T2]] nofold
+
+  //      CHECK: %[[T5:.*]] = linalg.matmul
+  // CHECK-SAME:              ins(%[[PAD0]], %[[PAD1]] : tensor<4x7xf32>, tensor<7x5xf32>)
+  // CHECK-SAME:              outs(%[[PAD2]] : tensor<4x5xf32>)
+
+  // Get unpadded result (no-op in this example).
+  //      CHECK: %[[T6:.*]] = tensor.extract_slice %[[T5]]
+  // Copy back result to the original buffer, so that the destination of the
+  // computation does not change.
+  //      CHECK: %[[T7:.*]] = bufferization.copy_tensor %[[T6]], %[[T2]]
+  %4 = linalg.matmul ins(%1, %2 : tensor<4x?xf32>, tensor<?x5xf32>) outs(%3 : tensor<4x5xf32>) -> tensor<4x5xf32>
+
+  //      CHECK: %[[T8:.*]] = tensor.insert_slice %[[T7]] into %{{.*}}
+  %5 = tensor.insert_slice %4 into %arg2[%iv0, %iv1] [4, 5] [1, 1] : tensor<4x5xf32> into tensor<24x25xf32>
+  func.return %5 : tensor<24x25xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %1 = transform.structured.pad %0 {
+    padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32],
+    padding_dimensions=[0, 1, 2],
+    pack_paddings=[1, 1, 1]
+  } : (!transform.any_op) -> !transform.any_op
+}