946 lines · plain
1// RUN: mlir-opt -split-input-file -test-linalg-transform-patterns=test-fold-into-pack-and-unpack %s | FileCheck %s2// RUN: mlir-opt -split-input-file -test-linalg-transform-patterns=test-fold-into-pack-and-unpack-control %s | FileCheck %s --check-prefix=CONTROL3 4func.func @fold_extract_slice_into_unpack_slicing_trailing_dim(%arg0 : tensor<28x2x1x16x16xf32>) -> tensor<28x28x10xf32> {5 %empty = tensor.empty() : tensor<28x28x15xf32>6 %unpack = linalg.unpack %arg07 inner_dims_pos = [1, 2]8 inner_tiles = [16, 16]9 into %empty : tensor<28x2x1x16x16xf32> -> tensor<28x28x15xf32>10 %extracted_slice = tensor.extract_slice %unpack11 [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_dim15// 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 25func.func @fold_extract_slice_into_unpack_slicing_dim_1(%arg0 : tensor<28x2x1x16x16xf32>) -> tensor<28x17x15xf32> {26 %empty = tensor.empty() : tensor<28x28x15xf32>27 %unpack = linalg.unpack %arg028 inner_dims_pos = [1, 2]29 inner_tiles = [16, 16]30 into %empty : tensor<28x2x1x16x16xf32> -> tensor<28x28x15xf32>31 %extracted_slice = tensor.extract_slice %unpack32 [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 46func.func @no_fold_extract_slice_into_unpack_artificial_padding(%arg0 : tensor<28x2x1x16x16xf32>) -> tensor<28x16x15xf32> {47 %empty = tensor.empty() : tensor<28x28x15xf32>48 %unpack = linalg.unpack %arg049 inner_dims_pos = [1, 2]50 inner_tiles = [16, 16]51 into %empty : tensor<28x2x1x16x16xf32> -> tensor<28x28x15xf32>52 %extracted_slice = tensor.extract_slice %unpack53 [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_padding57// CHECK: linalg.unpack58// CHECK: tensor.extract_slice59 60// -----61 62func.func @no_fold_extract_slice_into_unpack_dynamic(63 %src : tensor<28x2x?x16x16xf32>, %dest : tensor<28x32x?xf32>, %size : index64) -> tensor<28x28x?xf32> {65 %unpack = linalg.unpack %src66 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 %unpack71 [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_dynamic75// CHECK: linalg.unpack76// CHECK: tensor.extract_slice77 78// -----79 80func.func @nofold_dynamic_unpack_slice(%arg0 : tensor<?x?x8x4xf32>, %arg1 : tensor<?x?xf32>,81 %arg2 : index, %arg3 : index) -> tensor<?x?xf32> {82 %0 = linalg.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [8, 4] into %arg183 : tensor<?x?x8x4xf32> -> tensor<?x?xf32>84 %1 = tensor.extract_slice %0[0, 0] [%arg2, %arg3] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>85 return %1 : tensor<?x?xf32>86}87// CHECK-LABEL: func @nofold_dynamic_unpack_slice(88// CHECK: linalg.unpack89// CHECK: tensor.extract_slice90 91// -----92 93func.func @nofold_unpack_slice_non_zero_offset(%arg0 : tensor<?x?x8x4xf32>, %arg1 : tensor<?x?xf32>,94 %arg2 : index, %arg3 : index, %arg4 : index) -> tensor<?x?xf32> {95 %0 = linalg.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [8, 4] into %arg196 : tensor<?x?x8x4xf32> -> tensor<?x?xf32>97 %1 = tensor.extract_slice %0[0, %arg4] [%arg2, %arg3] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>98 return %1 : tensor<?x?xf32>99}100// CHECK-LABEL: func @nofold_unpack_slice_non_zero_offset(101// CHECK: %[[UNPACK:.+]] = linalg.unpack102// CHECK: tensor.extract_slice %[[UNPACK]]103 104// -----105 106func.func @nofold_unpack_slice_non_unit_stride(%arg0 : tensor<?x?x8x4xf32>, %arg1 : tensor<?x?xf32>,107 %arg2 : index, %arg3 : index, %arg4 : index) -> tensor<?x?xf32> {108 %0 = linalg.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [8, 4] into %arg1109 : tensor<?x?x8x4xf32> -> tensor<?x?xf32>110 %1 = tensor.extract_slice %0[0, 0] [%arg2, %arg3] [%arg4, 1] : tensor<?x?xf32> to tensor<?x?xf32>111 return %1 : tensor<?x?xf32>112}113// CHECK-LABEL: func @nofold_unpack_slice_non_unit_stride(114// CHECK: %[[UNPACK:.+]] = linalg.unpack115// CHECK: tensor.extract_slice %[[UNPACK]]116 117// -----118 119func.func @nofold_unpack_slice_rank_reduced(%arg0 : tensor<?x?x8x4xf32>, %arg1 : tensor<?x?xf32>,120 %arg2 : index, %arg3 : index) -> tensor<f32> {121 %0 = linalg.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [8, 4] into %arg1122 : tensor<?x?x8x4xf32> -> tensor<?x?xf32>123 %1 = tensor.extract_slice %0[0, 0] [1, 1] [1, 1] : tensor<?x?xf32> to tensor<f32>124 return %1 : tensor<f32>125}126// CHECK-LABEL: func @nofold_unpack_slice_rank_reduced(127// CHECK: %[[UNPACK:.+]] = linalg.unpack128// CHECK: tensor.extract_slice %[[UNPACK]]129 130// -----131 132func.func @fold_pad_pack(%src: tensor<9x16xf32>) -> tensor<2x1x8x32xf32> {133 %cst = arith.constant 0.000000e+00 : f32134 %padded = tensor.pad %src low[0, 0] high[7, 0] {135 ^bb0(%arg0: index, %arg1: index):136 tensor.yield %cst : f32137 } : tensor<9x16xf32> to tensor<16x16xf32>138 %empty = tensor.empty() : tensor<2x1x8x32xf32>139 %pack = linalg.pack %padded padding_value(%cst : f32) inner_dims_pos = [0, 1] inner_tiles = [8, 32] into %empty140 : tensor<16x16xf32> -> tensor<2x1x8x32xf32>141 return %pack : tensor<2x1x8x32xf32>142}143// CHECK-LABEL: func.func @fold_pad_pack144// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]145// CHECK: %[[PAD_VAL:.+]] = arith.constant 0.000000e+00 : f32146// CHECK: %[[DEST:.+]] = tensor.empty() : tensor<2x1x8x32xf32>147// CHECK: %[[PACK:.+]] = linalg.pack %[[SRC]]148// CHECK-SAME: padding_value(%[[PAD_VAL]] : f32)149// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [8, 32] into %[[DEST]]150 151// -----152 153func.func @nofold_pad_pack_artificial_padding(%src: tensor<9x16xf32>) -> tensor<3x1x8x32xf32> {154 %cst = arith.constant 0.000000e+00 : f32155 %padded = tensor.pad %src low[0, 0] high[8, 0] {156 ^bb0(%arg0: index, %arg1: index):157 tensor.yield %cst : f32158 } : tensor<9x16xf32> to tensor<17x16xf32>159 %empty = tensor.empty() : tensor<3x1x8x32xf32>160 %pack = linalg.pack %padded padding_value(%cst : f32) inner_dims_pos = [0, 1] inner_tiles = [8, 32] into %empty161 : tensor<17x16xf32> -> tensor<3x1x8x32xf32>162 return %pack : tensor<3x1x8x32xf32>163}164// CHECK-LABLE: func.func @nofold_pad_pack_artificial_padding(165// CHECK: tensor.pad166// CHECK: linalg.pack167 168// -----169 170func.func @nofold_pad_pack_with_nofold_attribute(%src: tensor<16649x16xf32>) -> tensor<2082x1x8x32xf32> {171 %cst = arith.constant 0.000000e+00 : f32172 %padded = tensor.pad %src nofold low[0, 0] high[7, 0] {173 ^bb0(%arg0: index, %arg1: index):174 tensor.yield %cst : f32175 } : tensor<16649x16xf32> to tensor<16656x16xf32>176 %empty = tensor.empty() : tensor<2082x1x8x32xf32>177 %pack = linalg.pack %padded padding_value(%cst : f32) inner_dims_pos = [0, 1] inner_tiles = [8, 32] into %empty178 : tensor<16656x16xf32> -> tensor<2082x1x8x32xf32>179 return %pack : tensor<2082x1x8x32xf32>180}181// CHECK-LABEL: func.func @nofold_pad_pack_with_nofold_attribute(182// CHECK: tensor.pad183// CHECK: linalg.pack184 185// -----186 187func.func @pad_pack_different_padding_value(%src: tensor<16641x16xf32>) -> tensor<2082x1x8x32xf32> {188 %cst0 = arith.constant 0.000000e+00 : f32189 %cst1 = arith.constant 1.000000e+00 : f32190 %padded = tensor.pad %src low[0, 0] high[15, 0] {191 ^bb0(%arg0: index, %arg1: index):192 tensor.yield %cst0 : f32193 } : tensor<16641x16xf32> to tensor<16656x16xf32>194 %empty = tensor.empty() : tensor<2082x1x8x32xf32>195 %pack = linalg.pack %padded padding_value(%cst1 : f32) inner_dims_pos = [0, 1] inner_tiles = [8, 32] into %empty196 : tensor<16656x16xf32> -> tensor<2082x1x8x32xf32>197 return %pack : tensor<2082x1x8x32xf32>198}199// CHECK-LABEL: func.func @pad_pack_different_padding_value200// CHECK: tensor.pad201// CHECK: linalg.pack202 203// -----204 205func.func @linalg.pack_linalg_transpose_fold(%arg0: tensor<56x57x1x64xf32>) -> tensor<1x57x56x2x32xf32> {206 %0 = tensor.empty() : tensor<56x2x1x57x32xf32>207 %pack = linalg.pack %arg0208 outer_dims_perm = [0, 3, 2, 1]209 inner_dims_pos = [3]210 inner_tiles = [32]211 into %0 : tensor<56x57x1x64xf32> -> tensor<56x2x1x57x32xf32>212 213 %1 = tensor.empty() : tensor<1x57x56x2x32xf32>214 %transposed = linalg.transpose215 ins(%pack : tensor<56x2x1x57x32xf32>)216 outs(%1 : tensor<1x57x56x2x32xf32>)217 permutation = [2, 3, 0, 1, 4]218 return %transposed : tensor<1x57x56x2x32xf32>219}220// CHECK: func @linalg.pack_linalg_transpose_fold(221// CHECK-SAME: %[[ARG0:.+]]: tensor<56x57x1x64xf32>)222// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x57x56x2x32xf32>223// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]224// CHECK-SAME: outer_dims_perm = [2, 1, 0, 3]225// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]226// CHECK-SAME: into %[[INIT]]227// CHECK: return %[[PACK]]228 229// -----230 231func.func @linalg.pack_linalg_transpose_fold_with_padding(%arg0: tensor<56x57x1x55xf32>, %padding: f32) -> tensor<1x57x56x2x32xf32> {232 %0 = tensor.empty() : tensor<56x2x1x57x32xf32>233 %pack = linalg.pack %arg0 padding_value(%padding : f32)234 outer_dims_perm = [0, 3, 2, 1]235 inner_dims_pos = [3]236 inner_tiles = [32]237 into %0 : tensor<56x57x1x55xf32> -> tensor<56x2x1x57x32xf32>238 239 %1 = tensor.empty() : tensor<1x57x56x2x32xf32>240 %transposed = linalg.transpose241 ins(%pack : tensor<56x2x1x57x32xf32>)242 outs(%1 : tensor<1x57x56x2x32xf32>)243 permutation = [2, 3, 0, 1, 4]244 return %transposed : tensor<1x57x56x2x32xf32>245}246// CHECK: func @linalg.pack_linalg_transpose_fold_with_padding(247// CHECK-SAME: %[[ARG0:.+]]: tensor<56x57x1x55xf32>, %[[PADDING:.+]]: f32)248// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x57x56x2x32xf32>249// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]] padding_value(%[[PADDING]] : f32)250// CHECK-SAME: outer_dims_perm = [2, 1, 0, 3]251// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]252// CHECK-SAME: into %[[INIT]]253// CHECK: return %[[PACK]]254 255// -----256 257func.func @linalg.pack_linalg_transpose_fold_no_outer_dims_perm(%arg0: tensor<56x57x1x64xf32>) -> tensor<1x2x56x57x32xf32> {258 %0 = tensor.empty() : tensor<56x57x1x2x32xf32>259 %pack = linalg.pack %arg0260 inner_dims_pos = [3]261 inner_tiles = [32]262 into %0 : tensor<56x57x1x64xf32> -> tensor<56x57x1x2x32xf32>263 264 %1 = tensor.empty() : tensor<1x2x56x57x32xf32>265 %transposed = linalg.transpose266 ins(%pack : tensor<56x57x1x2x32xf32>)267 outs(%1 : tensor<1x2x56x57x32xf32>)268 permutation = [2, 3, 0, 1, 4]269 return %transposed : tensor<1x2x56x57x32xf32>270}271// CHECK: func @linalg.pack_linalg_transpose_fold_no_outer_dims_perm(272// CHECK-SAME: %[[ARG0:.+]]: tensor<56x57x1x64xf32>)273// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x2x56x57x32xf32>274// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]275// CHECK-SAME: outer_dims_perm = [2, 3, 0, 1]276// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]277// CHECK-SAME: into %[[INIT]]278// CHECK: return %[[PACK]]279 280// -----281 282func.func @linalg.pack_linalg_transpose_fold_tile_dims_transpose(%arg0: tensor<56x72x24x128xf32>) -> tensor<12x56x4x9x32x8x2xf32> {283 %0 = tensor.empty() : tensor<4x9x12x56x8x2x32xf32>284 %pack = linalg.pack %arg0285 outer_dims_perm = [3, 1, 2, 0]286 inner_dims_pos = [1, 2, 3]287 inner_tiles = [8, 2, 32]288 into %0 : tensor<56x72x24x128xf32> -> tensor<4x9x12x56x8x2x32xf32>289 290 %1 = tensor.empty() : tensor<12x56x4x9x32x8x2xf32>291 %transposed = linalg.transpose292 ins(%pack : tensor<4x9x12x56x8x2x32xf32>)293 outs(%1 : tensor<12x56x4x9x32x8x2xf32>)294 permutation = [2, 3, 0, 1, 6, 4, 5]295 return %transposed : tensor<12x56x4x9x32x8x2xf32>296}297// CHECK: func @linalg.pack_linalg_transpose_fold_tile_dims_transpose(298// CHECK-SAME: %[[ARG0:.+]]: tensor<56x72x24x128xf32>)299// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<12x56x4x9x32x8x2xf32>300// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]301// CHECK-SAME: outer_dims_perm = [2, 0, 3, 1]302// CHECK-SAME: inner_dims_pos = [3, 1, 2] inner_tiles = [32, 8, 2]303// CHECK-SAME: into %[[INIT]]304// CHECK: return %[[PACK]]305 306// -----307 308func.func @linalg.pack_linalg_transpose_fold_tile_dims_outer_dims_transpose(%arg0: tensor<56x72x24x128xf32>) -> tensor<9x56x2x12x32x8x4xf32> {309 %0 = tensor.empty() : tensor<4x12x9x56x8x2x32xf32>310 %pack = linalg.pack %arg0311 outer_dims_perm = [3, 2, 1, 0]312 inner_dims_pos = [1, 2, 3]313 inner_tiles = [8, 2, 32]314 into %0 : tensor<56x72x24x128xf32> -> tensor<4x12x9x56x8x2x32xf32>315 316 %1 = tensor.empty() : tensor<9x56x2x12x32x8x4xf32>317 %transposed = linalg.transpose318 ins(%pack : tensor<4x12x9x56x8x2x32xf32>)319 outs(%1 : tensor<9x56x2x12x32x8x4xf32>)320 permutation = [2, 3, 5, 1, 6, 4, 0]321 return %transposed : tensor<9x56x2x12x32x8x4xf32>322}323// CHECK: func @linalg.pack_linalg_transpose_fold_tile_dims_outer_dims_transpose(324// CHECK-SAME: %[[ARG0:.+]]: tensor<56x72x24x128xf32>)325// CHECK: linalg.pack326// CHECK: linalg.transpose327 328// -----329 330func.func @linalg.pack_linalg_transpose_fold_dynamic_outer_dims(%arg0: tensor<56x?x?x64xf32>) -> tensor<?x?x56x2x32xf32> {331 %0 = tensor.empty() : tensor<56x2x1x57x32xf32>332 %pack = linalg.pack %arg0333 outer_dims_perm = [0, 3, 2, 1]334 inner_dims_pos = [3]335 inner_tiles = [32]336 into %0 : tensor<56x?x?x64xf32> -> tensor<56x2x1x57x32xf32>337 338 %1 = tensor.empty() : tensor<1x57x56x2x32xf32>339 %transposed = linalg.transpose340 ins(%pack : tensor<56x2x1x57x32xf32>)341 outs(%1 : tensor<1x57x56x2x32xf32>)342 permutation = [2, 3, 0, 1, 4]343 344 %return_value = tensor.cast %transposed : tensor<1x57x56x2x32xf32> to tensor<?x?x56x2x32xf32>345 return %return_value : tensor<?x?x56x2x32xf32>346}347// CHECK: func @linalg.pack_linalg_transpose_fold_dynamic_outer_dims(348// CHECK-SAME: %[[ARG0:.+]]: tensor<56x?x?x64xf32>)349// CHECK-DAG: %[[c1:.+]] = arith.constant 1 : index350// CHECK-DAG: %[[c2:.+]] = arith.constant 2 : index351// CHECK: %[[dim:.+]] = tensor.dim %[[ARG0]], %[[c1]] : tensor<56x?x?x64xf32>352// CHECK: %[[dim_0:.+]] = tensor.dim %[[ARG0]], %[[c2]] : tensor<56x?x?x64xf32>353// CHECK: %[[INIT:.+]] = tensor.empty(%[[dim_0]], %[[dim]]) : tensor<?x?x56x2x32xf32>354// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]355// CHECK-SAME: outer_dims_perm = [2, 1, 0, 3]356// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]357// CHECK-SAME: into %[[INIT]]358// CHECK: return %[[PACK]]359 360// -----361 362func.func @linalg.pack_linalg_transpose_fold_dynamic_outer_and_tile_dims(%arg0: tensor<56x?x?x128xf32>) -> tensor<?x?x56x9x32x8x2xf32> {363 %0 = tensor.empty() : tensor<56x9x12x4x8x2x32xf32>364 %pack = linalg.pack %arg0365 inner_dims_pos = [1, 2, 3]366 inner_tiles = [8, 2, 32]367 into %0 : tensor<56x?x?x128xf32> -> tensor<56x9x12x4x8x2x32xf32>368 369 %1 = tensor.empty() : tensor<12x4x56x9x32x8x2xf32>370 %transposed = linalg.transpose371 ins(%pack : tensor<56x9x12x4x8x2x32xf32>)372 outs(%1 : tensor<12x4x56x9x32x8x2xf32>)373 permutation = [2, 3, 0, 1, 6, 4, 5]374 375 %return_value = tensor.cast %transposed : tensor<12x4x56x9x32x8x2xf32> to tensor<?x?x56x9x32x8x2xf32>376 return %return_value : tensor<?x?x56x9x32x8x2xf32>377}378// CHECK-DAG: #[[$MAP0:.+]] = affine_map<()[s0] -> (s0 ceildiv 8)>379// CHECK-DAG: #[[$MAP1:.+]] = affine_map<()[s0] -> (s0 ceildiv 2)>380// CHECK-LABEL: func.func @linalg.pack_linalg_transpose_fold_dynamic_outer_and_tile_dims(381// CHECK-SAME: %[[ARG0:.+]]: tensor<56x?x?x128xf32>)382// CHECK-DAG: %[[c1:.+]] = arith.constant 1 : index383// CHECK-DAG: %[[c2:.+]] = arith.constant 2 : index384// CHECK: %[[dim:.+]] = tensor.dim %[[ARG0]], %[[c1]] : tensor<56x?x?x128xf32>385// CHECK: %[[dim_0:.+]] = tensor.dim %[[ARG0]], %[[c2]] : tensor<56x?x?x128xf32>386// CHECK: %[[mapped_dim1:.+]] = affine.apply #[[$MAP0]]()[%[[dim]]]387// CHECK: %[[mapped_dim2:.+]] = affine.apply #[[$MAP1]]()[%[[dim_0]]]388// CHECK: %[[INIT:.+]] = tensor.empty(%[[mapped_dim2]], %[[mapped_dim1]]) : tensor<?x4x56x?x32x8x2xf32>389// CHECK: %[[PACK:.+]] = linalg.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>390// CHECK: %[[CAST:.+]] = tensor.cast %[[PACK]] : tensor<?x4x56x?x32x8x2xf32> to tensor<?x?x56x9x32x8x2xf32>391// CHECK: return %[[CAST]] : tensor<?x?x56x9x32x8x2xf32>392// CHECK: }393 394// -----395 396func.func @linalg.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> {397 %pack = linalg.pack %arg0398 outer_dims_perm = [3, 0, 2, 1]399 inner_dims_pos = [1, 2, 3]400 inner_tiles = [%tile_p, %tile_q, %tile_r]401 into %pack_dest : tensor<?x?x?x?xf32> -> tensor<?x?x?x?x?x?x?xf32>402 403 %transposed = linalg.transpose404 ins(%pack : tensor<?x?x?x?x?x?x?xf32>)405 outs(%transpose_dest : tensor<?x?x?x?x?x?x?xf32>)406 permutation = [2, 3, 0, 1, 6, 4, 5]407 408 return %transposed : tensor<?x?x?x?x?x?x?xf32>409}410// CHECK: #[[$MAP:.+]] = affine_map<()[s0, s1] -> (s0 ceildiv s1)>411// CHECK: module {412// CHECK: func.func @linalg.pack_linalg_transpose_fold_dynamic_outer_dims_tile_dims_tile_sizes(413// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?xf32>,414// CHECK-SAME: %[[PACK_DEST:.+]]: tensor<?x?x?x?x?x?x?xf32>, %[[TRANSPOSE_DEST:.+]]: tensor<?x?x?x?x?x?x?xf32>,415// CHECK-SAME: %[[ARG1:.+]]: index, %[[ARG2:.+]]: index,416// CHECK-SAME: %[[ARG3:.+]]: index)417// CHECK-DAG: %[[c0:.+]] = arith.constant 0 : index418// CHECK-DAG: %[[c1:.+]] = arith.constant 1 : index419// CHECK-DAG: %[[c2:.+]] = arith.constant 2 : index420// CHECK-DAG: %[[c3:.+]] = arith.constant 3 : index421// CHECK: %[[dim:.+]] = tensor.dim %[[ARG0]], %[[c0]] : tensor<?x?x?x?xf32>422// CHECK: %[[dim_0:.+]] = tensor.dim %[[ARG0]], %[[c1]] : tensor<?x?x?x?xf32>423// CHECK: %[[dim_1:.+]] = tensor.dim %[[ARG0]], %[[c2]] : tensor<?x?x?x?xf32>424// CHECK: %[[dim_2:.+]] = tensor.dim %[[ARG0]], %[[c3]] : tensor<?x?x?x?xf32>425// CHECK: %[[mapped_dim0:.+]] = affine.apply #[[$MAP]]()[%[[dim_2]], %[[ARG3]]]426// CHECK: %[[mapped_dim1:.+]] = affine.apply #[[$MAP]]()[%[[dim_0]], %[[ARG1]]]427// CHECK: %[[mapped_dim2:.+]] = affine.apply #[[$MAP]]()[%[[dim_1]], %[[ARG2]]]428// CHECK: %[[INIT:.+]] = tensor.empty(%[[mapped_dim2]], %[[mapped_dim1]], %[[mapped_dim0]], %[[dim]], %[[ARG3]], %[[ARG1]], %[[ARG2]]) : tensor<?x?x?x?x?x?x?xf32>429// CHECK: %[[PACK:.+]] = linalg.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>430// CHECK: return %[[PACK]] : tensor<?x?x?x?x?x?x?xf32>431// CHECK: }432 433// -----434 435func.func @linalg_transpose_linalg.pack_fold(%arg0: tensor<56x57x1x64xf32>) -> tensor<1x57x56x2x32xf32> {436 %0 = tensor.empty() : tensor<1x56x57x64xf32>437 %transposed = linalg.transpose438 ins(%arg0 : tensor<56x57x1x64xf32>)439 outs(%0 : tensor<1x56x57x64xf32>)440 permutation = [2, 0, 1, 3]441 442 %1 = tensor.empty() : tensor<1x57x56x2x32xf32>443 %pack = linalg.pack %transposed444 outer_dims_perm = [0, 2, 1, 3]445 inner_dims_pos = [3]446 inner_tiles = [32]447 into %1 : tensor<1x56x57x64xf32> -> tensor<1x57x56x2x32xf32>448 return %pack : tensor<1x57x56x2x32xf32>449}450//CHECK-LABEL: func @linalg_transpose_linalg.pack_fold(451// CHECK-SAME: %[[ARG0:.+]]: tensor<56x57x1x64xf32>)452// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x57x56x2x32xf32>453// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]454// CHECK-SAME: outer_dims_perm = [2, 1, 0, 3]455// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]456// CHECK-SAME: into %[[INIT]]457// CHECK: return %[[PACK]]458 459// -----460 461func.func @linalg_transpose_linalg.pack_fold_multi_result(%arg0: tensor<56x57x1x64xf32>) -> (tensor<1x56x57x64xf32>, tensor<1x57x56x2x32xf32>) {462 %0 = tensor.empty() : tensor<1x56x57x64xf32>463 %transposed = linalg.transpose464 ins(%arg0 : tensor<56x57x1x64xf32>)465 outs(%0 : tensor<1x56x57x64xf32>)466 permutation = [2, 0, 1, 3]467 468 %1 = tensor.empty() : tensor<1x57x56x2x32xf32>469 %pack = linalg.pack %transposed470 outer_dims_perm = [0, 2, 1, 3]471 inner_dims_pos = [3]472 inner_tiles = [32]473 into %1 : tensor<1x56x57x64xf32> -> tensor<1x57x56x2x32xf32>474 return %transposed, %pack : tensor<1x56x57x64xf32>, tensor<1x57x56x2x32xf32>475}476// CHECK-LABEL: func @linalg_transpose_linalg.pack_fold_multi_result(477// CHECK-SAME: %[[ARG0:.+]]: tensor<56x57x1x64xf32>)478// CHECK: %[[TRANSPOSE:.+]] = linalg.transpose479// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]480// CHECK-SAME: outer_dims_perm = [2, 1, 0, 3]481// CHECK: return %[[TRANSPOSE]], %[[PACK]]482 483// CONTROL-LABEL: func @linalg_transpose_linalg.pack_fold_multi_result(484// CONTROL: %[[TRANSPOSE:.+]] = linalg.transpose485// CONTROL: %[[PACK:.+]] = linalg.pack %[[TRANSPOSE]]486// CONTROL-SAME: outer_dims_perm = [0, 2, 1, 3]487// CONTROL: return %[[TRANSPOSE]], %[[PACK]]488 489// -----490 491func.func @linalg_transpose_linalg.pack_fold_with_padding(%arg0: tensor<56x57x1x55xf32>, %padding: f32) -> tensor<1x57x56x2x32xf32> {492 %0 = tensor.empty() : tensor<1x56x57x55xf32>493 %transpose = linalg.transpose494 ins(%arg0 : tensor<56x57x1x55xf32>)495 outs(%0 : tensor<1x56x57x55xf32>)496 permutation = [2, 0, 1, 3]497 498 %1 = tensor.empty() : tensor<1x57x56x2x32xf32>499 %pack = linalg.pack %transpose padding_value(%padding : f32)500 outer_dims_perm = [0, 2, 1, 3]501 inner_dims_pos = [3]502 inner_tiles = [32]503 into %1 : tensor<1x56x57x55xf32> -> tensor<1x57x56x2x32xf32>504 return %pack : tensor<1x57x56x2x32xf32>505}506//CHECK-LABEL: func @linalg_transpose_linalg.pack_fold_with_padding(507// CHECK-SAME: %[[ARG0:.+]]: tensor<56x57x1x55xf32>, %[[PADDING:.+]]: f32)508// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x57x56x2x32xf32>509// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]] padding_value(%[[PADDING]] : f32)510// CHECK-SAME: outer_dims_perm = [2, 1, 0, 3]511// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]512// CHECK-SAME: into %[[INIT]]513// CHECK: return %[[PACK]]514 515// -----516 517func.func @linalg_transpose_linalg.pack_fold_no_outer_dims_perm(%arg0: tensor<56x57x1x64xf32>) -> tensor<1x56x57x2x32xf32> {518 %0 = tensor.empty() : tensor<1x56x57x64xf32>519 %transposed = linalg.transpose520 ins(%arg0 : tensor<56x57x1x64xf32>)521 outs(%0 : tensor<1x56x57x64xf32>)522 permutation = [2, 0, 1, 3]523 524 %1 = tensor.empty() : tensor<1x56x57x2x32xf32>525 %pack = linalg.pack %transposed526 inner_dims_pos = [3]527 inner_tiles = [32]528 into %1 : tensor<1x56x57x64xf32> -> tensor<1x56x57x2x32xf32>529 return %pack : tensor<1x56x57x2x32xf32>530}531//CHECK-LABEL: func @linalg_transpose_linalg.pack_fold_no_outer_dims_perm(532// CHECK-SAME: %[[ARG0:.+]]: tensor<56x57x1x64xf32>)533// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x56x57x2x32xf32>534// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]535// CHECK-SAME: outer_dims_perm = [2, 0, 1, 3]536// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]537// CHECK-SAME: into %[[INIT]]538// CHECK: return %[[PACK]]539 540// -----541 542func.func @linalg_transpose_linalg.pack_fold_complex_inner_dims_change(%arg0: tensor<25x30x35x40xf32>, %transpose_dest: tensor<35x40x25x30xf32>, %pack_dest: tensor<3x35x5x8x5x10x5xf32>) -> tensor<3x35x5x8x5x10x5xf32> {543 %transposed = linalg.transpose544 ins(%arg0 : tensor<25x30x35x40xf32>)545 outs(%transpose_dest : tensor<35x40x25x30xf32>)546 permutation = [2, 3, 0, 1]547 548 %pack = linalg.pack %transposed549 outer_dims_perm = [3, 0, 2, 1]550 inner_dims_pos = [1, 3, 2]551 inner_tiles = [5, 10, 5]552 into %pack_dest : tensor<35x40x25x30xf32> -> tensor<3x35x5x8x5x10x5xf32>553 return %pack : tensor<3x35x5x8x5x10x5xf32>554}555//CHECK-LABEL: func.func @linalg_transpose_linalg.pack_fold_complex_inner_dims_change(556// CHECK-SAME: %[[ARG0:.+]]: tensor<25x30x35x40xf32>,557// CHECK-SAME: %[[ARG1:.+]]: tensor<35x40x25x30xf32>,558// CHECK-SAME: %[[ARG2:.+]]: tensor<3x35x5x8x5x10x5xf32>) -> tensor<3x35x5x8x5x10x5xf32> {559// CHECK: %[[VAL0:.+]] = tensor.empty() : tensor<3x35x5x8x5x10x5xf32>560// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]561// CHECK-SAME: outer_dims_perm = [1, 2, 0, 3]562// CHECK-SAME: inner_dims_pos = [3, 1, 0]563// CHECK-SAME: inner_tiles = [5, 10, 5]564// CHECK-SAME: into %[[VAL0]]565// CHECK: return %[[PACK]]566 567// -----568 569func.func @linalg_transpose_linalg.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> {570 %transposed = linalg.transpose571 ins(%arg0 : tensor<?x?x?x?xf32>)572 outs(%transpose_dest : tensor<?x?x?x?xf32>)573 permutation = [2, 3, 0, 1]574 575 %pack = linalg.pack %transposed576 outer_dims_perm = [3, 0, 2, 1]577 inner_dims_pos = [1, 3, 2]578 inner_tiles = [%tile_p, %tile_q, %tile_r]579 into %pack_dest : tensor<?x?x?x?xf32> -> tensor<?x?x?x?x?x?x?xf32>580 return %pack : tensor<?x?x?x?x?x?x?xf32>581}582// CHECK: #[[$MAP:.+]] = affine_map<()[s0, s1] -> (s0 ceildiv s1)>583//CHECK-LABEL: func.func @linalg_transpose_linalg.pack_fold_dynamic_outer_dims_tile_dims_tile_sizes(584// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?xf32>, %[[ARG1:.+]]: tensor<?x?x?x?xf32>,585// 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> {586// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index587// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index588// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index589// CHECK-DAG: %[[C3:.+]] = arith.constant 3 : index590// CHECK: %[[DIM:.+]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x?x?x?xf32>591// CHECK: %[[DIM0:.+]] = tensor.dim %[[ARG0]], %[[C1]] : tensor<?x?x?x?xf32>592// CHECK: %[[DIM1:.+]] = tensor.dim %[[ARG0]], %[[C2]] : tensor<?x?x?x?xf32>593// CHECK: %[[DIM2:.+]] = tensor.dim %[[ARG0]], %[[C3]] : tensor<?x?x?x?xf32>594// CHECK: %[[VAL0:.+]] = affine.apply #[[$MAP]]()[%[[DIM2]], %[[ARG3]]]595// CHECK: %[[VAL1:.+]] = affine.apply #[[$MAP]]()[%[[DIM0]], %[[ARG4]]]596// CHECK: %[[VAL2:.+]] = affine.apply #[[$MAP]]()[%[[DIM]], %[[ARG5]]]597// CHECK: %[[VAL3:.+]] = tensor.empty(%[[VAL1]], %[[DIM1]], %[[VAL2]], %[[VAL0]], %[[ARG3]], %[[ARG4]], %[[ARG5]]) : tensor<?x?x?x?x?x?x?xf32>598// CHECK: %[[PACK:.+]] = linalg.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>599// CHECK: return %[[PACK]] : tensor<?x?x?x?x?x?x?xf32>600 601// -----602 603func.func @linalg_transpose_linalg.pack_multiple_tiles(%arg0: tensor<?x32x128xbf16>) -> tensor<32x?x64x16x2xbf16> {604 %c0 = arith.constant 0 : index605 %cst = arith.constant 0.000000e+00 : bf16606 %dim = tensor.dim %arg0, %c0 : tensor<?x32x128xbf16>607 608 %0 = tensor.empty(%dim) : tensor<32x128x?xbf16>609 %transposed = linalg.transpose610 ins(%arg0 : tensor<?x32x128xbf16>)611 outs(%0 : tensor<32x128x?xbf16>)612 permutation = [1, 2, 0]613 614 %2 = tensor.empty(%dim) : tensor<32x?x64x16x2xbf16>615 %pack = linalg.pack %transposed616 padding_value(%cst : bf16)617 outer_dims_perm = [0, 2, 1]618 inner_dims_pos = [2, 1]619 inner_tiles = [16, 2]620 into %2 : tensor<32x128x?xbf16> -> tensor<32x?x64x16x2xbf16>621 return %pack : tensor<32x?x64x16x2xbf16>622}623// CHECK: #[[$MAP:.+]] = affine_map<()[s0] -> (s0 ceildiv 16)>624//CHECK-LABEL: func.func @linalg_transpose_linalg.pack_multiple_tiles(625// CHECK-SAME: %[[ARG0:.+]]: tensor<?x32x128xbf16>) -> tensor<32x?x64x16x2xbf16> {626// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index627// CHECK-DAG: %[[CST:.+]] = arith.constant 0.000000e+00 : bf16628// CHECK: %[[DIM:.+]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x32x128xbf16>629// CHECK: %[[VAL0:.+]] = affine.apply #[[$MAP]]()[%[[DIM]]]630// CHECK: %[[VAL1:.+]] = tensor.empty(%[[VAL0]]) : tensor<32x?x64x16x2xbf16>631// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]632// CHECK-SAME: padding_value(%[[CST]] : bf16)633// CHECK-SAME: outer_dims_perm = [1, 0, 2]634// CHECK-SAME: inner_dims_pos = [0, 2]635// CHECK-SAME: inner_tiles = [16, 2]636// CHECK-SAME: into %[[VAL1]] : tensor<?x32x128xbf16> -> tensor<32x?x64x16x2xbf16>637// CHECK: return %[[PACK]] : tensor<32x?x64x16x2xbf16>638// CHECK: }639 640// -----641 642func.func @linalg_transpose_linalg.unpack_fold(%arg0: tensor<1x1x4x16xi32>) -> tensor<16x4xi32> {643 %0 = tensor.empty() : tensor<1x1x16x4xi32>644 %transposed = linalg.transpose ins(%arg0 : tensor<1x1x4x16xi32>)645 outs(%0 : tensor<1x1x16x4xi32>)646 permutation = [1, 0, 3, 2]647 %1 = tensor.empty() : tensor<16x4xi32>648 %unpack = linalg.unpack %transposed649 outer_dims_perm = [0, 1]650 inner_dims_pos = [0, 1]651 inner_tiles = [16, 4] into652 %1 : tensor<1x1x16x4xi32> -> tensor<16x4xi32>653 return %unpack : tensor<16x4xi32>654}655//CHECK-LABEL: func.func @linalg_transpose_linalg.unpack_fold(656// CHECK-SAME: %[[ARG0:.+]]: tensor<1x1x4x16xi32>) -> tensor<16x4xi32> {657// CHECK: %[[OUT:.+]] = tensor.empty() : tensor<16x4xi32>658// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[ARG0]]659// CHECK-SAME: outer_dims_perm = [1, 0]660// CHECK-SAME: inner_dims_pos = [1, 0]661// CHECK-SAME: inner_tiles = [4, 16]662// CHECK-SAME: into %[[OUT]] : tensor<1x1x4x16xi32> -> tensor<16x4xi32>663// CHECK: return %[[UNPACK]] : tensor<16x4xi32>664// CHECK: }665 666// -----667 668func.func @linalg_transpose_linalg.unpack_fold_multi_result(%arg0: tensor<1x1x4x16xi32>) -> (tensor<1x1x16x4xi32>, tensor<16x4xi32>) {669 %0 = tensor.empty() : tensor<1x1x16x4xi32>670 %transposed = linalg.transpose ins(%arg0 : tensor<1x1x4x16xi32>)671 outs(%0 : tensor<1x1x16x4xi32>)672 permutation = [1, 0, 3, 2]673 %1 = tensor.empty() : tensor<16x4xi32>674 %unpack = linalg.unpack %transposed675 outer_dims_perm = [0, 1]676 inner_dims_pos = [0, 1]677 inner_tiles = [16, 4] into678 %1 : tensor<1x1x16x4xi32> -> tensor<16x4xi32>679 return %transposed, %unpack : tensor<1x1x16x4xi32>, tensor<16x4xi32>680}681//CHECK-LABEL: func.func @linalg_transpose_linalg.unpack_fold_multi_result(682// CHECK-SAME: %[[ARG0:.+]]: tensor<1x1x4x16xi32>) 683// CHECK: %[[TRANSPOSE:.+]] = linalg.transpose684// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[ARG0]]685// CHECK-SAME: outer_dims_perm = [1, 0]686// CHECK: return %[[TRANSPOSE]], %[[UNPACK]]687// CHECK: }688 689//CONTROL-LABEL: func.func @linalg_transpose_linalg.unpack_fold_multi_result(690// CONTROL: %[[TRANSPOSE:.+]] = linalg.transpose691// CONTROL: %[[UNPACK:.+]] = linalg.unpack %[[TRANSPOSE]]692// CONTROL-SAME: outer_dims_perm = [0, 1]693// CONTROL: return %[[TRANSPOSE]], %[[UNPACK]]694// CONTROL: }695 696// -----697 698func.func @linalg_transpose_linalg.unpack_fold_partial_tile(%arg0: tensor<1x1x4x16xi32>) -> tensor<15x3xi32> {699 %0 = tensor.empty() : tensor<1x1x16x4xi32>700 %transposed = linalg.transpose ins(%arg0 : tensor<1x1x4x16xi32>)701 outs(%0 : tensor<1x1x16x4xi32>)702 permutation = [1, 0, 3, 2]703 %1 = tensor.empty() : tensor<15x3xi32>704 %unpack = linalg.unpack %transposed705 outer_dims_perm = [0, 1]706 inner_dims_pos = [0, 1]707 inner_tiles = [16, 4] into708 %1 : tensor<1x1x16x4xi32> -> tensor<15x3xi32>709 return %unpack : tensor<15x3xi32>710}711//CHECK-LABEL: func.func @linalg_transpose_linalg.unpack_fold_partial_tile(712// CHECK-SAME: %[[ARG0:.+]]: tensor<1x1x4x16xi32>) -> tensor<15x3xi32> {713// CHECK: %[[OUT:.+]] = tensor.empty() : tensor<15x3xi32>714// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[ARG0]]715// CHECK-SAME: outer_dims_perm = [1, 0]716// CHECK-SAME: inner_dims_pos = [1, 0]717// CHECK-SAME: inner_tiles = [4, 16]718// CHECK-SAME: into %[[OUT]] : tensor<1x1x4x16xi32> -> tensor<15x3xi32>719// CHECK: return %[[UNPACK]] : tensor<15x3xi32>720// CHECK: }721 722// -----723 724func.func @linalg_transpose_linalg.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> {725 %transposed = linalg.transpose726 ins(%arg0 : tensor<?x?x?x?xf32>)727 outs(%transpose_dest : tensor<?x?x?x?xf32>)728 permutation = [1, 0, 3, 2]729 730 %unpack = linalg.unpack %transposed731 outer_dims_perm = [1, 0]732 inner_dims_pos = [0, 1]733 inner_tiles = [%tile_p, %tile_q]734 into %unpack_dest : tensor<?x?x?x?xf32> -> tensor<?x?xf32>735 return %unpack : tensor<?x?xf32>736}737// CHECK-LABEL: func.func @linalg_transpose_linalg.unpack_fold_dynamic_outer_dims_tile_dims_tile_sizes(738// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?xf32>, %[[ARG1:.+]]: tensor<?x?x?x?xf32>, %[[ARG2:.+]]: tensor<?x?xf32>,739// CHECK-SAME: %[[IDX1:.+]]: index, %[[IDX2:.+]]: index) -> tensor<?x?xf32> {740// CHECK-DAG: %[[CST1:.+]] = arith.constant 1 : index741// CHECK-DAG: %[[CST0:.+]] = arith.constant 0 : index742// CHECK-DAG: %[[DIM0:.+]] = tensor.dim %[[ARG2]], %[[CST0]] : tensor<?x?xf32>743// CHECK-DAG: %[[DIM1:.+]] = tensor.dim %[[ARG2]], %[[CST1]] : tensor<?x?xf32>744// CHECK: %[[OUT:.+]] = tensor.empty(%[[DIM0]], %[[DIM1]]) : tensor<?x?xf32>745// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[ARG0]]746// CHECK-SAME: outer_dims_perm = [0, 1]747// CHECK-SAME: inner_dims_pos = [1, 0]748// CHECK-SAME: inner_tiles = [%[[IDX2]], %[[IDX1]]]749// CHECK-SAME: into %[[OUT]] : tensor<?x?x?x?xf32> -> tensor<?x?xf32>750// CHECK: return %[[UNPACK]] : tensor<?x?xf32>751// CHECK: }752 753// -----754 755func.func @linalg.unpack_linalg_transpose_fold(%arg0: tensor<56x57x1x64xf32>) -> tensor<3648x56xf32> {756 %0 = tensor.empty() : tensor<56x3648xf32>757 %pack = linalg.unpack %arg0758 outer_dims_perm = [0, 1]759 inner_dims_pos = [0, 1]760 inner_tiles = [1, 64]761 into %0 : tensor<56x57x1x64xf32> -> tensor<56x3648xf32>762 763 %1 = tensor.empty() : tensor<3648x56xf32>764 %transposed = linalg.transpose765 ins(%pack : tensor<56x3648xf32>)766 outs(%1 : tensor<3648x56xf32>)767 permutation = [1,0]768 return %transposed : tensor<3648x56xf32>769}770// CHECK-LABEL: func.func @linalg.unpack_linalg_transpose_fold(771// CHECK-SAME: %[[ARG0:.+]]: tensor<56x57x1x64xf32>) -> tensor<3648x56xf32> {772// CHECK: %[[OUT:.+]] = tensor.empty() : tensor<3648x56xf32>773// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[ARG0]]774// CHECK-SAME: outer_dims_perm = [1, 0]775// CHECK-SAME: inner_dims_pos = [1, 0]776// CHECK-SAME: inner_tiles = [1, 64]777// CHECK-SAME: into %[[OUT:.+]] : tensor<56x57x1x64xf32> -> tensor<3648x56xf32>778// CHECK: return %[[UNPACK]] : tensor<3648x56xf32>779// CHECK: }780 781// -----782 783func.func @tensor_padded_unpack_linalg_transpose_fold(%arg0: tensor<71x7x4x16x16xf32>) -> tensor<100x71x64xf32> {784 %0 = tensor.empty() : tensor<71x100x64xf32>785 %pack = linalg.unpack %arg0786 inner_dims_pos = [1, 2]787 inner_tiles = [16, 16]788 into %0 : tensor<71x7x4x16x16xf32> -> tensor<71x100x64xf32>789 790 %1 = tensor.empty() : tensor<100x71x64xf32>791 %transposed = linalg.transpose792 ins(%pack : tensor<71x100x64xf32>)793 outs(%1 : tensor<100x71x64xf32>)794 permutation = [1, 0, 2]795 return %transposed : tensor<100x71x64xf32>796}797// CHECK-LABEL: func.func @tensor_padded_unpack_linalg_transpose_fold(798// CHECK-SAME: %[[ARG0:.+]]: tensor<71x7x4x16x16xf32>) -> tensor<100x71x64xf32> {799// CHECK: %[[OUT:.+]] = tensor.empty() : tensor<100x71x64xf32>800// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[ARG0]]801// CHECK-SAME: outer_dims_perm = [1, 0, 2]802// CHECK-SAME: inner_dims_pos = [0, 2]803// CHECK-SAME: inner_tiles = [16, 16]804// CHECK-SAME: into %[[OUT:.+]] : tensor<71x7x4x16x16xf32> -> tensor<100x71x64xf32>805// CHECK: return %[[UNPACK]] : tensor<100x71x64xf32>806// CHECK: }807 808// -----809 810func.func @non_involution_transpose_unpack_fold(%arg0: tensor<2x3x5x4x16xi32>) -> tensor<5x48x8xi32> {811 %0 = tensor.empty() : tensor<5x2x3x16x4xi32>812 %transposed = linalg.transpose ins(%arg0 : tensor<2x3x5x4x16xi32>)813 outs(%0 : tensor<5x2x3x16x4xi32>)814 permutation = [2, 0, 1, 4, 3]815 %1 = tensor.empty() : tensor<5x48x8xi32>816 %unpack = linalg.unpack %transposed817 outer_dims_perm = [0, 2, 1]818 inner_dims_pos = [1, 2]819 inner_tiles = [16, 4] into820 %1 : tensor<5x2x3x16x4xi32> -> tensor<5x48x8xi32>821 return %unpack : tensor<5x48x8xi32>822}823//CHECK-LABEL: func.func @non_involution_transpose_unpack_fold(824// CHECK-SAME: %[[ARG0:.+]]: tensor<2x3x5x4x16xi32>) -> tensor<5x48x8xi32> {825// CHECK: %[[OUT:.+]] = tensor.empty() : tensor<5x48x8xi32>826// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[ARG0]]827// CHECK-SAME: outer_dims_perm = [2, 1, 0]828// CHECK-SAME: inner_dims_pos = [2, 1]829// CHECK-SAME: inner_tiles = [4, 16]830// CHEKC-SAME: into %[[OUT]] : tensor<2x3x5x4x16xi32> -> tensor<5x48x8xi32>831// CHECK: return %[[UNPACK]] : tensor<5x48x8xi32>832// CHECK: }833 834// -----835 836func.func @unpack_non_involution_transpose_fold(%arg0: tensor<57x3x56x1x64xf32>) -> tensor<3648x3x56xf32> {837 %0 = tensor.empty() : tensor<3x56x3648xf32>838 %unpack = linalg.unpack %arg0839 outer_dims_perm = [2, 0, 1]840 inner_dims_pos = [1, 2]841 inner_tiles = [1, 64]842 into %0 : tensor<57x3x56x1x64xf32> -> tensor<3x56x3648xf32>843 844 %1 = tensor.empty() : tensor<3648x3x56xf32>845 %transposed = linalg.transpose846 ins(%unpack : tensor<3x56x3648xf32>)847 outs(%1 : tensor<3648x3x56xf32>)848 permutation = [2, 0, 1]849 return %transposed : tensor<3648x3x56xf32>850}851// CHECK-LABEL: func.func @unpack_non_involution_transpose_fold(852// CHECK-SAME: %[[ARG0:.+]]: tensor<57x3x56x1x64xf32>) -> tensor<3648x3x56xf32> {853// CHECK: %[[OUT:.+]] = tensor.empty() : tensor<3648x3x56xf32>854// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[ARG0]]855// CHECK-SAME: outer_dims_perm = [0, 1, 2]856// CHECK-SAME: inner_dims_pos = [2, 0]857// CHECK-SAME: inner_tiles = [1, 64]858// CHECK-SAME: into %[[OUT:.+]] : tensor<57x3x56x1x64xf32> -> tensor<3648x3x56xf32>859// CHECK: return %[[UNPACK]] : tensor<3648x3x56xf32>860// CHECK: }861 862// -----863 864func.func @transpose_unpacked_dims_no_fold(%arg0: tensor<2x16x5x4x3xi32>) -> tensor<5x32x12xi32> {865 %0 = tensor.empty() : tensor<5x2x3x16x4xi32>866 %transposed = linalg.transpose ins(%arg0 : tensor<2x16x5x4x3xi32>)867 outs(%0 : tensor<5x2x3x16x4xi32>)868 permutation = [2, 0, 4, 1, 3]869 %1 = tensor.empty() : tensor<5x32x12xi32>870 %unpack = linalg.unpack %transposed871 inner_dims_pos = [1, 2]872 inner_tiles = [16, 4] into873 %1 : tensor<5x2x3x16x4xi32> -> tensor<5x32x12xi32>874 return %unpack : tensor<5x32x12xi32>875}876//CHECK-LABEL: func.func @transpose_unpacked_dims_no_fold(877// CHECK: linalg.transpose878// CHECK: linalg.unpack879 880// -----881 882#map = affine_map<(d0, d1, d2, d3, d4)->(d1, d2, d0, d4, d3)>883#map1 = affine_map<(d0, d1, d2, d3, d4)->(d0, d1, d2, d3, d4)>884func.func @generic_transpose_unpack_fold(%arg0: tensor<2x3x5x4x16xi32>) -> tensor<5x48x8xi32> {885 %0 = tensor.empty() : tensor<5x2x3x16x4xi32>886 %transposed = linalg.generic {887 iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"],888 indexing_maps = [#map, #map1]}889 ins(%arg0 : tensor<2x3x5x4x16xi32>)890 outs(%0 : tensor<5x2x3x16x4xi32>) {891 ^bb0(%in : i32, %out : i32):892 linalg.yield %in : i32893 } -> tensor<5x2x3x16x4xi32>894 %1 = tensor.empty() : tensor<5x48x8xi32>895 %unpack = linalg.unpack %transposed896 outer_dims_perm = [0, 2, 1]897 inner_dims_pos = [1, 2]898 inner_tiles = [16, 4] into899 %1 : tensor<5x2x3x16x4xi32> -> tensor<5x48x8xi32>900 return %unpack : tensor<5x48x8xi32>901}902//CHECK-LABEL: func.func @generic_transpose_unpack_fold(903// CHECK-SAME: %[[ARG0:.+]]: tensor<2x3x5x4x16xi32>) -> tensor<5x48x8xi32> {904// CHECK: %[[OUT:.+]] = tensor.empty() : tensor<5x48x8xi32>905// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[ARG0]]906// CHECK-SAME: outer_dims_perm = [2, 1, 0]907// CHECK-SAME: inner_dims_pos = [2, 1]908// CHECK-SAME: inner_tiles = [4, 16]909// CHEKC-SAME: into %[[OUT]] : tensor<2x3x5x4x16xi32> -> tensor<5x48x8xi32>910// CHECK: return %[[UNPACK]] : tensor<5x48x8xi32>911// CHECK: }912 913// -----914 915#map = affine_map<(d0, d1, d2)->(d1, d2, d0)>916#map1 = affine_map<(d0, d1, d2)->(d0, d1, d2)>917func.func @unpack_generic_transpose_fold(%arg0: tensor<57x3x56x1x64xf32>) -> tensor<3648x3x56xf32> {918 %0 = tensor.empty() : tensor<3x56x3648xf32>919 %unpack = linalg.unpack %arg0920 outer_dims_perm = [2, 0, 1]921 inner_dims_pos = [1, 2]922 inner_tiles = [1, 64]923 into %0 : tensor<57x3x56x1x64xf32> -> tensor<3x56x3648xf32>924 925 %1 = tensor.empty() : tensor<3648x3x56xf32>926 %transposed = linalg.generic {927 iterator_types = ["parallel", "parallel", "parallel"],928 indexing_maps = [#map, #map1]}929 ins(%unpack : tensor<3x56x3648xf32>)930 outs(%1 : tensor<3648x3x56xf32>) {931 ^bb0(%in : f32, %out : f32):932 linalg.yield %in : f32933 } -> tensor<3648x3x56xf32>934 return %transposed : tensor<3648x3x56xf32>935}936// CHECK-LABEL: func.func @unpack_generic_transpose_fold(937// CHECK-SAME: %[[ARG0:.+]]: tensor<57x3x56x1x64xf32>) -> tensor<3648x3x56xf32> {938// CHECK: %[[OUT:.+]] = tensor.empty() : tensor<3648x3x56xf32>939// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[ARG0]]940// CHECK-SAME: outer_dims_perm = [0, 1, 2]941// CHECK-SAME: inner_dims_pos = [2, 0]942// CHECK-SAME: inner_tiles = [1, 64]943// CHECK-SAME: into %[[OUT:.+]] : tensor<57x3x56x1x64xf32> -> tensor<3648x3x56xf32>944// CHECK: return %[[UNPACK]] : tensor<3648x3x56xf32>945// CHECK: }946