[Mlir-commits] [mlir] 93a735a - [mlir][Linalg] Add a structured transform to materialize a tensor.insert_slice via a linalg.copy
Nicolas Vasilache
llvmlistbot at llvm.org
Wed Apr 12 04:05:06 PDT 2023
Author: Nicolas Vasilache
Date: 2023-04-12T04:05:00-07:00
New Revision: 93a735ac789d92a7e495408cc202140ea24db3ee
URL: https://github.com/llvm/llvm-project/commit/93a735ac789d92a7e495408cc202140ea24db3ee
DIFF: https://github.com/llvm/llvm-project/commit/93a735ac789d92a7e495408cc202140ea24db3ee.diff
LOG: [mlir][Linalg] Add a structured transform to materialize a tensor.insert_slice via a linalg.copy
This is useful to materialize copies explicitly before bufferization and
transform them, avoiding the need to rediscover them after bufferization.
Differential Revision: https://reviews.llvm.org/D148108
Added:
mlir/test/Dialect/Linalg/transform-op-insert-slice-to-copy.mlir
Modified:
mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h
mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h
index 6d7c802c3e533..091c419f1774a 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h
@@ -26,6 +26,7 @@ class LinalgOp;
} // namespace linalg
namespace tensor {
+class InsertSliceOp;
class PackOp;
class PadOp;
class UnPackOp;
diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index 415c7992b6138..f60e6bdc79abb 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
@@ -2005,4 +2005,42 @@ def HoistRedundantTensorSubsetsOp :
}];
}
+//===----------------------------------------------------------------------===//
+// 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
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index a92ffa196a40f..53a8fae8560b5 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -3232,6 +3232,36 @@ transform::HoistRedundantTensorSubsetsOp::applyToOne(
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
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Linalg/transform-op-insert-slice-to-copy.mlir b/mlir/test/Dialect/Linalg/transform-op-insert-slice-to-copy.mlir
new file mode 100644
index 0000000000000..7c2461c52b2a2
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/transform-op-insert-slice-to-copy.mlir
@@ -0,0 +1,110 @@
+// 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: !transform.any_op):
+ %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1 = transform.structured.insert_slice_to_copy %0 : (!transform.any_op) -> !transform.any_op
+ transform.cast %1 : !transform.any_op 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: !transform.any_op):
+ %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1 = transform.structured.insert_slice_to_copy %0 : (!transform.any_op) -> !transform.any_op
+ transform.cast %1 : !transform.any_op 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: !transform.any_op):
+ %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+ %1 = transform.structured.insert_slice_to_copy %0 : (!transform.any_op) -> !transform.any_op
+ transform.cast %1 : !transform.any_op to !transform.op<"linalg.copy">
+}
+
More information about the Mlir-commits
mailing list