Lines Matching full:tensor
1 // RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-fold-into-pack-and-unpack %s | FileCheck %s
3 func.func @fold_unpack_slice(%arg0 : tensor<?x?x8x4xf32>, %arg1 : tensor<?x?xf32>,
4 %arg2 : index, %arg3 : index) -> tensor<?x?xf32> {
5 %0 = tensor.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [8, 4] into %arg1
6 : tensor<?x?x8x4xf32> -> tensor<?x?xf32>
7 %1 = tensor.extract_slice %0[0, 0] [%arg2, %arg3] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
8 return %1 : tensor<?x?xf32>
11 // CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x8x4xf32>
12 // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32>
15 // CHECK: %[[INIT:.+]] = tensor.empty(%[[ARG2]], %[[ARG3]]) : tensor<?x?xf32>
16 // CHECK: %[[UNPACK:.+]] = tensor.unpack %[[ARG0]] inner_dims_pos = [0, 1] inner_tiles = [8, 4]
22 func.func @nofold_unpack_slice_non_zero_offset(%arg0 : tensor<?x?x8x4xf32>, %arg1 : tensor<?x?xf32>,
23 %arg2 : index, %arg3 : index, %arg4 : index) -> tensor<?x?xf32> {
24 %0 = tensor.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [8, 4] into %arg1
25 : tensor<?x?x8x4xf32> -> tensor<?x?xf32>
26 %1 = tensor.extract_slice %0[0, %arg4] [%arg2, %arg3] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
27 return %1 : tensor<?x?xf32>
30 // CHECK: %[[UNPACK:.+]] = tensor.unpack
31 // CHECK: tensor.extract_slice %[[UNPACK]]
35 func.func @nofold_unpack_slice_non_unit_stride(%arg0 : tensor<?x?x8x4xf32>, %arg1 : tensor<?x?xf32>,
36 %arg2 : index, %arg3 : index, %arg4 : index) -> tensor<?x?xf32> {
37 %0 = tensor.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [8, 4] into %arg1
38 : tensor<?x?x8x4xf32> -> tensor<?x?xf32>
39 %1 = tensor.extract_slice %0[0, 0] [%arg2, %arg3] [%arg4, 1] : tensor<?x?xf32> to tensor<?x?xf32>
40 return %1 : tensor<?x?xf32>
43 // CHECK: %[[UNPACK:.+]] = tensor.unpack
44 // CHECK: tensor.extract_slice %[[UNPACK]]
48 func.func @nofold_unpack_slice_rank_reduced(%arg0 : tensor<?x?x8x4xf32>, %arg1 : tensor<?x?xf32>,
49 %arg2 : index, %arg3 : index) -> tensor<f32> {
50 %0 = tensor.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [8, 4] into %arg1
51 : tensor<?x?x8x4xf32> -> tensor<?x?xf32>
52 %1 = tensor.extract_slice %0[0, 0] [1, 1] [1, 1] : tensor<?x?xf32> to tensor<f32>
53 return %1 : tensor<f32>
56 // CHECK: %[[UNPACK:.+]] = tensor.unpack
57 // CHECK: tensor.extract_slice %[[UNPACK]]
61 func.func @pad_pack(%src: tensor<16641x16xf32>) -> tensor<2082x1x8x32xf32> {
64 %padded = tensor.pad %src low[0, 0] high[15, 0] {
66 tensor.yield %cst : f32
67 } : tensor<16641x16xf32> to tensor<16656x16xf32>
68 %empty = tensor.empty() : tensor<2082x1x8x32xf32>
69 %pack = tensor.pack %padded padding_value(%cst : f32) inner_dims_pos = [0, 1] inner_tiles = [8, 32] into %empty
70 : tensor<16656x16xf32> -> tensor<2082x1x8x32xf32>
71 return %pack : tensor<2082x1x8x32xf32>
76 // CHECK: %[[DEST:.+]] = tensor.empty() : tensor<2082x1x8x32xf32>
77 // CHECK: %[[PACK:.+]] = tensor.pack %[[SRC]]
83 func.func @nofold_pad_pack(%src: tensor<16641x16xf32>) -> tensor<2082x1x8x32xf32> {
86 %padded = tensor.pad %src nofold low[0, 0] high[15, 0] {
88 tensor.yield %cst : f32
89 } : tensor<16641x16xf32> to tensor<16656x16xf32>
90 %empty = tensor.empty() : tensor<2082x1x8x32xf32>
91 %pack = tensor.pack %padded padding_value(%cst : f32) inner_dims_pos = [0, 1] inner_tiles = [8, 32] into %empty
92 : tensor<16656x16xf32> -> tensor<2082x1x8x32xf32>
93 return %pack : tensor<2082x1x8x32xf32>
96 // CHECK: tensor.pad
97 // CHECK: tensor.pack
101 func.func @pad_pack_different_padding_value(%src: tensor<16641x16xf32>) -> tensor<2082x1x8x32xf32> {
105 %padded = tensor.pad %src low[0, 0] high[15, 0] {
107 tensor.yield %cst0 : f32
108 } : tensor<16641x16xf32> to tensor<16656x16xf32>
109 %empty = tensor.empty() : tensor<2082x1x8x32xf32>
110 %pack = tensor.pack %padded padding_value(%cst1 : f32) inner_dims_pos = [0, 1] inner_tiles = [8, 32] into %empty
111 : tensor<16656x16xf32> -> tensor<2082x1x8x32xf32>
112 return %pack : tensor<2082x1x8x32xf32>
115 // CHECK: tensor.pad
116 // CHECK: tensor.pack
120 func.func @tensor_pack_linalg_transpose_fold(%arg0: tensor<56x57x1x64xf32>) -> tensor<1x57x56x2x32xf32> {
121 %0 = tensor.empty() : tensor<56x2x1x57x32xf32>
122 %pack = tensor.pack %arg0
126 into %0 : tensor<56x57x1x64xf32> -> tensor<56x2x1x57x32xf32>
128 %1 = tensor.empty() : tensor<1x57x56x2x32xf32>
130 ins(%pack : tensor<56x2x1x57x32xf32>)
131 outs(%1 : tensor<1x57x56x2x32xf32>)
133 return %transposed : tensor<1x57x56x2x32xf32>
136 // CHECK-SAME: %[[ARG0:.+]]: tensor<56x57x1x64xf32>)
137 // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x57x56x2x32xf32>
138 // CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]]
146 func.func @tensor_pack_linalg_transpose_fold_with_padding(%arg0: tensor<56x57x1x55xf32>, %padding: f32) -> tensor<1x57x56x2x32xf32> {
147 %0 = tensor.empty() : tensor<56x2x1x57x32xf32>
148 %pack = tensor.pack %arg0 padding_value(%padding : f32)
152 into %0 : tensor<56x57x1x55xf32> -> tensor<56x2x1x57x32xf32>
154 %1 = tensor.empty() : tensor<1x57x56x2x32xf32>
156 ins(%pack : tensor<56x2x1x57x32xf32>)
157 outs(%1 : tensor<1x57x56x2x32xf32>)
159 return %transposed : tensor<1x57x56x2x32xf32>
162 // CHECK-SAME: %[[ARG0:.+]]: tensor<56x57x1x55xf32>, %[[PADDING:.+]]: f32)
163 // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x57x56x2x32xf32>
164 // CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]] padding_value(%[[PADDING]] : f32)
172 func.func @tensor_pack_linalg_transpose_fold_no_outer_dims_perm(%arg0: tensor<56x57x1x64xf32>) -> tensor<1x2x56x57x32xf32> {
173 %0 = tensor.empty() : tensor<56x57x1x2x32xf32>
174 %pack = tensor.pack %arg0
177 into %0 : tensor<56x57x1x64xf32> -> tensor<56x57x1x2x32xf32>
179 %1 = tensor.empty() : tensor<1x2x56x57x32xf32>
181 ins(%pack : tensor<56x57x1x2x32xf32>)
182 outs(%1 : tensor<1x2x56x57x32xf32>)
184 return %transposed : tensor<1x2x56x57x32xf32>
187 // CHECK-SAME: %[[ARG0:.+]]: tensor<56x57x1x64xf32>)
188 // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x2x56x57x32xf32>
189 // CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]]
197 func.func @tensor_pack_linalg_transpose_fold_tile_dims_transpose(%arg0: tensor<56x72x24x128xf32>) -> tensor<12x56x4x9x32x8x2xf32> {
198 %0 = tensor.empty() : tensor<4x9x12x56x8x2x32xf32>
199 %pack = tensor.pack %arg0
203 into %0 : tensor<56x72x24x128xf32> -> tensor<4x9x12x56x8x2x32xf32>
205 %1 = tensor.empty() : tensor<12x56x4x9x32x8x2xf32>
207 ins(%pack : tensor<4x9x12x56x8x2x32xf32>)
208 outs(%1 : tensor<12x56x4x9x32x8x2xf32>)
210 return %transposed : tensor<12x56x4x9x32x8x2xf32>
213 // CHECK-SAME: %[[ARG0:.+]]: tensor<56x72x24x128xf32>)
214 // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<12x56x4x9x32x8x2xf32>
215 // CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]]
223 func.func @tensor_pack_linalg_transpose_fold_tile_dims_outer_dims_transpose(%arg0: tensor<56x72x24x128xf32>) -> tensor<9x56x2x12x32x8x4xf32> {
224 %0 = tensor.empty() : tensor<4x12x9x56x8x2x32xf32>
225 %pack = tensor.pack %arg0
229 into %0 : tensor<56x72x24x128xf32> -> tensor<4x12x9x56x8x2x32xf32>
231 %1 = tensor.empty() : tensor<9x56x2x12x32x8x4xf32>
233 ins(%pack : tensor<4x12x9x56x8x2x32xf32>)
234 outs(%1 : tensor<9x56x2x12x32x8x4xf32>)
236 return %transposed : tensor<9x56x2x12x32x8x4xf32>
239 // CHECK-SAME: %[[ARG0:.+]]: tensor<56x72x24x128xf32>)
240 // CHECK: tensor.pack
245 func.func @tensor_pack_linalg_transpose_fold_dynamic_outer_dims(%arg0: tensor<56x?x?x64xf32>) -> tensor<?x?x56x2x32xf32> {
246 %0 = tensor.empty() : tensor<56x2x1x57x32xf32>
247 %pack = tensor.pack %arg0
251 into %0 : tensor<56x?x?x64xf32> -> tensor<56x2x1x57x32xf32>
253 %1 = tensor.empty() : tensor<1x57x56x2x32xf32>
255 ins(%pack : tensor<56x2x1x57x32xf32>)
256 outs(%1 : tensor<1x57x56x2x32xf32>)
259 %return_value = tensor.cast %transposed : tensor<1x57x56x2x32xf32> to tensor<?x?x56x2x32xf32>
260 return %return_value : tensor<?x?x56x2x32xf32>
263 // CHECK-SAME: %[[ARG0:.+]]: tensor<56x?x?x64xf32>)
266 // CHECK: %[[dim:.+]] = tensor.dim %[[ARG0]], %[[c1]] : tensor<56x?x?x64xf32>
267 // CHECK: %[[dim_0:.+]] = tensor.dim %[[ARG0]], %[[c2]] : tensor<56x?x?x64xf32>
268 // CHECK: %[[INIT:.+]] = tensor.empty(%[[dim_0]], %[[dim]]) : tensor<?x?x56x2x32xf32>
269 // CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]]
277 func.func @tensor_pack_linalg_transpose_fold_dynamic_outer_and_tile_dims(%arg0: tensor<56x?x?x128xf32>) -> tensor<?x?x56x9x32x8x2xf32> {
278 %0 = tensor.empty() : tensor<56x9x12x4x8x2x32xf32>
279 %pack = tensor.pack %arg0
282 into %0 : tensor<56x?x?x128xf32> -> tensor<56x9x12x4x8x2x32xf32>
284 %1 = tensor.empty() : tensor<12x4x56x9x32x8x2xf32>
286 ins(%pack : tensor<56x9x12x4x8x2x32xf32>)
287 outs(%1 : tensor<12x4x56x9x32x8x2xf32>)
290 %return_value = tensor.cast %transposed : tensor<12x4x56x9x32x8x2xf32> to tensor<?x?x56x9x32x8x2xf32>
291 return %return_value : tensor<?x?x56x9x32x8x2xf32>
296 // CHECK-SAME: %[[ARG0:.+]]: tensor<56x?x?x128xf32>)
299 // CHECK: %[[dim:.+]] = tensor.dim %[[ARG0]], %[[c1]] : tensor<56x?x?x128xf32>
300 // CHECK: %[[dim_0:.+]] = tensor.dim %[[ARG0]], %[[c2]] : tensor<56x?x?x128xf32>
303 // CHECK: %[[INIT:.+]] = tensor.empty(%[[mapped_dim2]], %[[mapped_dim1]]) : tensor<?x4x56x?x32x8x2xf32>
304 // CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]] outer_dims_perm = [2, 3, 0, 1] inner_dims_pos = [3, 1, 2] inner_tiles = [32, 8, 2] into %[[INIT]] : tensor<56x?x?x128xf32> -> tensor<?x4x56x?x32x8x2xf32>
305 // CHECK: %[[CAST:.+]] = tensor.cast %[[PACK]] : tensor<?x4x56x?x32x8x2xf32> to tensor<?x?x56x9x32x8x2xf32>
306 // CHECK: return %[[CAST]] : tensor<?x?x56x9x32x8x2xf32>
311 func.func @tensor_pack_linalg_transpose_fold_dynamic_outer_dims_tile_dims_tile_sizes(%arg0: tensor<?x?x?x?xf32>, %pack_dest: tensor<?x?x?x?x?x?x?xf32>, %transpose_dest: tensor<?x?x?x?x?x?x?xf32>, %tile_p : index, %tile_q : index, %tile_r : index) -> tensor<?x?x?x?x?x?x?xf32> {
312 %pack = tensor.pack %arg0
316 into %pack_dest : tensor<?x?x?x?xf32> -> tensor<?x?x?x?x?x?x?xf32>
319 ins(%pack : tensor<?x?x?x?x?x?x?xf32>)
320 outs(%transpose_dest : tensor<?x?x?x?x?x?x?xf32>)
323 return %transposed : tensor<?x?x?x?x?x?x?xf32>
328 // CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?xf32>,
329 // CHECK-SAME: %[[PACK_DEST:.+]]: tensor<?x?x?x?x?x?x?xf32>, %[[TRANSPOSE_DEST:.+]]: tensor<?x?x?x?x?x?x?xf32>,
336 // CHECK: %[[dim:.+]] = tensor.dim %[[ARG0]], %[[c0]] : tensor<?x?x?x?xf32>
337 // CHECK: %[[dim_0:.+]] = tensor.dim %[[ARG0]], %[[c1]] : tensor<?x?x?x?xf32>
338 // CHECK: %[[dim_1:.+]] = tensor.dim %[[ARG0]], %[[c2]] : tensor<?x?x?x?xf32>
339 // CHECK: %[[dim_2:.+]] = tensor.dim %[[ARG0]], %[[c3]] : tensor<?x?x?x?xf32>
343 // CHECK: %[[INIT:.+]] = tensor.empty(%[[mapped_dim2]], %[[mapped_dim1]], %[[mapped_dim0]], %[[dim]], %[[ARG3]], %[[ARG1]], %[[ARG2]]) : tensor<?x?x?x?x?x?x?xf32>
344 // CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]] outer_dims_perm = [2, 1, 3, 0] inner_dims_pos = [3, 1, 2] inner_tiles = [%[[ARG3]], %[[ARG1]], %[[ARG2]]] into %[[INIT]] : tensor<?x?x?x?xf32> -> tensor<?x?x?x?x?x?x?xf32>
345 // CHECK: return %[[PACK]] : tensor<?x?x?x?x?x?x?xf32>
350 func.func @linalg_transpose_tensor_pack_fold(%arg0: tensor<56x57x1x64xf32>) -> tensor<1x57x56x2x32xf32> {
351 %0 = tensor.empty() : tensor<1x56x57x64xf32>
353 ins(%arg0 : tensor<56x57x1x64xf32>)
354 outs(%0 : tensor<1x56x57x64xf32>)
357 %1 = tensor.empty() : tensor<1x57x56x2x32xf32>
358 %pack = tensor.pack %transposed
362 into %1 : tensor<1x56x57x64xf32> -> tensor<1x57x56x2x32xf32>
363 return %pack : tensor<1x57x56x2x32xf32>
366 // CHECK-SAME: %[[ARG0:.+]]: tensor<56x57x1x64xf32>)
367 // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x57x56x2x32xf32>
368 // CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]]
376 func.func @linalg_transpose_tensor_pack_fold_with_padding(%arg0: tensor<56x57x1x55xf32>, %padding: f32) -> tensor<1x57x56x2x32xf32> {
377 %0 = tensor.empty() : tensor<1x56x57x55xf32>
379 ins(%arg0 : tensor<56x57x1x55xf32>)
380 outs(%0 : tensor<1x56x57x55xf32>)
383 %1 = tensor.empty() : tensor<1x57x56x2x32xf32>
384 %pack = tensor.pack %transpose padding_value(%padding : f32)
388 into %1 : tensor<1x56x57x55xf32> -> tensor<1x57x56x2x32xf32>
389 return %pack : tensor<1x57x56x2x32xf32>
392 // CHECK-SAME: %[[ARG0:.+]]: tensor<56x57x1x55xf32>, %[[PADDING:.+]]: f32)
393 // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x57x56x2x32xf32>
394 // CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]] padding_value(%[[PADDING]] : f32)
402 func.func @linalg_transpose_tensor_pack_fold_no_outer_dims_perm(%arg0: tensor<56x57x1x64xf32>) -> tensor<1x56x57x2x32xf32> {
403 %0 = tensor.empty() : tensor<1x56x57x64xf32>
405 ins(%arg0 : tensor<56x57x1x64xf32>)
406 outs(%0 : tensor<1x56x57x64xf32>)
409 %1 = tensor.empty() : tensor<1x56x57x2x32xf32>
410 %pack = tensor.pack %transposed
413 into %1 : tensor<1x56x57x64xf32> -> tensor<1x56x57x2x32xf32>
414 return %pack : tensor<1x56x57x2x32xf32>
417 // CHECK-SAME: %[[ARG0:.+]]: tensor<56x57x1x64xf32>)
418 // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x56x57x2x32xf32>
419 // CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]]
427 func.func @linalg_transpose_tensor_pack_fold_complex_inner_dims_change(%arg0: tensor<25x30x35x40xf32>, %transpose_dest: tensor<35x40x25x30xf32>, %pack_dest: tensor<3x35x5x8x5x10x5xf32>) -> tensor<3x35x5x8x5x10x5xf32> {
429 ins(%arg0 : tensor<25x30x35x40xf32>)
430 outs(%transpose_dest : tensor<35x40x25x30xf32>)
433 %pack = tensor.pack %transposed
437 into %pack_dest : tensor<35x40x25x30xf32> -> tensor<3x35x5x8x5x10x5xf32>
438 return %pack : tensor<3x35x5x8x5x10x5xf32>
441 // CHECK-SAME: %[[ARG0:.+]]: tensor<25x30x35x40xf32>,
442 // CHECK-SAME: %[[ARG1:.+]]: tensor<35x40x25x30xf32>,
443 // CHECK-SAME: %[[ARG2:.+]]: tensor<3x35x5x8x5x10x5xf32>) -> tensor<3x35x5x8x5x10x5xf32> {
444 // CHECK: %[[VAL0:.+]] = tensor.empty() : tensor<3x35x5x8x5x10x5xf32>
445 // CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]]
454 func.func @linalg_transpose_tensor_pack_fold_dynamic_outer_dims_tile_dims_tile_sizes(%arg0: tensor<?x?x?x?xf32>, %transpose_dest: tensor<?x?x?x?xf32>, %pack_dest: tensor<?x?x?x?x?x?x?xf32>, %tile_p : index, %tile_q : index, %tile_r : index) -> tensor<?x?x?x?x?x?x?xf32> {
456 ins(%arg0 : tensor<?x?x?x?xf32>)
457 outs(%transpose_dest : tensor<?x?x?x?xf32>)
460 %pack = tensor.pack %transposed
464 into %pack_dest : tensor<?x?x?x?xf32> -> tensor<?x?x?x?x?x?x?xf32>
465 return %pack : tensor<?x?x?x?x?x?x?xf32>
469 // CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?xf32>, %[[ARG1:.+]]: tensor<?x?x?x?xf32>,
470 // CHECK-SAME: %[[ARG2:.+]]: tensor<?x?x?x?x?x?x?xf32>, %[[ARG3:.+]]: index, %[[ARG4:.+]]: index, %[[ARG5:.+]]: index) -> tensor<?x?x?x?x?x?x?xf32> {
475 // CHECK: %[[DIM:.+]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x?x?x?xf32>
476 // CHECK: %[[DIM0:.+]] = tensor.dim %[[ARG0]], %[[C1]] : tensor<?x?x?x?xf32>
477 // CHECK: %[[DIM1:.+]] = tensor.dim %[[ARG0]], %[[C2]] : tensor<?x?x?x?xf32>
478 // CHECK: %[[DIM2:.+]] = tensor.dim %[[ARG0]], %[[C3]] : tensor<?x?x?x?xf32>
482 // CHECK: %[[VAL3:.+]] = tensor.empty(%[[VAL1]], %[[DIM1]], %[[VAL2]], %[[VAL0]], %[[ARG3]], %[[ARG4]], %[[ARG5]]) : tensor<?x?x?x?x?x?x?xf32>
483 // CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]] outer_dims_perm = [1, 2, 0, 3] inner_dims_pos = [3, 1, 0] inner_tiles = [%[[ARG3]], %[[ARG4]], %[[ARG5]]] into %[[VAL3]] : tensor<?x?x?x?xf32> -> tensor<?x?x?x?x?x?x?xf32>
484 // CHECK: return %[[PACK]] : tensor<?x?x?x?x?x?x?xf32>
488 func.func @linalg_transpose_tensor_pack_multiple_tiles(%arg0: tensor<?x32x128xbf16>) -> tensor<32x?x64x16x2xbf16> {
491 %dim = tensor.dim %arg0, %c0 : tensor<?x32x128xbf16>
493 %0 = tensor.empty(%dim) : tensor<32x128x?xbf16>
495 ins(%arg0 : tensor<?x32x128xbf16>)
496 outs(%0 : tensor<32x128x?xbf16>)
499 %2 = tensor.empty(%dim) : tensor<32x?x64x16x2xbf16>
500 %pack = tensor.pack %transposed
505 into %2 : tensor<32x128x?xbf16> -> tensor<32x?x64x16x2xbf16>
506 return %pack : tensor<32x?x64x16x2xbf16>
510 // CHECK-SAME: %[[ARG0:.+]]: tensor<?x32x128xbf16>) -> tensor<32x?x64x16x2xbf16> {
513 // CHECK: %[[DIM:.+]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x32x128xbf16>
515 // CHECK: %[[VAL1:.+]] = tensor.empty(%[[VAL0]]) : tensor<32x?x64x16x2xbf16>
516 // CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]]
521 // CHECK-SAME: into %[[VAL1]] : tensor<?x32x128xbf16> -> tensor<32x?x64x16x2xbf16>
522 // CHECK: return %[[PACK]] : tensor<32x?x64x16x2xbf16>
527 func.func @linalg_transpose_tensor_unpack_fold(%arg0: tensor<1x1x4x16xi32>) -> tensor<16x4xi32> {
528 %0 = tensor.empty() : tensor<1x1x16x4xi32>
529 %transposed = linalg.transpose ins(%arg0 : tensor<1x1x4x16xi32>)
530 outs(%0 : tensor<1x1x16x4xi32>)
532 %1 = tensor.empty() : tensor<16x4xi32>
533 %unpack = tensor.unpack %transposed
537 %1 : tensor<1x1x16x4xi32> -> tensor<16x4xi32>
538 return %unpack : tensor<16x4xi32>
541 // CHECK-SAME: %[[ARG0:.+]]: tensor<1x1x4x16xi32>) -> tensor<16x4xi32> {
542 // CHECK: %[[OUT:.+]] = tensor.empty() : tensor<16x4xi32>
543 // CHECK: %[[UNPACK:.+]] = tensor.unpack %[[ARG0]]
547 // CHECK-SAME: into %[[OUT]] : tensor<1x1x4x16xi32> -> tensor<16x4xi32>
548 // CHECK: return %[[UNPACK]] : tensor<16x4xi32>
553 func.func @linalg_transpose_tensor_unpack_fold_partial_tile(%arg0: tensor<1x1x4x16xi32>) -> tensor<15x3xi32> {
554 %0 = tensor.empty() : tensor<1x1x16x4xi32>
555 %transposed = linalg.transpose ins(%arg0 : tensor<1x1x4x16xi32>)
556 outs(%0 : tensor<1x1x16x4xi32>)
558 %1 = tensor.empty() : tensor<15x3xi32>
559 %unpack = tensor.unpack %transposed
563 %1 : tensor<1x1x16x4xi32> -> tensor<15x3xi32>
564 return %unpack : tensor<15x3xi32>
567 // CHECK-SAME: %[[ARG0:.+]]: tensor<1x1x4x16xi32>) -> tensor<15x3xi32> {
568 // CHECK: %[[OUT:.+]] = tensor.empty() : tensor<15x3xi32>
569 // CHECK: %[[UNPACK:.+]] = tensor.unpack %[[ARG0]]
573 // CHECK-SAME: into %[[OUT]] : tensor<1x1x4x16xi32> -> tensor<15x3xi32>
574 // CHECK: return %[[UNPACK]] : tensor<15x3xi32>
579 func.func @linalg_transpose_tensor_unpack_fold_dynamic_outer_dims_tile_dims_tile_sizes(%arg0: tensor<?x?x?x?xf32>, %transpose_dest: tensor<?x?x?x?xf32>, %unpack_dest: tensor<?x?xf32>, %tile_p : index, %tile_q : index) -> tensor<?x?xf32> {
581 ins(%arg0 : tensor<?x?x?x?xf32>)
582 outs(%transpose_dest : tensor<?x?x?x?xf32>)
585 %unpack = tensor.unpack %transposed
589 into %unpack_dest : tensor<?x?x?x?xf32> -> tensor<?x?xf32>
590 return %unpack : tensor<?x?xf32>
593 // CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?xf32>, %[[ARG1:.+]]: tensor<?x?x?x?xf32>, %[[ARG2:.+]]: tensor<?x?xf32>,
594 // CHECK-SAME: %[[IDX1:.+]]: index, %[[IDX2:.+]]: index) -> tensor<?x?xf32> {
597 // CHECK-DAG: %[[DIM0:.+]] = tensor.dim %[[ARG2]], %[[CST0]] : tensor<?x?xf32>
598 // CHECK-DAG: %[[DIM1:.+]] = tensor.dim %[[ARG2]], %[[CST1]] : tensor<?x?xf32>
599 // CHECK: %[[OUT:.+]] = tensor.empty(%[[DIM0]], %[[DIM1]]) : tensor<?x?xf32>
600 // CHECK: %[[UNPACK:.+]] = tensor.unpack %[[ARG0]]
604 // CHECK-SAME: into %[[OUT]] : tensor<?x?x?x?xf32> -> tensor<?x?xf32>
605 // CHECK: return %[[UNPACK]] : tensor<?x?xf32>
610 func.func @tensor_unpack_linalg_transpose_fold(%arg0: tensor<56x57x1x64xf32>) -> tensor<3648x56xf32> {
611 %0 = tensor.empty() : tensor<56x3648xf32>
612 %pack = tensor.unpack %arg0
616 into %0 : tensor<56x57x1x64xf32> -> tensor<56x3648xf32>
618 %1 = tensor.empty() : tensor<3648x56xf32>
620 ins(%pack : tensor<56x3648xf32>)
621 outs(%1 : tensor<3648x56xf32>)
623 return %transposed : tensor<3648x56xf32>
626 // CHECK-SAME: %[[ARG0:.+]]: tensor<56x57x1x64xf32>) -> tensor<3648x56xf32> {
627 // CHECK: %[[OUT:.+]] = tensor.empty() : tensor<3648x56xf32>
628 // CHECK: %[[UNPACK:.+]] = tensor.unpack %[[ARG0]]
632 // CHECK-SAME: into %[[OUT:.+]] : tensor<56x57x1x64xf32> -> tensor<3648x56xf32>
633 // CHECK: return %[[UNPACK]] : tensor<3648x56xf32>
638 func.func @tensor_padded_unpack_linalg_transpose_fold(%arg0: tensor<71x7x4x16x16xf32>) -> tensor<100x71x64xf32> {
639 %0 = tensor.empty() : tensor<71x100x64xf32>
640 %pack = tensor.unpack %arg0
643 into %0 : tensor<71x7x4x16x16xf32> -> tensor<71x100x64xf32>
645 %1 = tensor.empty() : tensor<100x71x64xf32>
647 ins(%pack : tensor<71x100x64xf32>)
648 outs(%1 : tensor<100x71x64xf32>)
650 return %transposed : tensor<100x71x64xf32>
653 // CHECK-SAME: %[[ARG0:.+]]: tensor<71x7x4x16x16xf32>) -> tensor<100x71x64xf32> {
654 // CHECK: %[[OUT:.+]] = tensor.empty() : tensor<100x71x64xf32>
655 // CHECK: %[[UNPACK:.+]] = tensor.unpack %[[ARG0]]
659 // CHECK-SAME: into %[[OUT:.+]] : tensor<71x7x4x16x16xf32> -> tensor<100x71x64xf32>
660 // CHECK: return %[[UNPACK]] : tensor<100x71x64xf32>
665 func.func @non_involution_transpose_unpack_fold(%arg0: tensor<2x3x5x4x16xi32>) -> tensor<5x48x8xi32> {
666 %0 = tensor.empty() : tensor<5x2x3x16x4xi32>
667 %transposed = linalg.transpose ins(%arg0 : tensor<2x3x5x4x16xi32>)
668 outs(%0 : tensor<5x2x3x16x4xi32>)
670 %1 = tensor.empty() : tensor<5x48x8xi32>
671 %unpack = tensor.unpack %transposed
675 %1 : tensor<5x2x3x16x4xi32> -> tensor<5x48x8xi32>
676 return %unpack : tensor<5x48x8xi32>
679 // CHECK-SAME: %[[ARG0:.+]]: tensor<2x3x5x4x16xi32>) -> tensor<5x48x8xi32> {
680 // CHECK: %[[OUT:.+]] = tensor.empty() : tensor<5x48x8xi32>
681 // CHECK: %[[UNPACK:.+]] = tensor.unpack %[[ARG0]]
685 // CHEKC-SAME: into %[[OUT]] : tensor<2x3x5x4x16xi32> -> tensor<5x48x8xi32>
686 // CHECK: return %[[UNPACK]] : tensor<5x48x8xi32>
691 func.func @unpack_non_involution_transpose_fold(%arg0: tensor<57x3x56x1x64xf32>) -> tensor<3648x3x56xf32> {
692 %0 = tensor.empty() : tensor<3x56x3648xf32>
693 %unpack = tensor.unpack %arg0
697 into %0 : tensor<57x3x56x1x64xf32> -> tensor<3x56x3648xf32>
699 %1 = tensor.empty() : tensor<3648x3x56xf32>
701 ins(%unpack : tensor<3x56x3648xf32>)
702 outs(%1 : tensor<3648x3x56xf32>)
704 return %transposed : tensor<3648x3x56xf32>
707 // CHECK-SAME: %[[ARG0:.+]]: tensor<57x3x56x1x64xf32>) -> tensor<3648x3x56xf32> {
708 // CHECK: %[[OUT:.+]] = tensor.empty() : tensor<3648x3x56xf32>
709 // CHECK: %[[UNPACK:.+]] = tensor.unpack %[[ARG0]]
713 // CHECK-SAME: into %[[OUT:.+]] : tensor<57x3x56x1x64xf32> -> tensor<3648x3x56xf32>
714 // CHECK: return %[[UNPACK]] : tensor<3648x3x56xf32>
719 func.func @transpose_unpacked_dims_no_fold(%arg0: tensor<2x16x5x4x3xi32>) -> tensor<5x32x12xi32> {
720 %0 = tensor.empty() : tensor<5x2x3x16x4xi32>
721 %transposed = linalg.transpose ins(%arg0 : tensor<2x16x5x4x3xi32>)
722 outs(%0 : tensor<5x2x3x16x4xi32>)
724 %1 = tensor.empty() : tensor<5x32x12xi32>
725 %unpack = tensor.unpack %transposed
728 %1 : tensor<5x2x3x16x4xi32> -> tensor<5x32x12xi32>
729 return %unpack : tensor<5x32x12xi32>
733 // CHECK: tensor.unpack
739 func.func @generic_transpose_unpack_fold(%arg0: tensor<2x3x5x4x16xi32>) -> tensor<5x48x8xi32> {
740 %0 = tensor.empty() : tensor<5x2x3x16x4xi32>
744 ins(%arg0 : tensor<2x3x5x4x16xi32>)
745 outs(%0 : tensor<5x2x3x16x4xi32>) {
748 } -> tensor<5x2x3x16x4xi32>
749 %1 = tensor.empty() : tensor<5x48x8xi32>
750 %unpack = tensor.unpack %transposed
754 %1 : tensor<5x2x3x16x4xi32> -> tensor<5x48x8xi32>
755 return %unpack : tensor<5x48x8xi32>
758 // CHECK-SAME: %[[ARG0:.+]]: tensor<2x3x5x4x16xi32>) -> tensor<5x48x8xi32> {
759 // CHECK: %[[OUT:.+]] = tensor.empty() : tensor<5x48x8xi32>
760 // CHECK: %[[UNPACK:.+]] = tensor.unpack %[[ARG0]]
764 // CHEKC-SAME: into %[[OUT]] : tensor<2x3x5x4x16xi32> -> tensor<5x48x8xi32>
765 // CHECK: return %[[UNPACK]] : tensor<5x48x8xi32>
772 func.func @unpack_generic_transpose_fold(%arg0: tensor<57x3x56x1x64xf32>) -> tensor<3648x3x56xf32> {
773 %0 = tensor.empty() : tensor<3x56x3648xf32>
774 %unpack = tensor.unpack %arg0
778 into %0 : tensor<57x3x56x1x64xf32> -> tensor<3x56x3648xf32>
780 %1 = tensor.empty() : tensor<3648x3x56xf32>
784 ins(%unpack : tensor<3x56x3648xf32>)
785 outs(%1 : tensor<3648x3x56xf32>) {
788 } -> tensor<3648x3x56xf32>
789 return %transposed : tensor<3648x3x56xf32>
792 // CHECK-SAME: %[[ARG0:.+]]: tensor<57x3x56x1x64xf32>) -> tensor<3648x3x56xf32> {
793 // CHECK: %[[OUT:.+]] = tensor.empty() : tensor<3648x3x56xf32>
794 // CHECK: %[[UNPACK:.+]] = tensor.unpack %[[ARG0]]
798 // CHECK-SAME: into %[[OUT:.+]] : tensor<57x3x56x1x64xf32> -> tensor<3648x3x56xf32>
799 // CHECK: return %[[UNPACK]] : tensor<3648x3x56xf32>