Skip to content

Commit bab30f2

Browse files
hanhanWhanhanW
committed
Add tests and delete unused arguments
Signed-off-by: hanhanW <[email protected]>
1 parent 3aac143 commit bab30f2

File tree

2 files changed

+76
-25
lines changed

2 files changed

+76
-25
lines changed

‎mlir/test/Dialect/Linalg/canonicalize.mlir‎

Lines changed: 3 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1889,9 +1889,7 @@ func.func @fold_cast_unpack_dynamic_tile_size(
18891889
// linalg.unpack + tensor.extract_slice
18901890
//===----------------------------------------------------------------------===//
18911891

1892-
func.func @fold_extract_slice_into_unpack_slicing_trailing_dim(
1893-
%src : tensor<28x2x1x16x16xf32>, %dest : tensor<28x28x15xf32>, %size : index
1894-
) -> tensor<28x28x10xf32> {
1892+
func.func @fold_extract_slice_into_unpack_slicing_trailing_dim(%src : tensor<28x2x1x16x16xf32>, %dest : tensor<28x28x15xf32>) -> tensor<28x28x10xf32> {
18951893
%unpack = linalg.unpack %src
18961894
outer_dims_perm = [0, 1, 2]
18971895
inner_dims_pos = [1, 2]
@@ -1904,7 +1902,6 @@ func.func @fold_extract_slice_into_unpack_slicing_trailing_dim(
19041902
// CHECK-LABEL: func @fold_extract_slice_into_unpack_slicing_trailing_dim
19051903
// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]
19061904
// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]
1907-
// CHECK-SAME: %[[SIZE:[a-zA-Z0-9]+]]
19081905
// CHECK: %[[DEST_SLICE:.+]] = tensor.extract_slice %[[DEST]]
19091906
// CHECK-SAME: [0, 0, 0] [28, 28, 10] [1, 1, 1]
19101907
// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[SRC]]
@@ -1915,10 +1912,7 @@ func.func @fold_extract_slice_into_unpack_slicing_trailing_dim(
19151912

19161913
// The available dimension size is [17, 32], because CeilDiv(%d1, 16) == 2.
19171914

1918-
1919-
func.func @fold_extract_slice_into_unpack_slicing_dim_1(
1920-
%src : tensor<28x2x1x16x16xf32>, %dest : tensor<28x28x15xf32>, %size : index
1921-
) -> tensor<28x17x15xf32> {
1915+
func.func @fold_extract_slice_into_unpack_slicing_dim_1(%src : tensor<28x2x1x16x16xf32>, %dest : tensor<28x28x15xf32>) -> tensor<28x17x15xf32> {
19221916
%unpack = linalg.unpack %src
19231917
inner_dims_pos = [1, 2]
19241918
inner_tiles = [16, 16]
@@ -1930,7 +1924,6 @@ func.func @fold_extract_slice_into_unpack_slicing_dim_1(
19301924
// CHECK-LABEL: func @fold_extract_slice_into_unpack_slicing_dim_1(
19311925
// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]
19321926
// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]
1933-
// CHECK-SAME: %[[SIZE:[a-zA-Z0-9]+]]
19341927
// CHECK: %[[DEST_SLICE:.+]] = tensor.extract_slice %[[DEST]]
19351928
// CHECK-SAME: [0, 0, 0] [28, 17, 15] [1, 1, 1]
19361929
// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[SRC]]
@@ -1941,9 +1934,7 @@ func.func @fold_extract_slice_into_unpack_slicing_dim_1(
19411934

19421935
// The available dimension size is [17, 32], because CeilDiv(%d1, 16) == 2.
19431936

1944-
func.func @no_fold_extract_slice_into_unpack_artificial_padding(
1945-
%src : tensor<28x2x1x16x16xf32>, %dest : tensor<28x28x15xf32>, %size : index
1946-
) -> tensor<28x16x15xf32> {
1937+
func.func @no_fold_extract_slice_into_unpack_artificial_padding(%src : tensor<28x2x1x16x16xf32>, %dest : tensor<28x28x15xf32>) -> tensor<28x16x15xf32> {
19471938
%unpack = linalg.unpack %src
19481939
inner_dims_pos = [1, 2]
19491940
inner_tiles = [16, 16]

‎mlir/test/Dialect/Tensor/fold-into-pack-and-unpack.mlir‎

Lines changed: 73 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,79 @@
11
// RUN: mlir-opt -split-input-file -test-linalg-transform-patterns=test-fold-into-pack-and-unpack %s | FileCheck %s
22
// RUN: mlir-opt -split-input-file -test-linalg-transform-patterns=test-fold-into-pack-and-unpack-control %s | FileCheck %s --check-prefix=CONTROL
33

4-
func.func @fold_unpack_slice(%arg0 : tensor<2082x1x8x32xf32>) -> tensor<16649x16xf32> {
5-
%empty = tensor.empty() : tensor<16656x16xf32>
6-
%0 = linalg.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [8, 32] into %empty
7-
: tensor<2082x1x8x32xf32> -> tensor<16656x16xf32>
8-
%1 = tensor.extract_slice %0[0, 0] [16649, 16] [1, 1] : tensor<16656x16xf32> to tensor<16649x16xf32>
9-
return %1 : tensor<16649x16xf32>
10-
}
11-
// CHECK-LABEL: func @fold_unpack_slice(
12-
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
13-
// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<16649x16xf32>
14-
// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[ARG0]] inner_dims_pos = [0, 1] inner_tiles = [8, 32]
15-
// CHECK-SAME: into %[[INIT]]
16-
// CHECK: return %[[UNPACK]]
4+
func.func @fold_extract_slice_into_unpack_slicing_trailing_dim(%arg0 : tensor<28x2x1x16x16xf32>) -> tensor<28x28x10xf32> {
5+
%empty = tensor.empty() : tensor<28x28x15xf32>
6+
%unpack = linalg.unpack %arg0
7+
inner_dims_pos = [1, 2]
8+
inner_tiles = [16, 16]
9+
into %empty : tensor<28x2x1x16x16xf32> -> tensor<28x28x15xf32>
10+
%extracted_slice = tensor.extract_slice %unpack
11+
[0, 0, 0] [28, 28, 10] [1, 1, 1] : tensor<28x28x15xf32> to tensor<28x28x10xf32>
12+
return %extracted_slice : tensor<28x28x10xf32>
13+
}
14+
// CHECK-LABEL: func @fold_extract_slice_into_unpack_slicing_trailing_dim
15+
// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]
16+
// CHECK: %[[DEST_SLICE:.+]] = tensor.empty() : tensor<28x28x10xf32>
17+
// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[SRC]]
18+
// CHECK-SAME: into %[[DEST_SLICE]]
19+
// CHECK: return %[[UNPACK]]
20+
21+
// -----
22+
23+
// The available dimension size is [17, 32], because CeilDiv(%d1, 16) == 2.
24+
25+
func.func @fold_extract_slice_into_unpack_slicing_dim_1(%arg0 : tensor<28x2x1x16x16xf32>) -> tensor<28x17x15xf32> {
26+
%empty = tensor.empty() : tensor<28x28x15xf32>
27+
%unpack = linalg.unpack %arg0
28+
inner_dims_pos = [1, 2]
29+
inner_tiles = [16, 16]
30+
into %empty : tensor<28x2x1x16x16xf32> -> tensor<28x28x15xf32>
31+
%extracted_slice = tensor.extract_slice %unpack
32+
[0, 0, 0] [28, 17, 15] [1, 1, 1] : tensor<28x28x15xf32> to tensor<28x17x15xf32>
33+
return %extracted_slice : tensor<28x17x15xf32>
34+
}
35+
// CHECK-LABEL: func @fold_extract_slice_into_unpack_slicing_dim_1(
36+
// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]
37+
// CHECK: %[[DEST_SLICE:.+]] = tensor.empty() : tensor<28x17x15xf32>
38+
// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[SRC]]
39+
// CHECK-SAME: into %[[DEST_SLICE]]
40+
// CHECK: return %[[UNPACK]]
41+
42+
// -----
43+
44+
// The available dimension size is [17, 32], because CeilDiv(%d1, 16) == 2.
45+
46+
func.func @no_fold_extract_slice_into_unpack_artificial_padding(%arg0 : tensor<28x2x1x16x16xf32>) -> tensor<28x16x15xf32> {
47+
%empty = tensor.empty() : tensor<28x28x15xf32>
48+
%unpack = linalg.unpack %arg0
49+
inner_dims_pos = [1, 2]
50+
inner_tiles = [16, 16]
51+
into %empty : tensor<28x2x1x16x16xf32> -> tensor<28x28x15xf32>
52+
%extracted_slice = tensor.extract_slice %unpack
53+
[0, 0, 0] [28, 16, 15] [1, 1, 1] : tensor<28x28x15xf32> to tensor<28x16x15xf32>
54+
return %extracted_slice : tensor<28x16x15xf32>
55+
}
56+
// CHECK-LABEL: func @no_fold_extract_slice_into_unpack_artificial_padding
57+
// CHECK: linalg.unpack
58+
// CHECK: tensor.extract_slice
59+
60+
// -----
61+
62+
func.func @no_fold_extract_slice_into_unpack_dynamic(
63+
%src : tensor<28x2x?x16x16xf32>, %dest : tensor<28x32x?xf32>, %size : index
64+
) -> tensor<28x28x?xf32> {
65+
%unpack = linalg.unpack %src
66+
outer_dims_perm = [0, 1, 2]
67+
inner_dims_pos = [1, 2]
68+
inner_tiles = [16, 16]
69+
into %dest : tensor<28x2x?x16x16xf32> -> tensor<28x32x?xf32>
70+
%extracted_slice = tensor.extract_slice %unpack
71+
[0, 0, 0] [28, 28, %size] [1, 1, 1] : tensor<28x32x?xf32> to tensor<28x28x?xf32>
72+
return %extracted_slice : tensor<28x28x?xf32>
73+
}
74+
// CHECK-LABEL: func @no_fold_extract_slice_into_unpack_dynamic
75+
// CHECK: linalg.unpack
76+
// CHECK: tensor.extract_slice
1777

1878
// -----
1979

0 commit comments

Comments
 (0)