[mlir][Linalg] Add a structured transform to materialize a tensor.insert_slice via a linalg.copy

[nicolasvasilache]

Summary

This is useful to materialize copies explicitly before bufferization and
transform them, avoiding the need to rediscover them after bufferization.

[ftynse]

mlir/test/Dialect/Linalg/transform-op-insert-slice-to-copy.mlir

29–33

Nit: can we use transform.any_op instead of pdl here?

65–69

Ditto

[nicolasvasilache]

sed

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

 class RewriterBase;
 
 namespace linalg {
 class GenericOp;
 class LinalgOp;
 } // namespace linalg
 
 namespace tensor {
+class InsertSliceOp;
 class PackOp;
 class PadOp;
 class UnPackOp;
 } // namespace tensor
 
 namespace transform {
 class TransformHandleTypeInterface;
 // Types needed for builders.

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

def HoistRedundantTensorSubsetsOp :
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::Operation *target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// InsertSliceToCopyOp
+//===----------------------------------------------------------------------===//
+
+def InsertSliceToCopyOp :
+  Op<Transform_Dialect, "structured.insert_slice_to_copy",
+    [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
+     TransformEachOpTrait, TransformOpInterface]> {
+  let description = [{
+    Targeted rewrite of an tensor.insert_slice to linalg.copy.
+    This is useful to materialize copies explicitly before bufferization and 
+    transform them, avoiding the need to rediscover them after bufferization.
+
+    If the insert_slice source is already a linalg.copy, only return the source
+    op (i.e. do not create an additional linalg.copy op).
+
+    #### Return modes:
+
+    The operation always succeeds and returns a handle to the relevant 
+    linalg.copy op.
+  }];
+
+  let arguments = (ins TransformHandleTypeInterface:$target);
+  let results = (outs TransformHandleTypeInterface:$transformed);
+
+  let assemblyFormat = "$target attr-dict `:` functional-type(operands, results) ";
+
+  let builders = [
+    OpBuilder<(ins "Value":$target)>,
+  ];
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::tensor::InsertSliceOp target,
+        ::mlir::transform::ApplyToEachResultList &results,
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
 #endif // LINALG_TRANSFORM_OPS

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

transform::HoistRedundantTensorSubsetsOp::applyToOne(
   target->walk([&](scf::ForOp forOp) {
     hoistRedundantSubsetExtractInsert(rewriter, forOp);
   });
   results.push_back(target);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
+// InsertSliceToCopyOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure transform::InsertSliceToCopyOp::applyToOne(
+    tensor::InsertSliceOp target, transform::ApplyToEachResultList &results,
+    transform::TransformState &state) {
+  if (auto copySource = target.getSource().getDefiningOp<linalg::CopyOp>()) {
+    results.push_back(copySource);
+    return DiagnosedSilenceableFailure::success();
+  }
+
+  TrackingListener listener(state, *this);
+  IRRewriter rewriter(target->getContext(), &listener);
+  rewriter.setInsertionPoint(target);
+  Value extracted = rewriter.create<tensor::ExtractSliceOp>(
+      target.getLoc(), target.getDest(), target.getMixedOffsets(),
+      target.getMixedSizes(), target.getMixedStrides());
+  Value copied = rewriter
+                     .create<linalg::CopyOp>(target.getLoc(),
+                                             target.getSource(), extracted)
+                     .getResult(0);
+  rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
+      target, copied, target.getDest(), target.getMixedOffsets(),
+      target.getMixedSizes(), target.getMixedStrides());
+
+  results.push_back(copied.getDefiningOp());
+  return DiagnosedSilenceableFailure::success();
+}
+
+//===----------------------------------------------------------------------===//
 // Transform op registration
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// Registers new ops and declares PDL as dependent dialect since the
 /// additional ops are using PDL types for operands and results.
 class LinalgTransformDialectExtension
     : public transform::TransformDialectExtension<

mlir/test/Dialect/Linalg/transform-op-insert-slice-to-copy.mlir

+// RUN: mlir-opt -test-transform-dialect-interpreter %s --split-input-file | FileCheck %s
+
+// CHECK-LABEL: func @insert_slice_to_copy
+    // CHECK-SAME: %[[I:.*]]: tensor<2x3xf32>
+    // CHECK-SAME: %[[O:.*]]: tensor<?x?xf32>, 
+    // CHECK-SAME: %[[OFF0:[0-9a-zA-Z]+]]: index,
+    // CHECK-SAME: %[[OFF1:[0-9a-zA-Z]+]]: index,
+    // CHECK-SAME: %[[SZ0:[0-9a-zA-Z]+]]: index,
+    // CHECK-SAME: %[[SZ1:[0-9a-zA-Z]+]]: index,
+    // CHECK-SAME: %[[ST0:[0-9a-zA-Z]+]]: index,
+    // CHECK-SAME: %[[ST1:[0-9a-zA-Z]+]]: index)
+func.func @insert_slice_to_copy(
+    %I : tensor<2x3xf32>, %O : tensor<?x?xf32>, 
+    %off0 : index, %off1 : index,
+    %sz0 : index, %sz1 : index,
+    %st0 : index, %st1 : index) -> tensor<?x?xf32> {
+
+  //      CHECK: %[[EXTRACTED_SLICE:.*]] = tensor.extract_slice %[[O]][%[[OFF0]], %[[OFF1]]] [2, 3] [%[[ST0]], %[[ST1]]] 
+  // CHECK-SAME:   : tensor<?x?xf32> to tensor<2x3xf32>
+  //      CHECK: linalg.copy ins(%[[I]] : tensor<2x3xf32>) outs(%[[EXTRACTED_SLICE]] : tensor<2x3xf32>) -> tensor<2x3xf32>
+  //      CHECK: tensor.insert_slice %{{.*}} into %[[O]][%[[OFF0]], %[[OFF1]]] [2, 3] [%[[ST0]], %[[ST1]]]
+  // CHECK-SAME:   : tensor<2x3xf32> into tensor<?x?xf32>
+
+  %0 = tensor.insert_slice %I into %O[%off0, %off1] [2, 3] [%st0, %st1] 
+    : tensor<2x3xf32> into tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+  %1 = transform.structured.insert_slice_to_copy %0 : (!pdl.operation) -> !pdl.operation
+  transform.cast %1 : !pdl.operation to !transform.op<"linalg.copy">
+}
+
+// -----
+
+// CHECK-LABEL: func @insert_slice_to_copy
+    // CHECK-SAME: %[[I:[0-9a-zA-Z]+]]: tensor<?x?xf32>
+    // CHECK-SAME: %[[O:[0-9a-zA-Z]+]]: tensor<?x?xf32>, 
+    // CHECK-SAME: %[[OFF0:[0-9a-zA-Z]+]]: index,
+    // CHECK-SAME: %[[OFF1:[0-9a-zA-Z]+]]: index,
+    // CHECK-SAME: %[[SZ0:[0-9a-zA-Z]+]]: index,
+    // CHECK-SAME: %[[SZ1:[0-9a-zA-Z]+]]: index,
+    // CHECK-SAME: %[[ST0:[0-9a-zA-Z]+]]: index,
+    // CHECK-SAME: %[[ST1:[0-9a-zA-Z]+]]: index)
+func.func @insert_slice_to_copy(
+    %I : tensor<?x?xf32>, %O : tensor<?x?xf32>, 
+    %off0 : index, %off1 : index,
+    %sz0 : index, %sz1 : index,
+    %st0 : index, %st1 : index) -> tensor<?x?xf32> {
+
+  //      CHECK: %[[EXTRACTED_SLICE:.*]] = tensor.extract_slice %[[O]][%[[OFF0]], %[[OFF1]]] [%[[SZ0]], %[[SZ1]]] [1, 1] 
+  // CHECK-SAME:   : tensor<?x?xf32> to tensor<?x?xf32>
+  //      CHECK: linalg.copy ins(%[[I]] : tensor<?x?xf32>) outs(%[[EXTRACTED_SLICE]] : tensor<?x?xf32>) -> tensor<?x?xf32>
+  //      CHECK: tensor.insert_slice %{{.*}} into %[[O]][%[[OFF0]], %[[OFF1]]] [%[[SZ0]], %[[SZ1]]] [1, 1]
+  // CHECK-SAME:   : tensor<?x?xf32> into tensor<?x?xf32>
+
+  %0 = tensor.insert_slice %I into %O[%off0, %off1] [%sz0, %sz1] [1, 1] 
+    : tensor<?x?xf32> into tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+  %1 = transform.structured.insert_slice_to_copy %0 : (!pdl.operation) -> !pdl.operation
+  transform.cast %1 : !pdl.operation to !transform.op<"linalg.copy">
+}
+
+// -----
+// CHECK-LABEL: func @insert_slice_to_copy
+    // CHECK-SAME: %[[I:.*]]: tensor<2x3xf32>
+    // CHECK-SAME: %[[O:.*]]: tensor<?x?xf32>, 
+    // CHECK-SAME: %[[OFF0:[0-9a-zA-Z]+]]: index,
+    // CHECK-SAME: %[[OFF1:[0-9a-zA-Z]+]]: index,
+    // CHECK-SAME: %[[SZ0:[0-9a-zA-Z]+]]: index,
+    // CHECK-SAME: %[[SZ1:[0-9a-zA-Z]+]]: index,
+    // CHECK-SAME: %[[ST0:[0-9a-zA-Z]+]]: index,
+    // CHECK-SAME: %[[ST1:[0-9a-zA-Z]+]]: index)
+func.func @insert_slice_to_copy(
+    %I : tensor<2x3xf32>, %O : tensor<?x?xf32>, 
+    %off0 : index, %off1 : index,
+    %sz0 : index, %sz1 : index,
+    %st0 : index, %st1 : index) -> tensor<?x?xf32> {
+
+  //      CHECK: %[[EXTRACTED_SLICE:.*]] = tensor.extract_slice %[[O]][%[[OFF0]], %[[OFF1]]] [2, 3] [%[[ST0]], %[[ST1]]] 
+  // CHECK-SAME:   : tensor<?x?xf32> to tensor<2x3xf32>
+  //      CHECK: linalg.copy ins(%[[I]] : tensor<2x3xf32>) outs(%[[EXTRACTED_SLICE]] : tensor<2x3xf32>) -> tensor<2x3xf32>
+  //  CHECK-NOT: linalg.copy
+  //      CHECK: tensor.insert_slice %{{.*}} into %[[O]][%[[OFF0]], %[[OFF1]]] [2, 3] [%[[ST0]], %[[ST1]]]
+  // CHECK-SAME:   : tensor<2x3xf32> into tensor<?x?xf32>
+
+  %extracted_slice = tensor.extract_slice %O[%off0, %off1] [2, 3] [%st0, %st1] 
+    : tensor<?x?xf32> to tensor<2x3xf32>
+  %0 = linalg.copy ins(%I : tensor<2x3xf32>) outs(%extracted_slice 
+    : tensor<2x3xf32>) -> tensor<2x3xf32>
+  %inserted_slice = tensor.insert_slice %0 into %O[%off0, %off1] [2, 3] [%st0, %st1] 
+    : tensor<2x3xf32> into tensor<?x?xf32>
+
+  return %inserted_slice : tensor<?x?xf32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+  %1 = transform.structured.insert_slice_to_copy %0 : (!pdl.operation) -> !pdl.operation
+  transform.cast %1 : !pdl.operation to !transform.op<"linalg.copy">
+}
+