[mlir][linalg] Update vectorization tests for tensor.pad and tensor.insert_slice #138267
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Adds `-canonicalization -cse` to the vectorization tests for: * `tensor.pad` + `tensor.insert_slice`. This significantly simplies the expected output to match and hence improves maintainability of these tests. For for context/discussion: * llvm#138265 Other Ops/tests will be refactor in follow-up patches.
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-linalg Author: Andrzej Warzyński (banach-space) ChangesAdds
This significantly simplies the expected output to match and hence Other Ops/tests will be refactor in follow-up patches. Patch is 30.76 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/138267.diff 3 Files Affected:
diff --git a/mlir/test/Dialect/Linalg/vectorization.mlir b/mlir/test/Dialect/Linalg/vectorization.mlir
index 299be1296aa66..a970f927e32f5 100644
--- a/mlir/test/Dialect/Linalg/vectorization.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization.mlir
@@ -582,132 +582,6 @@ module attributes {transform.with_named_sequence} {
// -----
-// CHECK-LABEL: func @test_masked_vectorize_pad
-func.func @test_masked_vectorize_pad(
- %0 : tensor<?x?xf32>, %h0 : index, %h1 : index)
- -> tensor<2x4xf32>
-{
- // CHECK-DAG: %[[c42:.*]] = arith.constant 4.243000e+01 : f32
- // CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
- // CHECK-DAG: %[[c0_0:.*]] = arith.constant 0 : index
- // CHECK: %[[d0:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>
- // CHECK: %[[d1:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>
- // CHECK: %[[mask:.*]] = vector.create_mask %[[d0]], %[[d1]] : vector<2x4xi1>
- // CHECK: %[[masked_read:.*]] = vector.mask %[[mask]] {
- // CHECK-SAME: vector.transfer_read %{{.*}}[%[[c0_0]], %[[c0_0]]], %[[c42]]
- // CHECK-SAME: {in_bounds = [true, true]} : tensor<?x?xf32>, vector<2x4xf32>
- // CHECK-SAME: } : vector<2x4xi1> -> vector<2x4xf32>
- // CHECK-DAG: %[[c0_1:.*]] = arith.constant 0 : index
- // CHECK-DAG: %[[empty:.*]] = tensor.empty() : tensor<2x4xf32>
- // CHECK: vector.transfer_write %[[masked_read]], %[[empty]][%[[c0_1]], %[[c0_1]]]
- // CHECK-SAME: {in_bounds = [true, true]} : vector<2x4xf32>, tensor<2x4xf32>
- %cst = arith.constant 42.43 : f32
- %c0 = arith.constant 0 : index
- %1 = tensor.pad %0 low[0, %c0] high[%h0, %h1] {
- ^bb0(%hh1: index, %hh2: index):
- tensor.yield %cst : f32
- } : tensor<?x?xf32> to tensor<2x4xf32>
- return %1: tensor<2x4xf32>
-}
-
-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
- transform.structured.vectorize %0 vector_sizes [2, 4] : !transform.any_op
- transform.yield
- }
-}
-
-// -----
-
-// CHECK: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>
-// CHECK: func @test_masked_vectorize_dynamic_pad
-func.func @test_masked_vectorize_dynamic_pad(
- %0 : tensor<?x?xf32>, %h0 : index, %h1 : index)
- -> tensor<?x?xf32>
-{
- // CHECK-DAG: %[[c42:.*]] = arith.constant 4.243000e+01 : f32
- // CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
- // CHECK-DAG: %[[res_d0:.+]] = affine.apply #[[MAP]]()
- // CHECK-DAG: %[[res_d1:.+]] = affine.apply #[[MAP]]()
- // CHECK: %[[c0_2:.*]] = arith.constant 0 : index
- // CHECK: %[[d0:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>
- // CHECK: %[[d1:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>
- // CHECK: %[[mask:.*]] = vector.create_mask %[[d0]], %[[d1]] : vector<2x4xi1>
- // CHECK: %[[masked_read:.*]] = vector.mask %[[mask]] {
- // CHECK-SAME: vector.transfer_read %{{.*}}[%[[c0_2]], %[[c0_2]]], %[[c42]]
- // CHECK-SAME: {in_bounds = [true, true]} : tensor<?x?xf32>, vector<2x4xf32>
- // CHECK-SAME: } : vector<2x4xi1> -> vector<2x4xf32>
- // CHECK-DAG: %[[empty:.*]] = tensor.empty(%[[res_d0]], %[[res_d1]]) : tensor<?x?xf32>
- // CHECK-DAG: %[[c0_3:.*]] = arith.constant 0 : index
- // CHECK: %[[mask_2:.*]] = vector.create_mask %[[res_d0]], %[[res_d1]] : vector<2x4xi1>
- // CHECK: %[[masked_write:.*]] = vector.mask %[[mask_2]] {
- // CHECK-SAME: vector.transfer_write %[[masked_read]], %[[empty]][%[[c0_3]], %[[c0_3]]]
- // CHECK-SAME: {in_bounds = [true, true]} : vector<2x4xf32>, tensor<?x?xf32>
- // CHECK: return %[[masked_write]] : tensor<?x?xf32>
- %cst = arith.constant 42.43 : f32
- %c0 = arith.constant 0 : index
- %1 = tensor.pad %0 low[0, %c0] high[%h0, %h1] {
- ^bb0(%hh1: index, %hh2: index):
- tensor.yield %cst : f32
- } : tensor<?x?xf32> to tensor<?x?xf32>
- return %1: tensor<?x?xf32>
-}
-
-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
- transform.structured.vectorize %0 vector_sizes [2, 4] : !transform.any_op
- transform.yield
- }
-}
-
-// -----
-// This case is supported because low padding `%l0` is applied on
-// a unit dimension which is supported, non unit result dimension low
-// padding is currently unsupported.
-// CHECK-LABEL: func @test_masked_vectorize_non_zero_low_pad_unit_res_dim
-func.func @test_masked_vectorize_non_zero_low_pad_unit_res_dim(
- %0 : tensor<?x?xf32>, %h0 : index, %h1 : index, %l0 : index)
- -> tensor<1x4xf32>
-{
- // CHECK-DAG: %[[C42:.*]] = arith.constant 4.243000e+01 : f32
- // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
- // CHECK: %[[C0_1:.*]] = arith.constant 0 : index
- // CHECK-DAG: %[[D0:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>
- // CHECK-DAG: %[[D1:.*]] = tensor.dim {{.*}} : tensor<?x?xf32>
- // CHECK: %[[MASK:.*]] = vector.create_mask %[[D0]], %[[D1]] : vector<1x4xi1>
- // CHECK: %[[MASKED_READ:.*]] = vector.mask %[[MASK]] {
- // CHECK-SAME: vector.transfer_read %{{.*}}[%[[C0_1]], %[[C0_1]]], %[[C42]]
- // CHECK-SAME: {in_bounds = [true, true]} : tensor<?x?xf32>, vector<1x4xf32>
- // CHECK-SAME: } : vector<1x4xi1> -> vector<1x4xf32>
- // CHECK-DAG: %[[EMPTY:.*]] = tensor.empty() : tensor<1x4xf32>
- // CHECK-DAG: %[[C0_2:.*]] = arith.constant 0 : index
- // CHECK: %[[MASKED_WRITE:.*]] = vector.transfer_write %[[MASKED_READ]], %[[EMPTY]][%[[C0_2]], %[[C0_2]]]
- // CHECK-SAME: {in_bounds = [true, true]} : vector<1x4xf32>, tensor<1x4xf32>
- // CHECK: return %[[MASKED_WRITE]] : tensor<1x4xf32>
- %cst = arith.constant 42.43 : f32
- %c0 = arith.constant 0 : index
- %1 = tensor.pad %0 low[%l0, %c0] high[%h0, %h1] {
- ^bb0(%hh1: index, %hh2: index):
- tensor.yield %cst : f32
- } : tensor<?x?xf32> to tensor<1x4xf32>
- return %1: tensor<1x4xf32>
-}
-
-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
- transform.structured.vectorize %0 vector_sizes [1, 4] : !transform.any_op
- transform.yield
- }
-}
-
-// -----
-
// Input identical as the test in vectorization-with-patterns.mlir. Output is
// different - vector sizes are inferred (rather than user-specified) and hence
// masking was used.
@@ -1150,154 +1024,3 @@ func.func @test_vectorize_unpack_no_vector_sizes_permute(%source: tensor<4x7x4xf
transform.yield
}
}
-
-// -----
-
-///----------------------------------------------------------------------------------------
-/// tensor.insert_slice
-///----------------------------------------------------------------------------------------
-
-func.func private @insert_slice_static_sizes(%source: tensor<?x3x?x1xi32>) -> tensor<5x3xi32> {
- %c2 = arith.constant 2 : index
- %init = tensor.empty() : tensor<5x3xi32>
-
- %source_slice = tensor.extract_slice %source[0, %c2, 0, 0] [1, 1, 5, 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<5x1xi32>
- %res = tensor.insert_slice %source_slice into %init[0, %c2] [5, 1] [1, 1] : tensor<5x1xi32> into tensor<5x3xi32>
-
- return %res : tensor<5x3xi32>
-}
-
-// CHECK-LABEL: func.func private @insert_slice_static_sizes(
-// CHECK-SAME: %[[SEC:.*]]: tensor<?x3x?x1xi32>) -> tensor<5x3xi32> {
-// CHECK: %[[C_2:.*]] = arith.constant 2 : index
-// CHECK: %[[INIT:.*]] = tensor.empty() : tensor<5x3xi32>
-// CHECK: %[[SRC_SLICE:.*]] = tensor.extract_slice %[[SEC]][0, %[[C_2]], 0, 0] [1, 1, 5, 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<5x1xi32>
-// CHECK-DAG: %[[PAD:.*]] = arith.constant 0 : i32
-// CHECK-DAG: %[[C_5:.*]] = arith.constant 5 : index
-// CHECK-DAG: %[[C_1:.*]] = arith.constant 1 : index
-// CHECK: %[[MASK:.*]] = vector.create_mask %[[C_5]], %[[C_1]] : vector<8x1xi1>
-// CHECK: %[[C0:.*]] = arith.constant 0 : index
-// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[SRC_SLICE]][%[[C0]], %[[C0]]], %[[PAD]] : tensor<5x1xi32>, vector<8x1xi32> } : vector<8x1xi1> -> vector<8x1xi32>
-// CHECK: %[[C_0:.*]] = arith.constant 0 : index
-// CHECK: %[[RES:.*]] = vector.mask %[[MASK]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0]], %[[C_2]]] : vector<8x1xi32>, tensor<5x3xi32> } : vector<8x1xi1> -> tensor<5x3xi32>
-// CHECK: return %[[RES]] : tensor<5x3xi32>
-
- module attributes {transform.with_named_sequence} {
- transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
- %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg0 : (!transform.any_op) -> !transform.any_op
- transform.structured.vectorize %0 vector_sizes [8, 1] : !transform.any_op
- transform.yield
- }
- }
-
-// -----
-
-// One of the _source_ dimensions is dynamic (but _destination_ dimensions are static).
-
-func.func private @insert_slice_dynamic_src_dim(%source: tensor<?x3x?x1xi32>, %size: index) -> tensor<5x3xi32> {
- %c2 = arith.constant 2 : index
- %init = tensor.empty() : tensor<5x3xi32>
-
- %source_slice = tensor.extract_slice %source[0, %c2, 0, 0] [1, 1, %size, 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<?x1xi32>
- %res = tensor.insert_slice %source_slice into %init[0, %c2] [%size, 1] [1, 1] : tensor<?x1xi32> into tensor<5x3xi32>
-
- return %res : tensor<5x3xi32>
-}
-
-// CHECK-LABEL: func.func private @insert_slice_dynamic_src_dim(
-// CHECK-SAME: %[[SRC:.*]]: tensor<?x3x?x1xi32>,
-// CHECK-SAME: %[[SIZE:.*]]: index) -> tensor<5x3xi32> {
-// CHECK: %[[C_2:.*]] = arith.constant 2 : index
-// CHECK: %[[INIT:.*]] = tensor.empty() : tensor<5x3xi32>
-// CHECK: %[[SRC_SLICE:.*]] = tensor.extract_slice %[[SRC]][0, %[[C_2]], 0, 0] [1, 1, %[[SIZE]], 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<?x1xi32>
-// CHECK-DAG: %[[PAD:.*]] = arith.constant 0 : i32
-// CHECK-DAG: %[[C_1:.*]] = arith.constant 1 : index
-// CHECK: %[[MASK:.*]] = vector.create_mask %[[SIZE]], %[[C_1]] : vector<8x1xi1>
-// CHECK: %[[C_0:.*]] = arith.constant 0 : index
-// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[SRC_SLICE]][%[[C_0]], %[[C_0]]], %[[PAD]] : tensor<?x1xi32>, vector<8x1xi32> } : vector<8x1xi1> -> vector<8x1xi32>
-// CHECK: %[[C_0_1:.*]] = arith.constant 0 : index
-// CHECK: %[[RES:.*]] = vector.mask %[[MASK]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0_1]], %[[C_2]]] : vector<8x1xi32>, tensor<5x3xi32> } : vector<8x1xi1> -> tensor<5x3xi32>
-// CHECK: return %[[RES]] : tensor<5x3xi32>
-
- module attributes {transform.with_named_sequence} {
- transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
- %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg0 : (!transform.any_op) -> !transform.any_op
- transform.structured.vectorize %0 vector_sizes [8, 1] : !transform.any_op
- transform.yield
- }
- }
-
-// -----
-
-// One of the _destination_ dimensions is dynamic (but _source_ dimensions are static).
-
-func.func private @insert_slice_dynamic_dest_dim(%source: tensor<?x3x?x1xi32>, %size: index) -> tensor<?x3xi32> {
- %c2 = arith.constant 2 : index
- %init = tensor.empty(%size) : tensor<?x3xi32>
-
- %source_slice = tensor.extract_slice %source[0, %c2, 0, 0] [1, 1, 5, 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<5x1xi32>
- %res = tensor.insert_slice %source_slice into %init[0, %c2] [5, 1] [1, 1] : tensor<5x1xi32> into tensor<?x3xi32>
-
- return %res : tensor<?x3xi32>
-}
-
-// CHECK-LABEL: func.func private @insert_slice_dynamic_dest_dim(
-// CHECK-SAME: %[[SRC:.*]]: tensor<?x3x?x1xi32>,
-// CHECK-SAME: %[[SIZE:.*]]: index) -> tensor<?x3xi32> {
-// CHECK: %[[C_2:.*]] = arith.constant 2 : index
-// CHECK: %[[INIT:.*]] = tensor.empty(%[[SIZE]]) : tensor<?x3xi32>
-// CHECK: %[[SRC_SLICE:.*]] = tensor.extract_slice %[[SRC]][0, %[[C_2]], 0, 0] [1, 1, 5, 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<5x1xi32>
-// CHECK: %[[PAD:.*]] = arith.constant 0 : i32
-// CHECK: %[[C_5:.*]] = arith.constant 5 : index
-// CHECK: %[[C_1:.*]] = arith.constant 1 : index
-// CHECK: %[[MASK:.*]] = vector.create_mask %[[C_5]], %[[C_1]] : vector<8x1xi1>
-// CHECK: %[[C_0:.*]] = arith.constant 0 : index
-// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[SRC_SLICE]][%[[C_0]], %[[C_0]]], %[[PAD]] : tensor<5x1xi32>, vector<8x1xi32> } : vector<8x1xi1> -> vector<8x1xi32>
-// CHECK: %[[C_0_1:.*]] = arith.constant 0 : index
-// CHECK: %[[WRITE:.*]] = vector.mask %[[MASK]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0_1]], %[[C_2]]] : vector<8x1xi32>, tensor<?x3xi32> } : vector<8x1xi1> -> tensor<?x3xi32>
-// CHECK: return %[[WRITE]] : tensor<?x3xi32>
-
- module attributes {transform.with_named_sequence} {
- transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
- %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg0 : (!transform.any_op) -> !transform.any_op
- transform.structured.vectorize %0 vector_sizes [8, 1] : !transform.any_op
- transform.yield
- }
- }
-
-// -----
-
-// At least one _source_ and one _destination_ dimensions are dynamic.
-
-func.func private @insert_slice_dynamic_source_and_dest_dim(%source: tensor<?x3x?x1xi32>, %size: index) -> tensor<?x3xi32> {
- %c2 = arith.constant 2 : index
- %init = tensor.empty(%size) : tensor<?x3xi32>
-
- %source_slice = tensor.extract_slice %source[0, %c2, 0, 0] [1, 1, %size, 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<?x1xi32>
- %res = tensor.insert_slice %source_slice into %init[0, %c2] [%size, 1] [1, 1] : tensor<?x1xi32> into tensor<?x3xi32>
-
- return %res : tensor<?x3xi32>
-}
-
-// CHECK-LABEL: func.func private @insert_slice_dynamic_source_and_dest_dim(
-// CHECK-SAME: %[[SRC:.*]]: tensor<?x3x?x1xi32>,
-// CHECK-SAME: %[[SIZE:.*]]: index) -> tensor<?x3xi32> {
-// CHECK: %[[C_2:.*]] = arith.constant 2 : index
-// CHECK: %[[INIT:.*]] = tensor.empty(%[[SIZE]]) : tensor<?x3xi32>
-// CHECK: %[[SRC_SIZE:.*]] = tensor.extract_slice %[[SRC]][0, %[[C_2]], 0, 0] [1, 1, %[[SIZE]], 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<?x1xi32>
-// CHECK: %[[PAD:.*]] = arith.constant 0 : i32
-// CHECK: %[[C1:.*]] = arith.constant 1 : index
-// CHECK: %[[MASK:.*]] = vector.create_mask %[[SIZE]], %[[C1]] : vector<8x1xi1>
-// CHECK: %[[C0:.*]] = arith.constant 0 : index
-// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[SRC_SIZE]]{{\[}}%[[C0]], %[[C0]]], %[[PAD]] : tensor<?x1xi32>, vector<8x1xi32> } : vector<8x1xi1> -> vector<8x1xi32>
-// CHECK: %[[C_0_1:.*]] = arith.constant 0 : index
-// CHECK: %[[WRITE:.*]] = vector.mask %[[MASK]] { vector.transfer_write %[[READ]], %[[INIT]]{{\[}}%[[C_0_1]], %[[C_2]]] : vector<8x1xi32>, tensor<?x3xi32> } : vector<8x1xi1> -> tensor<?x3xi32>
-// CHECK: return %[[WRITE]] : tensor<?x3xi32>
-
- module attributes {transform.with_named_sequence} {
- transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
- %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg0 : (!transform.any_op) -> !transform.any_op
- transform.structured.vectorize %0 vector_sizes [8, 1] : !transform.any_op
- transform.yield
- }
- }
diff --git a/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir b/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
new file mode 100644
index 0000000000000..14522cb970ab5
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/vectorization/insert-slice.mlir
@@ -0,0 +1,150 @@
+// RUN: mlir-opt %s -transform-interpreter -split-input-file -cse -canonicalize | FileCheck %s
+
+func.func private @insert_slice_static_sizes(%source: tensor<?x3x?x1xi32>) -> tensor<5x3xi32> {
+ %c2 = arith.constant 2 : index
+ %init = tensor.empty() : tensor<5x3xi32>
+
+ %source_slice = tensor.extract_slice %source[0, %c2, 0, 0] [1, 1, 5, 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<5x1xi32>
+ %res = tensor.insert_slice %source_slice into %init[0, %c2] [5, 1] [1, 1] : tensor<5x1xi32> into tensor<5x3xi32>
+
+ return %res : tensor<5x3xi32>
+}
+
+// CHECK-LABEL: func.func private @insert_slice_static_sizes(
+// CHECK-SAME: %[[SEC:.*]]: tensor<?x3x?x1xi32>) -> tensor<5x3xi32> {
+// CHECK-DAG: %[[C_2:.*]] = arith.constant 2 : index
+// CHECK-DAG: %[[PAD:.*]] = arith.constant 0 : i32
+// CHECK-DAG: %[[C_0:.*]] = arith.constant 0 : index
+// CHECK: %[[INIT:.*]] = tensor.empty() : tensor<5x3xi32>
+// CHECK: %[[SRC_SLICE:.*]] = tensor.extract_slice %[[SEC]][0, 2, 0, 0] [1, 1, 5, 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<5x1xi32>
+// CHECK: %[[MASK:.*]] = vector.constant_mask [5, 1] : vector<8x1xi1>
+// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[SRC_SLICE]][%[[C_0]], %[[C_0]]], %[[PAD]] {{.*}}: tensor<5x1xi32>, vector<8x1xi32> } : vector<8x1xi1> -> vector<8x1xi32>
+/// The mask is vector.constant_mask [5, 1] rather than e.g.
+/// * vector.constant_mask [5, %c3]
+/// as the trailing dim of the mask is "1" anyway.
+// CHECK: %[[RES:.*]] = vector.mask %[[MASK]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0]], %[[C_2]]] {{.*}}: vector<8x1xi32>, tensor<5x3xi32> } : vector<8x1xi1> -> tensor<5x3xi32>
+// CHECK: return %[[RES]] : tensor<5x3xi32>
+
+ module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["tensor.insert_slice"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ transform.structured.vectorize %0 vector_sizes [8, 1] : !transform.any_op
+ transform.yield
+ }
+ }
+
+// -----
+
+// One of the _source_ dimensions is dynamic (but _destination_ dimensions are static).
+
+func.func private @insert_slice_dynamic_src_dim(%source: tensor<?x3x?x1xi32>, %size: index) -> tensor<5x3xi32> {
+ %c2 = arith.constant 2 : index
+ %init = tensor.empty() : tensor<5x3xi32>
+
+ %source_slice = tensor.extract_slice %source[0, %c2, 0, 0] [1, 1, %size, 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<?x1xi32>
+ %res = tensor.insert_slice %source_slice into %init[0, %c2] [%size, 1] [1, 1] : tensor<?x1xi32> into tensor<5x3xi32>
+
+ return %res : tensor<5x3xi32>
+}
+
+// CHECK-LABEL: func.func private @insert_slice_dynamic_src_dim(
+// CHECK-SAME: %[[SRC:.*]]: tensor<?x3x?x1xi32>,
+// CHECK-SAME: %[[SIZE:.*]]: index) -> tensor<5x3xi32> {
+// CHECK-DAG: %[[PAD:.*]] = arith.constant 0 : i32
+// CHECK-DAG: %[[C_1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C_2:.*]] = arith.constant 2 : index
+// CHECK-DAG: %[[C_0:.*]] = arith.constant 0 : index
+// CHECK: %[[INIT:.*]] = tensor.empty() : tensor<5x3xi32>
+// CHECK: %[[SRC_SLICE:.*]] = tensor.extract_slice %[[SRC]][0, 2, 0, 0] [1, 1, %[[SIZE]], 1] [1, 1, 1, 1] : tensor<?x3x?x1xi32> to tensor<?x1xi32>
+// CHECK: %[[MASK:.*]] = vector.create_mask %[[SIZE]], %[[C_1]] : vector<8x1xi1>
+// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[SRC_SLICE]][%[[C_0]], %[[C_0]]], %[[PAD]] {{.*}}: tensor<?x1xi32>, vector<8x1xi32> } : vector<8x1xi1> -> vector<8x1xi32>
+/// The mask is vector.constant_mask [5, 1] rather than e.g.
+/// * vector.constant_mask [5, %c3]
+/// as the trailing dim of the mask is "1" anyway.
+// CHECK: %[[RES:.*]] = vector.mask %[[MASK]] { vector.transfer_write %[[READ]], %[[INIT]][%[[C_0]], %[[C_2]]] {{.*}}: vector<8x1xi32>, tensor<5x3xi32> } : vector<8x1xi1> -> tensor<5x3xi32>
+// CHECK: return %[[RES]] : tensor<5x3xi32>
+
+ module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op...
[truncated]
|
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This patch adds
-canonicalize -cseto the vectorization tests for:tensor.pad+tensor.insert_sliceThis significantly simplifies the expected output, making the tests easier to maintain. For context and discussion, see:
-canonicalize -cse#138265This change is implemented by moving the affected tests to dedicated files in a newly created subdirectory:
Other ops/tests will be refactored in follow-up patches.