From 17bc7a9aafe277b584874158e94f42a8a5429afe Mon Sep 17 00:00:00 2001 From: Andrzej Warzynski Date: Wed, 16 Oct 2024 09:44:49 +0100 Subject: [PATCH 1/2] [mlir] Add apply_patterns.linalg.pad_vectorization TD Op This PR simply wraps `populatePadOpVectorizationPatterns` into a new Transform Dialect Op: `apply_patterns.linalg.pad_vectorization`. This change makes it possible to run (and test) the corresponding patterns without: `transform.structured.vectorize_children_and_apply_patterns`. Note that the Op above only supports non-masked vectorisation (i.e. when the inputs are static), so, effectively, only fixed-width vectorisation (as opposed to scalable vectorisation). As such, this change is required to construct vectorization pipelines for tensor.pad targeting scalable vectors. To test the new Op and the corresponding patterns, I added "vectorization-pad-patterns.mlir" - most tests have been extracted from "vectorization-with-patterns.mlir". As a side note, I feel that we should move `GenericPadOpVectorizationPattern` out of `populatePadOpVectorizationPatterns` as that's a "lower tensor.pad" rather than a "vectorize tensor.pad" pattern. I am leaving that as a TODO. --- .../Linalg/TransformOps/LinalgTransformOps.td | 17 ++ .../TransformOps/LinalgTransformOps.cpp | 5 + .../Linalg/Transforms/Vectorization.cpp | 2 +- .../Linalg/vectorization-pad-patterns.mlir | 272 ++++++++++++++++++ .../Linalg/vectorization-with-patterns.mlir | 143 --------- 5 files changed, 295 insertions(+), 144 deletions(-) create mode 100644 mlir/test/Dialect/Linalg/vectorization-pad-patterns.mlir diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td index 0915bbde3072b..af09708e9c794 100644 --- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td +++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td @@ -84,6 +84,23 @@ def ApplyFoldAddIntoDestPatternsOp : Op]> { + let description = [{ + Apply patterns that take tensor.pad and rewrites it as + vector.transfer_read/vector.transfer_write Ops. + + These patterns will either fold tensor.pad with an existing + vector.transfer_read or vector.transfer_write producer/consumers (requires + other surrounding Ops to be already vectorised) or rewrite it, together + with tensor.insert_slice consumer, as a vector.transfer_read + + vector.transfer_write pair. + }]; + + let assemblyFormat = "attr-dict"; +} + //===----------------------------------------------------------------------===// // BufferizeToAllocationOp //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp index ad72b5d7beccd..55e9f88a0046f 100644 --- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp +++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp @@ -253,6 +253,11 @@ void transform::ApplyFoldAddIntoDestPatternsOp::populatePatterns( linalg::populateFoldAddIntoDestPatterns(patterns); } +void transform::ApplyPadVectorizationPatternsOp::populatePatterns( + RewritePatternSet &patterns) { + linalg::populatePadOpVectorizationPatterns(patterns); +} + //===----------------------------------------------------------------------===// // BufferizeToAllocationOp //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp index e1b97fbf985df..c139ff5d82bf2 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp @@ -2304,7 +2304,7 @@ static SmallVector ofrToIndexValues(RewriterBase &rewriter, Location loc, return result; } -/// Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp and +/// Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp/GenerateOp and /// InsertSliceOp. For now, only constant padding values are supported. /// If there is enough static type information, TransferReadOps and /// TransferWriteOps may be generated instead of InsertSliceOps. diff --git a/mlir/test/Dialect/Linalg/vectorization-pad-patterns.mlir b/mlir/test/Dialect/Linalg/vectorization-pad-patterns.mlir new file mode 100644 index 0000000000000..76cd7d7b0721f --- /dev/null +++ b/mlir/test/Dialect/Linalg/vectorization-pad-patterns.mlir @@ -0,0 +1,272 @@ +// RUN: mlir-opt %s -transform-interpreter -split-input-file | FileCheck %s + +///---------------------------------------------------------------------------------------- +/// [Pattern: PadOpVectorizationWithTransferReadPattern] +///---------------------------------------------------------------------------------------- +// CHECK-LABEL: func @pad_and_transfer_read +// CHECK-SAME: %[[ARG0:.*]]: tensor<5x6xf32> +// CHECK-NOT: tensor.pad +// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index +// CHECK-DAG: %[[C5:.*]] = arith.constant 5.0 +// CHECK: %[[RESULT:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]]], %[[C5]] : tensor<5x6xf32>, vector<7x9xf32> +// CHECK: return %[[RESULT]] +func.func @pad_and_transfer_read(%arg0: tensor<5x6xf32>) -> vector<7x9xf32> { + %c0 = arith.constant 0 : index + %c5 = arith.constant 5.0 : f32 + %c6 = arith.constant 6.0 : f32 + %0 = tensor.pad %arg0 low[0, 0] high[5, 7] { + ^bb0(%arg1: index, %arg2: index): + tensor.yield %c5 : f32 + } : tensor<5x6xf32> to tensor<10x13xf32> + %1 = vector.transfer_read %0[%c0, %c0], %c6 + : tensor<10x13xf32>, vector<7x9xf32> + return %1 : vector<7x9xf32> +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { + %func_op = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.op<"func.func"> + + transform.apply_patterns to %func_op { + transform.apply_patterns.linalg.pad_vectorization + } : !transform.op<"func.func"> + transform.yield + } +} + +// ----- + +///---------------------------------------------------------------------------------------- +/// [Pattern: PadOpVectorizationWithTransferReadPattern +///---------------------------------------------------------------------------------------- +func.func private @make_vector() -> vector<7x9xf32> + +// CHECK-LABEL: func @pad_and_transfer_write_static +// CHECK-SAME: %[[ARG0:.*]]: tensor<5x6xf32> +// CHECK-NOT: tensor.pad +// CHECK: %[[C0:.*]] = arith.constant 0 : index +// CHECK: %[[VEC0:.*]] = call @make_vector() : () -> vector<7x9xf32> +// CHECK: %[[RESULT:.*]] = vector.transfer_write %[[VEC0]], %[[ARG0]][%[[C0]], %[[C0]]] : vector<7x9xf32>, tensor<5x6xf32> +// CHECK: return %[[RESULT]] +func.func @pad_and_transfer_write_static( + %arg0: tensor<5x6xf32>) -> tensor<5x6xf32> { + %c0 = arith.constant 0 : index + %c5 = arith.constant 5.0 : f32 + %0 = tensor.pad %arg0 low[0, 0] high[5, 7] { + ^bb0(%arg2: index, %arg3: index): + tensor.yield %c5 : f32 + } : tensor<5x6xf32> to tensor<10x13xf32> + %1 = call @make_vector() : () -> vector<7x9xf32> + %2 = vector.transfer_write %1, %0[%c0, %c0] + : vector<7x9xf32>, tensor<10x13xf32> + %3 = tensor.extract_slice %2[0, 0] [5, 6] [1, 1] : tensor<10x13xf32> to tensor<5x6xf32> + return %3 : tensor<5x6xf32> +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { + %func_op = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.op<"func.func"> + + transform.apply_patterns to %func_op { + transform.apply_patterns.linalg.pad_vectorization + } : !transform.op<"func.func"> + transform.yield + } +} + +// ----- + +func.func private @make_vector() -> vector<7x9xf32> + +// CHECK-LABEL: func @pad_and_transfer_write_dynamic_static +// CHECK-SAME: %[[ARG0:.*]]: tensor, %[[SIZE:.*]]: index, %[[PADDING:.*]]: index +// CHECK-NOT: tensor.pad +// CHECK: %[[C0:.*]] = arith.constant 0 : index +// CHECK: %[[SUB:.*]] = tensor.extract_slice %[[ARG0]][0, 0] [%[[SIZE]], 6] [1, 1] : tensor to tensor +// CHECK: %[[VEC0:.*]] = call @make_vector() : () -> vector<7x9xf32> +// CHECK: %[[RESULT:.*]] = vector.transfer_write %[[VEC0]], %[[SUB]][%[[C0]], %[[C0]]] : vector<7x9xf32>, tensor +// CHECK: return %[[RESULT]] +func.func @pad_and_transfer_write_dynamic_static( + %arg0: tensor, %size: index, %padding: index) -> tensor { + %c0 = arith.constant 0 : index + %c5 = arith.constant 5.0 : f32 + %s = tensor.extract_slice %arg0[0, 0] [%size, 6] [1, 1] + : tensor to tensor + %0 = tensor.pad %s low[0, 0] high[%padding, 7] { + ^bb0(%arg2: index, %arg3: index): + tensor.yield %c5 : f32 + } : tensor to tensor + %1 = call @make_vector() : () -> vector<7x9xf32> + %2 = vector.transfer_write %1, %0[%c0, %c0] + : vector<7x9xf32>, tensor + %3 = tensor.extract_slice %2[0, 0] [%size, 6] [1, 1] : tensor to tensor + return %3 : tensor +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { + %func_op = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.op<"func.func"> + + transform.apply_patterns to %func_op { + transform.apply_patterns.linalg.pad_vectorization + } : !transform.op<"func.func"> + transform.yield + } +} + + +// ----- + +///---------------------------------------------------------------------------------------- +/// [Pattern: PadOpVectorizationWithInsertSlicePattern] +///---------------------------------------------------------------------------------------- + +func.func private @make_vector() -> tensor<12x13xf32> + +// CHECK-LABEL: func @pad_and_insert_slice_source +// CHECK-SAME: %[[ARG0:.*]]: tensor<5x6xf32> +// CHECK-NOT: tensor.pad +// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index +// CHECK-DAG: %[[C5:.*]] = arith.constant 5.0 +// CHECK: %[[VEC0:.*]] = call @make_vector() : () -> tensor<12x13xf32> +// CHECK: %[[READ:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]]], %[[C5]] : tensor<5x6xf32>, vector<7x9xf32> +// CHECK: %[[WRITE:.*]] = vector.transfer_write %[[READ]], %[[VEC0]][%[[C0]], %[[C0]]] {in_bounds = [true, true]} : vector<7x9xf32>, tensor<12x13xf32> +// CHECK: return %[[WRITE]] +func.func @pad_and_insert_slice_source( + %arg0: tensor<5x6xf32>) -> tensor<12x13xf32> { + %c0 = arith.constant 0 : index + %c5 = arith.constant 5.0 : f32 + %0 = tensor.pad %arg0 low[0, 0] high[2, 3] { + ^bb0(%arg2: index, %arg3: index): + tensor.yield %c5 : f32 + } : tensor<5x6xf32> to tensor<7x9xf32> + %1 = call @make_vector() : () -> tensor<12x13xf32> + %r = tensor.insert_slice %0 into %1[0, 0][7, 9][1, 1] : tensor<7x9xf32> into tensor<12x13xf32> + return %r : tensor<12x13xf32> +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { + %func_op = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.op<"func.func"> + + transform.apply_patterns to %func_op { + transform.apply_patterns.linalg.pad_vectorization + } : !transform.op<"func.func"> + transform.yield + } +} + + +// ----- + +///---------------------------------------------------------------------------------------- +/// tensor::PadOp -> tensor::EmptyOp + linalg::FillOp/tensor::GenerateOp + tensor::InsertSliceOp +/// [Pattern: GenericPadOpVectorizationPattern] +///---------------------------------------------------------------------------------------- + +func.func private @make_vector() -> tensor<12x13xf32> + +// Same as @pad_and_insert_slice_dest in vectorization-wit-patterns.mlir, but +// CHECK-LABEL: func.func @pad_and_insert_slice_dest( +// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x5x6xf32>) -> tensor<1x12x13xf32> { +// CHECK-NOT: tensor.pad +// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index +// CHECK-DAG: %[[PAD:.*]] = arith.constant 5.000000e+00 : f32 +// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<1x12x13xf32> +// CHECK: %[[FILL:.*]] = linalg.fill ins(%[[PAD]] : f32) outs(%[[EMPTY]] : tensor<1x12x13xf32>) -> tensor<1x12x13xf32> +// CHECK: %[[READ:.*]] = vector.transfer_read %[[ARG_0]]{{\[}}%[[C0]], %[[C0]], %[[C0]]], %[[PAD]] {in_bounds = [true, true, true]} : tensor<1x5x6xf32>, vector<1x5x6xf32> +// CHECK: %[[WRITE:.*]] = vector.transfer_write %[[READ]], %[[FILL]]{{\[}}%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true]} : vector<1x5x6xf32>, tensor<1x12x13xf32> +// CHECK: %[[VEC:.*]] = call @make_vector() : () -> tensor<12x13xf32> +// CHECK: %[[RES:.*]] = tensor.insert_slice %[[VEC]] into %[[WRITE]][0, 0, 0] [1, 12, 13] [1, 1, 1] : tensor<12x13xf32> into tensor<1x12x13xf32> +// CHECK: return %[[RES]] : tensor<1x12x13xf32> + +func.func @pad_and_insert_slice_dest( + %arg0: tensor<1x5x6xf32>) -> tensor<1x12x13xf32> { + %c5 = arith.constant 5.0 : f32 + %0 = tensor.pad %arg0 low[0, 0, 0] high[0, 7, 7] { + ^bb0(%arg2: index, %arg3: index, %arg4: index): + tensor.yield %c5 : f32 + } : tensor<1x5x6xf32> to tensor<1x12x13xf32> + %1 = call @make_vector() : () -> tensor<12x13xf32> + %r = tensor.insert_slice %1 into %0[0, 0, 0][1, 12, 13][1, 1, 1] : tensor<12x13xf32> into tensor<1x12x13xf32> + return %r : tensor<1x12x13xf32> +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { + %func_op = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.op<"func.func"> + + transform.apply_patterns to %func_op { + transform.apply_patterns.linalg.pad_vectorization + } : !transform.op<"func.func"> + transform.yield + } +} + +// ----- +func.func private @make_vector() -> vector<7x9xf32> + +// Variant of @pad_and_transfer_write_static + +// CHECK-LABEL: func @pad_and_transfer_write_static_non_zero_low_pad +// CHECK-NOT: tensor.pad +// CHECK: linalg.fill +func.func @pad_and_transfer_write_static_non_zero_low_pad( + %arg0: tensor<5x6xf32>) -> tensor<5x6xf32> { + %c0 = arith.constant 0 : index + %c5 = arith.constant 5.0 : f32 + %0 = tensor.pad %arg0 low[0, 1] high[5, 6] { + ^bb0(%arg2: index, %arg3: index): + tensor.yield %c5 : f32 + } : tensor<5x6xf32> to tensor<10x13xf32> + %1 = call @make_vector() : () -> vector<7x9xf32> + %2 = vector.transfer_write %1, %0[%c0, %c0] + : vector<7x9xf32>, tensor<10x13xf32> + %3 = tensor.extract_slice %2[0, 0] [5, 6] [1, 1] : tensor<10x13xf32> to tensor<5x6xf32> + return %3 : tensor<5x6xf32> +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { + %func_op = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.op<"func.func"> + + transform.apply_patterns to %func_op { + transform.apply_patterns.linalg.pad_vectorization + } : !transform.op<"func.func"> + transform.yield + } +} + +// ----- +func.func private @make_vector() -> vector<7x9xf32> + +// Variant of @pad_and_transfer_write_static + +// CHECK-LABEL: func @pad_and_transfer_write_static_non_zero_offset +// CHECK-NOT: tensor.pad +// CHECK: linalg.fill +func.func @pad_and_transfer_write_static_non_zero_offset( + %arg0: tensor<5x6xf32>) -> tensor<5x6xf32> { + %c0 = arith.constant 0 : index + %c5 = arith.constant 5.0 : f32 + %0 = tensor.pad %arg0 low[0, 1] high[5, 6] { + ^bb0(%arg2: index, %arg3: index): + tensor.yield %c5 : f32 + } : tensor<5x6xf32> to tensor<10x13xf32> + %1 = call @make_vector() : () -> vector<7x9xf32> + %2 = vector.transfer_write %1, %0[%c0, %c0] + : vector<7x9xf32>, tensor<10x13xf32> + %3 = tensor.extract_slice %2[0, 1] [5, 6] [1, 1] : tensor<10x13xf32> to tensor<5x6xf32> + return %3 : tensor<5x6xf32> +} + +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { + %func_op = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.op<"func.func"> + + transform.apply_patterns to %func_op { + transform.apply_patterns.linalg.pad_vectorization + } : !transform.op<"func.func"> + transform.yield + } +} diff --git a/mlir/test/Dialect/Linalg/vectorization-with-patterns.mlir b/mlir/test/Dialect/Linalg/vectorization-with-patterns.mlir index 1c6a786bfa436..189507d97d6dc 100644 --- a/mlir/test/Dialect/Linalg/vectorization-with-patterns.mlir +++ b/mlir/test/Dialect/Linalg/vectorization-with-patterns.mlir @@ -935,149 +935,6 @@ module attributes {transform.with_named_sequence} { } } -// ----- - -// CHECK-LABEL: func @pad_and_transfer_read -// CHECK-SAME: %[[ARG0:.*]]: tensor<5x6xf32> -// CHECK-NOT: tensor.pad -// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index -// CHECK-DAG: %[[C5:.*]] = arith.constant 5.0 -// CHECK: %[[RESULT:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]]], %[[C5]] : tensor<5x6xf32>, vector<7x9xf32> -// CHECK: return %[[RESULT]] -func.func @pad_and_transfer_read(%arg0: tensor<5x6xf32>) -> vector<7x9xf32> { - %c0 = arith.constant 0 : index - %c5 = arith.constant 5.0 : f32 - %c6 = arith.constant 6.0 : f32 - %0 = tensor.pad %arg0 low[0, 0] high[5, 7] { - ^bb0(%arg1: index, %arg2: index): - tensor.yield %c5 : f32 - } : tensor<5x6xf32> to tensor<10x13xf32> - %1 = vector.transfer_read %0[%c0, %c0], %c6 - : tensor<10x13xf32>, vector<7x9xf32> - return %1 : vector<7x9xf32> -} - - -module attributes {transform.with_named_sequence} { - transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { - %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op - %1 = transform.get_parent_op %0 {isolated_from_above} : (!transform.any_op) -> !transform.any_op - %2 = transform.structured.vectorize_children_and_apply_patterns %1 { vectorize_padding } : (!transform.any_op) -> !transform.any_op - transform.yield - } -} - -// ----- - -func.func private @make_vector() -> vector<7x9xf32> - -// CHECK-LABEL: func @pad_and_transfer_write_static -// CHECK-SAME: %[[ARG0:.*]]: tensor<5x6xf32> -// CHECK-NOT: tensor.pad -// CHECK: %[[C0:.*]] = arith.constant 0 : index -// CHECK: %[[VEC0:.*]] = call @make_vector() : () -> vector<7x9xf32> -// CHECK: %[[RESULT:.*]] = vector.transfer_write %[[VEC0]], %[[ARG0]][%[[C0]], %[[C0]]] : vector<7x9xf32>, tensor<5x6xf32> -// CHECK: return %[[RESULT]] -func.func @pad_and_transfer_write_static( - %arg0: tensor<5x6xf32>) -> tensor<5x6xf32> { - %c0 = arith.constant 0 : index - %c5 = arith.constant 5.0 : f32 - %0 = tensor.pad %arg0 low[0, 0] high[5, 7] { - ^bb0(%arg2: index, %arg3: index): - tensor.yield %c5 : f32 - } : tensor<5x6xf32> to tensor<10x13xf32> - %1 = call @make_vector() : () -> vector<7x9xf32> - %2 = vector.transfer_write %1, %0[%c0, %c0] - : vector<7x9xf32>, tensor<10x13xf32> - %3 = tensor.extract_slice %2[0, 0] [5, 6] [1, 1] : tensor<10x13xf32> to tensor<5x6xf32> - return %3 : tensor<5x6xf32> -} - -module attributes {transform.with_named_sequence} { - transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { - %3 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op - %4 = transform.get_parent_op %3 {isolated_from_above} : (!transform.any_op) -> !transform.any_op - %5 = transform.structured.vectorize_children_and_apply_patterns %4 { vectorize_padding } : (!transform.any_op) -> !transform.any_op - transform.yield - } -} - - -// ----- - -func.func private @make_vector() -> vector<7x9xf32> - -// CHECK-LABEL: func @pad_and_transfer_write_dynamic_static -// CHECK-SAME: %[[ARG0:.*]]: tensor, %[[SIZE:.*]]: index, %[[PADDING:.*]]: index -// CHECK-NOT: tensor.pad -// CHECK: %[[C0:.*]] = arith.constant 0 : index -// CHECK: %[[SUB:.*]] = tensor.extract_slice %[[ARG0]][0, 0] [%[[SIZE]], 6] [1, 1] : tensor to tensor -// CHECK: %[[VEC0:.*]] = call @make_vector() : () -> vector<7x9xf32> -// CHECK: %[[RESULT:.*]] = vector.transfer_write %[[VEC0]], %[[SUB]][%[[C0]], %[[C0]]] : vector<7x9xf32>, tensor -// CHECK: return %[[RESULT]] -func.func @pad_and_transfer_write_dynamic_static( - %arg0: tensor, %size: index, %padding: index) -> tensor { - %c0 = arith.constant 0 : index - %c5 = arith.constant 5.0 : f32 - %s = tensor.extract_slice %arg0[0, 0] [%size, 6] [1, 1] - : tensor to tensor - %0 = tensor.pad %s low[0, 0] high[%padding, 7] { - ^bb0(%arg2: index, %arg3: index): - tensor.yield %c5 : f32 - } : tensor to tensor - %1 = call @make_vector() : () -> vector<7x9xf32> - %2 = vector.transfer_write %1, %0[%c0, %c0] - : vector<7x9xf32>, tensor - %3 = tensor.extract_slice %2[0, 0] [%size, 6] [1, 1] : tensor to tensor - return %3 : tensor -} - -module attributes {transform.with_named_sequence} { - transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { - %3 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op - %4 = transform.get_parent_op %3 {isolated_from_above} : (!transform.any_op) -> !transform.any_op - %5 = transform.structured.vectorize_children_and_apply_patterns %4 { vectorize_padding } : (!transform.any_op) -> !transform.any_op - transform.yield - } -} - - -// ----- - -func.func private @make_vector() -> tensor<12x13xf32> - -// CHECK-LABEL: func @pad_and_insert_slice_source -// CHECK-SAME: %[[ARG0:.*]]: tensor<5x6xf32> -// CHECK-NOT: tensor.pad -// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index -// CHECK-DAG: %[[C5:.*]] = arith.constant 5.0 -// CHECK: %[[VEC0:.*]] = call @make_vector() : () -> tensor<12x13xf32> -// CHECK: %[[READ:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]]], %[[C5]] : tensor<5x6xf32>, vector<7x9xf32> -// CHECK: %[[WRITE:.*]] = vector.transfer_write %[[READ]], %[[VEC0]][%[[C0]], %[[C0]]] {in_bounds = [true, true]} : vector<7x9xf32>, tensor<12x13xf32> -// CHECK: return %[[WRITE]] -func.func @pad_and_insert_slice_source( - %arg0: tensor<5x6xf32>) -> tensor<12x13xf32> { - %c0 = arith.constant 0 : index - %c5 = arith.constant 5.0 : f32 - %0 = tensor.pad %arg0 low[0, 0] high[2, 3] { - ^bb0(%arg2: index, %arg3: index): - tensor.yield %c5 : f32 - } : tensor<5x6xf32> to tensor<7x9xf32> - %1 = call @make_vector() : () -> tensor<12x13xf32> - %r = tensor.insert_slice %0 into %1[0, 0][7, 9][1, 1] : tensor<7x9xf32> into tensor<12x13xf32> - return %r : tensor<12x13xf32> -} - -module attributes {transform.with_named_sequence} { - transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { - %3 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op - %4 = transform.get_parent_op %3 {isolated_from_above} : (!transform.any_op) -> !transform.any_op - %5 = transform.structured.vectorize_children_and_apply_patterns %4 { vectorize_padding } : (!transform.any_op) -> !transform.any_op - transform.yield - } -} - - // ----- func.func private @make_vector() -> tensor<12x13xf32> From c9efeda1b9f534d73598fd0333ff202875313beb Mon Sep 17 00:00:00 2001 From: Andrzej Warzynski Date: Sat, 19 Oct 2024 12:57:15 +0100 Subject: [PATCH 2/2] fixup! [mlir] Add apply_patterns.linalg.pad_vectorization TD Op Address comment from Javed --- .../Linalg/TransformOps/LinalgTransformOps.td | 19 +++++++++++-------- .../Linalg/Transforms/Vectorization.cpp | 5 ++++- .../Linalg/vectorization-pad-patterns.mlir | 14 ++++++++------ 3 files changed, 23 insertions(+), 15 deletions(-) diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td index af09708e9c794..10ab65047d646 100644 --- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td +++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td @@ -88,14 +88,17 @@ def ApplyPadVectorizationPatternsOp : Op]> { let description = [{ - Apply patterns that take tensor.pad and rewrites it as - vector.transfer_read/vector.transfer_write Ops. - - These patterns will either fold tensor.pad with an existing - vector.transfer_read or vector.transfer_write producer/consumers (requires - other surrounding Ops to be already vectorised) or rewrite it, together - with tensor.insert_slice consumer, as a vector.transfer_read + - vector.transfer_write pair. + Apply patterns that vectorize tensor.pad. + + These patterns rewrite tensor.pad Ops using vector.transfer_read and + vector.transfer_write operations. This is done either by: + 1. Folding tensor.pad with an existing vector.transfer_read / + vector.transfer_write Op (generated prior to running these patterns). + 2. Rewriting it (when matched together with q tensor.insert_slice + consumer Op) as a vector.transfer_read + vector.transfer_write pair. + + In both cases, these patterns look at producers and consumers for the + matched tensor.pad Op to find opportunities for vectorization. }]; let assemblyFormat = "attr-dict"; diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp index c139ff5d82bf2..0a2457176a1d4 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp @@ -2304,7 +2304,7 @@ static SmallVector ofrToIndexValues(RewriterBase &rewriter, Location loc, return result; } -/// Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp/GenerateOp and +/// Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp and /// InsertSliceOp. For now, only constant padding values are supported. /// If there is enough static type information, TransferReadOps and /// TransferWriteOps may be generated instead of InsertSliceOps. @@ -2712,6 +2712,9 @@ struct PadOpVectorizationWithInsertSlicePattern void mlir::linalg::populatePadOpVectorizationPatterns( RewritePatternSet &patterns, PatternBenefit baseBenefit) { + // TODO: The following pattern implements "decomposition" and + // optional "vectorization". Seperate "decomposition" into a sepereate + // pre-processing pattern group. patterns.add(patterns.getContext(), baseBenefit); // Try these specialized patterns first before resorting to the generic one. diff --git a/mlir/test/Dialect/Linalg/vectorization-pad-patterns.mlir b/mlir/test/Dialect/Linalg/vectorization-pad-patterns.mlir index 76cd7d7b0721f..2aa4638af3f0f 100644 --- a/mlir/test/Dialect/Linalg/vectorization-pad-patterns.mlir +++ b/mlir/test/Dialect/Linalg/vectorization-pad-patterns.mlir @@ -37,18 +37,18 @@ module attributes {transform.with_named_sequence} { // ----- ///---------------------------------------------------------------------------------------- -/// [Pattern: PadOpVectorizationWithTransferReadPattern +/// [Pattern: PadOpVectorizationWithTransferWritePattern] ///---------------------------------------------------------------------------------------- func.func private @make_vector() -> vector<7x9xf32> -// CHECK-LABEL: func @pad_and_transfer_write_static +// CHECK-LABEL: func @pad_and_transfer_write_static_low_and_high // CHECK-SAME: %[[ARG0:.*]]: tensor<5x6xf32> // CHECK-NOT: tensor.pad // CHECK: %[[C0:.*]] = arith.constant 0 : index // CHECK: %[[VEC0:.*]] = call @make_vector() : () -> vector<7x9xf32> // CHECK: %[[RESULT:.*]] = vector.transfer_write %[[VEC0]], %[[ARG0]][%[[C0]], %[[C0]]] : vector<7x9xf32>, tensor<5x6xf32> // CHECK: return %[[RESULT]] -func.func @pad_and_transfer_write_static( +func.func @pad_and_transfer_write_static_low_and_high( %arg0: tensor<5x6xf32>) -> tensor<5x6xf32> { %c0 = arith.constant 0 : index %c5 = arith.constant 5.0 : f32 @@ -78,7 +78,7 @@ module attributes {transform.with_named_sequence} { func.func private @make_vector() -> vector<7x9xf32> -// CHECK-LABEL: func @pad_and_transfer_write_dynamic_static +// CHECK-LABEL: func @pad_and_transfer_write_static_low_dynamic_high // CHECK-SAME: %[[ARG0:.*]]: tensor, %[[SIZE:.*]]: index, %[[PADDING:.*]]: index // CHECK-NOT: tensor.pad // CHECK: %[[C0:.*]] = arith.constant 0 : index @@ -86,7 +86,7 @@ func.func private @make_vector() -> vector<7x9xf32> // CHECK: %[[VEC0:.*]] = call @make_vector() : () -> vector<7x9xf32> // CHECK: %[[RESULT:.*]] = vector.transfer_write %[[VEC0]], %[[SUB]][%[[C0]], %[[C0]]] : vector<7x9xf32>, tensor // CHECK: return %[[RESULT]] -func.func @pad_and_transfer_write_dynamic_static( +func.func @pad_and_transfer_write_static_low_dynamic_high( %arg0: tensor, %size: index, %padding: index) -> tensor { %c0 = arith.constant 0 : index %c5 = arith.constant 5.0 : f32 @@ -166,7 +166,9 @@ module attributes {transform.with_named_sequence} { func.func private @make_vector() -> tensor<12x13xf32> -// Same as @pad_and_insert_slice_dest in vectorization-wit-patterns.mlir, but +// Same as @pad_and_insert_slice_dest in vectorization-with-patterns.mlir, but +// over here linalg::fill is not vectorized (patterns for linalg.fill are not +// included here) // CHECK-LABEL: func.func @pad_and_insert_slice_dest( // CHECK-SAME: %[[ARG_0:.*]]: tensor<1x5x6xf32>) -> tensor<1x12x13xf32> { // CHECK-NOT: tensor.pad