Lines Matching full:tensor
1 // RUN: mlir-opt -split-input-file -test-tensor-transform-patterns="test-simplify-pack-unpack-patte…
4 // CHECK-SAME: %[[ARG0:.+]]: tensor<256xf32>)
5 …: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1]] output_shape [8, 32] : te…
6 // CHECK: return %[[EXPANDED]] : tensor<8x32xf32>
7 func.func @single_dim_packing(%arg0: tensor<256xf32>) -> tensor<8x32xf32> {
8 %empty = tensor.empty() : tensor<8x32xf32>
9 …%0 = tensor.pack %arg0 inner_dims_pos = [0] inner_tiles = [32] into %empty : tensor<256xf32> -> te…
10 return %0 : tensor<8x32xf32>
16 // CHECK-SAME: %[[ARG0:.+]]: tensor<255xf32>)
17 // CHECK-NOT: tensor.expand_shape
18 // CHECK: tensor.pack
19 func.func @single_dim_packing_with_padding(%arg0: tensor<255xf32>) -> tensor<8x32xf32> {
20 %empty = tensor.empty() : tensor<8x32xf32>
22 …%0 = tensor.pack %arg0 padding_value(%cst : f32) inner_dims_pos = [0] inner_tiles = [32] into %emp…
23 return %0 : tensor<8x32xf32>
29 // CHECK-SAME: %[[ARG0:.+]]: tensor<5x256xf32>)
30 … %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2]] output_shape [5, 8, 32] :…
31 // CHECK: return %[[EXPANDED]] : tensor<5x8x32xf32>
32 func.func @single_last_inner_dim_packing(%arg0: tensor<5x256xf32>) -> tensor<5x8x32xf32> {
33 %empty = tensor.empty() : tensor<5x8x32xf32>
34 …%0 = tensor.pack %arg0 inner_dims_pos = [1] inner_tiles = [32] into %empty : tensor<5x256xf32> -> …
35 return %0 : tensor<5x8x32xf32>
41 // CHECK-SAME: %[[ARG0:.+]]: tensor<64xf32>)
42 …K: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1]] output_shape [2, 32] : t…
43 // CHECK: return %[[EXPANDED]] : tensor<2x32xf32>
44 func.func @pack_1d_with_outer_dims_perm(%arg0: tensor<64xf32>) -> tensor<2x32xf32> {
45 %empty = tensor.empty() : tensor<2x32xf32>
46 …%pack = tensor.pack %arg0 outer_dims_perm = [0] inner_dims_pos = [0] inner_tiles = [32] into %empt…
47 return %pack : tensor<2x32xf32>
53 // CHECK-SAME: %[[ARG0:.+]]: tensor<5x256xf32>)
54 … %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2]] output_shape [5, 8, 32] :…
55 // CHECK: return %[[EXPANDED]] : tensor<5x8x32xf32>
56 …e_last_inner_dim_packing_with_identity_outer_dims_perm(%arg0: tensor<5x256xf32>) -> tensor<5x8x32x…
57 %empty = tensor.empty() : tensor<5x8x32xf32>
58 …%0 = tensor.pack %arg0 outer_dims_perm = [0, 1] inner_dims_pos = [1] inner_tiles = [32] into %empt…
59 return %0 : tensor<5x8x32xf32>
65 // CHECK-NOT: tensor.expand_shape
66 // CHECK: tensor.pack
67 func.func @packing_with_outer_dims_perm(%arg0: tensor<5x256xf32>) -> tensor<8x5x32xf32> {
68 %empty = tensor.empty() : tensor<8x5x32xf32>
69 …%0 = tensor.pack %arg0 outer_dims_perm = [1, 0] inner_dims_pos = [1] inner_tiles = [32] into %empt…
70 return %0 : tensor<8x5x32xf32>
76 // CHECK-NOT: tensor.expand_shape
77 // CHECK: tensor.pack
78 func.func @single_first_inner_dim_packing(%arg0: tensor<256x5xf32>) -> tensor<8x5x32xf32> {
79 %empty = tensor.empty() : tensor<8x5x32xf32>
80 …%0 = tensor.pack %arg0 inner_dims_pos = [0] inner_tiles = [32] into %empty : tensor<256x5xf32> -> …
81 return %0 : tensor<8x5x32xf32>
88 // CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2, 3]] output_sha…
90 func.func @pack_1x32_to_1x32x1x1(%arg0 : tensor<1x32xf32>) -> tensor<1x32x1x1xf32> {
91 %empty = tensor.empty() : tensor<1x32x1x1xf32>
92 %pack = tensor.pack %arg0 inner_dims_pos = [0, 1] inner_tiles = [1, 1] into %empty
93 : tensor<1x32xf32> -> tensor<1x32x1x1xf32>
94 return %pack : tensor<1x32x1x1xf32>
101 // CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2, 3]] output_sha…
103 func.func @pack_1x32_to_1x16x1x2(%arg0 : tensor<1x32xf32>) -> tensor<1x16x1x2xf32> {
104 %empty = tensor.empty() : tensor<1x16x1x2xf32>
105 %pack = tensor.pack %arg0 inner_dims_pos = [0, 1] inner_tiles = [1, 2] into %empty
106 : tensor<1x32xf32> -> tensor<1x16x1x2xf32>
107 return %pack : tensor<1x16x1x2xf32>
114 // CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1, 2], [3]] output_sha…
116 func.func @pack_32x1_to_16x1x2x1(%arg0 : tensor<32x1xf32>) -> tensor<1x16x2x1xf32> {
117 %empty = tensor.empty() : tensor<1x16x2x1xf32>
118 …%pack = tensor.pack %arg0 outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [2, 1] in…
119 : tensor<32x1xf32> -> tensor<1x16x2x1xf32>
120 return %pack : tensor<1x16x2x1xf32>
126 // CHECK-NOT: tensor.expand_shape
127 // CHECK: tensor.pack
128 func.func @pack_32x1_to_16x1x1x2(%arg0 : tensor<32x1xf32>) -> tensor<16x1x1x2xf32> {
129 %empty = tensor.empty() : tensor<16x1x1x2xf32>
130 %pack = tensor.pack %arg0 inner_dims_pos = [1, 0] inner_tiles = [1, 2] into %empty
131 : tensor<32x1xf32> -> tensor<16x1x1x2xf32>
132 return %pack : tensor<16x1x1x2xf32>
138 // CHECK-SAME: %[[ARG0:.+]]: tensor<8x32xf32>)
139 // CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1]] : tensor<8x32xf…
141 func.func @unpack_1d_to_collapse(%arg0: tensor<8x32xf32>) -> tensor<256xf32> {
142 %empty = tensor.empty() : tensor<256xf32>
143 …%0 = tensor.unpack %arg0 inner_dims_pos = [0] inner_tiles = [32] into %empty : tensor<8x32xf32> ->…
144 return %0 : tensor<256xf32>
150 // CHECK-NOT: tensor.collapse
151 // CHECK: tensor.unpack
152 func.func @unpack_to_partial_slice(%arg0: tensor<8x32xf32>) -> tensor<255xf32> {
153 %empty = tensor.empty() : tensor<255xf32>
154 …%0 = tensor.unpack %arg0 inner_dims_pos = [0] inner_tiles = [32] into %empty : tensor<8x32xf32> ->…
155 return %0 : tensor<255xf32>
161 // CHECK-NOT: tensor.collapse
162 // CHECK: tensor.unpack
163 func.func @unpack_dynamic(%arg0: tensor<?x32xf32>) -> tensor<?xf32> {
166 %d0 = tensor.dim %arg0, %c0 : tensor<?x32xf32>
168 %empty = tensor.empty(%size) : tensor<?xf32>
169 …%0 = tensor.unpack %arg0 inner_dims_pos = [0] inner_tiles = [32] into %empty : tensor<?x32xf32> ->…
170 return %0 : tensor<?xf32>
176 // CHECK-SAME: %[[ARG0:.+]]: tensor<5x8x32xf32>)
177 …CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2]] : tensor<5x8…
178 // CHECK: return %[[COLLAPSED]] : tensor<5x256xf32>
179 func.func @single_last_inner_dim_unpacking(%arg0: tensor<5x8x32xf32>) -> tensor<5x256xf32> {
180 %empty = tensor.empty() : tensor<5x256xf32>
181 …%0 = tensor.unpack %arg0 inner_dims_pos = [1] inner_tiles = [32] into %empty : tensor<5x8x32xf32> …
182 return %0 : tensor<5x256xf32>
188 // CHECK-SAME: %[[ARG0:.+]]: tensor<5x8x32xf32>)
189 …CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2]] : tensor<5x8…
190 // CHECK: return %[[COLLAPSED]] : tensor<5x256xf32>
191 …last_inner_dim_unpacking_with_identity_outer_dims_perm(%arg0: tensor<5x8x32xf32>) -> tensor<5x256x…
192 %empty = tensor.empty() : tensor<5x256xf32>
193 …%0 = tensor.unpack %arg0 outer_dims_perm = [0, 1] inner_dims_pos = [1] inner_tiles = [32] into %em…
194 return %0 : tensor<5x256xf32>
200 // CHECK-NOT: tensor.collpase_shape
201 // CHECK: tensor.unpack
202 func.func @unpacking_with_outer_dims_perm(%arg0: tensor<8x5x32xf32>) -> tensor<5x256xf32> {
203 %empty = tensor.empty() : tensor<5x256xf32>
204 …%0 = tensor.unpack %arg0 outer_dims_perm = [1, 0] inner_dims_pos = [1] inner_tiles = [32] into %em…
205 return %0 : tensor<5x256xf32>
211 // CHECK-NOT: tensor.collapse_shape
212 // CHECK: tensor.unpack
213 func.func @single_first_inner_dim_unpacking(%arg0: tensor<8x5x32xf32>) -> tensor<256x5xf32> {
214 %empty = tensor.empty() : tensor<256x5xf32>
215 …%0 = tensor.unpack %arg0 inner_dims_pos = [0] inner_tiles = [32] into %empty : tensor<8x5x32xf32> …
216 return %0 : tensor<256x5xf32>
223 // CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2, 3]]
225 func.func @unpack_1x32x1x1_to_1x32(%arg0 : tensor<1x32x1x1xf32>) -> tensor<1x32xf32> {
226 %empty = tensor.empty() : tensor<1x32xf32>
227 %unpack = tensor.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [1, 1] into %empty
228 : tensor<1x32x1x1xf32> -> tensor<1x32xf32>
229 return %unpack : tensor<1x32xf32>
236 // CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2, 3]]
238 func.func @unpack_1x2x1x16_to_1x32(%arg0 : tensor<1x2x1x16xf32>) -> tensor<1x32xf32> {
239 %empty = tensor.empty() : tensor<1x32xf32>
240 …%unpack = tensor.unpack %arg0 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [1, 1…
241 : tensor<1x2x1x16xf32> -> tensor<1x32xf32>
242 return %unpack : tensor<1x32xf32>
249 // CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1, 2], [3]]
251 func.func @unpack_16x1x2x1_to_32x1(%arg0 : tensor<1x16x2x1xf32>) -> tensor<32x1xf32> {
252 %empty = tensor.empty() : tensor<32x1xf32>
253 …%unpack = tensor.unpack %arg0 outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [2, 1…
254 : tensor<1x16x2x1xf32> -> tensor<32x1xf32>
255 return %unpack : tensor<32x1xf32>
261 // CHECK-NOT: tensor.collapse_shape
262 // CHECK: tensor.unpack
263 func.func @unpack_16x1x1x2_to_32x1(%arg0 : tensor<16x1x1x2xf32>) -> tensor<32x1xf32> {
264 %empty = tensor.empty() : tensor<32x1xf32>
265 %unpack = tensor.unpack %arg0 inner_dims_pos = [1, 0] inner_tiles = [1, 2] into %empty
266 : tensor<16x1x1x2xf32> -> tensor<32x1xf32>
267 return %unpack : tensor<32x1xf32>
273 // CHECK-SAME: %[[ARG0:.+]]: tensor<32x64xf32>)
274 … %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1, 2], [3]] output_shape [1, 1, 32, 6…
275 // CHECK: return %[[EXPANDED]] : tensor<1x1x32x64xf32>
276 func.func @pad_like_pack(%arg0: tensor<32x64xf32>) -> tensor<1x1x32x64xf32> {
277 %empty = tensor.empty() : tensor<1x1x32x64xf32>
278 …%0 = tensor.pack %arg0 inner_dims_pos = [0, 1] inner_tiles = [32, 64] into %empty : tensor<32x64xf…
279 return %0 : tensor<1x1x32x64xf32>
285 // CHECK-SAME: %[[ARG0:.+]]: tensor<32x64xf32>)
286 … %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1, 2], [3]] output_shape [1, 1, 32, 6…
287 // CHECK: return %[[EXPANDED]] : tensor<1x1x32x64xf32>
288 func.func @pad_like_pack_with_outer_dims_perm(%arg0: tensor<32x64xf32>) -> tensor<1x1x32x64xf32> {
289 %empty = tensor.empty() : tensor<1x1x32x64xf32>
290 …%0 = tensor.pack %arg0 outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [32, 64] int…
291 return %0 : tensor<1x1x32x64xf32>
297 // CHECK-SAME: %[[ARG0:.+]]: tensor<32x64xf32>)
298 … %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2]] output_shape [32, 1, 64] …
299 // CHECK: return %[[EXPANDED]] : tensor<32x1x64xf32>
300 func.func @inner_pad_like_pack(%arg0: tensor<32x64xf32>) -> tensor<32x1x64xf32> {
301 %empty = tensor.empty() : tensor<32x1x64xf32>
302 …%0 = tensor.pack %arg0 inner_dims_pos = [1] inner_tiles = [64] into %empty : tensor<32x64xf32> -> …
303 return %0 : tensor<32x1x64xf32>
310 // CHECK-SAME: %[[ARG0:.+]]: tensor<32x64xf32>)
311 // CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<1x1x64x32xf32>
312 … %[[PACK:.+]] = tensor.pack %[[ARG0]] inner_dims_pos = [1, 0] inner_tiles = [64, 32] into %[[EM…
313 // CHECK: return %[[PACK]] : tensor<1x1x64x32xf32>
314 func.func @pad_and_inner_dim_shuffle_pack(%arg0: tensor<32x64xf32>) -> tensor<1x1x64x32xf32> {
315 %empty = tensor.empty() : tensor<1x1x64x32xf32>
316 …%0 = tensor.pack %arg0 inner_dims_pos = [1, 0] inner_tiles = [64, 32] into %empty : tensor<32x64xf…
317 return %0 : tensor<1x1x64x32xf32>
324 // CHECK-SAME: %[[ARG0:.+]]: tensor<32x64x16xf32>)
325 // CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<32x1x16x64xf32>
326 … %[[PACK:.+]] = tensor.pack %[[ARG0]] inner_dims_pos = [1] inner_tiles = [64] into %[[EMPTY]]…
327 // CHECK: return %[[PACK]] : tensor<32x1x16x64xf32>
328 func.func @pad_like_pack_with_transpose(%arg0: tensor<32x64x16xf32>) -> tensor<32x1x16x64xf32> {
329 %empty = tensor.empty() : tensor<32x1x16x64xf32>
330 …%0 = tensor.pack %arg0 inner_dims_pos = [1] inner_tiles = [64] into %empty : tensor<32x64x16xf32> …
331 return %0 : tensor<32x1x16x64xf32>
337 // CHECK-SAME: %[[ARG0:.+]]: tensor<1x1x32x64xf32>)
338 …CK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1, 2], [3]] : tensor<1x1…
339 // CHECK: return %[[COLLAPSED]] : tensor<32x64xf32>
340 func.func @unpad_like_unpack(%arg0: tensor<1x1x32x64xf32>) -> tensor<32x64xf32> {
341 %empty = tensor.empty() : tensor<32x64xf32>
342 …%0 = tensor.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [32, 64] into %empty : tensor<1x1x3…
343 return %0 : tensor<32x64xf32>
349 // CHECK-SAME: %[[ARG0:.+]]: tensor<1x1x32x64xf32>)
350 …CK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1, 2], [3]] : tensor<1x1…
351 // CHECK: return %[[COLLAPSED]] : tensor<32x64xf32>
352 func.func @unpad_like_unpack_with_outer_dims_perm(%arg0: tensor<1x1x32x64xf32>) -> tensor<32x64xf32…
353 %empty = tensor.empty() : tensor<32x64xf32>
354 …%0 = tensor.unpack %arg0 outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [32, 64] i…
355 return %0 : tensor<32x64xf32>
361 // CHECK-SAME: %[[ARG0:.+]]: tensor<32x1x64xf32>)
362 …HECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2]] : tensor<32x1…
363 // CHECK: return %[[COLLAPSED]] : tensor<32x64xf32>
364 func.func @inner_unpad_like_unpack(%arg0: tensor<32x1x64xf32>) -> tensor<32x64xf32> {
365 %empty = tensor.empty() : tensor<32x64xf32>
366 …%0 = tensor.unpack %arg0 inner_dims_pos = [1] inner_tiles = [64] into %empty : tensor<32x1x64xf32>…
367 return %0 : tensor<32x64xf32>
374 // CHECK-SAME: %[[ARG0:.+]]: tensor<1x1x32x64xf32>)
375 // CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<64x32xf32>
376 …[[UNPACK:.+]] = tensor.unpack %[[ARG0]] inner_dims_pos = [1, 0] inner_tiles = [32, 64] into %[[EMP…
377 // CHECK: return %[[UNPACK]] : tensor<64x32xf32>
378 func.func @unpad_and_inner_dim_shuffle_pack(%arg0: tensor<1x1x32x64xf32>) -> tensor<64x32xf32> {
379 %empty = tensor.empty() : tensor<64x32xf32>
380 …%0 = tensor.unpack %arg0 inner_dims_pos = [1, 0] inner_tiles = [32, 64] into %empty : tensor<1x1x3…
381 return %0 : tensor<64x32xf32>
388 // CHECK-SAME: %[[ARG0:.+]]: tensor<32x1x16x64xf32>)
389 // CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<32x64x16xf32>
390 … %[[UNPACK:.+]] = tensor.unpack %[[ARG0]] inner_dims_pos = [1] inner_tiles = [64] into %[[EMPTY]]…
391 // CHECK: return %[[UNPACK]] : tensor<32x64x16xf32>
392 func.func @unpad_like_unpack_with_transpose(%arg0: tensor<32x1x16x64xf32>) -> tensor<32x64x16xf32> {
393 %empty = tensor.empty() : tensor<32x64x16xf32>
394 …%0 = tensor.unpack %arg0 inner_dims_pos = [1] inner_tiles = [64] into %empty : tensor<32x1x16x64xf…
395 return %0 : tensor<32x64x16xf32>