1637 lines · plain
1// RUN: mlir-opt %s -test-linalg-data-layout-propagation -split-input-file | FileCheck %s2 3#map0 = affine_map<(d0, d1) -> (d0, d1)>4func.func @dynamic_elem_pack(%arg0: tensor<?x?xf32>, %dest: tensor<?x?x8x2xf32>) -> tensor<?x?x8x2xf32>5{6 %c0 = arith.constant 0 : index7 %c1 = arith.constant 1 : index8 %0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>9 %1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>10 %2 = tensor.empty(%0, %1) : tensor<?x?xf32>11 %3 = linalg.generic {indexing_maps = [#map0, #map0], iterator_types = ["parallel", "parallel"]}12 ins(%arg0 : tensor<?x?xf32>)13 outs(%2 : tensor<?x?xf32>) {14 ^bb0(%arg3: f32, %arg4: f32):15 %4 = arith.addf %arg3, %arg3 : f3216 linalg.yield %4 : f3217 } -> tensor<?x?xf32>18 %4 = linalg.pack %319 inner_dims_pos = [0, 1]20 inner_tiles = [8, 2]21 into %dest : tensor<?x?xf32> -> tensor<?x?x8x2xf32>22 return %4 : tensor<?x?x8x2xf32>23}24// CHECK-DAG: #[[$MAP0:.+]] = affine_map<()[s0] -> (s0 ceildiv 8)>25// CHECK-DAG: #[[$MAP1:.+]] = affine_map<()[s0] -> (s0 ceildiv 2)>26// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>27// CHECK-LABEL: func.func @dynamic_elem_pack28// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]29// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]30// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index31// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index32// CHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]]33// CHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG0]], %[[C1]]34// CHECK-DAG: %[[OUTER_D0:.+]] = affine.apply #[[$MAP0]]()[%[[D0]]]35// CHECK-DAG: %[[OUTER_D1:.+]] = affine.apply #[[$MAP1]]()[%[[D1]]]36// CHECK: %[[ARG0_EMPTY:.+]] = tensor.empty(%[[OUTER_D0]], %[[OUTER_D1]]) : tensor<?x?x8x2xf32>37// CHECK: %[[PACK_ARG0:.+]] = linalg.pack %[[ARG0]]38// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [8, 2]39// CHECK-SAME: into %[[ARG0_EMPTY]]40// CHECK: %[[ELEM:.+]] = linalg.generic41// CHECK-SAME: indexing_maps = [#[[$MAP2]], #[[$MAP2]]]42// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]43// CHECK-SAME: ins(%[[PACK_ARG0]]44// CHECK-SAME: outs(%[[DEST]]45// CHECK: return %[[ELEM]] : tensor<?x?x8x2xf32>46 47// -----48 49#map0 = affine_map<(d0, d1) -> (d0, d1)>50func.func @dynamic_elem_pack_padding_value(%arg0: tensor<?x?xf32>, %dest: tensor<?x?x8x2xf32>) -> tensor<?x?x8x2xf32>51{52 %c0 = arith.constant 0 : index53 %c1 = arith.constant 1 : index54 %cst = arith.constant 3.000000e+00 : f3255 %0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>56 %1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>57 %2 = tensor.empty(%0, %1) : tensor<?x?xf32>58 %3 = linalg.generic {indexing_maps = [#map0, #map0], iterator_types = ["parallel", "parallel"]}59 ins(%arg0 : tensor<?x?xf32>)60 outs(%2 : tensor<?x?xf32>) {61 ^bb0(%arg3: f32, %arg4: f32):62 %4 = arith.addf %arg3, %arg3 : f3263 linalg.yield %4 : f3264 } -> tensor<?x?xf32>65 %4 = linalg.pack %3 padding_value(%cst : f32)66 inner_dims_pos = [0, 1]67 inner_tiles = [8, 2]68 into %dest : tensor<?x?xf32> -> tensor<?x?x8x2xf32>69 return %4 : tensor<?x?x8x2xf32>70}71// CHECK-LABEL: func.func @dynamic_elem_pack_padding_value72// CHECK: %[[GENERIC:.+]] = linalg.generic73// CHECK: linalg.pack %[[GENERIC]]74 75// -----76 77#map0 = affine_map<(d0, d1) -> (d0, d1)>78func.func @elem_pack_transpose_inner_dims(%arg0: tensor<128x256xi32>, %dest: tensor<4x16x16x32xi32>) -> tensor<4x16x16x32xi32>{79 %init = tensor.empty() : tensor<128x256xi32>80 %elem = linalg.generic {indexing_maps = [#map0, #map0], iterator_types = ["parallel", "parallel"]}81 ins(%arg0 : tensor<128x256xi32>)82 outs(%init : tensor<128x256xi32>) {83 ^bb0(%arg3: i32, %arg4: i32):84 %4 = arith.addi %arg3, %arg3 : i3285 linalg.yield %4 : i3286 } -> tensor<128x256xi32>87 %pack = linalg.pack %elem88 inner_dims_pos = [1, 0]89 inner_tiles = [16, 32]90 into %dest : tensor<128x256xi32> -> tensor<4x16x16x32xi32>91 return %pack : tensor<4x16x16x32xi32>92}93// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>94// CHECK-LABEL: func.func @elem_pack_transpose_inner_dims95// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]96// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]97// CHECK: %[[ARG0_EMPTY:.+]] = tensor.empty() : tensor<4x16x16x32xi32>98// CHECK: %[[PACK_ARG0:.+]] = linalg.pack %[[ARG0]]99// CHECK-SAME: inner_dims_pos = [1, 0] inner_tiles = [16, 32]100// CHECK-SAME: into %[[ARG0_EMPTY]]101// CHECK: %[[ELEM:.+]] = linalg.generic102// CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP]]]103// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]104// CHECK-SAME: ins(%[[PACK_ARG0]]105// CHECK-SAME: outs(%[[DEST]]106// CHECK: return %[[ELEM]] : tensor<4x16x16x32xi32>107 108// -----109 110#map0 = affine_map<(d0, d1) -> (d0, d1)>111func.func @elem_pack_transpose_outer_dims(%arg0: tensor<128x256xi32>, %dest: tensor<16x4x32x16xi32>) -> tensor<16x4x32x16xi32>{112 %init = tensor.empty() : tensor<128x256xi32>113 %elem = linalg.generic {indexing_maps = [#map0, #map0], iterator_types = ["parallel", "parallel"]}114 ins(%arg0 : tensor<128x256xi32>)115 outs(%init : tensor<128x256xi32>) {116 ^bb0(%arg3: i32, %arg4: i32):117 %4 = arith.addi %arg3, %arg3 : i32118 linalg.yield %4 : i32119 } -> tensor<128x256xi32>120 %pack = linalg.pack %elem121 outer_dims_perm = [1, 0]122 inner_dims_pos = [0, 1]123 inner_tiles = [32, 16]124 into %dest : tensor<128x256xi32> -> tensor<16x4x32x16xi32>125 return %pack : tensor<16x4x32x16xi32>126}127// CHECK-DAG: #[[$MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>128// CHECK-LABEL: func.func @elem_pack_transpose_outer_dims129// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]130// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]131// CHECK: %[[ARG0_EMPTY:.+]] = tensor.empty() : tensor<16x4x32x16xi32>132// CHECK: %[[PACK_ARG0:.+]] = linalg.pack %[[ARG0]]133// CHECK-SAME: outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [32, 16]134// CHECK-SAME: into %[[ARG0_EMPTY]] : tensor<128x256xi32> -> tensor<16x4x32x16xi32>135// CHECK: %[[ELEM:.+]] = linalg.generic136// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP0]]]137// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]138// CHECK-SAME: ins(%[[PACK_ARG0]]139// CHECK-SAME: outs(%[[DEST]]140// CHECK: return %[[ELEM]] : tensor<16x4x32x16xi32>141 142// -----143 144#map0 = affine_map<(d0, d1) -> (d0, d1)>145func.func @elem_pack_transpose_inner_and_outer_dims(%arg0: tensor<128x256xi32>, %dest: tensor<16x4x16x32xi32>) -> tensor<16x4x16x32xi32>{146 %init = tensor.empty() : tensor<128x256xi32>147 %elem = linalg.generic {indexing_maps = [#map0, #map0], iterator_types = ["parallel", "parallel"]}148 ins(%arg0 : tensor<128x256xi32>)149 outs(%init : tensor<128x256xi32>) {150 ^bb0(%arg3: i32, %arg4: i32):151 %4 = arith.addi %arg3, %arg3 : i32152 linalg.yield %4 : i32153 } -> tensor<128x256xi32>154 %pack = linalg.pack %elem155 outer_dims_perm = [1, 0]156 inner_dims_pos = [1, 0]157 inner_tiles = [16, 32]158 into %dest : tensor<128x256xi32> -> tensor<16x4x16x32xi32>159 return %pack : tensor<16x4x16x32xi32>160}161// CHECK-DAG: #[[$MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>162// CHECK-LABEL: func.func @elem_pack_transpose_inner_and_outer_dims163// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]164// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]165// CHECK: %[[ARG0_EMPTY:.+]] = tensor.empty() : tensor<16x4x16x32xi32>166// CHECK: %[[PACK_ARG0:.+]] = linalg.pack %[[ARG0]]167// CHECK-SAME: outer_dims_perm = [1, 0] inner_dims_pos = [1, 0] inner_tiles = [16, 32]168// CHECK-SAME: into %[[ARG0_EMPTY]]169// CHECK: %[[ELEM:.+]] = linalg.generic170// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP0]]]171// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]172// CHECK-SAME: ins(%[[PACK_ARG0]]173// CHECK-SAME: outs(%[[DEST]]174// CHECK: return %[[ELEM]] : tensor<16x4x16x32xi32>175 176// -----177 178#map0 = affine_map<(d0, d1) -> (d0, d1)>179#map1 = affine_map<(d0, d1) -> (d0)>180#map2 = affine_map<(d0, d1) -> (d1)>181func.func @dynamic_broadcast_pack(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>, %dest: tensor<?x?x8x2xf32>) -> tensor<?x?x8x2xf32>182{183 %c0 = arith.constant 0 : index184 %0 = tensor.dim %arg0, %c0 : tensor<?xf32>185 %1 = tensor.dim %arg1, %c0 : tensor<?xf32>186 %2 = tensor.empty(%0, %1) : tensor<?x?xf32>187 %3 = linalg.generic {indexing_maps = [#map1, #map2, #map0], iterator_types = ["parallel", "parallel"]}188 ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>)189 outs(%2 : tensor<?x?xf32>) {190 ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):191 %4 = arith.addf %arg3, %arg4 : f32192 linalg.yield %4 : f32193 } -> tensor<?x?xf32>194 %4 = linalg.pack %3195 inner_dims_pos = [0, 1]196 inner_tiles = [8, 2]197 into %dest : tensor<?x?xf32> -> tensor<?x?x8x2xf32>198 return %4 : tensor<?x?x8x2xf32>199}200// CHECK-DAG: #[[$MAP0:.+]] = affine_map<()[s0] -> (s0 ceildiv 8)>201// CHECK-DAG: #[[$MAP1:.+]] = affine_map<()[s0] -> (s0 ceildiv 2)>202// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2)>203// CHECK-DAG: #[[$MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d3)>204// CHECK-DAG: #[[$MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>205// CHECK-LABEL: func.func @dynamic_broadcast_pack206// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]207// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]208// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]209// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index210// CHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]]211// CHECK-DAG: %[[OUTER_D0:.+]] = affine.apply #[[$MAP0]]()[%[[D0]]]212// CHECK: %[[ARG0_EMPTY:.+]] = tensor.empty(%[[OUTER_D0]]) : tensor<?x8xf32>213// CHECK: %[[PACK_ARG0:.+]] = linalg.pack %[[ARG0]]214// CHECK-SAME: inner_dims_pos = [0] inner_tiles = [8]215// CHECK-SAME: into %[[ARG0_EMPTY]]216// CHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG1]], %[[C0]]217// CHECK-DAG: %[[OUTER_D1:.+]] = affine.apply #[[$MAP1]]()[%[[D1]]]218// CHECK: %[[ARG1_EMPTY:.+]] = tensor.empty(%[[OUTER_D1]]) : tensor<?x2xf32>219// CHECK: %[[PACK_ARG1:.+]] = linalg.pack %[[ARG1]]220// CHECK-SAME: inner_dims_pos = [0] inner_tiles = [2]221// CHECK-SAME: into %[[ARG1_EMPTY]]222// CHECK: %[[ELEM:.+]] = linalg.generic223// CHECK-SAME: indexing_maps = [#[[$MAP2]], #[[$MAP3]], #[[$MAP4]]]224// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]225// CHECK-SAME: ins(%[[PACK_ARG0]], %[[PACK_ARG0]]226// CHECK-SAME: outs(%[[DEST]]227// CHECK: return %[[ELEM]] : tensor<?x?x8x2xf32>228 229// -----230 231#map = affine_map<(d0, d1, d2, d3) -> (d3)>232#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>233func.func @elem_pack_transpose_inner_and_outer_dims2(%arg0: tensor<64xf32>, %dest: tensor<1x2x56x57x32xf32>) -> tensor<1x2x56x57x32xf32> {234 %0 = tensor.empty() : tensor<1x56x57x64xf32>235 %1 = linalg.generic {236 indexing_maps = [#map, #map1],237 iterator_types = ["parallel", "parallel", "parallel", "parallel"]}238 ins(%arg0 : tensor<64xf32>)239 outs(%0 : tensor<1x56x57x64xf32>) {240 ^bb0(%in: f32, %out: f32):241 linalg.yield %in : f32242 } -> tensor<1x56x57x64xf32>243 %2 = linalg.pack %1 outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32] into %dest : tensor<1x56x57x64xf32> -> tensor<1x2x56x57x32xf32>244 return %2 : tensor<1x2x56x57x32xf32>245}246// CHECK-DAG: #[[$MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d1, d4)>247// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>248// CHECK-LABEL: func.func @elem_pack_transpose_inner_and_outer_dims2249// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]250// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]251// CHECK: %[[ARG0_EMPTY:.+]] = tensor.empty() : tensor<2x32xf32>252// CHECK: %[[PACKED_ARG0:.+]] = linalg.pack %[[ARG0]]253// CHECK-SAME: inner_dims_pos = [0] inner_tiles = [32]254// CHECK-SAME: into %[[ARG0_EMPTY]]255// CHECK: %[[RES:.+]] = linalg.generic256// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]]]257// CHECK-SAME: ins(%[[PACKED_ARG0]]258// CHECK-SAME: outs(%[[DEST]]259 260// -----261 262func.func @transpose_pack(%arg0: tensor<100x128x200x256xi32>, %arg1: tensor<100xi32>, %arg2: tensor<128xi32>, %dest: tensor<100x200x4x16x16x32xi32>) -> tensor<100x200x4x16x16x32xi32>263{264 %init_transpose = tensor.empty() : tensor<100x200x128x256xi32>265 %transpose = linalg.generic {266 indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>,267 affine_map<(d0, d1, d2, d3) -> (d0)>,268 affine_map<(d0, d1, d2, d3) -> (d1)>,269 affine_map<(d0, d1, d2, d3) -> (d0, d2, d1, d3)>],270 iterator_types = ["parallel", "parallel", "parallel", "parallel"]}271 ins(%arg0, %arg1, %arg2 : tensor<100x128x200x256xi32>, tensor<100xi32>, tensor<128xi32>)272 outs(%init_transpose : tensor<100x200x128x256xi32>) {273 ^bb0(%b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32):274 %0 = arith.addi %b0, %b1 : i32275 %1 = arith.addi %0, %b2 : i32276 linalg.yield %1 : i32277 } -> tensor<100x200x128x256xi32>278 %4 = linalg.pack %transpose279 inner_dims_pos = [3, 2]280 inner_tiles = [16, 32]281 into %dest : tensor<100x200x128x256xi32> -> tensor<100x200x4x16x16x32xi32>282 return %4 : tensor<100x200x4x16x16x32xi32>283}284// CHECK-DAG: #[[$MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3, d4, d5)>285// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0)>286// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d5)>287// CHECK-DAG: #[[$MAP3:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d2, d1, d3, d4, d5)>288// CHECK-LABEL: func.func @transpose_pack289// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]290// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]291// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]292// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]293// CHECK: %[[ARG0_EMPTY:.+]] = tensor.empty() : tensor<100x4x200x16x16x32xi32>294// CHECK: %[[PACKED_ARG0:.+]] = linalg.pack %[[ARG0]]295// CHECK-SAME: inner_dims_pos = [3, 1] inner_tiles = [16, 32]296// CHECK-SAME: into %[[ARG0_EMPTY]]297// CHECK: %[[ARG2_EMPTY:.+]] = tensor.empty() : tensor<4x32xi32>298// CHECK: %[[PACKED_ARG2:.+]] = linalg.pack %[[ARG2]]299// CHECK-SAME: inner_dims_pos = [0] inner_tiles = [32]300// CHECK-SAME: into %[[ARG2_EMPTY]]301// CHECK: %[[RES:.+]] = linalg.generic302// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]], #[[$MAP3]]]303// CHECK-SAME: ins(%[[PACKED_ARG0]], %[[ARG1]], %[[PACKED_ARG2]]304// CHECK-SAME: outs(%[[DEST]]305 306// -----307 308func.func @affine_constant_expr_pack(%arg0: tensor<100x128x200x256xi32>, %arg1: tensor<100x1x1x1xi32>, %arg2: tensor<1x128x1x1xi32>, %dest: tensor<100x200x4x16x16x32xi32>) -> tensor<100x200x4x16x16x32xi32>309{310 %init_transpose = tensor.empty() : tensor<100x200x128x256xi32>311 %transpose = linalg.generic {312 indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>,313 affine_map<(d0, d1, d2, d3) -> (d0, 0, 0, 0)>,314 affine_map<(d0, d1, d2, d3) -> (0, d1, 0, 0)>,315 affine_map<(d0, d1, d2, d3) -> (d0, d2, d1, d3)>],316 iterator_types = ["parallel", "parallel", "parallel", "parallel"]}317 ins(%arg0, %arg1, %arg2 : tensor<100x128x200x256xi32>, tensor<100x1x1x1xi32>, tensor<1x128x1x1xi32>)318 outs(%init_transpose : tensor<100x200x128x256xi32>) {319 ^bb0(%b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32):320 %0 = arith.addi %b0, %b1 : i32321 %1 = arith.addi %0, %b2 : i32322 linalg.yield %1 : i32323 } -> tensor<100x200x128x256xi32>324 %4 = linalg.pack %transpose325 inner_dims_pos = [3, 2]326 inner_tiles = [16, 32]327 into %dest : tensor<100x200x128x256xi32> -> tensor<100x200x4x16x16x32xi32>328 return %4 : tensor<100x200x4x16x16x32xi32>329}330// CHECK-DAG: #[[$MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3, d4, d5)>331// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, 0, 0, 0)>332// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (0, d1, 0, 0, d5)>333// CHECK-DAG: #[[$MAP3:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d2, d1, d3, d4, d5)>334// CHECK-LABEL: func.func @affine_constant_expr_pack335// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]336// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]337// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]338// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]339// CHECK: %[[ARG0_EMPTY:.+]] = tensor.empty() : tensor<100x4x200x16x16x32xi32>340// CHECK: %[[PACKED_ARG0:.+]] = linalg.pack %[[ARG0]]341// CHECK-SAME: inner_dims_pos = [3, 1] inner_tiles = [16, 32]342// CHECK-SAME: into %[[ARG0_EMPTY]]343// CHECK: %[[ARG2_EMPTY:.+]] = tensor.empty() : tensor<1x4x1x1x32xi32>344// CHECK: %[[PACKED_ARG2:.+]] = linalg.pack %[[ARG2]]345// CHECK-SAME: inner_dims_pos = [1] inner_tiles = [32]346// CHECK-SAME: into %[[ARG2_EMPTY]]347// CHECK: %[[RES:.+]] = linalg.generic348// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]], #[[$MAP3]]]349// CHECK-SAME: ins(%[[PACKED_ARG0]], %[[ARG1]], %[[PACKED_ARG2]]350// CHECK-SAME: outs(%[[DEST]]351 352// -----353 354func.func @transpose_pack_with_outer_dims(%arg0: tensor<100x128x200x256xi32>, %arg1: tensor<100xi32>, %arg2: tensor<128xi32>, %dest: tensor<200x4x16x100x16x32xi32>) -> tensor<200x4x16x100x16x32xi32>355{356 %init_transpose = tensor.empty() : tensor<100x200x128x256xi32>357 %transpose = linalg.generic {358 indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>,359 affine_map<(d0, d1, d2, d3) -> (d0)>,360 affine_map<(d0, d1, d2, d3) -> (d1)>,361 affine_map<(d0, d1, d2, d3) -> (d0, d2, d1, d3)>],362 iterator_types = ["parallel", "parallel", "parallel", "parallel"]}363 ins(%arg0, %arg1, %arg2 : tensor<100x128x200x256xi32>, tensor<100xi32>, tensor<128xi32>)364 outs(%init_transpose : tensor<100x200x128x256xi32>) {365 ^bb0(%b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32):366 %0 = arith.addi %b0, %b1 : i32367 %1 = arith.addi %0, %b2 : i32368 linalg.yield %1 : i32369 } -> tensor<100x200x128x256xi32>370 %4 = linalg.pack %transpose371 outer_dims_perm = [1, 2, 3, 0]372 inner_dims_pos = [3, 2]373 inner_tiles = [16, 32]374 into %dest : tensor<100x200x128x256xi32> -> tensor<200x4x16x100x16x32xi32>375 return %4 : tensor<200x4x16x100x16x32xi32>376}377 378// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3, d4, d5)>379// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d3)>380// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d5)>381// CHECK-LABEL: func.func @transpose_pack_with_outer_dims382// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]383// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]384// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]385// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]386// CHECK: %[[ARG0_EMPTY:.+]] = tensor.empty() : tensor<200x4x16x100x16x32xi32>387// CHECK: %[[PACKED_ARG0:.+]] = linalg.pack %[[ARG0]]388// CHECK-SAME: outer_dims_perm = [2, 1, 3, 0] inner_dims_pos = [3, 1] inner_tiles = [16, 32]389// CHECK-SAME: into %[[ARG0_EMPTY]]390// CHECK: %[[ARG2_EMPTY:.+]] = tensor.empty() : tensor<4x32xi32>391// CHECK: %[[PACKED_ARG2:.+]] = linalg.pack %[[ARG2]]392// CHECK-SAME: inner_dims_pos = [0] inner_tiles = [32]393// CHECK-SAME: into %[[ARG2_EMPTY]]394// CHECK: %[[RES:.+]] = linalg.generic395// CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP1]], #[[$MAP2]], #[[$MAP]]]396// CHECK-SAME: ins(%[[PACKED_ARG0]], %[[ARG1]], %[[PACKED_ARG2]]397// CHECK-SAME: outs(%[[DEST]]398 399// -----400 401#map0 = affine_map<(d0, d1) -> (d0, d1)>402func.func @elem_pack_transpose_outer_dims(%arg0: tensor<128x256xi32>, %init: tensor<128x256xi32>) -> tensor<16x4x32x16xi32>{403 %elem = linalg.generic {indexing_maps = [#map0, #map0], iterator_types = ["parallel", "parallel"]}404 ins(%arg0 : tensor<128x256xi32>)405 outs(%init : tensor<128x256xi32>) {406 ^bb0(%arg3: i32, %arg4: i32):407 %4 = arith.addi %arg3, %arg4 : i32408 linalg.yield %4 : i32409 } -> tensor<128x256xi32>410 %empty = tensor.empty() : tensor<16x4x32x16xi32>411 %pack = linalg.pack %elem412 outer_dims_perm = [1, 0]413 inner_dims_pos = [0, 1]414 inner_tiles = [32, 16]415 into %empty : tensor<128x256xi32> -> tensor<16x4x32x16xi32>416 return %pack : tensor<16x4x32x16xi32>417}418 419// CHECK: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>420// CHECK-LABEL: func.func @elem_pack_transpose_outer_dims421// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]422// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]423// CHECK: %[[ARG1_EMPTY:.+]] = tensor.empty() : tensor<16x4x32x16xi32>424// CHECK: %[[PACKED_ARG1:.+]] = linalg.pack %[[ARG1]]425// CHECK-SAME: outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [32, 16]426// CHECK-SAME: into %[[ARG1_EMPTY]]427// CHECK: %[[ARG0_EMPTY:.+]] = tensor.empty() : tensor<16x4x32x16xi32>428// CHECK: %[[PACKED_ARG0:.+]] = linalg.pack %[[ARG0]]429// CHECK-SAME: outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [32, 16]430// CHECK-SAME: into %[[ARG0_EMPTY]]431// CHECK: %[[RES:.+]] = linalg.generic432// CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP]]]433// CHECK-SAME: ins(%[[PACKED_ARG0]]434// CHECK-SAME: outs(%[[PACKED_ARG1]]435 436// -----437 438#map0 = affine_map<(d0, d1) -> (d0, d1)>439func.func @elem_pack_transpose_outer_dims_unused_init(%arg0: tensor<128x256xi32>, %init: tensor<128x256xi32>) -> tensor<16x4x32x16xi32>{440 %elem = linalg.generic {indexing_maps = [#map0, #map0], iterator_types = ["parallel", "parallel"]}441 ins(%arg0 : tensor<128x256xi32>)442 outs(%init : tensor<128x256xi32>) {443 ^bb0(%arg3: i32, %arg4: i32):444 %4 = arith.addi %arg3, %arg3 : i32445 linalg.yield %4 : i32446 } -> tensor<128x256xi32>447 %empty = tensor.empty() : tensor<16x4x32x16xi32>448 %pack = linalg.pack %elem449 outer_dims_perm = [1, 0]450 inner_dims_pos = [0, 1]451 inner_tiles = [32, 16]452 into %empty : tensor<128x256xi32> -> tensor<16x4x32x16xi32>453 return %pack : tensor<16x4x32x16xi32>454}455 456// CHECK: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>457// CHECK-LABEL: func.func @elem_pack_transpose_outer_dims458// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]459// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]460// CHECK: %[[ARG1_EMPTY:.+]] = tensor.empty() : tensor<16x4x32x16xi32>461// CHECK: %[[ARG0_EMPTY:.+]] = tensor.empty() : tensor<16x4x32x16xi32>462// CHECK: %[[PACKED_ARG0:.+]] = linalg.pack %[[ARG0]]463// CHECK-SAME: outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [32, 16]464// CHECK-SAME: into %[[ARG0_EMPTY]]465// CHECK: %[[RES:.+]] = linalg.generic466// CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP]]]467// CHECK-SAME: ins(%[[PACKED_ARG0]]468// CHECK-SAME: outs(%[[ARG1_EMPTY]]469 470// -----471 472#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>473 474func.func @unpack_on_output(%arg0: tensor<12x2x56x56x32xf32>) -> tensor<12x56x56x64xf32> {475 %0 = tensor.empty() : tensor<12x56x56x64xf32>476 %1 = linalg.unpack %arg0 outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32] into %0 : tensor<12x2x56x56x32xf32> -> tensor<12x56x56x64xf32>477 %2 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} outs(%1 : tensor<12x56x56x64xf32>) {478 ^bb0(%out: f32):479 %3 = arith.addf %out, %out : f32480 linalg.yield %3 : f32481 } -> tensor<12x56x56x64xf32>482 return %2 : tensor<12x56x56x64xf32>483}484 485// CHECK: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>486// CHECK-LABEL: func.func @unpack_on_output487// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]488// CHECK: %[[ARG0_EMPTY_UNPACK:.+]] = tensor.empty() : tensor<12x56x56x64xf32>489// CHECK: %[[UNPACKED_ARG0:.+]] = linalg.unpack %[[ARG0]]490// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32]491// CHECK-SAME: into %[[ARG0_EMPTY_UNPACK]]492// CHECK: %[[RES:.+]] = linalg.generic493// CHECK-SAME: indexing_maps = [#[[$MAP]]]494// CHECK-SAME: outs(%[[ARG0]]495// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[RES]]496// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32]497// CHECK-SAME: into %[[UNPACKED_ARG0]]498 499// -----500 501#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>502 503func.func @unpack_on_input(%arg0: tensor<12x2x56x56x32xf32>, %init: tensor<12x56x56x64xf32>) -> tensor<12x56x56x64xf32> {504 %0 = tensor.empty() : tensor<12x56x56x64xf32>505 %1 = linalg.unpack %arg0 outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32] into %0 : tensor<12x2x56x56x32xf32> -> tensor<12x56x56x64xf32>506 %2 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%1: tensor<12x56x56x64xf32>) outs(%init : tensor<12x56x56x64xf32>) {507 ^bb0(%in: f32, %out: f32):508 %3 = arith.addf %in, %out : f32509 linalg.yield %3 : f32510 } -> tensor<12x56x56x64xf32>511 return %2 : tensor<12x56x56x64xf32>512}513 514// CHECK: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>515// CHECK-LABEL: func.func @unpack_on_input516// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]517// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]518// CHECK: %[[ARG1_PACK_EMPTY:.+]] = tensor.empty() : tensor<12x2x56x56x32xf32>519// CHECK: %[[ARG1_PACK:.+]] = linalg.pack %[[ARG1]]520// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32]521// CHECK-SAME: into %[[ARG1_PACK_EMPTY]]522// CHECK: %[[RES:.+]] = linalg.generic523// CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP]]]524// CHECK-SAME: ins(%[[ARG0]]525// CHECK-SAME: outs(%[[ARG1_PACK]]526// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[RES]]527// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32]528// CHECK-SAME: into %[[ARG1]]529 530// -----531 532#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>533 534func.func @unpack_element_type_change_no_use(%arg0: tensor<12x2x56x56x32xf32>, %init: tensor<12x56x56x64xf16>) -> tensor<12x56x56x64xf16> {535 %0 = tensor.empty() : tensor<12x56x56x64xf32>536 %1 = linalg.unpack %arg0 outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32] into %0 : tensor<12x2x56x56x32xf32> -> tensor<12x56x56x64xf32>537 %2 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%1: tensor<12x56x56x64xf32>) outs(%init : tensor<12x56x56x64xf16>) {538 ^bb0(%in: f32, %out: f16):539 %3 = arith.truncf %in : f32 to f16540 linalg.yield %3 : f16541 } -> tensor<12x56x56x64xf16>542 return %2 : tensor<12x56x56x64xf16>543}544 545// CHECK: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>546// CHECK-LABEL: func.func @unpack_element_type_change_no_use547// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]548// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]549// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<12x2x56x56x32xf16>550// CHECK: %[[RES:.+]] = linalg.generic551// CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP]]]552// CHECK-SAME: ins(%[[ARG0]]553// CHECK-SAME: outs(%[[EMPTY]]554// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[RES]]555// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32]556// CHECK-SAME: into %[[ARG1]]557 558// -----559 560#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>561 562func.func @forward_tensor_empty(%arg0: tensor<12x2x56x56x32xf32>) -> tensor<12x56x56x64xf32> {563 %init = tensor.empty() : tensor<12x56x56x64xf32>564 %0 = tensor.empty() : tensor<12x56x56x64xf32>565 %1 = linalg.unpack %arg0 outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32] into %0 : tensor<12x2x56x56x32xf32> -> tensor<12x56x56x64xf32>566 %2 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%1: tensor<12x56x56x64xf32>) outs(%init : tensor<12x56x56x64xf32>) {567 ^bb0(%in: f32, %out: f32):568 %3 = arith.addf %in, %in : f32569 linalg.yield %3 : f32570 } -> tensor<12x56x56x64xf32>571 return %2 : tensor<12x56x56x64xf32>572}573 574// CHECK: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>575// CHECK-LABEL: func.func @forward_tensor_empty576// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]577// CHECK: %[[FINAL_RES:.+]] = tensor.empty() : tensor<12x56x56x64xf32>578// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<12x2x56x56x32xf32>579// CHECK: %[[RES:.+]] = linalg.generic580// CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP]]]581// CHECK-SAME: ins(%[[ARG0]]582// CHECK-SAME: outs(%[[EMPTY]]583// CHECK: %[[UNPACKED:.+]] = linalg.unpack %[[RES]]584// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32]585// CHECK-SAME: into %[[FINAL_RES]]586 587// -----588 589func.func @pad_valid_unpack_propagation(%arg0: tensor<1x2x56x56x32xf32>) -> tensor<1x58x58x64xf32> {590 %cst = arith.constant 0.000000e+00 : f32591 %0 = tensor.empty() : tensor<1x56x56x64xf32>592 %1 = linalg.unpack %arg0 outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32] into %0 : tensor<1x2x56x56x32xf32> -> tensor<1x56x56x64xf32>593 %padded = tensor.pad %1 low[0, 1, 1, 0] high[0, 1, 1, 0] {594 ^bb0(%arg3: index, %arg4: index, %arg5: index, %arg6: index):595 tensor.yield %cst : f32596 } : tensor<1x56x56x64xf32> to tensor<1x58x58x64xf32>597 return %padded : tensor<1x58x58x64xf32>598}599 600// CHECK-LABEL: func.func @pad_valid_unpack_propagation(601// CHECK-SAME: %[[ARG0:.+]]: tensor<1x2x56x56x32xf32>)602// CHECK: %[[CST:.+]] = arith.constant 0.000000e+00 : f32603// CHECK: %[[PADDED:.+]] = tensor.pad %[[ARG0]] low[0, 0, 1, 1, 0] high[0, 0, 1, 1, 0]604// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<1x58x58x64xf32>605// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[PADDED]]606// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32]607// CHECK-SAME: into %[[EMPTY]] : tensor<1x2x58x58x32xf32> -> tensor<1x58x58x64xf32>608 609// -----610 611func.func @pad_valid_unpack_propagation(%arg0: tensor<1x2x56x56x32xf32>) -> tensor<2x58x58x64xf32> {612 %cst = arith.constant 0.000000e+00 : f32613 %0 = tensor.empty() : tensor<1x56x56x64xf32>614 %1 = linalg.unpack %arg0 outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32] into %0 : tensor<1x2x56x56x32xf32> -> tensor<1x56x56x64xf32>615 %padded = tensor.pad %1 low[1, 1, 1, 0] high[0, 1, 1, 0] {616 ^bb0(%arg3: index, %arg4: index, %arg5: index, %arg6: index):617 tensor.yield %cst : f32618 } : tensor<1x56x56x64xf32> to tensor<2x58x58x64xf32>619 return %padded : tensor<2x58x58x64xf32>620}621 622// CHECK-LABEL: func.func @pad_valid_unpack_propagation(623// CHECK-SAME: %[[ARG0:.+]]: tensor<1x2x56x56x32xf32>)624// CHECK: %[[CST:.+]] = arith.constant 0.000000e+00 : f32625// CHECK: %[[PADDED:.+]] = tensor.pad %[[ARG0]] low[1, 0, 1, 1, 0] high[0, 0, 1, 1, 0]626// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<2x58x58x64xf32>627// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[PADDED]]628// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32]629// CHECK-SAME: into %[[EMPTY]] : tensor<2x2x58x58x32xf32> -> tensor<2x58x58x64xf32>630 631// -----632 633func.func @pad_along_unpacked_dim(%arg0: tensor<1x2x56x56x32xf32>) -> tensor<1x58x58x66xf32> {634 %cst = arith.constant 0.000000e+00 : f32635 %0 = tensor.empty() : tensor<1x56x56x64xf32>636 %1 = linalg.unpack %arg0 outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32] into %0 : tensor<1x2x56x56x32xf32> -> tensor<1x56x56x64xf32>637 %padded = tensor.pad %1 low[0, 1, 1, 1] high[0, 1, 1, 1] {638 ^bb0(%arg3: index, %arg4: index, %arg5: index, %arg6: index):639 tensor.yield %cst : f32640 } : tensor<1x56x56x64xf32> to tensor<1x58x58x66xf32>641 return %padded : tensor<1x58x58x66xf32>642}643 644// CHECK-LABEL: func.func @pad_along_unpacked_dim(645// CHECK: %[[ARG0:.+]]: tensor<1x2x56x56x32xf32>)646// CHECK: %[[CST:.+]] = arith.constant 0.000000e+00 : f32647// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<1x56x56x64xf32>648// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[ARG0]]649// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32]650// CHECK-SAME: into %[[EMPTY]] : tensor<1x2x56x56x32xf32> -> tensor<1x56x56x64xf32>651// CHECK: %[[PADDED:.+]] = tensor.pad %[[UNPACK]] low[0, 1, 1, 1] high[0, 1, 1, 1]652 653// -----654 655func.func @pad_valid_pack_propagation(%arg0: tensor<1x64x56x56xf32>) -> tensor<1x2x58x58x32xf32> {656 %cst = arith.constant 0.000000e+00 : f32657 %padded = tensor.pad %arg0 low[0, 0, 1, 1] high[0, 0, 1, 1] {658 ^bb0(%arg3: index, %arg4: index, %arg5: index, %arg6: index):659 tensor.yield %cst : f32660 } : tensor<1x64x56x56xf32> to tensor<1x64x58x58xf32>661 %0 = tensor.empty() : tensor<1x2x58x58x32xf32>662 %1 = linalg.pack %padded inner_dims_pos = [1] inner_tiles = [32] into %0 : tensor<1x64x58x58xf32> -> tensor<1x2x58x58x32xf32>663 return %1 : tensor<1x2x58x58x32xf32>664}665 666// CHECK-LABEL: func.func @pad_valid_pack_propagation(667// CHECK-SAME: %[[ARG0:.+]]: tensor<1x64x56x56xf32>)668// CHECK: %[[CST:.+]] = arith.constant 0.000000e+00 : f32669// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<1x2x56x56x32xf32>670// CHECK: %[[PACKED:.+]] = linalg.pack %[[ARG0]] inner_dims_pos = [1] inner_tiles = [32]671// CHECK-SAME: into %[[EMPTY]] : tensor<1x64x56x56xf32> -> tensor<1x2x56x56x32xf32>672// CHECK: %[[PADDED:.+]] = tensor.pad %[[PACKED]] low[0, 0, 1, 1, 0] high[0, 0, 1, 1, 0]673// CHECK: return %[[PADDED]]674 675// -----676 677func.func @pad_valid_outer_dims_pack_propagation(%arg0: tensor<1x64x56x56xf32>) -> tensor<1x58x58x2x32xf32> {678 %cst = arith.constant 0.000000e+00 : f32679 %padded = tensor.pad %arg0 low[0, 0, 1, 1] high[0, 0, 1, 1] {680 ^bb0(%arg3: index, %arg4: index, %arg5: index, %arg6: index):681 tensor.yield %cst : f32682 } : tensor<1x64x56x56xf32> to tensor<1x64x58x58xf32>683 %0 = tensor.empty() : tensor<1x58x58x2x32xf32>684 %1 = linalg.pack %padded outer_dims_perm = [0, 3, 2, 1] inner_dims_pos = [1] inner_tiles = [32] into %0 : tensor<1x64x58x58xf32> -> tensor<1x58x58x2x32xf32>685 return %1 : tensor<1x58x58x2x32xf32>686}687 688// CHECK-LABEL: func.func @pad_valid_outer_dims_pack_propagation(689// CHECK-SAME: %[[ARG0:.+]]: tensor<1x64x56x56xf32>)690// CHECK: %[[CST:.+]] = arith.constant 0.000000e+00 : f32691// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<1x56x56x2x32xf32>692// CHECK: %[[PACKED:.+]] = linalg.pack %[[ARG0]]693// CHECK-SAME: outer_dims_perm = [0, 3, 2, 1] inner_dims_pos = [1] inner_tiles = [32]694// CHECK-SAME: into %[[EMPTY]] : tensor<1x64x56x56xf32> -> tensor<1x56x56x2x32xf32>695// CHECK: %[[PADDED:.+]] = tensor.pad %[[PACKED]] low[0, 1, 1, 0, 0] high[0, 1, 1, 0, 0]696// CHECK: return %[[PADDED]]697 698// -----699 700func.func @pad_along_packed_dim(%arg0: tensor<1x60x56x56xf32>) -> tensor<1x2x58x58x32xf32> {701 %cst = arith.constant 0.000000e+00 : f32702 %padded = tensor.pad %arg0 low[0, 2, 1, 1] high[0, 2, 1, 1] {703 ^bb0(%arg3: index, %arg4: index, %arg5: index, %arg6: index):704 tensor.yield %cst : f32705 } : tensor<1x60x56x56xf32> to tensor<1x64x58x58xf32>706 %0 = tensor.empty() : tensor<1x2x58x58x32xf32>707 %1 = linalg.pack %padded inner_dims_pos = [1] inner_tiles = [32] into %0 : tensor<1x64x58x58xf32> -> tensor<1x2x58x58x32xf32>708 return %1 : tensor<1x2x58x58x32xf32>709}710 711// CHECK-LABEL: func.func @pad_along_packed_dim(712// CHECK-SAME: %[[ARG0:.+]]: tensor<1x60x56x56xf32>)713// CHECK: %[[CST:.+]] = arith.constant 0.000000e+00 : f32714// CHECK: %[[PADDED:.+]] = tensor.pad %[[ARG0]] low[0, 2, 1, 1] high[0, 2, 1, 1]715// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<1x2x58x58x32xf32>716// CHECK: linalg.pack %[[PADDED]] inner_dims_pos = [1] inner_tiles = [32]717// CHECK-SAME: into %[[EMPTY]] : tensor<1x64x58x58xf32> -> tensor<1x2x58x58x32xf32>718 719// -----720 721func.func @multi_use_pad_pack_propagation(%arg0: tensor<1x64x56x56xf32>) -> (tensor<1x64x58x58xf32>, tensor<1x2x58x58x32xf32>) {722 %cst = arith.constant 0.000000e+00 : f32723 %padded = tensor.pad %arg0 low[0, 0, 1, 1] high[0, 0, 1, 1] {724 ^bb0(%arg3: index, %arg4: index, %arg5: index, %arg6: index):725 tensor.yield %cst : f32726 } : tensor<1x64x56x56xf32> to tensor<1x64x58x58xf32>727 %0 = tensor.empty() : tensor<1x2x58x58x32xf32>728 %1 = linalg.pack %padded inner_dims_pos = [1] inner_tiles = [32] into %0 : tensor<1x64x58x58xf32> -> tensor<1x2x58x58x32xf32>729 return %padded, %1 : tensor<1x64x58x58xf32>, tensor<1x2x58x58x32xf32>730}731 732// CHECK-LABEL: func.func @multi_use_pad_pack_propagation(733// CHECK-SAME: %[[ARG0:.+]]: tensor<1x64x56x56xf32>)734// CHECK: %[[CST:.+]] = arith.constant 0.000000e+00 : f32735// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<1x2x56x56x32xf32>736// CHECK: %[[PACKED:.+]] = linalg.pack %[[ARG0]] inner_dims_pos = [1] inner_tiles = [32]737// CHECK-SAME: into %[[EMPTY]] : tensor<1x64x56x56xf32> -> tensor<1x2x56x56x32xf32>738// CHECK: %[[PADDED:.+]] = tensor.pad %[[PACKED]] low[0, 0, 1, 1, 0] high[0, 0, 1, 1, 0]739// CHECK: %[[UNPACKED:.+]] = linalg.unpack %[[PADDED]] inner_dims_pos = [1] inner_tiles = [32]740// CHECK: return %[[UNPACKED]], %[[PADDED]]741 742// -----743 744#map0 = affine_map<(d0, d1) -> (d0, d1)>745func.func @would_break_dominance(%arg0: tensor<128x256xi32>) -> tensor<4x16x16x32xi32>{746 %init = tensor.empty() : tensor<128x256xi32>747 %elem = linalg.generic {indexing_maps = [#map0, #map0], iterator_types = ["parallel", "parallel"]}748 ins(%arg0 : tensor<128x256xi32>)749 outs(%init : tensor<128x256xi32>) {750 ^bb0(%arg3: i32, %arg4: i32):751 %4 = arith.addi %arg3, %arg3 : i32752 linalg.yield %4 : i32753 } -> tensor<128x256xi32>754 %dest = bufferization.alloc_tensor() : tensor<4x16x16x32xi32>755 %pack = linalg.pack %elem756 inner_dims_pos = [1, 0]757 inner_tiles = [16, 32]758 into %dest : tensor<128x256xi32> -> tensor<4x16x16x32xi32>759 return %pack : tensor<4x16x16x32xi32>760}761 762// CHECK-LABEL: func.func @would_break_dominance(763// CHECK-SAME: %[[ARG0:.+]]: tensor<128x256xi32>)764// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<128x256xi32>765// CHECK-NEXT: %[[GEN:.+]] = linalg.generic766// CHECK-SAME: ins(%[[ARG0]]767// CHECK-SAME: outs(%[[EMPTY]]768// CHECK: %[[ALLOC:.+]] = bufferization.alloc_tensor() : tensor<4x16x16x32xi32>769// CHECK-NEXT: %{{.+}} = linalg.pack %[[GEN]]770// CHECK-SAME: inner_dims_pos = [1, 0] inner_tiles = [16, 32]771// CHECK-SAME: into %[[ALLOC]]772 773// -----774 775#map0 = affine_map<(d0, d1, d2, d3) -> ()>776#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>777 778func.func @scalar_tensor(%arg0 : tensor<f32>) -> tensor<1x32x7x7x32xf32> {779 %empty_gen = tensor.empty() : tensor<1x7x7x1024xf32>780 %gen = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<f32>) outs(%empty_gen : tensor<1x7x7x1024xf32>) {781 ^bb0(%in: f32, %out: f32):782 linalg.yield %in : f32783 } -> tensor<1x7x7x1024xf32>784 %empty_pack = tensor.empty() : tensor<1x32x7x7x32xf32>785 %pack = linalg.pack %gen outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32] into %empty_pack : tensor<1x7x7x1024xf32> -> tensor<1x32x7x7x32xf32>786 return %pack : tensor<1x32x7x7x32xf32>787}788 789// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3, d4) -> ()>790// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>791// CHECK-LABEL: func.func @scalar_tensor792// CHECK-SAME: %[[ARG0:.+]]: tensor<f32>)793// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<1x32x7x7x32xf32>794// CHECK: linalg.generic795// CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP1]]]796// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]797// CHECK-SAME: ins(%[[ARG0]]798// CHECK-SAME: outs(%[[EMPTY]]799 800// -----801 802#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>803func.func @unpack_empty_inner_dims(%arg0: tensor<12x64x56x56xf32>) -> tensor<12x56x56x64xf32> {804 %init = tensor.empty() : tensor<12x56x56x64xf32>805 %0 = tensor.empty() : tensor<12x56x56x64xf32>806 %1 = linalg.unpack %arg0 outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [] inner_tiles = [] into %0 : tensor<12x64x56x56xf32> -> tensor<12x56x56x64xf32>807 %2 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%1: tensor<12x56x56x64xf32>) outs(%init : tensor<12x56x56x64xf32>) {808 ^bb0(%in: f32, %out: f32):809 %3 = arith.addf %in, %in : f32810 linalg.yield %3 : f32811 } -> tensor<12x56x56x64xf32>812 return %2 : tensor<12x56x56x64xf32>813}814 815// CHECK-LABEL: func.func @unpack_empty_inner_dims816// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<12x64x56x56xf32>)817// CHECK: %[[RES:.+]] = linalg.generic818// CHECK-SAME: ins(%[[ARG0]]819// CHECK: %[[UNPACKED:.+]] = linalg.unpack %[[RES]]820// CHECK-SAME: outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [] inner_tiles = []821 822// -----823 824#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>825#map1 = affine_map<(d0, d1, d2) -> (d0, d1)>826func.func @reduction_pack_transpose_inner_dims(%arg0: tensor<128x256x32xi32>,827 %arg1: tensor<128x256xi32>) -> tensor<4x16x16x32xi32>{828 %elem = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "parallel", "reduction"]}829 ins(%arg0 : tensor<128x256x32xi32>)830 outs(%arg1 : tensor<128x256xi32>) {831 ^bb0(%arg3: i32, %arg4: i32):832 %4 = arith.addi %arg3, %arg4 : i32833 linalg.yield %4 : i32834 } -> tensor<128x256xi32>835 %dest = tensor.empty() : tensor<4x16x16x32xi32>836 %pack = linalg.pack %elem837 inner_dims_pos = [1, 0]838 inner_tiles = [16, 32]839 into %dest : tensor<128x256xi32> -> tensor<4x16x16x32xi32>840 return %pack : tensor<4x16x16x32xi32>841}842// CHECK-DAG: #[[$MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>843// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d3, d4)>844// CHECK-LABEL: func.func @reduction_pack_transpose_inner_dims845// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]846// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]847// CHECK: %[[ARG1_EMPTY:.+]] = tensor.empty() : tensor<4x16x16x32xi32>848// CHECK: %[[PACK_ARG1:.+]] = linalg.pack %[[ARG1]]849// CHECK-SAME: inner_dims_pos = [1, 0] inner_tiles = [16, 32]850// CHECK-SAME: into %[[ARG1_EMPTY]]851// CHECK: %[[ARG0_EMPTY:.+]] = tensor.empty() : tensor<4x16x32x16x32xi32>852// CHECK: %[[PACK_ARG0:.+]] = linalg.pack %[[ARG0]]853// CHECK-SAME: inner_dims_pos = [1, 0] inner_tiles = [16, 32]854// CHECK-SAME: into %[[ARG0_EMPTY]]855// CHECK: %[[RED:.+]] = linalg.generic856// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]]]857// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction", "parallel", "parallel"]858// CHECK-SAME: ins(%[[PACK_ARG0]]859// CHECK-SAME: outs(%[[PACK_ARG1]]860// CHECK: return %[[RED]] : tensor<4x16x16x32xi32>861 862// -----863 864func.func @reduction_pack_with_outer_dims(%arg0: tensor<100x128x200x256xi32>, %arg1: tensor<100xi32>,865 %arg2: tensor<128xi32>, %init_reduction: tensor<100x128x256xi32>) -> tensor<4x16x100x16x32xi32>866{867 %reduction = linalg.generic {868 indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>,869 affine_map<(d0, d1, d2, d3) -> (d0)>,870 affine_map<(d0, d1, d2, d3) -> (d1)>,871 affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>],872 iterator_types = ["parallel", "parallel", "reduction", "parallel"]}873 ins(%arg0, %arg1, %arg2 : tensor<100x128x200x256xi32>, tensor<100xi32>, tensor<128xi32>)874 outs(%init_reduction : tensor<100x128x256xi32>) {875 ^bb0(%b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32):876 %0 = arith.addi %b0, %b1 : i32877 %1 = arith.addi %0, %b2 : i32878 %2 = arith.addi %1, %b3 : i32879 linalg.yield %2 : i32880 } -> tensor<100x128x256xi32>881 %init_pack = tensor.empty() : tensor<4x16x100x16x32xi32>882 %4 = linalg.pack %reduction883 outer_dims_perm = [1, 2, 0]884 inner_dims_pos = [2, 1]885 inner_tiles = [16, 32]886 into %init_pack : tensor<100x128x256xi32> -> tensor<4x16x100x16x32xi32>887 return %4 : tensor<4x16x100x16x32xi32>888}889 890// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3, d4, d5)>891// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d3)>892// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d5)>893// CHECK-DAG: #[[$MAP3:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d3, d4, d5)>894// CHECK-LABEL: func.func @reduction_pack_with_outer_dims895// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]896// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]897// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]898// CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]899// CHECK: %[[ARG3_EMPTY:.+]] = tensor.empty() : tensor<4x16x100x16x32xi32>900// CHECK: %[[PACKED_ARG3:.+]] = linalg.pack %[[ARG3]]901// CHECK-SAME: outer_dims_perm = [1, 2, 0] inner_dims_pos = [2, 1] inner_tiles = [16, 32]902// CHECK-SAME: into %[[ARG3_EMPTY]]903// CHECK: %[[ARG0_EMPTY:.+]] = tensor.empty() : tensor<4x16x200x100x16x32xi32>904// CHECK: %[[PACKED_ARG0:.+]] = linalg.pack %[[ARG0]]905// CHECK-SAME: outer_dims_perm = [1, 3, 2, 0] inner_dims_pos = [3, 1] inner_tiles = [16, 32]906// CHECK-SAME: into %[[ARG0_EMPTY]]907// CHECK: %[[ARG2_EMPTY:.+]] = tensor.empty() : tensor<4x32xi32>908// CHECK: %[[PACKED_ARG2:.+]] = linalg.pack %[[ARG2]]909// CHECK-SAME: inner_dims_pos = [0] inner_tiles = [32]910// CHECK-SAME: into %[[ARG2_EMPTY]]911// CHECK: %[[RES:.+]] = linalg.generic912// CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP1]], #[[$MAP2]], #[[$MAP3]]]913// CHECK-SAME: ins(%[[PACKED_ARG0]], %[[ARG1]], %[[PACKED_ARG2]]914// CHECK-SAME: outs(%[[PACKED_ARG3]]915 916// -----917 918#map0 = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2 * 2 + d4, d3 * 2 + d5)>919#map1 = affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5)>920#map2 = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d2, d3)>921func.func @unpack_different_destination_shape(%arg0: tensor<1x1x1080x1920x16xi32>,922 %filter: tensor<2x2xi32>) -> tensor<16x540x960xi32>{923 %init = tensor.empty() : tensor<16x540x960xi32>924 %empty = tensor.empty() : tensor<1x16x1080x1920xi32>925 %unpack = linalg.unpack %arg0926 inner_dims_pos = [1]927 inner_tiles = [16]928 into %empty : tensor<1x1x1080x1920x16xi32> -> tensor<1x16x1080x1920xi32>929 %pool = linalg.generic {indexing_maps = [#map0, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}930 ins(%unpack, %filter : tensor<1x16x1080x1920xi32>, tensor<2x2xi32>)931 outs(%init : tensor<16x540x960xi32>) {932 ^bb0(%in: i32, %in_1: i32, %out: i32):933 %max = arith.maxui %in, %in_1 : i32934 linalg.yield %max : i32935 } -> tensor<16x540x960xi32>936 return %pool : tensor<16x540x960xi32>937}938// CHECK-DAG: #[[$MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2 * 2 + d4, d3 * 2 + d5, d6)>939// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5)>940// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d1, d2, d3, d6)>941// CHECK-LABEL: func.func @unpack_different_destination_shape942// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]943// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]944// CHECK: %[[FINAL_RES:.+]] = tensor.empty() : tensor<16x540x960xi32>945// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x540x960x16xi32>946// CHECK: %[[POOL:.+]] = linalg.generic947// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]]948// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "parallel"]949// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]]950// CHECK-SAME: outs(%[[INIT]]951// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[POOL]]952// CHECK-SAME: inner_dims_pos = [0] inner_tiles = [16]953// CHECK-SAME: into %[[FINAL_RES]]954// CHECK: return %[[UNPACK]] : tensor<16x540x960xi32>955 956// -----957 958func.func @bubble_up_pack_through_collapse(%1: tensor<?x16x4xf32>, %dim : index) -> tensor<?x4x8x1xf32> {959 %collapsed = tensor.collapse_shape %1 [[0, 1], [2]] : tensor<?x16x4xf32> into tensor<?x4xf32>960 %2 = tensor.empty(%dim) : tensor<?x4x8x1xf32>961 %pack = linalg.pack %collapsed outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [8, 1] into %2 : tensor<?x4xf32> -> tensor<?x4x8x1xf32>962 func.return %pack : tensor<?x4x8x1xf32>963}964// CHECK-LABEL: func.func @bubble_up_pack_through_collapse965// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]966// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]967// CHECK: %[[C0:.+]] = arith.constant 0 : index968// CHECK: %[[DIM:.+]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x16x4xf32>969// CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM]]) : tensor<?x2x4x8x1xf32>970// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]] outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %[[EMPTY]] : tensor<?x16x4xf32> -> tensor<?x2x4x8x1xf32>971// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[PACK]] {{\[}}[0, 1], [2], [3], [4]] : tensor<?x2x4x8x1xf32> into tensor<?x4x8x1xf32>972// CHECK: return %[[COLLAPSED]] : tensor<?x4x8x1xf32>973 974// -----975 976func.func @bubble_up_pack_through_collapse_empty_outer_dims_perm(%1: tensor<?x16x4xf32>, %dim : index) -> tensor<?x4x8x1xf32> {977 %collapsed = tensor.collapse_shape %1 [[0, 1], [2]] : tensor<?x16x4xf32> into tensor<?x4xf32>978 %2 = tensor.empty(%dim) : tensor<?x4x8x1xf32>979 %pack = linalg.pack %collapsed inner_dims_pos = [0, 1] inner_tiles = [8, 1] into %2 : tensor<?x4xf32> -> tensor<?x4x8x1xf32>980 func.return %pack : tensor<?x4x8x1xf32>981}982// CHECK-LABEL: func.func @bubble_up_pack_through_collapse_empty_outer_dims_perm983// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]984// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]985// CHECK: %[[C0:.+]] = arith.constant 0 : index986// CHECK: %[[DIM:.+]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x16x4xf32>987// CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM]]) : tensor<?x2x4x8x1xf32>988// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %[[EMPTY]] : tensor<?x16x4xf32> -> tensor<?x2x4x8x1xf32>989// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[PACK]] {{\[}}[0, 1], [2], [3], [4]] : tensor<?x2x4x8x1xf32> into tensor<?x4x8x1xf32>990// CHECK: return %[[COLLAPSED]] : tensor<?x4x8x1xf32>991 992// -----993 994func.func @bubble_up_permuted_pack_through_collapse(%1: tensor<4x192x16x256xf32>) -> tensor<4x32x3072x8x1xf32> {995 %collapsed = tensor.collapse_shape %1 [[0], [1, 2], [3]] : tensor<4x192x16x256xf32> into tensor<4x3072x256xf32>996 %2 = tensor.empty() : tensor<4x32x3072x8x1xf32>997 %pack = linalg.pack %collapsed outer_dims_perm = [0, 2, 1] inner_dims_pos = [2, 1] inner_tiles = [8, 1] into %2 : tensor<4x3072x256xf32> -> tensor<4x32x3072x8x1xf32>998 func.return %pack : tensor<4x32x3072x8x1xf32>999}1000// CHECK-LABEL: func.func @bubble_up_permuted_pack_through_collapse1001// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1002// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<4x32x192x16x8x1xf32>1003// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]] outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3, 2] inner_tiles = [8, 1] into %[[EMPTY]] : tensor<4x192x16x256xf32> -> tensor<4x32x192x16x8x1xf32>1004// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %pack {{\[}}[0], [1], [2, 3], [4], [5]] : tensor<4x32x192x16x8x1xf32> into tensor<4x32x3072x8x1xf32>1005// CHECK: return %[[COLLAPSED]] : tensor<4x32x3072x8x1xf32>1006 1007// -----1008 1009func.func @bubble_up_pack_through_unit_collapse(%1: tensor<1x64x1x4xf32>) -> tensor<8x4x8x1xf32> {1010 %collapsed = tensor.collapse_shape %1 [[0, 1, 2], [3]] : tensor<1x64x1x4xf32> into tensor<64x4xf32>1011 %2 = tensor.empty() : tensor<8x4x8x1xf32>1012 %pack = linalg.pack %collapsed outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [8, 1] into %2 : tensor<64x4xf32> -> tensor<8x4x8x1xf32>1013 func.return %pack : tensor<8x4x8x1xf32>1014}1015// CHECK-LABEL: func.func @bubble_up_pack_through_unit_collapse1016// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1017// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<1x8x1x4x8x1xf32>1018// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]] outer_dims_perm = [0, 1, 2, 3] inner_dims_pos = [1, 3] inner_tiles = [8, 1] into %[[EMPTY]] : tensor<1x64x1x4xf32> -> tensor<1x8x1x4x8x1xf32>1019// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[PACK]] {{\[}}[0, 1, 2], [3], [4], [5]] : tensor<1x8x1x4x8x1xf32> into tensor<8x4x8x1xf32>1020// CHECK: return %[[COLLAPSED]] : tensor<8x4x8x1xf32>1021 1022// -----1023 1024func.func @bubble_up_pack_through_collapse_on_outer_dims(%1: tensor<?x16x4xf32>, %dim : index) -> tensor<?x1x4xf32> {1025 %collapsed = tensor.collapse_shape %1 [[0, 1], [2]] : tensor<?x16x4xf32> into tensor<?x4xf32>1026 %2 = tensor.empty(%dim) : tensor<?x1x4xf32>1027 %pack = linalg.pack %collapsed outer_dims_perm = [0, 1] inner_dims_pos = [1] inner_tiles = [4] into %2 : tensor<?x4xf32> -> tensor<?x1x4xf32>1028 func.return %pack : tensor<?x1x4xf32>1029}1030// CHECK-LABEL: func.func @bubble_up_pack_through_collapse_on_outer_dims1031// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1032// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]1033// CHECK: %[[C0:.+]] = arith.constant 0 : index1034// CHECK: %[[DIM:.+]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x16x4xf32>1035// CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM]]) : tensor<?x16x1x4xf32>1036// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]] outer_dims_perm = [0, 1, 2] inner_dims_pos = [2] inner_tiles = [4] into %[[EMPTY]] : tensor<?x16x4xf32> -> tensor<?x16x1x4xf32>1037// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[PACK]] {{\[}}[0, 1], [2], [3]] : tensor<?x16x1x4xf32> into tensor<?x1x4xf32>1038// CHECK: return %[[COLLAPSED]] : tensor<?x1x4xf32>1039 1040// -----1041 1042func.func @no_bubble_up_pack_through_non_divisible_collapse(%1: tensor<3072x64x4xf32>) -> tensor<384x32x8x8xf32> {1043 %collapsed = tensor.collapse_shape %1 [[0], [1, 2]] : tensor<3072x64x4xf32> into tensor<3072x256xf32>1044 %2 = tensor.empty() : tensor<384x32x8x8xf32>1045 %pack = linalg.pack %collapsed outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [8, 8] into %2 : tensor<3072x256xf32> -> tensor<384x32x8x8xf32>1046 func.return %pack : tensor<384x32x8x8xf32>1047}1048// CHECK-LABEL: func.func @no_bubble_up_pack_through_non_divisible_collapse1049// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1050// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2]] : tensor<3072x64x4xf32> into tensor<3072x256xf32>1051// CHECK: %[[PACK:.+]] = linalg.pack %[[COLLAPSED]]1052// CHECK: return %[[PACK]] : tensor<384x32x8x8xf32>1053 1054// -----1055 1056func.func @bubble_up_pack_outer_expanded_through_expand(%arg0: tensor<32x64xf32>) -> tensor<4x2x64x4xf32> {1057 %empty = tensor.empty() : tensor<4x2x64x4xf32>1058 %expanded = tensor.expand_shape %arg0 [[0, 1], [2]] output_shape [4, 8, 64] : tensor<32x64xf32> into tensor<4x8x64xf32>1059 %pack = linalg.pack %expanded inner_dims_pos = [1] inner_tiles = [4] into %empty : tensor<4x8x64xf32> -> tensor<4x2x64x4xf32>1060 return %pack : tensor<4x2x64x4xf32>1061}1062// CHECK-LABEL: func.func @bubble_up_pack_outer_expanded_through_expand(1063// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1064// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<8x64x4xf32>1065// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]1066// CHECK-SAME: inner_dims_pos = [0] inner_tiles = [4] into %[[EMPTY]] : tensor<32x64xf32> -> tensor<8x64x4xf32>1067// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[PACK]] {{\[}}[0, 1], [2], [3]]1068// CHECK-SAME: output_shape [4, 2, 64, 4] : tensor<8x64x4xf32> into tensor<4x2x64x4xf32>1069// CHECK: return %[[EXPANDED]] : tensor<4x2x64x4xf32>1070 1071// -----1072 1073func.func @bubble_up_pack_inner_expanded_through_expand(%arg0: tensor<32x64xf32>) -> tensor<32x4x4x4xf32> {1074 %empty = tensor.empty() : tensor<32x4x4x4xf32>1075 %expanded = tensor.expand_shape %arg0 [[0], [1, 2]] output_shape [32, 4, 16] : tensor<32x64xf32> into tensor<32x4x16xf32>1076 %pack = linalg.pack %expanded inner_dims_pos = [2] inner_tiles = [4] into %empty : tensor<32x4x16xf32> -> tensor<32x4x4x4xf32>1077 return %pack : tensor<32x4x4x4xf32>1078}1079// CHECK-LABEL: func.func @bubble_up_pack_inner_expanded_through_expand(1080// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1081// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<32x16x4xf32>1082// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]1083// CHECK-SAME: inner_dims_pos = [1] inner_tiles = [4] into %[[EMPTY]]1084// CHECK-SAME: : tensor<32x64xf32> -> tensor<32x16x4xf32>1085// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[PACK]] {{\[}}[0], [1, 2], [3]]1086// CHECK-SAME: output_shape [32, 4, 4, 4] : tensor<32x16x4xf32> into tensor<32x4x4x4xf32>1087// CHECK: return %[[EXPANDED]] : tensor<32x4x4x4xf32>1088 1089// -----1090 1091func.func @bubble_up_pack_non_expanded_dims_through_expand(%arg0: tensor<32x64x16xf32>) -> tensor<8x2x32x16x4xf32> {1092 %empty = tensor.empty() : tensor<8x2x32x16x4xf32>1093 %expanded = tensor.expand_shape %arg0 [[0], [1, 2], [3]] output_shape [32, 2, 32, 16] : tensor<32x64x16xf32> into tensor<32x2x32x16xf32>1094 %pack = linalg.pack %expanded inner_dims_pos = [0] inner_tiles = [4] into %empty : tensor<32x2x32x16xf32> -> tensor<8x2x32x16x4xf32>1095 return %pack : tensor<8x2x32x16x4xf32>1096}1097// CHECK-LABEL: func.func @bubble_up_pack_non_expanded_dims_through_expand(1098// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1099// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<8x64x16x4xf32>1100// CHECK: %[[PACK:.+]] = linalg.pack1101// CHECK-SAME: %[[ARG0]] inner_dims_pos = [0] inner_tiles = [4] into %[[EMPTY]]1102// CHECK-SAME: : tensor<32x64x16xf32> -> tensor<8x64x16x4xf32>1103// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[PACK]] {{\[}}[0], [1, 2], [3], [4]]1104// CHECK-SAME: output_shape [8, 2, 32, 16, 4] : tensor<8x64x16x4xf32> into tensor<8x2x32x16x4xf32>1105// CHECK: return %[[EXPANDED]] : tensor<8x2x32x16x4xf32>1106 1107// -----1108 1109func.func @bubble_up_pack_through_expand_dynamic(%arg0: tensor<?x64xf32>) -> tensor<?x4x2x8xf32> {1110 %c0 = arith.constant 0 : index1111 %dim = tensor.dim %arg0, %c0 : tensor<?x64xf32>1112 %empty = tensor.empty(%dim) : tensor<?x4x2x8xf32>1113 %expanded = tensor.expand_shape %arg0 [[0], [1, 2]] output_shape [%dim, 4, 16] : tensor<?x64xf32> into tensor<?x4x16xf32>1114 %pack = linalg.pack %expanded inner_dims_pos = [2] inner_tiles = [8] into %empty : tensor<?x4x16xf32> -> tensor<?x4x2x8xf32>1115 return %pack : tensor<?x4x2x8xf32>1116}1117// CHECK-LABEL: func.func @bubble_up_pack_through_expand_dynamic(1118// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1119// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index1120// CHECK: %[[DIM_INPUT:.+]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x64xf32>1121// CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM_INPUT]]) : tensor<?x8x8xf32>1122// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]1123// CHECK-SAME: inner_dims_pos = [1] inner_tiles = [8] into %[[EMPTY]]1124// CHECK-SAME: : tensor<?x64xf32> -> tensor<?x8x8xf32>1125// CHECK: %[[DIM_PACK:.+]] = tensor.dim %[[PACK]], %[[C0]] : tensor<?x8x8xf32>1126// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[PACK]] {{\[}}[0], [1, 2], [3]]1127// CHECK-SAME: output_shape [%[[DIM_PACK]], 4, 2, 8] : tensor<?x8x8xf32> into tensor<?x4x2x8xf32>1128// CHECK: return %[[EXPANDED]] : tensor<?x4x2x8xf32>1129 1130// -----1131 1132func.func @bubble_up_pack_non_expanded_padding_through_expand(%arg0: tensor<32x60xf32>) -> tensor<4x2x8x4x8xf32> {1133 %cst = arith.constant 3.000000e+00 : f321134 %empty = tensor.empty() : tensor<4x2x8x4x8xf32>1135 %expanded = tensor.expand_shape %arg0 [[0, 1], [2]] output_shape [4, 8, 64] : tensor<32x60xf32> into tensor<4x8x60xf32>1136 %pack = linalg.pack %expanded padding_value(%cst : f32) inner_dims_pos = [1, 2] inner_tiles = [4, 8] into %empty : tensor<4x8x60xf32> -> tensor<4x2x8x4x8xf32>1137 return %pack : tensor<4x2x8x4x8xf32>1138}1139// CHECK-LABEL: func.func @bubble_up_pack_non_expanded_padding_through_expand(1140// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1141// CHECK-DAG: %[[CST:.+]] = arith.constant 3.000000e+00 : f321142// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<8x8x4x8xf32>1143// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]] padding_value(%[[CST]] : f32)1144// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [4, 8] into %[[EMPTY]]1145// CHECK-SAME: : tensor<32x60xf32> -> tensor<8x8x4x8xf32>1146// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[PACK]] {{\[}}[0, 1], [2], [3], [4]]1147// CHECK-SAME: output_shape [4, 2, 8, 4, 8] : tensor<8x8x4x8xf32> into tensor<4x2x8x4x8xf32>1148// CHECK: return %[[EXPANDED]] : tensor<4x2x8x4x8xf32>1149 1150// -----1151 1152func.func @bubble_up_pack_outer_dims_perm_identity_through_expand(%arg0: tensor<32x64xf32>) -> tensor<4x2x32x4x2xf32> {1153 %empty = tensor.empty() : tensor<4x2x32x4x2xf32>1154 %expanded = tensor.expand_shape %arg0 [[0, 1], [2]] output_shape [4, 8, 64] : tensor<32x64xf32> into tensor<4x8x64xf32>1155 %pack = linalg.pack %expanded outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [4, 2] into %empty : tensor<4x8x64xf32> -> tensor<4x2x32x4x2xf32>1156 return %pack : tensor<4x2x32x4x2xf32>1157}1158// CHECK-LABEL: func.func @bubble_up_pack_outer_dims_perm_identity_through_expand(1159// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1160// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<8x32x4x2xf32>1161// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]1162// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [4, 2] into %[[EMPTY]]1163// CHECK-SAME: : tensor<32x64xf32> -> tensor<8x32x4x2xf32>1164// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[PACK]] {{\[}}[0, 1], [2], [3], [4]]1165// CHECK-SAME: output_shape [4, 2, 32, 4, 2] : tensor<8x32x4x2xf32> into tensor<4x2x32x4x2xf32>1166// CHECK: return %[[EXPANDED]] : tensor<4x2x32x4x2xf32>1167 1168// -----1169 1170func.func @bubble_up_pack_multiple_dims_through_expand(%arg0: tensor<32x64x16xf32>) -> tensor<8x2x4x8x4x8x2xf32> {1171 %empty = tensor.empty() : tensor<8x2x4x8x4x8x2xf32>1172 %expanded = tensor.expand_shape %arg0 [[0], [1, 2], [3]] output_shape [32, 2, 32, 16] : tensor<32x64x16xf32> into tensor<32x2x32x16xf32>1173 %pack = linalg.pack %expanded inner_dims_pos = [0, 2, 3] inner_tiles = [4, 8, 2] into %empty : tensor<32x2x32x16xf32> -> tensor<8x2x4x8x4x8x2xf32>1174 return %pack : tensor<8x2x4x8x4x8x2xf32>1175}1176// CHECK-LABEL: func.func @bubble_up_pack_multiple_dims_through_expand(1177// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1178// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<8x8x8x4x8x2xf32>1179// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]1180// CHECK-SAME: inner_dims_pos = [0, 1, 2] inner_tiles = [4, 8, 2] into %[[EMPTY]]1181// CHECK-SAME: : tensor<32x64x16xf32> -> tensor<8x8x8x4x8x2xf32>1182// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[PACK]] {{\[}}[0], [1, 2], [3], [4], [5], [6]]1183// CHECK-SAME: output_shape [8, 2, 4, 8, 4, 8, 2] : tensor<8x8x8x4x8x2xf32> into tensor<8x2x4x8x4x8x2xf32>1184// CHECK: return %[[EXPANDED]] : tensor<8x2x4x8x4x8x2xf32>1185 1186// -----1187 1188func.func @bubble_up_pack_inner_dims_reorder_through_expand(%arg0: tensor<32x64xf32>) -> tensor<4x2x4x16x4xf32> {1189 %empty = tensor.empty() : tensor<4x2x4x16x4xf32>1190 %expanded = tensor.expand_shape %arg0 [[0, 1], [2]] output_shape [4, 8, 64] : tensor<32x64xf32> into tensor<4x8x64xf32>1191 %pack = linalg.pack %expanded inner_dims_pos = [2, 1] inner_tiles = [16, 4] into %empty : tensor<4x8x64xf32> -> tensor<4x2x4x16x4xf32>1192 return %pack : tensor<4x2x4x16x4xf32>1193}1194// CHECK-LABEL: func.func @bubble_up_pack_inner_dims_reorder_through_expand(1195// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1196// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<8x4x16x4xf32>1197// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]1198// CHECK-SAME: inner_dims_pos = [1, 0] inner_tiles = [16, 4] into %[[EMPTY]]1199// CHECK-SAME: : tensor<32x64xf32> -> tensor<8x4x16x4xf32>1200// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[PACK]] {{\[}}[0, 1], [2], [3], [4]]1201// CHECK-SAME: output_shape [4, 2, 4, 16, 4] : tensor<8x4x16x4xf32> into tensor<4x2x4x16x4xf32>1202// CHECK: return %[[EXPANDED]] : tensor<4x2x4x16x4xf32>1203 1204// -----1205 1206func.func @bubble_up_pack_multiple_different_expanded_dims_through_expand(%arg0: tensor<32x64x16xf32>) -> tensor<4x2x2x8x16x4x4xf32> {1207 %empty = tensor.empty() : tensor<4x2x2x8x16x4x4xf32>1208 %expanded = tensor.expand_shape %arg0 [[0, 1], [2, 3], [4]] output_shape [4, 8, 2, 32, 16] : tensor<32x64x16xf32> into tensor<4x8x2x32x16xf32>1209 %pack = linalg.pack %expanded inner_dims_pos = [1, 3] inner_tiles = [4, 4] into %empty : tensor<4x8x2x32x16xf32> -> tensor<4x2x2x8x16x4x4xf32>1210 return %pack : tensor<4x2x2x8x16x4x4xf32>1211}1212// CHECK-LABEL: func.func @bubble_up_pack_multiple_different_expanded_dims_through_expand(1213// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1214// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<8x16x16x4x4xf32>1215// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG0]]1216// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [4, 4] into %[[EMPTY]]1217// CHECK-SAME: : tensor<32x64x16xf32> -> tensor<8x16x16x4x4xf32>1218// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[PACK]] {{\[}}[0, 1], [2, 3], [4], [5], [6]]1219// CHECK-SAME: output_shape [4, 2, 2, 8, 16, 4, 4] : tensor<8x16x16x4x4xf32> into tensor<4x2x2x8x16x4x4xf32>1220// CHECK: return %[[EXPANDED]] : tensor<4x2x2x8x16x4x4xf32>1221 1222// -----1223 1224func.func @no_bubble_up_pack_outer_dims_permutation_through_expand(%arg0: tensor<32x64xf32>) -> tensor<32x4x2x4x2xf32> {1225 %empty = tensor.empty() : tensor<32x4x2x4x2xf32>1226 %expanded = tensor.expand_shape %arg0 [[0, 1], [2]] output_shape [4, 8, 64] : tensor<32x64xf32> into tensor<4x8x64xf32>1227 %pack = linalg.pack %expanded outer_dims_perm = [2, 0, 1] inner_dims_pos = [1, 2] inner_tiles = [4, 2] into %empty : tensor<4x8x64xf32> -> tensor<32x4x2x4x2xf32>1228 return %pack : tensor<32x4x2x4x2xf32>1229}1230// CHECK-LABEL: func.func @no_bubble_up_pack_outer_dims_permutation_through_expand(1231// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1232// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<32x4x2x4x2xf32>1233// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1], [2]]1234// CHECK-SAME: output_shape [4, 8, 64] : tensor<32x64xf32> into tensor<4x8x64xf32>1235// CHECK: %[[PACK:.+]] = linalg.pack %[[EXPANDED]]1236// CHECK-SAME: outer_dims_perm = [2, 0, 1] inner_dims_pos = [1, 2] inner_tiles = [4, 2] into %[[EMPTY]]1237// CHECK-SAME: : tensor<4x8x64xf32> -> tensor<32x4x2x4x2xf32>1238// CHECK: return %[[PACK]] : tensor<32x4x2x4x2xf32>1239 1240// -----1241 1242func.func @no_bubble_up_pack_multiple_same_expanded_dim_through_expand(%arg0: tensor<32x64xf32>) -> tensor<2x2x64x2x4xf32> {1243 %empty = tensor.empty() : tensor<2x2x64x2x4xf32>1244 %expanded = tensor.expand_shape %arg0 [[0, 1], [2]] output_shape [4, 8, 64] : tensor<32x64xf32> into tensor<4x8x64xf32>1245 %pack = linalg.pack %expanded inner_dims_pos = [0, 1] inner_tiles = [2, 4] into %empty : tensor<4x8x64xf32> -> tensor<2x2x64x2x4xf32>1246 return %pack : tensor<2x2x64x2x4xf32>1247}1248// CHECK-LABEL: func.func @no_bubble_up_pack_multiple_same_expanded_dim_through_expand(1249// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1250// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<2x2x64x2x4xf32>1251// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1], [2]]1252// CHECK-SAME: output_shape [4, 8, 64] : tensor<32x64xf32> into tensor<4x8x64xf32>1253// CHECK: %[[PACK:.+]] = linalg.pack %[[EXPANDED]]1254// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [2, 4] into %[[EMPTY]]1255// CHECK-SAME: : tensor<4x8x64xf32> -> tensor<2x2x64x2x4xf32>1256// CHECK: return %[[PACK]] : tensor<2x2x64x2x4xf32>1257 1258// -----1259 1260func.func @no_bubble_up_pack_non_innermost_expanded_dim_through_expand(%arg0: tensor<32x64xf32>) -> tensor<2x8x64x2xf32> {1261 %empty = tensor.empty() : tensor<2x8x64x2xf32>1262 %expanded = tensor.expand_shape %arg0 [[0, 1], [2]] output_shape [4, 8, 64] : tensor<32x64xf32> into tensor<4x8x64xf32>1263 %pack = linalg.pack %expanded inner_dims_pos = [0] inner_tiles = [2] into %empty : tensor<4x8x64xf32> -> tensor<2x8x64x2xf32>1264 return %pack : tensor<2x8x64x2xf32>1265}1266// CHECK-LABEL: func.func @no_bubble_up_pack_non_innermost_expanded_dim_through_expand(1267// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1268// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<2x8x64x2xf32>1269// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1], [2]]1270// CHECK-SAME: output_shape [4, 8, 64] : tensor<32x64xf32> into tensor<4x8x64xf32>1271// CHECK: %[[PACK:.+]] = linalg.pack %[[EXPANDED]]1272// CHECK-SAME: inner_dims_pos = [0] inner_tiles = [2] into %[[EMPTY]]1273// CHECK-SAME: : tensor<4x8x64xf32> -> tensor<2x8x64x2xf32>1274// CHECK: return %[[PACK]] : tensor<2x8x64x2xf32>1275 1276// -----1277 1278func.func @no_bubble_up_pack_expanded_padding_through_expand_cannot_reassociate(%arg0: tensor<30x60xf32>) -> tensor<3x2x60x8xf32> {1279 %cst = arith.constant 3.000000e+00 : f321280 %empty = tensor.empty() : tensor<3x2x60x8xf32>1281 %expanded = tensor.expand_shape %arg0 [[0, 1], [2]] output_shape [3, 10, 60] : tensor<30x60xf32> into tensor<3x10x60xf32>1282 %pack = linalg.pack %expanded padding_value(%cst : f32) inner_dims_pos = [1] inner_tiles = [8] into %empty : tensor<3x10x60xf32> -> tensor<3x2x60x8xf32>1283 return %pack : tensor<3x2x60x8xf32>1284}1285// CHECK-LABEL: func.func @no_bubble_up_pack_expanded_padding_through_expand_cannot_reassociate(1286// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1287// CHECK-DAG: %[[CST:.+]] = arith.constant 3.000000e+00 : f321288// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<3x2x60x8xf32>1289// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1], [2]]1290// CHECK-SAME: output_shape [3, 10, 60] : tensor<30x60xf32> into tensor<3x10x60xf32>1291// CHECK: %[[PACK:.+]] = linalg.pack %[[EXPANDED]] padding_value(%[[CST]] : f32)1292// CHECK-SAME: inner_dims_pos = [1] inner_tiles = [8] into %[[EMPTY]]1293// CHECK-SAME: : tensor<3x10x60xf32> -> tensor<3x2x60x8xf32>1294// CHECK: return %[[PACK]] : tensor<3x2x60x8xf32>1295 1296// -----1297 1298func.func @push_down_unpack_through_expand(%5: tensor<?x32x8x8xf32>, %dim: index, %sz0: index) -> tensor<?x256x256xf32> {1299 %6 = tensor.empty(%dim) : tensor<?x256xf32>1300 %unpack = linalg.unpack %5 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [8, 8] into %6 : tensor<?x32x8x8xf32> -> tensor<?x256xf32>1301 %expanded = tensor.expand_shape %unpack [[0, 1], [2]] output_shape [%sz0, 256, 256] : tensor<?x256xf32> into tensor<?x256x256xf32>1302 func.return %expanded : tensor<?x256x256xf32>1303}1304// CHECK-LABEL: func.func @push_down_unpack_through_expand1305// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1306// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]1307// CHECK: %[[C32:.+]] = arith.constant 32 : index1308// CHECK: %[[C0:.+]] = arith.constant 0 : index1309// CHECK: %[[DIM0:.+]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x32x8x8xf32>1310// CHECK: %[[SZ0:.+]] = arith.divsi %[[DIM0]], %[[C32]] : index1311// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1], [2], [3], [4]] output_shape [%[[SZ0]], 32, 32, 8, 8] : tensor<?x32x8x8xf32> into tensor<?x32x32x8x8xf32>1312// CHECK: %[[DIM:.+]] = tensor.dim %[[EXPANDED]], %[[C0]] : tensor<?x32x32x8x8xf32>1313// CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM]]) : tensor<?x256x256xf32>1314// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[EXPANDED:.+]] outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 8] into %[[EMPTY]] : tensor<?x32x32x8x8xf32> -> tensor<?x256x256xf32>1315// CHECK: return %[[UNPACK]] : tensor<?x256x256xf32>1316 1317// -----1318 1319func.func @push_down_unpack_through_expand_empty_outer_dims_perm(%5: tensor<?x32x8x8xf32>, %dim: index, %sz0: index) -> tensor<?x256x256xf32> {1320 %6 = tensor.empty(%dim) : tensor<?x256xf32>1321 %unpack = linalg.unpack %5 inner_dims_pos = [0, 1] inner_tiles = [8, 8] into %6 : tensor<?x32x8x8xf32> -> tensor<?x256xf32>1322 %expanded = tensor.expand_shape %unpack [[0, 1], [2]] output_shape [%sz0, 256, 256] : tensor<?x256xf32> into tensor<?x256x256xf32>1323 func.return %expanded : tensor<?x256x256xf32>1324}1325// CHECK-LABEL: func.func @push_down_unpack_through_expand_empty_outer_dims_perm1326// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1327// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]1328// CHECK: %[[C32:.+]] = arith.constant 32 : index1329// CHECK: %[[C0:.+]] = arith.constant 0 : index1330// CHECK: %[[DIM0:.+]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x32x8x8xf32>1331// CHECK: %[[SZ0:.+]] = arith.divsi %[[DIM0]], %[[C32]] : index1332// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1], [2], [3], [4]] output_shape [%[[SZ0]], 32, 32, 8, 8] : tensor<?x32x8x8xf32> into tensor<?x32x32x8x8xf32>1333// CHECK: %[[DIM:.+]] = tensor.dim %[[EXPANDED]], %[[C0]] : tensor<?x32x32x8x8xf32>1334// CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM]]) : tensor<?x256x256xf32>1335// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[EXPANDED:.+]] inner_dims_pos = [1, 2] inner_tiles = [8, 8] into %[[EMPTY]] : tensor<?x32x32x8x8xf32> -> tensor<?x256x256xf32>1336// CHECK: return %[[UNPACK]] : tensor<?x256x256xf32>1337 1338// -----1339 1340func.func @push_down_permuted_unpack_through_expand(%5: tensor<4x32x384x8x8xf32>) -> tensor<4x12x256x256xf32> {1341 %6 = tensor.empty() : tensor<4x3072x256xf32>1342 %unpack = linalg.unpack %5 outer_dims_perm = [0, 2, 1] inner_dims_pos = [2, 1] inner_tiles = [8, 8] into %6 : tensor<4x32x384x8x8xf32> -> tensor<4x3072x256xf32>1343 %expanded = tensor.expand_shape %unpack [[0], [1, 2], [3]] output_shape [4, 12, 256, 256] : tensor<4x3072x256xf32> into tensor<4x12x256x256xf32>1344 func.return %expanded : tensor<4x12x256x256xf32>1345}1346// CHECK-LABEL: @push_down_permuted_unpack_through_expand1347// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1348// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1], [2, 3], [4], [5]] output_shape [4, 32, 12, 32, 8, 8] : tensor<4x32x384x8x8xf32> into tensor<4x32x12x32x8x8xf32>1349// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<4x12x256x256xf32>1350// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[EXPANDED]] outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3, 2] inner_tiles = [8, 8] into %[[EMPTY]] : tensor<4x32x12x32x8x8xf32> -> tensor<4x12x256x256xf32>1351// CHECK: return %[[UNPACK]] : tensor<4x12x256x256xf32>1352 1353// -----1354 1355func.func @push_down_unpack_through_unit_expand(%5: tensor<6x32x8x8xf32>) -> tensor<3x16x1x256xf32> {1356 %6 = tensor.empty() : tensor<48x256xf32>1357 %unpack = linalg.unpack %5 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [8, 8] into %6 : tensor<6x32x8x8xf32> -> tensor<48x256xf32>1358 %expanded = tensor.expand_shape %unpack [[0, 1, 2], [3]] output_shape [3, 16, 1, 256] : tensor<48x256xf32> into tensor<3x16x1x256xf32>1359 func.return %expanded : tensor<3x16x1x256xf32>1360}1361// CHECK-LABEL: func.func @push_down_unpack_through_unit_expand1362// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1363// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1, 2], [3], [4], [5]] output_shape [3, 2, 1, 32, 8, 8] : tensor<6x32x8x8xf32> into tensor<3x2x1x32x8x8xf32>1364// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<3x16x1x256xf32>1365// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[EXPANDED]] outer_dims_perm = [0, 1, 2, 3] inner_dims_pos = [1, 3] inner_tiles = [8, 8] into %[[EMPTY]] : tensor<3x2x1x32x8x8xf32> -> tensor<3x16x1x256xf32>1366// CHECK: return %[[UNPACK]] : tensor<3x16x1x256xf32>1367 1368// -----1369 1370func.func @push_down_unpack_through_expand_on_outer_dims(%5: tensor<?x32x8xf32>, %dim: index, %sz0: index) -> tensor<?x256x256xf32> {1371 %6 = tensor.empty(%dim) : tensor<?x256xf32>1372 %unpack = linalg.unpack %5 outer_dims_perm = [0, 1] inner_dims_pos = [1] inner_tiles = [8] into %6 : tensor<?x32x8xf32> -> tensor<?x256xf32>1373 %expanded = tensor.expand_shape %unpack [[0, 1], [2]] output_shape [%sz0, 256, 256] : tensor<?x256xf32> into tensor<?x256x256xf32>1374 func.return %expanded : tensor<?x256x256xf32>1375}1376// CHECK-LABEL: func.func @push_down_unpack_through_expand_on_outer_dims1377// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1378// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]1379// CHECK: %[[C256:.+]] = arith.constant 256 : index1380// CHECK: %[[C0:.+]] = arith.constant 0 : index1381// CHECK: %[[DIM0:.+]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x32x8xf32>1382// CHECK: %[[SZ0:.+]] = arith.divsi %[[DIM0]], %[[C256]] : index1383// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1], [2], [3]] output_shape [%[[SZ0]], 256, 32, 8] : tensor<?x32x8xf32> into tensor<?x256x32x8xf32>1384// CHECK: %[[DIM:.+]] = tensor.dim %[[EXPANDED]], %[[C0]] : tensor<?x256x32x8xf32>1385// CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM]]) : tensor<?x256x256xf32>1386// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[EXPANDED:.+]] outer_dims_perm = [0, 1, 2] inner_dims_pos = [2] inner_tiles = [8] into %[[EMPTY]] : tensor<?x256x32x8xf32> -> tensor<?x256x256xf32>1387// CHECK: return %[[UNPACK]] : tensor<?x256x256xf32>1388 1389// -----1390 1391func.func @no_push_down_unpack_through_non_divisible_expand(%5: tensor<384x32x8x8xf32>) -> tensor<256x12x256xf32> {1392 %6 = tensor.empty() : tensor<3072x256xf32>1393 %unpack = linalg.unpack %5 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [8, 8] into %6 : tensor<384x32x8x8xf32> -> tensor<3072x256xf32>1394 %expanded = tensor.expand_shape %unpack [[0, 1], [2]] output_shape [256, 12, 256] : tensor<3072x256xf32> into tensor<256x12x256xf32>1395 func.return %expanded : tensor<256x12x256xf32>1396}1397// CHECK-LABEL: func.func @no_push_down_unpack_through_non_divisible_expand1398// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1399// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[ARG0]]1400// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[UNPACK]] {{\[}}[0, 1], [2]] output_shape [256, 12, 256] : tensor<3072x256xf32> into tensor<256x12x256xf32>1401// CHECK: return %[[EXPANDED]] : tensor<256x12x256xf32>1402 1403// -----1404 1405func.func @push_unpack_in_padded_domain_foldable(%arg0: tensor<8x8x4x8xf32>, %dest: tensor<?x64xf32>, %arg1: tensor<?x64xbf16>) -> tensor<?x64xbf16> {1406 %unpack = linalg.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [4, 8] into %dest : tensor<8x8x4x8xf32> -> tensor<?x64xf32>1407 %0 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} ins(%unpack : tensor<?x64xf32>) outs(%arg1 : tensor<?x64xbf16>) {1408 ^bb0(%in: f32, %out: bf16):1409 %1 = arith.truncf %in : f32 to bf161410 linalg.yield %1 : bf161411 } -> tensor<?x64xbf16>1412 return %0 : tensor<?x64xbf16>1413}1414// CHECK-LABEL: func.func @push_unpack_in_padded_domain_foldable1415// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1416// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]1417// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]1418// CHECK: %[[EMPTY:.+]] = tensor.empty1419// CHECK: %[[GENERIC:.+]] = linalg.generic1420// CHECK-SAME: ins(%[[ARG0]] : tensor<8x8x4x8xf32>)1421// CHECK-SAME: outs(%[[EMPTY]] : tensor<?x8x4x8xbf16>)1422// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[GENERIC]]1423// CHECK-SAME: into %[[ARG2]]1424// CHECK: return %[[UNPACK]] : tensor<?x64xbf16>1425 1426// -----1427 1428func.func @push_unpack_in_padded_domain_out_used(%arg0: tensor<8x8x4x8xf32>, %arg1: tensor<?x64xf32>) -> tensor<?x64xf32> {1429 %unpack = linalg.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [4, 8] into %arg1 : tensor<8x8x4x8xf32> -> tensor<?x64xf32>1430 %0 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} ins(%unpack : tensor<?x64xf32>) outs(%arg1 : tensor<?x64xf32>) {1431 ^bb0(%in: f32, %out: f32):1432 %1 = arith.addf %in, %out : f321433 linalg.yield %1 : f321434 } -> tensor<?x64xf32>1435 return %0 : tensor<?x64xf32>1436}1437// CHECK-LABEL: func.func @push_unpack_in_padded_domain_out_used1438// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1439// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]1440// CHECK: %[[ARG1_PACK_EMPTY:.+]] = tensor.empty1441// CHECK: %[[ARG1_PACK:.+]] = linalg.pack %[[ARG1]]1442// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [4, 8]1443// CHECK-SAME: into %[[ARG1_PACK_EMPTY]]1444// CHECK: %[[GENERIC:.+]] = linalg.generic1445// CHECK-SAME: ins(%[[ARG0]] : tensor<8x8x4x8xf32>)1446// CHECK-SAME: outs(%[[ARG1_PACK]] : tensor<?x8x4x8xf32>)1447// CHECK: %[[UNPACK2:.+]] = linalg.unpack %[[GENERIC]]1448// CHECK-SAME: into %[[ARG1]]1449// CHECK: return %[[UNPACK2]] : tensor<?x64xf32>1450 1451// -----1452 1453#map = affine_map<(d0, d1) -> (d0, d1)>1454func.func @push_unpack_in_padded_domain_multiple_inputs(%arg0: tensor<1x4x16x16xf32>, %arg1: tensor<8x64xf32>, %arg2: tensor<8x64xf32>) -> tensor<8x64xf32> {1455 %0 = tensor.empty() : tensor<8x64xf32>1456 %unpack = linalg.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [16, 16] into %0 : tensor<1x4x16x16xf32> -> tensor<8x64xf32>1457 %1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg1, %unpack : tensor<8x64xf32>, tensor<8x64xf32>) outs(%arg2 : tensor<8x64xf32>) {1458 ^bb0(%in: f32, %in_0: f32, %out: f32):1459 %2 = arith.addf %in, %in_0 : f321460 linalg.yield %2 : f321461 } -> tensor<8x64xf32>1462 return %1 : tensor<8x64xf32>1463}1464// CHECK-LABEL: func.func @push_unpack_in_padded_domain_multiple_inputs1465// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1466// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]1467// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]1468// CHECK-DAG: %[[POISON:.+]] = ub.poison : f321469// CHECK: %[[PACK:.+]] = linalg.pack %[[ARG1]] padding_value(%[[POISON]] : f32)1470// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [16, 16]1471// CHECK: %[[ELEM:.+]] = linalg.generic1472// CHECK: ins(%[[PACK]], %[[ARG0]]1473// CHECK: %[[UNPACK:.+]] = linalg.unpack %[[ELEM]]1474// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [16, 16]1475// CHECK-SAME: into %[[ARG2]]1476// CHECK: return %[[UNPACK]]1477 1478// -----1479 1480module {1481 func.func @push_extract_through_generic(%arg0: tensor<128x7x128xf32>, %arg1: tensor<?x5x3x128xf32>, %arg2: tensor<?x5x128xbf16>, %arg3: index) -> tensor<?x5x128xbf16> {1482 %extracted_slice = tensor.extract_slice %arg0[0, 0, %arg3] [128, 7, %arg3] [1, 1, 1] : tensor<128x7x128xf32> to tensor<128x7x?xf32>1483 %0 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d1, d2 + d3, d0)>, affine_map<(d0, d1, d2, d3) -> (d0, d2, d3, d1)>, affine_map<(d0, d1, d2, d3) -> (d0, d2, d1)>], iterator_types = ["parallel", "parallel", "parallel", "reduction"]} ins(%extracted_slice, %arg1 : tensor<128x7x?xf32>, tensor<?x5x3x128xf32>) outs(%arg2 : tensor<?x5x128xbf16>) {1484 ^bb0(%in: f32, %in_0: f32, %out: bf16):1485 %1 = arith.truncf %in : f32 to bf161486 linalg.yield %1 : bf161487 } -> tensor<?x5x128xbf16>1488 return %0 : tensor<?x5x128xbf16>1489 }1490}1491 1492// CHECK-LABEL: func.func @push_extract_through_generic1493// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1494// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]1495// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]1496// CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]1497// CHECK: %[[POISON:.+]] = ub.poison : f321498// CHECK: %[[PADDED:.+]] = tensor.pad %arg11499// CHECK: tensor.yield %[[POISON]] : f321500// CHECK: } : tensor<?x5x3x128xf32> to tensor<?x5x3x128xf32>1501// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<128x5x128xbf16>1502// CHECK: %[[GENERIC:.+]] = linalg.generic1503// CHECK-SAME: ins(%[[ARG0]], %[[PADDED]]1504// CHECK-SAME: outs(%[[EMPTY]]1505// CHECK: %[[EXTRACT:.+]] = tensor.extract_slice %3[%[[ARG3]], 0, 0] [%[[ARG3]], 5, 128] [1, 1, 1] : tensor<128x5x128xbf16> to tensor<?x5x128xbf16>1506// CHECK: return %[[EXTRACT]]1507 1508// -----1509 1510func.func @nopush_extract_through_generic_nodimexpr1(%arg0: tensor<128x7x128xf32>, %arg1: tensor<?x5x3x128xf32>, %arg2: tensor<?x5x128xbf16>, %arg3: index) -> tensor<?x5x128xbf16> {1511 %extracted_slice = tensor.extract_slice %arg0[0, %arg3, %arg3] [128, 7, %arg3] [1, 1, 1] : tensor<128x7x128xf32> to tensor<128x7x?xf32>1512 %0 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d1, d2 + d3, d0)>, affine_map<(d0, d1, d2, d3) -> (d0, d2, d3, d1)>, affine_map<(d0, d1, d2, d3) -> (d0, d2, d1)>], iterator_types = ["parallel", "parallel", "parallel", "reduction"]} ins(%extracted_slice, %arg1 : tensor<128x7x?xf32>, tensor<?x5x3x128xf32>) outs(%arg2 : tensor<?x5x128xbf16>) {1513 ^bb0(%in: f32, %in_0: f32, %out: bf16):1514 %1 = arith.truncf %in : f32 to bf161515 linalg.yield %1 : bf161516 } -> tensor<?x5x128xbf16>1517 return %0 : tensor<?x5x128xbf16>1518}1519 1520// CHECK-LABEL: func.func @nopush_extract_through_generic_nodimexpr11521// CHECK: %[[GENERIC:.+]] = linalg.generic1522// CHECK: return %[[GENERIC]]1523 1524// -----1525 1526func.func @nopush_extract_through_generic_nodimexpr2(%arg0: tensor<128x?x128xf32>, %arg1: tensor<128x5x3x128xf32>, %arg2: tensor<128x?x128xbf16>, %arg3: index) -> tensor<128x?x128xbf16> {1527 %extracted_slice = tensor.extract_slice %arg1[0, %arg3, 0, 0] [128, %arg3, 3, 128] [1, 1, 1, 1] : tensor<128x5x3x128xf32> to tensor<128x?x3x128xf32>1528 %0 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d1, d2 + d3, d0)>, affine_map<(d0, d1, d2, d3) -> (d0, d2, d3, d1)>, affine_map<(d0, d1, d2, d3) -> (d0, d2, d1)>], iterator_types = ["parallel", "parallel", "parallel", "reduction"]} ins(%arg0, %extracted_slice : tensor<128x?x128xf32>, tensor<128x?x3x128xf32>) outs(%arg2 : tensor<128x?x128xbf16>) {1529 ^bb0(%in: f32, %in_0: f32, %out: bf16):1530 %1 = arith.truncf %in : f32 to bf161531 linalg.yield %1 : bf161532 } -> tensor<128x?x128xbf16>1533 return %0 : tensor<128x?x128xbf16>1534}1535 1536// CHECK-LABEL: func.func @nopush_extract_through_generic_nodimexpr21537// CHECK: %[[GENERIC:.+]] = linalg.generic1538// CHECK: return %[[GENERIC]]1539 1540// -----1541 1542func.func @push_redcutionextract_through_generic_withoutsused_2(%arg0: tensor<128x128xf32>, %arg1: tensor<?xbf16>, %arg2: index) -> tensor<?xbf16> {1543 %extracted_slice = tensor.extract_slice %arg0[%arg2, %arg2] [%arg2, %arg2] [1, 1] : tensor<128x128xf32> to tensor<?x?xf32>1544 %0 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>], iterator_types = ["parallel", "reduction"]} ins(%extracted_slice : tensor<?x?xf32>) outs(%arg1 : tensor<?xbf16>) {1545 ^bb0(%in: f32, %out: bf16):1546 %1 = arith.truncf %in : f32 to bf161547 %2 = arith.addf %1, %out : bf161548 linalg.yield %2 : bf161549 } -> tensor<?xbf16>1550 return %0 : tensor<?xbf16>1551}1552 1553// CHECK-LABEL: func.func @push_redcutionextract_through_generic_withoutsused_21554// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1555// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]1556// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]1557// CHECK: %[[POISON_BF16:.+]] = ub.poison : bf161558// CHECK: %[[POISON_F32:.+]] = ub.poison : f321559// CHECK: %[[EXTRACT:.+]] = tensor.extract_slice %[[ARG0]][%[[ARG2]], %[[ARG2]]] [%[[ARG2]], %[[ARG2]]] [1, 1] : tensor<128x128xf32> to tensor<?x?xf32>1560// CHECK: %[[PADDED:.+]] = tensor.pad %[[EXTRACT]]1561// CHECK: tensor.yield %[[POISON_F32]] : f321562// CHECK: } : tensor<?x?xf32> to tensor<?x?xf32>1563// CHECK: %[[APPLY2:.+]] = affine.apply #map()[%[[ARG2]]]1564// CHECK: %[[PADDED1:.+]] = tensor.pad %[[ARG1]] low[%[[ARG2]]] high[%[[APPLY2]]]1565// CHECK: tensor.yield %[[POISON_BF16]] : bf161566// CHECK: } : tensor<?xbf16> to tensor<?xbf16>1567// CHECK: %[[GENERIC:.+]] = linalg.generic1568// CHECK-SAME: ins(%[[PADDED]]1569// CHECK-SAME: outs(%[[PADDED1]]1570// CHECK: %[[EXTRACT1:.+]] = tensor.extract_slice %[[GENERIC]][%[[ARG2]]] [%[[ARG2]]] [1] : tensor<?xbf16> to tensor<?xbf16>1571// CHECK: return %[[EXTRACT1]]1572 1573 1574// -----1575 1576func.func @nopush_rankreducingextract(%arg0: tensor<128x128x128xf32>, %arg1: tensor<?xbf16>, %arg2: index) -> tensor<?xbf16> {1577 %extracted_slice = tensor.extract_slice %arg0[0, %arg2, %arg2] [1, %arg2, %arg2] [1, 1, 1] : tensor<128x128x128xf32> to tensor<?x?xf32>1578 %0 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>], iterator_types = ["parallel", "reduction"]} ins(%extracted_slice : tensor<?x?xf32>) outs(%arg1 : tensor<?xbf16>) {1579 ^bb0(%in: f32, %out: bf16):1580 %1 = arith.truncf %in : f32 to bf161581 %2 = arith.addf %1, %out : bf161582 linalg.yield %2 : bf161583 } -> tensor<?xbf16>1584 return %0 : tensor<?xbf16>1585}1586 1587// CHECK-LABEL: func.func @nopush_rankreducingextract1588// CHECK: %[[GENERIC:.+]] = linalg.generic1589// CHECK: return %[[GENERIC]]1590 1591// -----1592 1593func.func @push_extract_through_generic_rank0_operand(%arg0: tensor<128x128xf32>, %arg1: tensor<?x?xbf16>, %arg2: index, %arg3 : f32) -> tensor<?x?xbf16> {1594 %extracted_slice = tensor.extract_slice %arg0[%arg2, %arg2] [%arg2, %arg2] [1, 1] : tensor<128x128xf32> to tensor<?x?xf32>1595 %0 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,affine_map<(d0, d1) -> ()> ,affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} ins(%extracted_slice, %arg3 : tensor<?x?xf32>, f32) outs(%arg1 : tensor<?x?xbf16>) {1596 ^bb0(%in: f32, %in_1 : f32, %out: bf16):1597 %1 = arith.truncf %in : f32 to bf161598 linalg.yield %1 : bf161599 } -> tensor<?x?xbf16>1600 return %0 : tensor<?x?xbf16>1601}1602 1603// CHECK-LABEL: func.func @push_extract_through_generic_rank0_operand1604// CHECK: %[[GENERIC:.+]] = linalg.generic1605// CHECK: %[[EXTRACT:.+]] = tensor.extract_slice %[[GENERIC]]1606// CHECK: return %[[EXTRACT]]1607 1608// -----1609// Test that if one extract doesnt pass the control function which in this case is set to1610// only allow extracts from the same block, then an extract from a later operand can still be pushed1611// down.1612func.func @push_extract_through_generic_secondextract(%arg0: tensor<128x128xf32>, %arg1: tensor<?x?xbf16>, %arg2: index) -> tensor<?x?xbf16> {1613 %c0 = arith.constant 0 : index1614 %c32 = arith.constant 32 : index1615 %extracted_slice1 = tensor.extract_slice %arg0[%arg2, %arg2] [%arg2, %arg2] [1, 1] : tensor<128x128xf32> to tensor<?x?xf32>1616 %for = scf.for %arg3 = %c0 to %c32 step %arg2 iter_args(%arg4 = %arg1) -> tensor<?x?xbf16> {1617 %extracted_slice = tensor.extract_slice %arg0[%arg2, %arg2] [%arg2, %arg2] [1, 1] : tensor<128x128xf32> to tensor<?x?xf32>1618 %0 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,affine_map<(d0, d1) -> (d0, d1)> ,affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} ins(%extracted_slice1, %extracted_slice : tensor<?x?xf32>, tensor<?x?xf32>) outs(%arg1 : tensor<?x?xbf16>) {1619 ^bb0(%in: f32, %in_1 : f32, %out: bf16):1620 %1 = arith.truncf %in : f32 to bf161621 linalg.yield %1 : bf161622 } -> tensor<?x?xbf16>1623 scf.yield %0 : tensor<?x?xbf16>1624 }1625 return %for : tensor<?x?xbf16>1626}1627 1628// CHECK-LABEL: func.func @push_extract_through_generic_secondextract1629// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]1630// CHECK: %[[EXTRACT:.+]] = tensor.extract_slice1631// CHECK: %[[FOR:.+]] = scf.for1632// CHECK: %[[PAD:.+]] = tensor.pad %[[EXTRACT]]1633// CHECK: %[[GENERIC:.+]] = linalg.generic1634// CHECK-SAME: ins(%[[PAD]], %[[ARG0]]1635// CHECK: %[[EXTRACT2:.+]] = tensor.extract_slice %[[GENERIC]]1636// CHECK: scf.yield %[[EXTRACT2]]1637