359 lines · plain
1// RUN: mlir-opt -split-input-file \2// RUN: -transform-preload-library='transform-library-paths=%p/td/decompose-pack.mlir' \3// RUN: -transform-interpreter=entry-point=decompose_pack %s | FileCheck %s4 5func.func @simple_KCRS_to_KCRSsr(%arg0: tensor<?x?xi32>, %arg1: tensor<1x1x?x1xi32>) -> tensor<1x1x?x1xi32> {6 %c8 = arith.constant 8 : index7 %c5 = arith.constant 5 : i328 %pack = linalg.pack %arg0 padding_value(%c5 : i32) inner_dims_pos = [0, 1] inner_tiles = [%c8, 1] into %arg1 : tensor<?x?xi32> -> tensor<1x1x?x1xi32>9 return %pack : tensor<1x1x?x1xi32>10}11 12// CHECK: #[[$ATTR_0:.+]] = affine_map<()[s0] -> (-s0 + 8)>13// CHECK: #[[$ATTR_1:.+]] = affine_map<()[s0] -> (-s0 + 1)>14 15// CHECK-LABEL: func.func @simple_KCRS_to_KCRSsr(16// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]: tensor<?x?xi32>,17// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]: tensor<1x1x?x1xi32>) -> tensor<1x1x?x1xi32>18// CHECK: %[[VAL_2:.*]] = arith.constant 1 : index19// CHECK: %[[VAL_3:.*]] = arith.constant 5 : i3220// CHECK: %[[VAL_4:.*]] = arith.constant 0 : index21// CHECK: %[[VAL_5:.*]] = tensor.dim %[[SRC]], %[[VAL_4]] : tensor<?x?xi32>22// CHECK: %[[VAL_6:.*]] = affine.apply #[[$ATTR_0]](){{\[}}%[[VAL_5]]]23// CHECK: %[[VAL_7:.*]] = tensor.dim %[[SRC]], %[[VAL_2]] : tensor<?x?xi32>24// CHECK: %[[VAL_8:.*]] = affine.apply #[[$ATTR_1]](){{\[}}%[[VAL_7]]]25// CHECK: %[[PAD:.*]] = tensor.pad %[[SRC]] low[0, 0] high{{\[}}%[[VAL_6]], %[[VAL_8]]] {26// CHECK: ^bb0(%[[VAL_10:.*]]: index, %[[VAL_11:.*]]: index):27// CHECK: tensor.yield %[[VAL_3]] : i3228// CHECK: } : tensor<?x?xi32> to tensor<8x1xi32>29// CHECK: %[[INSERT:.*]] = tensor.insert_slice %[[PAD:.*]] into %[[DEST]][0, 0, 0, 0] [1, 1, 8, 1] [1, 1, 1, 1] : tensor<8x1xi32> into tensor<1x1x?x1xi32>30// CHECK: return %[[INSERT]] : tensor<1x1x?x1xi32>31 32// -----33 34func.func @NCHW_to_NCHWc(%src: tensor<2x32x16x8xf32>, %dest: tensor<2x1x16x8x32xf32>) -> tensor<2x1x16x8x32xf32> {35 %pack = linalg.pack %src36 inner_dims_pos = [1]37 inner_tiles = [32] into %dest38 : tensor<2x32x16x8xf32> -> tensor<2x1x16x8x32xf32>39 return %pack : tensor<2x1x16x8x32xf32>40}41// CHECK-LABEL: func.func @NCHW_to_NCHWc(42// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]43// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]44// CHECK: %[[INIT:.*]] = tensor.empty() : tensor<2x16x8x32xf32>45// CHECK: %[[TR:.*]] = linalg.transpose ins(%[[SRC]] : tensor<2x32x16x8xf32>) outs(%[[INIT]] : tensor<2x16x8x32xf32>) permutation = [0, 2, 3, 1]46// CHECK: %[[RES:.*]] = tensor.insert_slice %[[TR]] into %[[DEST]]47// CHECK-SAME: [0, 0, 0, 0, 0] [2, 1, 16, 8, 32] [1, 1, 1, 1, 1]48// CHECK-SAME: : tensor<2x16x8x32xf32> into tensor<2x1x16x8x32xf32>49// CHECK: return %[[RES]] : tensor<2x1x16x8x32xf32>50 51// -----52 53func.func @simple_pad_and_pack_static_tiles(%input: tensor<5x1xf32>, %output: tensor<1x1x8x2xf32>, %pad: f32) -> tensor<1x1x8x2xf32> {54 %0 = linalg.pack %input padding_value(%pad : f32) inner_dims_pos = [0, 1] inner_tiles = [8, 2] into %output : tensor<5x1xf32> -> tensor<1x1x8x2xf32>55 return %0 : tensor<1x1x8x2xf32>56}57// CHECK: #[[$ATTR_0:.+]] = affine_map<()[s0] -> (s0 - 5)>58// CHECK: #[[$ATTR_1:.+]] = affine_map<()[s0] -> (s0 - 1)>59 60// CHECK-LABEL: func.func @simple_pad_and_pack_static_tiles61// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]62// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]63// CHECK-SAME: %[[PAD_VAL:[a-zA-Z0-9]+]]64// CHECK: %[[PAD:.+]] = tensor.pad %[[SRC]] low[0, 0] high[3, 1]65// CHECK: tensor.yield %[[PAD_VAL]]66// CHECK-NOT: linalg.transpose67// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[PAD]] into %[[DEST]]68// CHECK-SAME: [0, 0, 0, 0] [1, 1, 8, 2] [1, 1, 1, 1]69// CHECK: return %[[INSERT]]70 71/// Same as example above, but with 1 dynamic tile size.72 73func.func @simple_pad_and_pack_dynamic_tile(%input: tensor<5x1xf32>, %output: tensor<1x1x?x2xf32>, %pad: f32, %tile_dim_0: index) -> tensor<1x1x?x2xf32> {74 %0 = linalg.pack %input padding_value(%pad : f32) inner_dims_pos = [0, 1] inner_tiles = [%tile_dim_0, 2] into %output : tensor<5x1xf32> -> tensor<1x1x?x2xf32>75 return %0 : tensor<1x1x?x2xf32>76}77// CHECK-LABEL: func.func @simple_pad_and_pack_dynamic_tile(78// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]79// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]80// CHECK-SAME: %[[PAD_VAL:[a-zA-Z0-9]+]]81// CHECK-SAME: %[[TILE_DIM_0:[a-zA-Z0-9]+]]: index) -> tensor<1x1x?x2xf32> {82// CHECK: %[[PAD_HIGH:.*]] = affine.apply #[[$ATTR_0]](){{\[}}%[[TILE_DIM_0]]]83// CHECK: %[[PAD:.*]] = tensor.pad %[[SRC]] low[0, 0] high{{\[}}%[[PAD_HIGH]], 1] {84// CHECK: tensor.yield %[[PAD_VAL]] : f3285// CHECK-NOT: linalg.transpose86// CHECK: %[[RES:.*]] = tensor.insert_slice %[[PAD:.*]] into %[[DEST]][0, 0, 0, 0] [1, 1, %[[TILE_DIM_0]], 2] [1, 1, 1, 1] : tensor<?x2xf32> into tensor<1x1x?x2xf32>87// CHECK: return %[[RES]] : tensor<1x1x?x2xf32>88 89/// Same as example above, but the dynamic tile size is a compile-time constant90/// that's folded away.91 92func.func @simple_pad_and_pack_dynamic_tile_cst(%input: tensor<5x1xf32>, %output: tensor<1x1x?x2xf32>, %pad: f32) -> tensor<1x1x?x2xf32> {93 %tile_dim_0 = arith.constant 8 : index94 %0 = linalg.pack %input padding_value(%pad : f32) inner_dims_pos = [0, 1] inner_tiles = [%tile_dim_0, 2] into %output : tensor<5x1xf32> -> tensor<1x1x?x2xf32>95 return %0 : tensor<1x1x?x2xf32>96}97// CHECK-LABEL: func.func @simple_pad_and_pack_dynamic_tile_cst(98// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]99// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]100// CHECK-SAME: %[[PAD_VAL:[a-zA-Z0-9]+]]: f32) -> tensor<1x1x?x2xf32> {101// CHECK: %[[PAD:.*]] = tensor.pad %[[SRC]] low[0, 0] high[3, 1] {102// CHECK: tensor.yield %[[PAD_VAL]] : f32103// CHECK: } : tensor<5x1xf32> to tensor<8x2xf32>104// CHECK: %[[RES:.*]] = tensor.insert_slice %[[PAD:.*]] into %[[DEST]][0, 0, 0, 0] [1, 1, 8, 2] [1, 1, 1, 1] : tensor<8x2xf32> into tensor<1x1x?x2xf32>105// CHECK: return %[[RES]] : tensor<1x1x?x2xf32>106 107func.func @simple_pad_and_pack_dynamic_tile_transpose(%input: tensor<5x1xf32>, %output: tensor<1x1x2x?xf32>, %pad: f32, %tile_dim_1: index) -> tensor<1x1x2x?xf32> {108 %0 = linalg.pack %input padding_value(%pad : f32) inner_dims_pos = [1, 0] inner_tiles = [2, %tile_dim_1] into %output : tensor<5x1xf32> -> tensor<1x1x2x?xf32>109 return %0 : tensor<1x1x2x?xf32>110}111// CHECK-LABEL: func.func @simple_pad_and_pack_dynamic_tile_transpose(112// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]113// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]114// CHECK-SAME: %[[PAD_VAL:[a-zA-Z0-9]+]]115// CHECK-SAME: %[[TILE_DIM_1:[a-zA-Z0-9]+]]: index) -> tensor<1x1x2x?xf32> {116// CHECK: %[[PAD_HIGH:.*]] = affine.apply #[[$ATTR_0]](){{\[}}%[[TILE_DIM_1]]]117// CHECK: %[[PAD:.*]] = tensor.pad %[[SRC]] low[0, 0] high{{\[}}%[[PAD_HIGH]], 1] {118// CHECK: tensor.yield %[[PAD_VAL]] : f32119// CHECK-NEXT: } : tensor<5x1xf32> to tensor<?x2xf32>120// CHECK: %[[EMPTY:.*]] = tensor.empty(%[[TILE_DIM_1]]) : tensor<2x?xf32>121// CHECK: %[[TR:.*]] = linalg.transpose122// CHECK-SAME: ins(%[[PAD:.*]] : tensor<?x2xf32>)123// CHECK-SAME: outs(%[[EMPTY]] : tensor<2x?xf32>)124// CHECK-SAME: permutation = [1, 0]125// CHECK: %[[RES:.*]] = tensor.insert_slice %[[TR]] into %[[DEST]][0, 0, 0, 0] [1, 1, 2, %[[TILE_DIM_1]]] [1, 1, 1, 1] : tensor<2x?xf32> into tensor<1x1x2x?xf32>126// CHECK: return %[[RES]] : tensor<1x1x2x?xf32>127 128/// Same as example above, but with 1 scalable tile size.129 130/// NOTE: For this example to make sense in practice, the "?" in the output shape131/// should effectively be 8 * vector.vscale (and that's what tensor.dim132/// below should return).133 134func.func @simple_pad_and_pack_scalable_tile(%input: tensor<5x1xf32>, %output: tensor<1x1x?x2xf32>, %pad: f32) -> tensor<1x1x?x2xf32> {135 %c8 = arith.constant 8 : index136 %vscale = vector.vscale137 %c8_vscale = arith.muli %vscale, %c8 : index138 %0 = linalg.pack %input padding_value(%pad : f32) inner_dims_pos = [0, 1] inner_tiles = [%c8_vscale, 2] into %output : tensor<5x1xf32> -> tensor<1x1x?x2xf32>139 return %0 : tensor<1x1x?x2xf32>140}141 142// CHECK-LABEL: func.func @simple_pad_and_pack_scalable_tile(143// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]: tensor<5x1xf32>,144// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]: tensor<1x1x?x2xf32>,145// CHECK-SAME: %[[PAD_VAL:[a-zA-Z0-9]+]]: f32) -> tensor<1x1x?x2xf32> {146// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index147// CHECK-DAG: %[[VS:.+]] = vector.vscale148// CHECK: %[[C8_VS:.+]] = arith.muli %[[VS]], %[[C8]] : index149// CHECK: %[[PAD_HIGH:.*]] = affine.apply #[[$ATTR_0]](){{\[}}%[[C8_VS]]]150// CHECK: %[[PAD:.+]] = tensor.pad %[[SRC]] low[0, 0] high{{\[}}%[[PAD_HIGH]], 1] {151// CHECK: tensor.yield %[[PAD_VAL]] : f32152// CHECK-NOT: linalg.transpose153// CHECK: %[[RES:.*]] = tensor.insert_slice %[[PAD:.*]] into %[[DEST]][0, 0, 0, 0] [1, 1, %[[C8_VS]], 2] [1, 1, 1, 1] : tensor<?x2xf32> into tensor<1x1x?x2xf32>154// CHECK: return %[[RES]] : tensor<1x1x?x2xf32>155 156 157/// Same as example above, but with both tile sizes dynamic.158 159func.func @simple_pad_and_pack_dynamic_tiles(%input: tensor<5x1xf32>, %output: tensor<1x1x?x?xf32>, %pad: f32, %tile_dim_0: index, %tile_dim_1: index) -> tensor<1x1x?x?xf32> {160 %0 = linalg.pack %input padding_value(%pad : f32) inner_dims_pos = [0, 1] inner_tiles = [%tile_dim_0, %tile_dim_1] into %output : tensor<5x1xf32> -> tensor<1x1x?x?xf32>161 return %0 : tensor<1x1x?x?xf32>162}163// CHECK-LABEL: func.func @simple_pad_and_pack_dynamic_tiles(164// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]: tensor<5x1xf32>,165// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]: tensor<1x1x?x?xf32>,166// CHECK-SAME: %[[PAD_VAL:[a-zA-Z0-9]+]]: f32,167// CHECK-SAME: %[[TILE_DIM_0:[a-zA-Z0-9]+]]: index,168// CHECK-SAME: %[[TILE_DIM_1:[a-zA-Z0-9]+]]: index) -> tensor<1x1x?x?xf32> {169// CHECK: %[[PAD_HIGH_1:.*]] = affine.apply #[[$ATTR_0]](){{\[}}%[[TILE_DIM_0]]]170// CHECK: %[[PAD_HIGH_2:.*]] = affine.apply #[[$ATTR_1]](){{\[}}%[[TILE_DIM_1]]]171// CHECK: %[[PAD:.*]] = tensor.pad %[[SRC]] low[0, 0] high{{\[}}%[[PAD_HIGH_1]], %[[PAD_HIGH_2]]] {172// CHECK: tensor.yield %[[PAD_VAL]] : f32173// CHECK-NOT: linalg.transpose174// CHECK: %[[RES:.*]] = tensor.insert_slice %[[PAD]] into %[[DEST]][0, 0, 0, 0] [1, 1, %[[TILE_DIM_0]], %[[TILE_DIM_1]]] [1, 1, 1, 1] : tensor<?x?xf32> into tensor<1x1x?x?xf32>175// CHECK: return %[[RES]] : tensor<1x1x?x?xf32>176 177// -----178 179// Note - un-tiled outer dims are permueted. However, these are unit dims, which is supported.180 181func.func @simple_pad_and_pack_dynamic_tile_not_all_dims_tiled(%input: tensor<1x1x5x1xf32>, %output: tensor<1x1x1x1x2x?xf32>, %pad: f32, %high: index) -> tensor<1x1x1x1x2x?xf32> {182 %0 = linalg.pack %input padding_value(%pad : f32) outer_dims_perm = [1, 0, 2, 3] inner_dims_pos = [3, 2] inner_tiles = [2, %high] into %output : tensor<1x1x5x1xf32> -> tensor<1x1x1x1x2x?xf32>183 return %0 : tensor<1x1x1x1x2x?xf32>184}185// CHECK: #[[$ATTR_2:.+]] = affine_map<()[s0] -> (s0 - 5)>186// CHECK-LABEL: func.func @simple_pad_and_pack_dynamic_tile_not_all_dims_tiled187// CHECK-SAME: %[[VAL_0:.*]]: tensor<1x1x5x1xf32>,188// CHECK-SAME: %[[VAL_1:.*]]: tensor<1x1x1x1x2x?xf32>,189// CHECK-SAME: %[[VAL_2:.*]]: f32,190// CHECK-SAME: %[[VAL_3:.*]]: index) -> tensor<1x1x1x1x2x?xf32> {191// CHECK: %[[VAL_4:.*]] = affine.apply #[[$ATTR_2]](){{\[}}%[[VAL_3]]]192// CHECK: %[[VAL_5:.*]] = tensor.pad %[[VAL_0]] low[0, 0, 0, 0] high[0, 0, %[[VAL_4]], 1] {193// CHECK: ^bb0(%[[VAL_6:.*]]: index, %[[VAL_7:.*]]: index, %[[VAL_8:.*]]: index, %[[VAL_9:.*]]: index):194// CHECK: tensor.yield %[[VAL_2]] : f32195// CHECK: } : tensor<1x1x5x1xf32> to tensor<1x1x?x2xf32>196// CHECK: %[[VAL_10:.*]] = tensor.empty(%[[VAL_3]]) : tensor<1x1x2x?xf32>197// CHECK: %[[VAL_11:.*]] = linalg.transpose198// CHECK-SAME: ins(%[[VAL_12:.*]] : tensor<1x1x?x2xf32>)199// CHECK-SAME: outs(%[[VAL_10]] : tensor<1x1x2x?xf32>)200// CHECK-SAME: permutation = [0, 1, 3, 2]201// CHECK: %[[VAL_13:.*]] = tensor.insert_slice %[[VAL_11]] into %[[VAL_1]][0, 0, 0, 0, 0, 0] [1, 1, 1, 1, 2, %[[VAL_3]]] [1, 1, 1, 1, 1, 1] : tensor<1x1x2x?xf32> into tensor<1x1x1x1x2x?xf32>202// CHECK: return %[[VAL_13]] : tensor<1x1x1x1x2x?xf32>203 204// -----205 206// Similar as the example above, but one of the un-tiled outer dims that are permuted is non-unit: (7,1) -> (1, 7)207 208func.func @negative_not_all_dims_tiled_outer_dim_0_permuted(%input: tensor<7x1x5x1xf32>, %output: tensor<1x7x1x1x2x?xf32>, %pad: f32, %high: index) -> tensor<1x7x1x1x2x?xf32> {209 %0 = linalg.pack %input padding_value(%pad : f32) outer_dims_perm = [1, 0, 2, 3] inner_dims_pos = [3, 2] inner_tiles = [2, %high] into %output : tensor<7x1x5x1xf32> -> tensor<1x7x1x1x2x?xf32>210 return %0 : tensor<1x7x1x1x2x?xf32>211}212// CHECK-LABEL: func.func @negative_not_all_dims_tiled_outer_dim_0_permuted213// CHECK: linalg.pack214 215// -----216 217// Similar as the example above, but one of the un-tiled outer dims that are permuted is non-unit: (1, 7) -> (7, 1).218 219func.func @negative_not_all_dims_tiled_outer_dim_1_permuted(%input: tensor<1x7x5x1xf32>, %output: tensor<7x1x1x1x2x?xf32>, %pad: f32, %high: index) -> tensor<7x1x1x1x2x?xf32> {220 %0 = linalg.pack %input padding_value(%pad : f32) outer_dims_perm = [1, 0, 2, 3] inner_dims_pos = [3, 2] inner_tiles = [2, %high] into %output : tensor<1x7x5x1xf32> -> tensor<7x1x1x1x2x?xf32>221 return %0 : tensor<7x1x1x1x2x?xf32>222}223// CHECK-LABEL: func.func @negative_not_all_dims_tiled_outer_dim_1_permuted224// CHECK: linalg.pack225 226// -----227 228func.func @simple_NC_to_CNnc(%arg0: tensor<32x8xf32>, %arg1: tensor<1x1x32x8xf32>) -> tensor<1x1x32x8xf32>{229 %0 = linalg.pack %arg0 outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [32, 8] into %arg1 : tensor<32x8xf32> -> tensor<1x1x32x8xf32>230 return %0 : tensor<1x1x32x8xf32>231}232// CHECK-LABEL: func.func @simple_NC_to_CNnc233// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]234// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]235// CHECK-NOT: linalg.transpose236// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[SRC]] into %[[DEST]]237// CHECK-SAME: [0, 0, 0, 0] [1, 1, 32, 8] [1, 1, 1, 1]238// CHECK: return %[[INSERT]]239 240// -----241 242func.func @simple_CHW_to_CHWhwc(%arg0: tensor<3x5x7xf32>, %arg1: tensor<1x1x1x5x7x3xf32>) -> tensor<1x1x1x5x7x3xf32> {243 %0 = linalg.pack %arg0 inner_dims_pos = [1, 2, 0] inner_tiles = [5, 7, 3] into %arg1 : tensor<3x5x7xf32> -> tensor<1x1x1x5x7x3xf32>244 return %0 : tensor<1x1x1x5x7x3xf32>245}246// CHECK-LABEL: func.func @simple_CHW_to_CHWhwc247// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]248// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]249// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<5x7x3xf32>250// CHECK: %[[TRANSP:.+]] = linalg.transpose251// CHECK-SAME: ins(%[[SRC]] : tensor<3x5x7xf32>)252// CHECK-SAME: outs(%[[EMPTY]] : tensor<5x7x3xf32>)253// CHECK-SAME: permutation = [1, 2, 0]254// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[TRANSP]] into %[[DEST]]255// CHECK-SAME: [0, 0, 0, 0, 0, 0] [1, 1, 1, 5, 7, 3] [1, 1, 1, 1, 1, 1]256// CHECK: return %[[INSERT]]257 258// -----259 260func.func @simple_KCRS_to_KRSCsr(%arg0: tensor<1x1x32x8xf32>, %arg1: tensor<1x1x1x1x8x32xf32>) -> tensor<1x1x1x1x8x32xf32> {261 %0 = linalg.pack %arg0 outer_dims_perm = [0, 2, 3, 1] inner_dims_pos = [3, 2] inner_tiles = [8, 32] into %arg1 : tensor<1x1x32x8xf32> -> tensor<1x1x1x1x8x32xf32>262 return %0 : tensor<1x1x1x1x8x32xf32>263}264// CHECK-LABEL: func.func @simple_KCRS_to_KRSCsr265// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]266// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]267// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<1x1x8x32xf32>268// CHECK: %[[TRANSP:.+]] = linalg.transpose269// CHECK-SAME: ins(%[[SRC]] : tensor<1x1x32x8xf32>270// CHECK-SAME: outs(%[[EMPTY]] : tensor<1x1x8x32xf32>)271// CHECK-SAME: permutation = [0, 1, 3, 2]272// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[TRANSP]] into %[[DEST]]273// CHECK-SAME: [0, 0, 0, 0, 0, 0] [1, 1, 1, 1, 8, 32] [1, 1, 1, 1, 1, 1]274// CHECK: return %[[INSERT]]275 276// -----277 278// The following example shows a pack operation that is defined with inner279// dimension positions that are not adjacent, i.e. `[2, 0]`. And the outer280// dimensions of the packed tensor are of unit values, i.e. `1x1x1`.281func.func @pack_with_non_adjacent_inner_dims_pos_and_unit_outer(%arg0: tensor<1x1x4xf32>, %arg1: tensor<1x1x1x4x1xf32>) -> tensor<1x1x1x4x1xf32> {282 %pack = linalg.pack %arg0 outer_dims_perm = [1, 2, 0] inner_dims_pos = [2, 0] inner_tiles = [4, 1] into %arg1 : tensor<1x1x4xf32> -> tensor<1x1x1x4x1xf32>283 return %pack : tensor<1x1x1x4x1xf32>284}285// CHECK-LABEL: func.func @pack_with_non_adjacent_inner_dims_pos_and_unit_outer286// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]287// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]288// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<1x4x1xf32>289// CHECK: %[[TRANSP:.+]] = linalg.transpose290// CHECK-SAME: ins(%[[SRC]] : tensor<1x1x4xf32>)291// CHECK-SAME: outs(%[[EMPTY]] : tensor<1x4x1xf32>)292// CHECK-SAME: permutation = [1, 2, 0]293// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[TRANSP]] into %[[DEST]]294// CHECK-SAME: [0, 0, 0, 0, 0] [1, 1, 1, 4, 1] [1, 1, 1, 1, 1] : tensor<1x4x1xf32> into tensor<1x1x1x4x1xf32>295// CHECK: return %[[INSERT]]296 297// -----298 299// The following example shows a pack operation where the inner dimension300// positions are specified as [2, 1] which are termed adjacent trailing301// dimensions as they contain the last dimension of the source tensor with a302// neighboring dimension. [1, 2] would also be considered trailing adjacent.303// And the outer dimensions of the packed tensor are all set to unit values304// of `1x1x1`.305func.func @pack_with_adjacent_trailing_dimensions_inner_dims_pos_and_unit_outer(%arg0: tensor<1x1x4xf32>, %arg1: tensor<1x1x1x4x1xf32>) -> tensor<1x1x1x4x1xf32> {306 %pack = linalg.pack %arg0 outer_dims_perm = [1, 2, 0] inner_dims_pos = [2, 1] inner_tiles = [4, 1] into %arg1 : tensor<1x1x4xf32> -> tensor<1x1x1x4x1xf32>307 return %pack : tensor<1x1x1x4x1xf32>308}309// CHECK-LABEL: func.func @pack_with_adjacent_trailing_dimensions_inner_dims_pos_and_unit_outer310// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]311// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]312// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<1x4x1xf32>313// CHECK: %[[TRANSP:.+]] = linalg.transpose314// CHECK-SAME: ins(%[[SRC]] : tensor<1x1x4xf32>)315// CHECK-SAME: outs(%[[EMPTY]] : tensor<1x4x1xf32>)316// CHECK-SAME: permutation = [0, 2, 1]317// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[TRANSP]] into %[[DEST]]318// CHECK-SAME: [0, 0, 0, 0, 0] [1, 1, 1, 4, 1] [1, 1, 1, 1, 1] : tensor<1x4x1xf32> into tensor<1x1x1x4x1xf32>319// CHECK: return %[[INSERT]]320 321// -----322 323// The following example shows a pack operation where the inner dims324// positions are non-adjacent and non-permuted.325func.func @pack_with_non_adjacent_and_non_permuted_inner_dims(%arg0: tensor<8x1x1x1xf32>, %arg1:tensor<1x1x1x1x8x1xf32>) -> tensor<1x1x1x1x8x1xf32> {326 %pack = linalg.pack %arg0 outer_dims_perm = [0, 1, 2, 3] inner_dims_pos = [0, 3] inner_tiles = [8, 1] into %arg1: tensor<8x1x1x1xf32> -> tensor<1x1x1x1x8x1xf32>327 return %pack : tensor<1x1x1x1x8x1xf32>328}329 330// CHECK-LABEL: func.func @pack_with_non_adjacent_and_non_permuted_inner_dims331// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]332// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]333// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<1x1x8x1xf32>334// CHECK: %[[TRANSP:.+]] = linalg.transpose335// CHECK-SAME: ins(%[[SRC]] : tensor<8x1x1x1xf32>)336// CHECK-SAME: outs(%[[EMPTY]] : tensor<1x1x8x1xf32>)337// CHECK-SAME: permutation = [1, 2, 0, 3]338// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[TRANSP]] into %[[DEST]]339// CHECK-SAME: [0, 0, 0, 0, 0, 0] [1, 1, 1, 1, 8, 1] [1, 1, 1, 1, 1, 1] : tensor<1x1x8x1xf32> into tensor<1x1x1x1x8x1xf32>340// CHECK: return %[[INSERT]]341 342// -----343 344/// Note "126", which is a non-unit tiled-outer-dim. This is not supported.345 346func.func @negative_non_unit_tiled_outer_dim(%dest: tensor<1x126x1x1x8xf32>, %src: tensor<1x1x1x1001xf32>, %pad: f32) -> tensor<1x126x1x1x8xf32> {347 %pack = linalg.pack %src348 padding_value(%pad : f32)349 outer_dims_perm = [0, 3, 2, 1]350 inner_dims_pos = [3]351 inner_tiles = [8]352 into %dest353 : tensor<1x1x1x1001xf32> -> tensor<1x126x1x1x8xf32>354 355 return %pack : tensor<1x126x1x1x8xf32>356}357// CHECK-LABEL: @negative_non_unit_tiled_outer_dim(358// CHECK: linalg.pack359