544 lines · plain
1// RUN: mlir-opt %s --transform-interpreter --split-input-file -canonicalize | FileCheck %s2 3// CHECK-LABEL: func.func @fuse_unary4func.func @fuse_unary(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {5 6 // CHECK: %[[RES:.*]] = scf.for7 // CHECK: scf.for8 // CHECK: linalg.exp9 // CHECK: linalg.add10 // CHECK: return %[[RES]]11 %0 = linalg.exp ins(%arg0 : tensor<?x?xf32>)12 outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>13 %1 = linalg.add ins(%0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>)14 outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>15 return %1 : tensor<?x?xf32>16}17 18module attributes {transform.with_named_sequence} {19 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {20 %0 = transform.structured.match ops{["linalg.add"]} in %arg1 : (!transform.any_op) -> !transform.any_op21 %1, %loops:2 = transform.structured.fuse %0 tile_sizes [32, 32] interchange [0, 1]22 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)23 transform.yield24 }25}26 27// -----28 29// CHECK-LABEL: func.func @fuse_unary30func.func @fuse_unary(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {31 32 // CHECK: %[[PARTIAL_RES:.*]] = scf.for33 // CHECK: scf.for34 // CHECK: linalg.exp35 // CHECK: linalg.add36 // CHECK: %[[RES:.*]] = scf.for {{.*}}%[[PARTIAL_RES]]37 // CHECK: scf.for38 // CHECK: linalg.exp39 // CHECK: linalg.add40 // CHECK: return %[[RES]]41 %0 = linalg.exp ins(%arg0 : tensor<?x?xf32>)42 outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>43 %1 = linalg.add ins(%0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>)44 outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>45 return %1 : tensor<?x?xf32>46}47 48module attributes {transform.with_named_sequence} {49 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {50 %0 = transform.structured.match ops{["linalg.add"]} in %arg1 : (!transform.any_op) -> !transform.any_op51 %1, %loops:2 = transform.structured.fuse %0 tile_sizes [32, 32] interchange [0, 1]52 : (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">, !transform.any_op)53 transform.loop.peel %loops#0 : (!transform.op<"scf.for">) -> (!transform.any_op, !transform.any_op)54 transform.yield55 }56}57 58// -----59 60// CHECK-LABEL: func.func @fuse_unary_param61func.func @fuse_unary_param(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {62 63 // CHECK: %[[RES:.*]] = scf.for64 // CHECK: scf.for65 // CHECK: linalg.exp66 // CHECK: linalg.add67 // CHECK: return %[[RES]]68 %0 = linalg.exp ins(%arg0 : tensor<?x?xf32>)69 outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>70 %1 = linalg.add ins(%0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>)71 outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>72 return %1 : tensor<?x?xf32>73}74 75module attributes {transform.with_named_sequence} {76 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {77 %0 = transform.structured.match ops{["linalg.add"]} in %arg1 : (!transform.any_op) -> !transform.any_op78 %1 = transform.param.constant 32 : i32 -> !transform.param<i32>79 %2 = transform.param.constant 1 : i32 -> !transform.param<i32>80 %3, %loops:2 = transform.structured.fuse %0 tile_sizes [%1, 32] interchange [0, %2]81 : (!transform.any_op, !transform.param<i32>, !transform.param<i32>) ->82 (!transform.any_op, !transform.any_op, !transform.any_op)83 transform.yield84 }85}86 87// -----88 89// CHECK-LABEL: func.func @fuse_unary_forall90func.func @fuse_unary_forall(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {91 92 // CHECK: %[[RES:.*]] = scf.forall93 // CHECK: linalg.exp94 // CHECK: linalg.add95 // CHECK: return %[[RES]]96 %0 = linalg.exp ins(%arg0 : tensor<?x?xf32>)97 outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>98 %1 = linalg.add ins(%0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>)99 outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>100 return %1 : tensor<?x?xf32>101}102 103module attributes {transform.with_named_sequence} {104 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {105 %0 = transform.structured.match ops{["linalg.add"]} in %arg1 : (!transform.any_op) -> !transform.any_op106 %1, %loop = transform.structured.fuse %0 tile_sizes [32, 32] {use_forall}107 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)108 transform.yield109 }110}111 112// -----113 114// CHECK-LABEL: func.func @interchange_reduction115// CHECK-SAME: (%[[INPUT:.+]]: tensor<12x7x25xf32>)116func.func @interchange_reduction(%input: tensor<12x7x25xf32>) -> tensor<12x25xf32> {117 %five = arith.constant 5.0 : f32118 %init = tensor.empty() : tensor<12x25xf32>119 120// CHECK-DAG: %[[INIT:.+]] = tensor.empty()121// CHECK-DAG: %[[C5:.+]] = arith.constant 5 : index122// CHECK-DAG: %[[C7:.+]] = arith.constant 7 : index123// CHECK-DAG: %[[C4:.+]] = arith.constant 4 : index124// CHECK: %[[RES:.*]] = scf.for %[[IV0:.+]] = %{{.+}} to %{{.+}} step %[[C5]] iter_args(%[[FOR_ARG0:.+]] = %[[INIT]])125// CHECK: scf.for %[[IV1:.+]] = %{{.+}} to %{{.+}} step %[[C7]] iter_args(%[[FOR_ARG1:.+]] = %[[FOR_ARG0]])126// CHECK: %[[OUT_SLICE0:.+]] = tensor.extract_slice %[[INPUT]][%[[IV0]], 0, %[[IV1]]]127// CHECK: %[[OUT_SLICE1:.+]] = tensor.extract_slice %[[FOR_ARG1]][%[[IV0]], %[[IV1]]]128// CHECK: %[[FILL:.+]] = linalg.fill {{.+}} outs(%[[OUT_SLICE1]] : tensor<?x?xf32>)129// CHECK: scf.for %[[IV2:.+]] = %{{.+}} to %{{.+}} step %[[C4]] iter_args(%[[FOR_ARG2:.+]] = %[[FILL]])130// CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice %[[OUT_SLICE0]]131// CHECK: %[[OUT_SLICE2:.+]] = tensor.extract_slice %[[FOR_ARG2]][0, 0]132// CHECK: linalg.generic {{.+}} ins(%[[IN_SLICE]] : tensor<?x?x?xf32>) outs(%[[OUT_SLICE2]] : tensor<?x?xf32>)133// CHECK: return %[[RES]]134 135 %fill = linalg.fill ins(%five : f32) outs(%init : tensor<12x25xf32>) -> tensor<12x25xf32>136 %0 = linalg.generic {137 indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d2)>],138 iterator_types = ["parallel", "reduction", "parallel"]139 } ins(%input : tensor<12x7x25xf32>) outs(%fill : tensor<12x25xf32>) {140 ^bb0(%arg0: f32, %arg1: f32):141 %2 = arith.addf %arg0, %arg1 : f32142 linalg.yield %2 : f32143 } -> tensor<12x25xf32>144 func.return %0 : tensor<12x25xf32>145}146 147module attributes {transform.with_named_sequence} {148 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {149 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op150 %1, %loops:2 = transform.structured.fuse %0 tile_sizes [5, 0, 7] interchange [0, 2, 1]151 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)152 %2, %loops_2 = transform.structured.tile_using_for %1 tile_sizes [0, 4]153 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)154 transform.yield155 }156}157 158// -----159 160// CHECK-LABEL: func.func @unpack_elemwise161// CHECK: %[[RES:.*]] = scf.for162// CHECK: scf.for163// CHECK: linalg.unpack164// CHECK: linalg.exp165// CHECK: return %[[RES]]166func.func @unpack_elemwise(%arg0: tensor<16x48x8x8xf32>, %arg1: tensor<128x384xf32>) -> tensor<128x384xf32> {167 %0 = tensor.empty() : tensor<128x384xf32>168 %1 = linalg.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [8, 8] into %0169 : tensor<16x48x8x8xf32> -> tensor<128x384xf32>170 %2 = linalg.exp ins(%1: tensor<128x384xf32>)171 outs(%arg1: tensor<128x384xf32>) -> tensor<128x384xf32>172 return %2 : tensor<128x384xf32>173}174 175module attributes {transform.with_named_sequence} {176 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {177 %0 = transform.structured.match ops{["linalg.exp"]} in %arg1 : (!transform.any_op) -> !transform.any_op178 %1, %loops:2 = transform.structured.fuse %0 tile_sizes [16, 32] interchange [0, 1]179 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)180 transform.yield181 }182}183 184// -----185 186// CHECK-LABEL: func.func @pack_elemwise187// CHECK: %[[RES:.*]] = scf.for188// CHECK: scf.for189// CHECK: linalg.pack190// CHECK: linalg.exp191// CHECK: return %[[RES]]192func.func @pack_elemwise(%arg0: tensor<128x384xf32>, %arg1: tensor<16x48x8x8xf32>) -> tensor<16x48x8x8xf32> {193 %0 = tensor.empty() : tensor<16x48x8x8xf32>194 %1 = linalg.pack %arg0 inner_dims_pos = [0, 1] inner_tiles = [8, 8] into %0195 : tensor<128x384xf32> -> tensor<16x48x8x8xf32>196 %2 = linalg.exp ins(%1: tensor<16x48x8x8xf32>)197 outs(%arg1: tensor<16x48x8x8xf32>) -> tensor<16x48x8x8xf32>198 return %2 : tensor<16x48x8x8xf32>199}200 201module attributes {transform.with_named_sequence} {202 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {203 %0 = transform.structured.match ops{["linalg.exp"]} in %arg1 : (!transform.any_op) -> !transform.any_op204 %1, %loops:2 = transform.structured.fuse %0 tile_sizes [3, 5, 0, 0]205 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)206 transform.yield207 }208}209 210// -----211 212// CHECK-LABEL: func.func @nofuse_pack_elemwise213// CHECK: linalg.pack214// CHECK: %[[RES:.*]] = scf.for215// CHECK: scf.for216// CHECK: linalg.exp217// CHECK: return %[[RES]]218func.func @nofuse_pack_elemwise(%arg0: tensor<128x384xf32>, %arg1: tensor<16x48x8x8xf32>) -> tensor<16x48x8x8xf32> {219 %0 = tensor.empty() : tensor<16x48x8x8xf32>220 %1 = linalg.pack %arg0 inner_dims_pos = [0, 1] inner_tiles = [8, 8] into %0221 : tensor<128x384xf32> -> tensor<16x48x8x8xf32>222 %2 = linalg.exp ins(%1: tensor<16x48x8x8xf32>)223 outs(%arg1: tensor<16x48x8x8xf32>) -> tensor<16x48x8x8xf32>224 return %2 : tensor<16x48x8x8xf32>225}226 227module attributes {transform.with_named_sequence} {228 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {229 %0 = transform.structured.match ops{["linalg.exp"]} in %arg1 : (!transform.any_op) -> !transform.any_op230 %1, %loops:3 = transform.structured.fuse %0 tile_sizes [3, 5, 2, 0]231 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)232 transform.yield233 }234}235 236// -----237 238// CHECK-LABEL: func.func @fuse_through_slice239func.func @fuse_through_slice(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {240 241 // CHECK: %[[RES:.*]] = scf.for242 // CHECK: scf.for243 // CHECK: linalg.exp244 // CHECK: linalg.add245 // CHECK: return %[[RES]]246 %0 = linalg.exp ins(%arg0 : tensor<?x?xf32>)247 outs(%arg0: tensor<?x?xf32>) -> tensor<?x?xf32>248 %c0 = arith.constant 0 : index249 %c1 = arith.constant 1 : index250 %dim0 = tensor.dim %arg1, %c0 : tensor<?x?xf32>251 %dim1 = tensor.dim %arg1, %c1 : tensor<?x?xf32>252 %1 = tensor.extract_slice %0 [1, 1] [%dim0, %dim1] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>253 %2 = linalg.add ins(%1, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)254 outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>255 return %2 : tensor<?x?xf32>256}257 258module attributes {transform.with_named_sequence} {259 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {260 %0 = transform.structured.match ops{["linalg.add"]} in %arg1 : (!transform.any_op) -> !transform.any_op261 %1, %loops:2 = transform.structured.fuse %0 tile_sizes [32, 32] interchange [0, 1] {apply_cleanup}262 : (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">, !transform.any_op)263 transform.yield264 }265}266 267// -----268 269// CHECK-LABEL: func.func @fuse_through_slice_and_cast_chain270func.func @fuse_through_slice_and_cast_chain(%arg0: tensor<100x100xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {271 272 // CHECK: %[[RES:.*]] = scf.for273 // CHECK: scf.for274 // CHECK: linalg.exp275 // CHECK: linalg.add276 // CHECK: return %[[RES]]277 %0 = linalg.exp ins(%arg0 : tensor<100x100xf32>)278 outs(%arg0: tensor<100x100xf32>) -> tensor<100x100xf32>279 %1 = tensor.cast %0 : tensor<100x100xf32> to tensor<100x?xf32>280 %2 = tensor.extract_slice %1 [1, 1] [98, 98] [1, 1] : tensor<100x?xf32> to tensor<98x98xf32>281 %3 = tensor.cast %2 : tensor<98x98xf32> to tensor<?x?xf32>282 %c0 = arith.constant 0 : index283 %c1 = arith.constant 1 : index284 %dim0 = tensor.dim %arg1, %c0 : tensor<?x?xf32>285 %dim1 = tensor.dim %arg1, %c1 : tensor<?x?xf32>286 %4 = tensor.extract_slice %3 [1, 1] [%dim0, %dim1] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>287 %5 = linalg.add ins(%4, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)288 outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>289 return %5 : tensor<?x?xf32>290}291 292module attributes {transform.with_named_sequence} {293 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {294 %0 = transform.structured.match ops{["linalg.add"]} in %arg1 : (!transform.any_op) -> !transform.any_op295 %1, %loops:2 = transform.structured.fuse %0 tile_sizes [32, 32] interchange [0, 1] {apply_cleanup}296 : (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">, !transform.any_op)297 transform.yield298 }299}300 301// -----302 303// CHECK-LABEL: func.func @fuse_unrelated_slice304func.func @fuse_unrelated_slices(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> (tensor<?x?xf32>, tensor<10x10xf32>) {305 306 // CHECK: %[[SLICE1:.+]] = tensor.extract_slice307 // CHECK: %[[SLICE2:.+]] = tensor.extract_slice %[[SLICE1]]308 // CHECK: %[[RES:.*]] = scf.for309 // CHECK: scf.for310 // CHECK: linalg.exp311 // CHECK: linalg.add312 // CHECK: return %[[RES]], %[[SLICE2]]313 %c0 = arith.constant 0 : index314 %c1 = arith.constant 1 : index315 %dim0 = tensor.dim %arg1, %c0 : tensor<?x?xf32>316 %dim1 = tensor.dim %arg1, %c1 : tensor<?x?xf32>317 %slice1 = tensor.extract_slice %arg0 [1, 1] [%dim0, %dim1] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>318 %slice2 = tensor.extract_slice %slice1 [1, 1] [10, 10] [1, 1] : tensor<?x?xf32> to tensor<10x10xf32>319 %0 = linalg.exp ins(%arg0 : tensor<?x?xf32>)320 outs(%arg0: tensor<?x?xf32>) -> tensor<?x?xf32>321 %1 = tensor.extract_slice %0 [1, 1] [%dim0, %dim1] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>322 %2 = linalg.add ins(%1, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)323 outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>324 return %2, %slice2 : tensor<?x?xf32>, tensor<10x10xf32>325}326 327module attributes {transform.with_named_sequence} {328 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {329 %0 = transform.structured.match ops{["linalg.add"]} in %arg1 : (!transform.any_op) -> !transform.any_op330 %1, %loops:2 = transform.structured.fuse %0 tile_sizes [32, 32] interchange [0, 1] {apply_cleanup}331 : (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">, !transform.any_op)332 transform.yield333 }334}335 336// -----337 338// CHECK-LABEL: func.func @bubble_up_extract_slice_through_expand_shape339// CHECK: scf.for %[[X:[A-Za-z0-9]+]] = {{.*}}340// CHECK: scf.for %[[Y:[A-Za-z0-9]+]] = {{.*}}341// CHECK: scf.for %[[Z:[A-Za-z0-9]+]] = {{.*}}342// CHECK: %[[LINEAR_IDX:.+]] = affine.linearize_index disjoint [%[[X]], %[[Y]], %[[Z]]] by (2, 3, 10)343// CHECK: %[[SLICE:.+]] = tensor.extract_slice %{{.*}}[%[[LINEAR_IDX]]] [5] [1] : tensor<60xf32> to tensor<5xf32>344// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[SLICE]] {{\[\[}}0, 1, 2]] output_shape [1, 1, 5]345// CHECK: linalg.exp ins(%[[EXPAND]]346func.func @bubble_up_extract_slice_through_expand_shape(%0: tensor<60xf32>) -> tensor<2x3x10xf32> {347 %expand = tensor.expand_shape %0 [[0, 1, 2]] output_shape [2, 3, 10] : tensor<60xf32> into tensor<2x3x10xf32>348 %empty = tensor.empty() : tensor<2x3x10xf32>349 %exp = linalg.exp ins(%expand : tensor<2x3x10xf32>) outs(%empty : tensor<2x3x10xf32>) -> tensor<2x3x10xf32>350 return %exp : tensor<2x3x10xf32>351}352 353module attributes {transform.with_named_sequence} {354 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {355 %0 = transform.structured.match ops{["linalg.exp"]} in %arg0 : (!transform.any_op) -> !transform.any_op356 %transformed, %loops:3 = transform.structured.fuse %0 tile_sizes [1, 1, 5] interchange [0, 1, 2] {apply_cleanup} : 357 (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">, !transform.any_op, !transform.any_op)358 transform.yield 359 }360}361 362// -----363 364// CHECK-LABEL: func.func @bubble_up_extract_slice_through_expand_shape_full_inner_dim365// CHECK: scf.for %[[X:[A-Za-z0-9]+]] = {{.*}}366// CHECK: scf.for %[[Y:[A-Za-z0-9]+]] = {{.*}}367// CHECK: %[[LINEAR_IDX:.+]] = affine.linearize_index disjoint [%[[X]], %[[Y]]{{.*}} by (3, 4, 10)368// CHECK: %[[SLICE:.+]] = tensor.extract_slice %{{.*}}[%[[LINEAR_IDX]]] [20] [1] : tensor<120xf32> to tensor<20xf32>369// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[SLICE]] {{\[\[}}0, 1, 2]] output_shape [1, 2, 10]370// CHECK: linalg.exp ins(%[[EXPAND]]371func.func @bubble_up_extract_slice_through_expand_shape_full_inner_dim(%0: tensor<120xf32>) -> tensor<3x4x10xf32> {372 %expand = tensor.expand_shape %0 [[0, 1, 2]] output_shape [3, 4, 10] : tensor<120xf32> into tensor<3x4x10xf32>373 %empty = tensor.empty() : tensor<3x4x10xf32>374 %exp = linalg.exp ins(%expand : tensor<3x4x10xf32>) outs(%empty : tensor<3x4x10xf32>) -> tensor<3x4x10xf32>375 return %exp : tensor<3x4x10xf32>376}377 378module attributes {transform.with_named_sequence} {379 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {380 %0 = transform.structured.match ops{["linalg.exp"]} in %arg0 : (!transform.any_op) -> !transform.any_op381 %transformed, %loops:2 = transform.structured.fuse %0 tile_sizes [1, 2, 0] interchange [0, 1, 2] {apply_cleanup} :382 (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">, !transform.any_op)383 transform.yield 384 }385}386 387// -----388 389// CHECK-LABEL: func.func @no_bubble_up_extract_slice_through_expand_shape_non_contiguous390// CHECK: tensor.expand_shape391// CHECK: scf.for392// CHECK: scf.for393// CHECK: scf.for394// CHECK: linalg.exp395func.func @no_bubble_up_extract_slice_through_expand_shape_non_contiguous(%0: tensor<120xf32>) -> tensor<3x4x10xf32> {396 %expand = tensor.expand_shape %0 [[0, 1, 2]] output_shape [3, 4, 10] : tensor<120xf32> into tensor<3x4x10xf32>397 %empty = tensor.empty() : tensor<3x4x10xf32>398 %exp = linalg.exp ins(%expand : tensor<3x4x10xf32>) outs(%empty : tensor<3x4x10xf32>) -> tensor<3x4x10xf32>399 return %exp : tensor<3x4x10xf32>400}401 402module attributes {transform.with_named_sequence} {403 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {404 %0 = transform.structured.match ops{["linalg.exp"]} in %arg0 : (!transform.any_op) -> !transform.any_op405 %transformed, %loops:3 = transform.structured.fuse %0 tile_sizes [1, 2, 5] interchange [0, 1, 2] {apply_cleanup} :406 (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">, !transform.any_op, !transform.any_op)407 transform.yield 408 }409}410 411// -----412 413// CHECK-LABEL: func.func @bubble_up_extract_slice_through_expand_shape_multiple_expanded_dims414// CHECK: %[[C0:.+]] = arith.constant 0 : index415// CHECK: scf.for %[[X:[A-Za-z0-9]+]] = {{.*}}416// CHECK: scf.for %[[Y:[A-Za-z0-9]+]] = {{.*}}417// CHECK: scf.for %[[Z:[A-Za-z0-9]+]] = {{.*}}418// CHECK: scf.for %[[W:[A-Za-z0-9]+]] = {{.*}}419// CHECK: %[[LINEAR_IDX0:.+]] = affine.linearize_index disjoint [%[[X]], %[[Y]], %[[C0]]] by (3, 4, 10)420// CHECK: %[[LINEAR_IDX1:.+]] = affine.linearize_index disjoint [%[[Z]], %[[W]]] by (7, 8)421// CHECK: %[[SLICE:.+]] = tensor.extract_slice %{{.*}}[%[[LINEAR_IDX0]], %[[LINEAR_IDX1]]] [20, 4] [1, 1] : tensor<120x56xf32> to tensor<20x4xf32>422// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[SLICE]] {{\[\[}}0, 1, 2], [3, 4]] output_shape [1, 2, 10, 1, 4]423// CHECK: linalg.exp ins(%[[EXPAND]]424module {425 func.func @bubble_up_extract_slice_through_expand_shape_multiple_expanded_dims(%0: tensor<120x56xf32>) -> tensor<3x4x10x7x8xf32> {426 %expand = tensor.expand_shape %0 [[0, 1, 2], [3, 4]] output_shape [3, 4, 10, 7, 8] : tensor<120x56xf32> into tensor<3x4x10x7x8xf32>427 %empty = tensor.empty() : tensor<3x4x10x7x8xf32>428 %exp = linalg.exp ins(%expand : tensor<3x4x10x7x8xf32>) outs(%empty : tensor<3x4x10x7x8xf32>) -> tensor<3x4x10x7x8xf32>429 return %exp : tensor<3x4x10x7x8xf32>430 }431}432 433module attributes {transform.with_named_sequence} {434 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {435 %0 = transform.structured.match ops{["linalg.exp"]} in %arg0 : (!transform.any_op) -> !transform.any_op436 %transformed, %loops:4 = transform.structured.fuse %0 tile_sizes [1, 2, 0, 1, 4] interchange [0, 1, 2, 3, 4] {apply_cleanup} :437 (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">, !transform.any_op, !transform.any_op, !transform.any_op)438 transform.yield 439 }440}441 442// -----443 444// CHECK-LABEL: func.func @bubble_up_extract_slice_through_expand_shape_and_fuse_with_expand_producer445// CHECK: scf.for %[[X:[A-Za-z0-9]+]] = {{.*}}446// CHECK: %[[LINEAR_IDX:.+]] = affine.linearize_index disjoint [%[[X]], {{.*}} by (8, 32)447// CHECK: %[[SLICE:.+]] = tensor.extract_slice %{{.*}}[0, 0, %[[LINEAR_IDX]]] [1, 1800, 32] [1, 1, 1] : tensor<1x1800x256xf32> to tensor<1x1800x32xf32>448// CHECK: %[[ABS:.+]] = linalg.abs ins(%[[SLICE]]449// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[ABS]] {{\[\[}}0], [1], [2, 3]] output_shape [1, 1800, 1, 32]450// CHECK: linalg.exp ins(%[[EXPAND]]451module {452 func.func @bubble_up_extract_slice_through_expand_shape_and_fuse_with_expand_producer(%0: tensor<1x1800x256xf32>) -> tensor<1x1800x8x32xf32> {453 %empty1 = tensor.empty() : tensor<1x1800x256xf32>454 %exp1 = linalg.abs ins(%0 : tensor<1x1800x256xf32>) outs(%empty1 : tensor<1x1800x256xf32>) -> tensor<1x1800x256xf32>455 %expand = tensor.expand_shape %exp1 [[0], [1], [2, 3]] output_shape [1, 1800, 8, 32] : tensor<1x1800x256xf32> into tensor<1x1800x8x32xf32>456 %empty2 = tensor.empty() : tensor<1x1800x8x32xf32>457 %exp2 = linalg.exp ins(%expand : tensor<1x1800x8x32xf32>) outs(%empty2 : tensor<1x1800x8x32xf32>) -> tensor<1x1800x8x32xf32>458 return %exp2 : tensor<1x1800x8x32xf32>459 }460}461 462module attributes {transform.with_named_sequence} {463 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {464 %0 = transform.structured.match ops{["linalg.exp"]} in %arg0 : (!transform.any_op) -> !transform.any_op465 %transformed, %loops:1 = transform.structured.fuse %0 tile_sizes [0, 0, 1, 0] interchange [0, 1, 2, 3] {apply_cleanup} :466 (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">)467 transform.yield 468 }469}470 471// -----472 473// CHECK-LABEL: func.func @no_bubble_up_extract_slice_through_expand_shape_on_cleanup_false474// CHECK: %[[EXPAND:.+]] = tensor.expand_shape {{.*}} {{\[\[}}0, 1, 2]] output_shape [2, 3, 10]475// CHECK: scf.for %[[X:[A-Za-z0-9]+]] = {{.*}}476// CHECK: scf.for %[[Y:[A-Za-z0-9]+]] = {{.*}}477// CHECK: scf.for %[[Z:[A-Za-z0-9]+]] = {{.*}}478// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[EXPAND]]{{.*}} [1, 1, 5] [1, 1, 1] : tensor<2x3x10xf32> to tensor<1x1x5xf32>479// CHECK: linalg.exp ins(%[[SLICE]]480func.func @no_bubble_up_extract_slice_through_expand_shape_on_cleanup_false(%0: tensor<60xf32>) -> tensor<2x3x10xf32> {481 %expand = tensor.expand_shape %0 [[0, 1, 2]] output_shape [2, 3, 10] : tensor<60xf32> into tensor<2x3x10xf32>482 %empty = tensor.empty() : tensor<2x3x10xf32>483 %exp = linalg.exp ins(%expand : tensor<2x3x10xf32>) outs(%empty : tensor<2x3x10xf32>) -> tensor<2x3x10xf32>484 return %exp : tensor<2x3x10xf32>485}486 487module attributes {transform.with_named_sequence} {488 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {489 %0 = transform.structured.match ops{["linalg.exp"]} in %arg0 : (!transform.any_op) -> !transform.any_op490 %transformed, %loops:3 = transform.structured.fuse %0 tile_sizes [1, 1, 5] interchange [0, 1, 2] :491 (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">, !transform.any_op, !transform.any_op)492 transform.yield 493 }494}495 496// -----497 498// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape(499// CHECK: scf.for %[[X:[A-Za-z0-9]+]] = {{.*}} -> (tensor<8x1800x32xf32>) {500// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice501// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[EXTRACT]]502// CHECK: %[[EXP1:.*]] = linalg.exp ins(%[[COLLAPSE]]503func.func @bubble_up_extract_slice_through_collapse_shape(%0: tensor<1x8x1800x32xf32>) -> tensor<8x1800x32xf32> {504 %expand = tensor.collapse_shape %0 [[0, 1], [2], [3]] : tensor<1x8x1800x32xf32> into tensor<8x1800x32xf32>505 %empty = tensor.empty() : tensor<8x1800x32xf32>506 %exp = linalg.exp ins(%expand : tensor<8x1800x32xf32>) outs(%empty : tensor<8x1800x32xf32>) -> tensor<8x1800x32xf32>507 return %exp : tensor<8x1800x32xf32>508}509 510module attributes {transform.with_named_sequence} {511 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {512 %0 = transform.structured.match ops{["linalg.exp"]} in %arg0 : (!transform.any_op) -> !transform.any_op513 %transformed, %loops:1 = transform.structured.fuse %0 tile_sizes [1, 0, 0] interchange [0, 1, 2] {apply_cleanup} : 514 (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">)515 transform.yield 516 }517}518 519// -----520 521// CHECK-LABEL: func.func @bubble_up_extract_slice_through_collapse_shape_with_collapse_producer(522// CHECK: scf.for %[[X:[A-Za-z0-9]+]] = {{.*}}523// CHECK: %[[EXTRACT:.*]] = tensor.extract_slice524// CHECK: %[[ABS:.*]] = linalg.abs ins(%[[EXTRACT]]525// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[ABS]]526// CHECK: %[[EXP:.*]] = linalg.exp ins(%[[COLLAPSE]]527func.func @bubble_up_extract_slice_through_collapse_shape_with_collapse_producer(%0: tensor<1x8x1800x32xf32>) -> tensor<8x1800x32xf32> {528 %empty1 = tensor.empty() : tensor<1x8x1800x32xf32>529 %abs = linalg.abs ins(%0 : tensor<1x8x1800x32xf32>) outs(%empty1 : tensor<1x8x1800x32xf32>) -> tensor<1x8x1800x32xf32>530 %expand = tensor.collapse_shape %abs [[0, 1], [2], [3]] : tensor<1x8x1800x32xf32> into tensor<8x1800x32xf32>531 %empty2 = tensor.empty() : tensor<8x1800x32xf32>532 %exp = linalg.exp ins(%expand : tensor<8x1800x32xf32>) outs(%empty2 : tensor<8x1800x32xf32>) -> tensor<8x1800x32xf32>533 return %exp : tensor<8x1800x32xf32>534}535 536module attributes {transform.with_named_sequence} {537 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {538 %0 = transform.structured.match ops{["linalg.exp"]} in %arg0 : (!transform.any_op) -> !transform.any_op539 %transformed, %loops:1 = transform.structured.fuse %0 tile_sizes [1, 0, 0] interchange [0, 1, 2] {apply_cleanup} : 540 (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">)541 transform.yield 542 }543}544