Skip to content

Commit 33a9c26

Browse files
authored
[mlir][tensor] Add e2e test for tensor.pack with dynamic tile sizes (#115698)
Adds an end-to-end test for `tensor.pack` with dynamic inner tile sizes. While relatively simple (e.g., no vectorization), this example required a few non-trivial fixes in handling `tensor.pack`: * #114315, #114559, #113108. The end goal for this test is to incrementally increase its complexity and to work towards scalable tile sizes.
1 parent 980316e commit 33a9c26

File tree

1 file changed

+97
-0
lines changed

1 file changed

+97
-0
lines changed
Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
// DEFINE: %{compile} = mlir-opt %s \
2+
// DEFINE: -transform-interpreter -test-transform-dialect-erase-schedule |\
3+
// DEFINE: mlir-opt --test-linalg-transform-patterns="test-generalize-tensor-pack"\
4+
// DEFINE: --test-transform-dialect-erase-schedule \
5+
// DEFINE: -one-shot-bufferize="bufferize-function-boundaries" \
6+
// DEFINE: -buffer-deallocation-pipeline="private-function-dynamic-ownership" \
7+
// DEFINE: -cse -canonicalize -test-lower-to-llvm -o %t
8+
// DEFINE: %{entry_point} = main
9+
// DEFINE: %{run} = mlir-cpu-runner %t -e %{entry_point} -entry-point-result=void \
10+
// DEFINE: -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils
11+
12+
// RUN: rm -f %t && %{compile} && %{run} | FileCheck %s
13+
14+
/// End-to-end test for tensor.pack where one of the inner tile sizes is
15+
/// dynamic.
16+
///
17+
/// Note, ATM this is a relatively simple example, with no vectorization and
18+
/// the dynamic tile size being a compile-time constant. The intention is to
19+
/// incrementally expand the config to something much more complex.
20+
21+
func.func @main() {
22+
// Allocate and initialise the inputs
23+
%A_alloc = tensor.empty() : tensor<7x16xi32>
24+
25+
%A = arith.constant dense<[
26+
[ 1, 8, 15, 22, 29, 36, 43, 50, 57, 64, 71, 78, 85, 92, 99 , 106],
27+
[ 2, 9, 16, 23, 30, 37, 44, 51, 58, 65, 72, 79, 86, 93, 100, 107],
28+
[ 3, 10, 17, 24, 31, 38, 45, 52, 59, 66, 73, 80, 87, 94, 101, 108],
29+
[ 4, 11, 18, 25, 32, 39, 46, 53, 60, 67, 74, 81, 88, 95, 102, 109],
30+
[ 5, 12, 19, 26, 33, 40, 47, 54, 61, 68, 75, 82, 89, 96, 103, 110],
31+
[ 6, 13, 20, 27, 34, 41, 48, 55, 62, 69, 76, 83, 90, 97, 104, 111],
32+
[ 7, 14, 21, 28, 35, 42, 49, 56, 63, 70, 77, 84, 91, 98, 105, 112]
33+
]> : tensor<7x16xi32>
34+
35+
func.call @pack(%A) : (tensor<7x16xi32>) -> ()
36+
37+
return
38+
}
39+
40+
func.func private @pack(%A: tensor<7x16xi32>) {
41+
%c1 = arith.constant 1 : index
42+
%pad_val = arith.constant 123 : i32
43+
44+
// Dynamic tile size
45+
%tile_size = arith.constant 8 : index
46+
%A_pack_empty = tensor.empty(%c1, %tile_size) : tensor<?x16x?x1xi32>
47+
48+
%A_pack = tensor.pack %A
49+
padding_value(%pad_val : i32)
50+
inner_dims_pos = [0, 1]
51+
inner_tiles = [%tile_size, 1]
52+
into %A_pack_empty : tensor<7x16xi32> -> tensor<?x16x?x1xi32>
53+
%A_cast = tensor.cast %A_pack : tensor<?x16x?x1xi32> to tensor<*xi32>
54+
55+
// Print the results
56+
// CHECK: Unranked Memref base@ = 0{{.*}} rank = 4 offset = 0 sizes = [1, 16, 8, 1] strides = [128, 8, 1, 1] data =
57+
// Tile 1: (8 x 1)
58+
// CHECK-NEXT: 1
59+
// CHECK-NEXT: 2
60+
// CHECK-NEXT: 3
61+
// CHECK-NEXT: 4
62+
// CHECK-NEXT: 5
63+
// CHECK-NEXT: 6
64+
// CHECK-NEXT: 7
65+
// Expect pad value after 7 elements
66+
// CHECK-NEXT: 123
67+
// Tile 2: (8 x 1)
68+
// CHECK-NEXT: 8
69+
// CHECK-NEXT: 9
70+
// CHECK-NEXT: 10
71+
// CHECK-NEXT: 11
72+
// CHECK-NEXT: 12
73+
// CHECK-NEXT: 13
74+
// CHECK-NEXT: 14
75+
// Expect pad value after further 7 elements
76+
// CHECK-NEXT: 123
77+
// Tile 3: (8 x 1)
78+
// CHECK-NEXT: 15
79+
// CHECK-NEXT: 16
80+
// ...
81+
call @printMemrefI32(%A_cast) : (tensor<*xi32>) -> ()
82+
83+
return
84+
}
85+
86+
module @transforms attributes { transform.with_named_sequence } {
87+
transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) {
88+
%pack = transform.structured.match ops{["tensor.pack"]} in %module : (!transform.any_op) -> !transform.any_op
89+
90+
%tiled_linalg_op_p, %loops:2 = transform.structured.tile_using_for %pack tile_sizes [1, 1]
91+
: (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
92+
93+
transform.yield
94+
}
95+
}
96+
97+
func.func private @printMemrefI32(%ptr : tensor<*xi32>)

0 commit comments

Comments
 (0)