823 lines · plain
1// RUN: mlir-opt --transform-interpreter --split-input-file %s -verify-diagnostics | FileCheck %s2 3#map0 = affine_map<()[s0, s1] -> (s0 ceildiv s1)>4#map1 = affine_map<(d0)[s0] -> (d0 * s0)>5#map2 = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)>6 7module {8 // CHECK-LABEL: func.func @fuse_tileable_op9 // CHECK-SAME: %[[CHUNK_SIZE:[0-9a-z]+]]: index10 // CHECK-SAME: %[[IN:[0-9a-z]+]]: tensor<?xf32>11 // CHECK-SAME: %[[OUT:[0-9a-z]+]]: tensor<?xf32>12 func.func @fuse_tileable_op(%arg0: index, %arg1: tensor<?xf32>, %arg2: tensor<?xf32>) -> tensor<?xf32> {13 %cst = arith.constant 4.200000e+01 : f3214 %c0 = arith.constant 0 : index15 %0 = linalg.fill ins(%cst : f32) outs(%arg1 : tensor<?xf32>) -> tensor<?xf32>16 %d0 = tensor.dim %arg1, %c0 : tensor<?xf32>17 %1 = affine.apply #map0()[%d0, %arg0]18 19 // CHECK: scf.forall {{.*}} {20 %2 = scf.forall (%arg3) in (%1) shared_outs(%o = %arg2) -> (tensor<?xf32>) {21 %3 = affine.apply #map1(%arg3)[%arg0]22 %4 = affine.min #map2(%arg3)[%d0, %arg0]23 %5 = tensor.extract_slice %o[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32>24 25 // CHECK: %[[T0:.*]] = tensor.extract_slice %[[IN]][%{{.*}}] [%{{.*}}] [{{.*}}]26 // CHECK: %[[T1:.*]] = linalg.fill {{.*}} outs(%[[T0]]27 %6 = tensor.extract_slice %0[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32>28 29 // CHECK: %[[T2:.*]] = linalg.elementwise kind=#linalg.elementwise_kind<exp> ins(%[[T1]]30 %7 = linalg.elementwise kind=#linalg.elementwise_kind<exp> ins(%6 : tensor<?xf32>) outs(%5 : tensor<?xf32>) -> tensor<?xf32>31 scf.forall.in_parallel {32 tensor.parallel_insert_slice %7 into %o[%3] [%4] [1] : tensor<?xf32> into tensor<?xf32>33 }34 }35 // CHECK: }36 func.return %2 : tensor<?xf32>37 }38 39 // Check no failure when nothing happens.40 func.func @dummy1() { return }41 func.func @dummy2() { return }42 func.func @dummy3() { return }43 44 module attributes {transform.with_named_sequence} {45 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {46 %0 = transform.structured.match ops{["linalg.fill"]} in %arg1 : (!transform.any_op) -> !transform.op<"linalg.fill">47 %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.op<"scf.forall">48 49 // linalg.fill is tileable. The op is tiled and fused.50 transform.structured.fuse_into_containing_op %0 into %151 : (!transform.op<"linalg.fill">, !transform.op<"scf.forall">) -> (!transform.any_op, !transform.any_op)52 transform.yield53 }54 }55}56 57// -----58 59#map0 = affine_map<()[s0] -> (64 ceildiv s0)>60#map1 = affine_map<(d0)[s0] -> (d0 * s0)>61#map2 = affine_map<(d0)[s0] -> (-(d0 * s0) + 64, s0)>62 63module {64 // CHECK-LABEL: func.func @fuse_untileable_op65 // CHECK-SAME: %[[CHUNK_SIZE:[0-9a-z]+]]: index66 // CHECK-SAME: %[[IN:[0-9a-z]+]]: tensor<64xf32>67 // CHECK-SAME: %[[OUT:[0-9a-z]+]]: tensor<64xf32>68 func.func @fuse_untileable_op(%arg0: index, %arg1: tensor<64xf32>, %arg2: tensor<64xf32>) -> tensor<64xf32> {69 %0 = tensor.empty(%arg0) : tensor<?xf32>70 %1 = affine.apply #map0()[%arg0]71 72 // CHECK: scf.forall {{.*}} {73 %2 = scf.forall (%arg3) in (%1) shared_outs(%o = %arg2) -> (tensor<64xf32>) {74 // CHECK: %[[INIT_TENSOR:.*]] = tensor.empty75 %3 = affine.apply #map1(%arg3)[%arg0]76 %4 = affine.min #map2(%arg3)[%arg0]77 %5 = tensor.extract_slice %o[%3] [%4] [1] : tensor<64xf32> to tensor<?xf32>78 79 // CHECK: %[[T2:.*]] = linalg.exp ins(%[[INIT_TENSOR]]80 %7 = linalg.exp ins(%0 : tensor<?xf32>) outs(%5 : tensor<?xf32>) -> tensor<?xf32>81 scf.forall.in_parallel {82 tensor.parallel_insert_slice %7 into %o[%3] [%4] [1] : tensor<?xf32> into tensor<64xf32>83 }84 }85 // CHECK: }86 87 func.return %2 : tensor<64xf32>88 }89 90 module attributes {transform.with_named_sequence} {91 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {92 %0 = transform.structured.match ops{["tensor.empty"]} in %arg1 : (!transform.any_op) -> !transform.op<"tensor.empty">93 %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.op<"scf.forall">94 95 // tensor.empty is not tileable. The op is cloned and fused.96 transform.structured.fuse_into_containing_op %0 into %197 : (!transform.op<"tensor.empty">, !transform.op<"scf.forall">) -> (!transform.any_op, !transform.any_op)98 transform.yield99 }100 }101}102 103// -----104 105module {106 func.func @foo(%0: tensor<f32>) -> tensor<f32> {107 return %0: tensor<f32>108 }109 110 // CHECK-LABEL: func.func @fuse_tileable_op_rank_reducing111 // CHECK-SAME: %[[CHUNK_SIZE:[0-9a-z]+]]: index112 // CHECK-SAME: %[[IN:[0-9a-z]+]]: tensor<?xf32>113 // CHECK-SAME: %[[OUT:[0-9a-z]+]]: tensor<?xf32>114 func.func @fuse_tileable_op_rank_reducing(%arg0: index, %arg1: tensor<?xf32>, %arg2: tensor<?xf32>) -> tensor<?xf32> {115 %cst = arith.constant 4.200000e+01 : f32116 %c0 = arith.constant 0 : index117 %0 = linalg.fill ins(%cst : f32) outs(%arg2 : tensor<?xf32>) -> tensor<?xf32>118 %d0 = tensor.dim %arg1, %c0 : tensor<?xf32>119 120 // CHECK: scf.forall {{.*}} -> (tensor<?xf32>) {121 %2 = scf.forall (%arg3) in (%d0) shared_outs(%o = %0) -> (tensor<?xf32>) {122 %5 = tensor.extract_slice %o[%arg3] [1] [1] : tensor<?xf32> to tensor<f32>123 124 // CHECK: tensor.extract_slice %{{.*}}[%{{.*}}] [1] [1] : tensor<?xf32> to tensor<1xf32>125 // CHECK: linalg.fill ins(%{{.*}} : f32) outs(%{{.*}} : tensor<1xf32>) -> tensor<1xf32>126 // CHECK: tensor.extract_slice %{{.*}}[0] [1] [1] : tensor<1xf32> to tensor<f32>127 // CHECK: func.call @foo(%{{.*}}) : (tensor<f32>) -> tensor<f32>128 %7 = func.call @foo(%5) : (tensor<f32>) -> tensor<f32>129 130 scf.forall.in_parallel {131 // CHECK: tensor.parallel_insert_slice %{{.*}} into %{{.*}}[%{{.*}}] [1] [1] : tensor<f32> into tensor<?xf32>132 tensor.parallel_insert_slice %7 into %o[%arg3] [1] [1] : tensor<f32> into tensor<?xf32>133 }134 }135 // CHECK: }136 func.return %2 : tensor<?xf32>137 }138 139 module attributes {transform.with_named_sequence} {140 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {141 %0 = transform.structured.match ops{["linalg.fill"]} in %arg1 : (!transform.any_op) -> !transform.op<"linalg.fill">142 %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.op<"scf.forall">143 144 // linalg.fill is tileable. The op is tiled and fused.145 transform.structured.fuse_into_containing_op %0 into %1146 : (!transform.op<"linalg.fill">, !transform.op<"scf.forall">) -> (!transform.any_op, !transform.any_op)147 transform.yield148 }149 }150}151 152// -----153 154#map0 = affine_map<()[s0, s1] -> (s0 ceildiv s1)>155#map1 = affine_map<(d0)[s0] -> (d0 * s0)>156#map2 = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)>157 158module {159 // CHECK-LABEL: func.func @fuse_tileable_op_through_bbarg160 // CHECK-SAME: %[[CHUNK_SIZE:[0-9a-z]+]]: index161 // CHECK-SAME: %[[IN:[0-9a-z]+]]: tensor<?xf32>162 // CHECK-SAME: %[[OUT:[0-9a-z]+]]: tensor<?xf32>163 func.func @fuse_tileable_op_through_bbarg(%arg0: index, %arg1: tensor<?xf32>, %arg2: tensor<?xf32>) -> tensor<?xf32> {164 %cst = arith.constant 4.200000e+01 : f32165 %c0 = arith.constant 0 : index166 %0 = linalg.fill ins(%cst : f32) outs(%arg2 : tensor<?xf32>) -> tensor<?xf32>167 %d0 = tensor.dim %arg1, %c0 : tensor<?xf32>168 %1 = affine.apply #map0()[%d0, %arg0]169 170 // CHECK: scf.forall {{.*}} shared_outs(%[[BBARGOUT:.*]] = %[[OUT]]) -> (tensor<?xf32>) {171 %2 = scf.forall (%arg3) in (%1) shared_outs(%o = %0) -> (tensor<?xf32>) {172 %3 = affine.apply #map1(%arg3)[%arg0]173 %4 = affine.min #map2(%arg3)[%d0, %arg0]174 %5 = tensor.extract_slice %o[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32>175 176 // CHECK: %[[T0:.*]] = tensor.extract_slice %[[BBARGOUT]][%{{.*}}] [%{{.*}}] [{{.*}}]177 // CHECK: %[[T1:.*]] = linalg.fill {{.*}} outs(%[[T0]]178 %6 = tensor.extract_slice %arg1[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32>179 180 // CHECK: %[[T2:.*]] = linalg.exp {{.*}} outs(%[[T1]]181 %7 = linalg.exp ins(%6 : tensor<?xf32>) outs(%5 : tensor<?xf32>) -> tensor<?xf32>182 scf.forall.in_parallel {183 tensor.parallel_insert_slice %7 into %o[%3] [%4] [1] : tensor<?xf32> into tensor<?xf32>184 }185 }186 // CHECK: }187 func.return %2 : tensor<?xf32>188 }189 190 module attributes {transform.with_named_sequence} {191 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {192 %0 = transform.structured.match ops{["linalg.fill"]} in %arg1 : (!transform.any_op) -> !transform.any_op193 %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.any_op194 195 // linalg.fill is tileable. The op is tiled and fused.196 transform.structured.fuse_into_containing_op %0 into %1197 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)198 transform.yield199 }200 }201}202 203// -----204 205#map0 = affine_map<()[s0, s1] -> (s0 ceildiv s1)>206#map1 = affine_map<(d0)[s0] -> (d0 * s0)>207#map2 = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)>208 209module {210 // CHECK-LABEL: func.func @fuse_tileable_op_through_bbarg_inout211 // CHECK-SAME: %[[CHUNK_SIZE:[0-9a-z]+]]: index212 // CHECK-SAME: %[[INOUT:[0-9a-z]+]]: tensor<?xf32>213 func.func @fuse_tileable_op_through_bbarg_inout(%arg0: index, %arg1: tensor<?xf32>) -> tensor<?xf32> {214 %cst = arith.constant 4.200000e+01 : f32215 %c0 = arith.constant 0 : index216 %0 = linalg.fill ins(%cst : f32) outs(%arg1 : tensor<?xf32>) -> tensor<?xf32>217 %d0 = tensor.dim %arg1, %c0 : tensor<?xf32>218 %1 = affine.apply #map0()[%d0, %arg0]219 220 // CHECK: scf.forall {{.*}} shared_outs(%[[BBARGOUT:.*]] = %[[INOUT]]) -> (tensor<?xf32>) {221 %2 = scf.forall (%arg3) in (%1) shared_outs(%o = %arg1) -> (tensor<?xf32>) {222 %3 = affine.apply #map1(%arg3)[%arg0]223 %4 = affine.min #map2(%arg3)[%d0, %arg0]224 %5 = tensor.extract_slice %o[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32>225 226 // CHECK: %[[T0:.*]] = tensor.extract_slice %[[BBARGOUT]][%{{.*}}] [%{{.*}}] [{{.*}}]227 // CHECK: %[[T1:.*]] = tensor.extract_slice %[[BBARGOUT]][%{{.*}}] [%{{.*}}] [{{.*}}]228 // CHECK: %[[T2:.*]] = linalg.fill {{.*}} outs(%[[T1]]229 %6 = tensor.extract_slice %0[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32>230 231 // CHECK: %[[T3:.*]] = linalg.exp ins(%[[T2]] : tensor<?xf32>) outs(%[[T0]] : tensor<?xf32>)232 %7 = linalg.exp ins(%6 : tensor<?xf32>) outs(%5 : tensor<?xf32>) -> tensor<?xf32>233 scf.forall.in_parallel {234 tensor.parallel_insert_slice %7 into %o[%3] [%4] [1] : tensor<?xf32> into tensor<?xf32>235 }236 }237 // CHECK: }238 func.return %2 : tensor<?xf32>239 }240 241 module attributes {transform.with_named_sequence} {242 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {243 %0 = transform.structured.match ops{["linalg.fill"]} in %arg1 : (!transform.any_op) -> !transform.any_op244 %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.any_op245 246 // linalg.fill is tileable. The op is tiled and fused.247 transform.structured.fuse_into_containing_op %0 into %1248 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)249 transform.yield250 }251 }252}253 254// -----255 256#map = affine_map<(d0) -> (d0 * 2)>257#map1 = affine_map<(d0) -> (d0 * 4)>258module {259 // CHECK-LABEL: func.func @fuse_tileable_op_no_dps260 func.func @fuse_tileable_op_no_dps(%arg0: tensor<4x4x4xf32>, %arg1: tensor<4x4x4xf32>) -> tensor<4x4x4xf32> {261 %0 = "test.tiling_no_dps_op"(%arg0, %arg1) : (tensor<4x4x4xf32>, tensor<4x4x4xf32>) -> tensor<4x4x4xf32>262 %1 = tensor.empty() : tensor<4x4x4xf32>263 // CHECK: scf.forall264 %2 = scf.forall (%arg2, %arg3, %arg4) in (4, 2, 1) shared_outs(%arg5 = %1) -> (tensor<4x4x4xf32>) {265 %3 = affine.apply #map(%arg3)266 %4 = affine.apply #map1(%arg4)267 // CHECK: "test.tiling_no_dps_op"268 // CHECK: "test.unregistered_op"269 %extracted_slice = tensor.extract_slice %0[%arg2, %3, %4] [1, 2, 4] [1, 1, 1] : tensor<4x4x4xf32> to tensor<1x2x4xf32>270 %5 = "test.unregistered_op"(%extracted_slice, %extracted_slice) : (tensor<1x2x4xf32>, tensor<1x2x4xf32>) -> tensor<1x2x4xf32>271 scf.forall.in_parallel {272 tensor.parallel_insert_slice %5 into %arg5[%arg2, %3, %4] [1, 2, 4] [1, 1, 1] : tensor<1x2x4xf32> into tensor<4x4x4xf32>273 }274 }275 return %2 : tensor<4x4x4xf32>276 }277 278 module attributes {transform.with_named_sequence} {279 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {280 %op = transform.structured.match ops{["test.tiling_no_dps_op"]} in %arg0 : (!transform.any_op) -> !transform.any_op281 %forall = transform.structured.match ops{["scf.forall"]} in %arg0 : (!transform.any_op) -> !transform.any_op282 %fused, %new_containing = transform.structured.fuse_into_containing_op %op into %forall : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)283 transform.yield284 }285 }286}287 288// -----289 290module {291 // CHECK-LABEL: func.func @fuse_tileable_op_through_bbarg_inout_nested292 // CHECK-SAME: %[[ARG0:[0-9a-z]+]]: tensor<?x?x?xf32>293 // CHECK-SAME: %[[ARG1:[0-9a-z]+]]: tensor<?x?x?xf32>294 func.func @fuse_tileable_op_through_bbarg_inout_nested(%arg0: tensor<?x?x?xf32>, %arg1: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {295 %c2 = arith.constant 2 : index296 %c1 = arith.constant 1 : index297 %c0 = arith.constant 0 : index298 %0 = linalg.elementwise kind=#linalg.elementwise_kind<abs> ins(%arg0 : tensor<?x?x?xf32>) outs(%arg1 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>299 %dim = tensor.dim %arg1, %c0 : tensor<?x?x?xf32>300 %dim_0 = tensor.dim %arg1, %c1 : tensor<?x?x?xf32>301 %dim_1 = tensor.dim %arg1, %c2 : tensor<?x?x?xf32>302 // CHECK: scf.for {{.*}} iter_args(%[[BBARG0:.*]] = %[[ARG1]]) -> (tensor<?x?x?xf32>) {303 // CHECK: scf.for {{.*}} iter_args(%[[BBARG1:.*]] = %[[BBARG0]]) -> (tensor<?x?x?xf32>) {304 // CHECK: scf.for {{.*}} iter_args(%[[BBARG2:.*]] = %[[BBARG1]]) -> (tensor<?x?x?xf32>) {305 %1 = scf.for %arg2 = %c0 to %dim step %c1 iter_args(%arg3 = %arg1) -> (tensor<?x?x?xf32>) {306 %2 = scf.for %arg4 = %c0 to %dim_0 step %c1 iter_args(%arg5 = %arg3) -> (tensor<?x?x?xf32>) {307 %3 = scf.for %arg6 = %c0 to %dim_1 step %c1 iter_args(%arg7 = %arg5) -> (tensor<?x?x?xf32>) {308 // CHECK: %[[EX1:.*]] = tensor.extract_slice %[[BBARG2]]{{.*}}: tensor<?x?x?xf32> to tensor<1x1x1xf32>309 // CHECK: linalg.elementwise kind=#linalg.elementwise_kind<abs> ins({{.*}} : tensor<1x1x1xf32>) outs(%[[EX1]] : tensor<1x1x1xf32>) -> tensor<1x1x1xf32>310 // CHECK: %[[EX2:.*]] = tensor.extract_slice %[[BBARG2]]{{.*}} : tensor<?x?x?xf32> to tensor<1x1x1xf32>311 // CHECK: linalg.elementwise kind=#linalg.elementwise_kind<exp> ins({{.*}} : tensor<1x1x1xf32>) outs(%[[EX2]] : tensor<1x1x1xf32>) -> tensor<1x1x1xf32>312 %extracted_slice = tensor.extract_slice %0[%arg2, %arg4, %arg6] [1, 1, 1] [1, 1, 1] : tensor<?x?x?xf32> to tensor<1x1x1xf32>313 %extracted_slice_2 = tensor.extract_slice %arg7[%arg2, %arg4, %arg6] [1, 1, 1] [1, 1, 1] : tensor<?x?x?xf32> to tensor<1x1x1xf32>314 %4 = linalg.elementwise kind=#linalg.elementwise_kind<exp> ins(%extracted_slice : tensor<1x1x1xf32>) outs(%extracted_slice_2 : tensor<1x1x1xf32>) -> tensor<1x1x1xf32>315 %inserted_slice = tensor.insert_slice %4 into %arg7[%arg2, %arg4, %arg6] [1, 1, 1] [1, 1, 1] : tensor<1x1x1xf32> into tensor<?x?x?xf32>316 scf.yield %inserted_slice : tensor<?x?x?xf32>317 }318 scf.yield %3 : tensor<?x?x?xf32>319 }320 scf.yield %2 : tensor<?x?x?xf32>321 }322 return %1 : tensor<?x?x?xf32>323 }324 325 module attributes {transform.with_named_sequence} {326 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {327 %0 = transform.structured.match ops{["linalg.elementwise"]} in %arg0 : (!transform.any_op) -> !transform.any_op328 %1 = transform.structured.match ops{["scf.for"]} in %arg0 : (!transform.any_op) -> !transform.any_op329 %2:2 = transform.split_handle %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)330 %3:3 = transform.split_handle %1 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)331 transform.structured.fuse_into_containing_op %2#0 into %3#2 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)332 transform.yield333 }334 }335}336 337// -----338 339#map0 = affine_map<()[s0, s1] -> (s0 ceildiv s1)>340#map1 = affine_map<(d0)[s0] -> (d0 * s0)>341#map2 = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)>342 343module {344 // CHECK-LABEL: func.func @fuse_tileable_multi_output_op345 // CHECK-SAME: %[[CHUNK_SIZE:[0-9a-z]+]]: index346 // CHECK-SAME: %[[IN:[0-9a-z]+]]: tensor<?xf32>347 // CHECK-SAME: %[[OUT_1:[0-9a-z]+]]: tensor<?xf32>348 // CHECK-SAME: %[[OUT_2:[0-9a-z]+]]: tensor<?xf32>349 // CHECK-SAME: %[[OUT_3:[0-9a-z]+]]: tensor<?xf32>350 func.func @fuse_tileable_multi_output_op(%idx: index, %in: tensor<?xf32>, %out_1: tensor<?xf32>, %out_2: tensor<?xf32>, %out_3: tensor<?xf32>) -> tensor<?xf32> {351 %cst = arith.constant 4.200000e+01 : f32352 %c0 = arith.constant 0 : index353 354 %0:2 = linalg.generic {355 indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],356 iterator_types = ["parallel"]357 } ins(%in : tensor<?xf32>) outs(%out_1, %out_3 : tensor<?xf32>, tensor<?xf32>) {358 ^bb0(%a: f32, %b: f32, %c: f32):359 %d = arith.addf %a, %b : f32360 %e = arith.addf %d, %c : f32361 linalg.yield %d, %e : f32, f32362 } -> (tensor<?xf32>, tensor<?xf32>)363 %d0 = tensor.dim %out_1, %c0 : tensor<?xf32>364 365 %1 = affine.apply #map0()[%d0, %idx]366 367 // CHECK: scf.forall {{.*}} {368 %2 = scf.forall (%i) in (%1) shared_outs(%o = %out_2) -> (tensor<?xf32>) {369 %3 = affine.apply #map1(%i)[%idx]370 %4 = affine.min #map2(%i)[%d0, %idx]371 %5 = tensor.extract_slice %o[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32>372 373 // CHECK: %[[T0:.*]] = tensor.extract_slice %[[IN]][%{{.*}}] [%{{.*}}] [{{.*}}]374 // CHECK: %[[T1:.*]]:2 = linalg.generic {{.*}} ins(%[[T0]]375 %6 = tensor.extract_slice %0#0[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32>376 377 // CHECK: %[[T2:.*]] = linalg.exp ins(%[[T1]]#0378 %7 = linalg.exp ins(%6 : tensor<?xf32>) outs(%5 : tensor<?xf32>) -> tensor<?xf32>379 scf.forall.in_parallel {380 tensor.parallel_insert_slice %7 into %o[%3] [%4] [1] : tensor<?xf32> into tensor<?xf32>381 }382 }383 // CHECK: }384 func.return %2 : tensor<?xf32>385 }386 387 module attributes {transform.with_named_sequence} {388 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {389 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.op<"linalg.generic">390 %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.op<"scf.forall">391 392 // linalg.generic is tileable. The op is tiled and fused.393 transform.structured.fuse_into_containing_op %0 into %1394 : (!transform.op<"linalg.generic">, !transform.op<"scf.forall">) -> (!transform.any_op, !transform.any_op)395 transform.yield396 }397 }398}399 400// -----401 402module {403 // CHECK-LABEL: func.func @fuse_repeated404 func.func @fuse_repeated(%fill: tensor<2xf32>, %output: tensor<2xf32>) -> tensor<2xf32> {405 %c0 = arith.constant 0.0 : f32406 %0 = linalg.fill ins(%c0 : f32) outs(%fill : tensor<2xf32>) -> tensor<2xf32>407 408 // CHECK: scf.forall409 %1 = scf.forall (%i) in (2) shared_outs(%arg1 = %output) -> (tensor<2xf32>) {410 %2 = tensor.extract_slice %0[%i][1][1] : tensor<2xf32> to tensor<1xf32>411 %3 = tensor.extract_slice %arg1[%i][1][1] : tensor<2xf32> to tensor<1xf32>412 // CHECK: %[[FUSED:.+]] = linalg.fill413 // CHECK: exp ins(%[[FUSED]]414 %4 = linalg.exp ins(%2 : tensor<1xf32>) outs(%3 : tensor<1xf32>) -> tensor<1xf32>415 scf.forall.in_parallel {416 tensor.parallel_insert_slice %4 into %arg1[%i][1][1] : tensor<1xf32> into tensor<2xf32>417 }418 }419 420 return %1 : tensor<2xf32>421 }422 423 module attributes {transform.with_named_sequence} {424 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {425 %0 = transform.structured.match ops{["linalg.fill"]} in %arg1 : (!transform.any_op) -> !transform.any_op426 %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.any_op427 428 // Create a new handle that points to `linalg.fill` twice.429 %2 = transform.merge_handles %0, %0 : !transform.any_op430 431 // It shouldn't be a problem to fuse this handle.432 transform.structured.fuse_into_containing_op %2 into %1 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)433 transform.yield434 }435 }436}437 438// -----439 440#map0 = affine_map<()[s0, s1] -> (s0 ceildiv s1)>441#map1 = affine_map<(d0)[s0] -> (d0 * s0)>442#map2 = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)>443 444module {445 // CHECK-LABEL: func.func @fuse_tileable_multi_output_op_multi_use446 // CHECK-SAME: %[[CHUNK_SIZE:[0-9a-z]+]]: index447 // CHECK-SAME: %[[IN:[0-9a-z]+]]: tensor<?xf32>448 // CHECK-SAME: %[[OUT_1:[0-9a-z]+]]: tensor<?xf32>449 // CHECK-SAME: %[[OUT_2:[0-9a-z]+]]: tensor<?xf32>450 // CHECK-SAME: %[[OUT_3:[0-9a-z]+]]: tensor<?xf32>451 func.func @fuse_tileable_multi_output_op_multi_use(%idx: index, %in: tensor<?xf32>, %out_1: tensor<?xf32>, %out_2: tensor<?xf32>, %out_3: tensor<?xf32>)452 -> (tensor<?xf32>, tensor<?xf32>, tensor<?xf32>) {453 %cst = arith.constant 4.200000e+01 : f32454 %c0 = arith.constant 0 : index455 456 // CHECK: %[[G0:.*]]:2 = linalg.generic457 %0:2 = linalg.generic {458 indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],459 iterator_types = ["parallel"]460 } ins(%in : tensor<?xf32>) outs(%out_1, %out_3 : tensor<?xf32>, tensor<?xf32>) {461 ^bb0(%a: f32, %b: f32, %c: f32):462 %d = arith.addf %a, %b : f32463 %e = arith.addf %d, %c : f32464 linalg.yield %d, %e : f32, f32465 } -> (tensor<?xf32>, tensor<?xf32>)466 %d0 = tensor.dim %out_1, %c0 : tensor<?xf32>467 468 %1 = affine.apply #map0()[%d0, %idx]469 470 // CHECK: %[[R0:.*]]:2 = scf.forall (%[[ARG5:.*]]) in (%{{.*}}) shared_outs(%[[ARG6:.*]] = %[[OUT_2]], %[[ARG7:.*]] = %[[OUT_1]])471 // CHECK-SAME: -> (tensor<?xf32>, tensor<?xf32>) {472 // expected-remark @below{{new containing op}}473 %2 = scf.forall (%i) in (%1) shared_outs(%o = %out_2) -> (tensor<?xf32>) {474 // CHECK: %[[I0:.*]] = affine.apply {{.*}}475 %3 = affine.apply #map1(%i)[%idx]476 // CHECK: %[[I1:.*]] = affine.min {{.*}}477 %4 = affine.min #map2(%i)[%d0, %idx]478 %5 = tensor.extract_slice %o[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32>479 480 // CHECK: %[[T1:.*]]:2 = linalg.generic {{.*}}481 %6 = tensor.extract_slice %0#0[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32>482 483 %7 = linalg.exp ins(%6 : tensor<?xf32>) outs(%5 : tensor<?xf32>) -> tensor<?xf32>484 scf.forall.in_parallel {485 // CHECK: tensor.parallel_insert_slice %[[T1]]#0 into %[[ARG7]][%[[I0]]] [%[[I1]]] [1] : tensor<?xf32> into tensor<?xf32>486 tensor.parallel_insert_slice %7 into %o[%3] [%4] [1] : tensor<?xf32> into tensor<?xf32>487 }488 }489 // CHECK: return %[[R0]]#0, %[[R0]]#1, %[[G0]]#1490 func.return %2, %0#0, %0#1 : tensor<?xf32>, tensor<?xf32>, tensor<?xf32>491 // CHECK: }492 }493 494 module attributes {transform.with_named_sequence} {495 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {496 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.op<"linalg.generic">497 %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.op<"scf.forall">498 499 // linalg.generic is tileable. The op is tiled and fused.500 %fused, %containing = transform.structured.fuse_into_containing_op %0 into %1501 : (!transform.op<"linalg.generic">, !transform.op<"scf.forall">) -> (!transform.any_op, !transform.any_op)502 transform.debug.emit_remark_at %containing, "new containing op" : !transform.any_op503 transform.yield504 }505 }506}507 508// -----509 510#map0 = affine_map<()[s0, s1] -> (s0 ceildiv s1)>511#map1 = affine_map<(d0)[s0] -> (d0 * s0)>512#map2 = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)>513 514module {515 // CHECK-LABEL: func.func @fuse_tileable_mixed_dominating_uses516 // CHECK-SAME: %[[CHUNK_SIZE:[0-9a-z]+]]: index517 // CHECK-SAME: %[[IN:[0-9a-z]+]]: tensor<?xf32>518 // CHECK-SAME: %[[OUT_1:[0-9a-z]+]]: tensor<?xf32>519 // CHECK-SAME: %[[OUT_2:[0-9a-z]+]]: tensor<?xf32>520 // CHECK-SAME: %[[OUT_3:[0-9a-z]+]]: tensor<?xf32>521 func.func @fuse_tileable_mixed_dominating_uses(%idx: index, %in: tensor<?xf32>, %out_1: tensor<?xf32>, %out_2: tensor<?xf32>, %out_3: tensor<?xf32>)522 -> (tensor<?xf32>, tensor<?xf32>) {523 %cst = arith.constant 4.200000e+01 : f32524 %c0 = arith.constant 0 : index525 526 // CHECK: %[[G0:.*]] = linalg.generic527 %0 = linalg.generic {528 indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],529 iterator_types = ["parallel"]530 } ins(%in : tensor<?xf32>) outs(%out_1 : tensor<?xf32>) {531 ^bb0(%a: f32, %b: f32):532 %d = arith.addf %a, %b : f32533 linalg.yield %d : f32534 } -> tensor<?xf32>535 // CHECK: %[[D0:.*]] = tensor.dim %[[G0]]536 %d0 = tensor.dim %0, %c0 : tensor<?xf32>537 538 %1 = affine.apply #map0()[%d0, %idx]539 540 // CHECK: %[[R0:.*]]:2 = scf.forall (%[[ARG5:.*]]) in (%{{.*}}) shared_outs(%[[ARG6:.*]] = %[[OUT_2]], %[[ARG7:.*]] = %[[OUT_1]])541 // CHECK-SAME: -> (tensor<?xf32>, tensor<?xf32>) {542 %2 = scf.forall (%i) in (%1) shared_outs(%o = %out_2) -> (tensor<?xf32>) {543 // CHECK: %[[I0:.*]] = affine.apply {{.*}}544 %3 = affine.apply #map1(%i)[%idx]545 // CHECK: %[[I1:.*]] = affine.min {{.*}}546 %4 = affine.min #map2(%i)[%d0, %idx]547 %5 = tensor.extract_slice %o[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32>548 549 // CHECK: %[[T1:.*]] = linalg.generic {{.*}}550 %6 = tensor.extract_slice %0[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32>551 552 %7 = linalg.exp ins(%6 : tensor<?xf32>) outs(%5 : tensor<?xf32>) -> tensor<?xf32>553 scf.forall.in_parallel {554 // CHECK: tensor.parallel_insert_slice %[[T1]] into %[[ARG7]][%[[I0]]] [%[[I1]]] [1] : tensor<?xf32> into tensor<?xf32>555 tensor.parallel_insert_slice %7 into %o[%3] [%4] [1] : tensor<?xf32> into tensor<?xf32>556 }557 }558 // CHECK: return %[[R0]]#0, %[[R0]]#1559 func.return %2, %0 : tensor<?xf32>, tensor<?xf32>560 // CHECK: }561 }562 563 module attributes {transform.with_named_sequence} {564 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {565 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.op<"linalg.generic">566 %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.op<"scf.forall">567 568 // linalg.generic is tileable. The op is tiled and fused.569 transform.structured.fuse_into_containing_op %0 into %1570 : (!transform.op<"linalg.generic">, !transform.op<"scf.forall">) -> (!transform.any_op, !transform.any_op)571 transform.yield572 }573 }574}575 576// -----577 578#map0 = affine_map<()[s0, s1] -> (s0 ceildiv s1)>579#map1 = affine_map<(d0)[s0] -> (d0 * s0)>580#map2 = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)>581#map3 = affine_map<(d0, d1) -> (d0, d1)>582#map4 = affine_map<(d0, d1) -> (d0)>583 584module {585 // CHECK-LABEL: func.func @fuse_tileable_reductions586 // CHECK-SAME: %[[CHUNK_SIZE:[0-9a-z]+]]: index587 // CHECK-SAME: %[[IN:[0-9a-z]+]]: tensor<?x?xf32>588 // CHECK-SAME: %[[OUT_1:[0-9a-z]+]]: tensor<?xf32>589 // CHECK-SAME: %[[OUT_2:[0-9a-z]+]]: tensor<?xf32>590 // CHECK-SAME: %[[OUT_3:[0-9a-z]+]]: tensor<?xf32>591 func.func @fuse_tileable_reductions(%idx: index, %in: tensor<?x?xf32>, %out_1: tensor<?xf32>, %out_2: tensor<?xf32>, %out_3: tensor<?xf32>)592 -> (tensor<?xf32>, tensor<?xf32>) {593 %cst = arith.constant 4.200000e+01 : f32594 %c0 = arith.constant 0 : index595 596 %0 = linalg.generic {597 indexing_maps = [#map3, #map4], iterator_types = ["parallel", "reduction"]598 } ins(%in : tensor<?x?xf32>) outs(%out_1 : tensor<?xf32>) {599 ^bb0(%a: f32, %b: f32):600 %d = arith.maximumf %a, %b : f32601 linalg.yield %d : f32602 } -> tensor<?xf32>603 %d0 = tensor.dim %out_1, %c0 : tensor<?xf32>604 605 %1 = affine.apply #map0()[%d0, %idx]606 607 // CHECK: %[[R0:.*]]:2 = scf.forall (%[[ARG5:.*]]) in (%{{.*}}) shared_outs(%[[ARG6:.*]] = %[[OUT_2]], %[[ARG7:.*]] = %[[OUT_1]])608 // CHECK-SAME: -> (tensor<?xf32>, tensor<?xf32>) {609 %2 = scf.forall (%i) in (%1) shared_outs(%o = %out_2) -> (tensor<?xf32>) {610 // CHECK: %[[I0:.*]] = affine.apply {{.*}}611 %3 = affine.apply #map1(%i)[%idx]612 // CHECK: %[[I1:.*]] = affine.min {{.*}}613 %4 = affine.min #map2(%i)[%d0, %idx]614 %5 = tensor.extract_slice %o[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32>615 616 // CHECK: %[[T1:.*]] = linalg.generic {{.*}}617 %6 = tensor.extract_slice %0[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32>618 619 %7 = linalg.exp ins(%6 : tensor<?xf32>) outs(%5 : tensor<?xf32>) -> tensor<?xf32>620 scf.forall.in_parallel {621 // CHECK: tensor.parallel_insert_slice %[[T1]] into %[[ARG7]][%[[I0]]] [%[[I1]]] [1] : tensor<?xf32> into tensor<?xf32>622 tensor.parallel_insert_slice %7 into %o[%3] [%4] [1] : tensor<?xf32> into tensor<?xf32>623 }624 }625 // CHECK: return %[[R0]]#0, %[[R0]]#1626 func.return %2, %0 : tensor<?xf32>, tensor<?xf32>627 // CHECK: }628 }629 630 module attributes {transform.with_named_sequence} {631 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {632 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.op<"linalg.generic">633 %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.op<"scf.forall">634 635 // linalg.generic is tileable. The op is tiled and fused.636 transform.structured.fuse_into_containing_op %0 into %1637 : (!transform.op<"linalg.generic">, !transform.op<"scf.forall">) -> (!transform.any_op, !transform.any_op)638 transform.yield639 }640 }641}642 643// -----644 645#map0 = affine_map<()[s0, s1] -> (s0 ceildiv s1)>646#map1 = affine_map<(d0)[s0] -> (d0 * s0)>647#map2 = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)>648#map3 = affine_map<(d0) -> (d0)>649 650module {651 // CHECK-LABEL: func.func @fuse_tileable_using_new_handle652 // CHECK-SAME: %[[CHUNK_SIZE:[0-9a-z]+]]: index653 // CHECK-SAME: %[[IN:[0-9a-z]+]]: tensor<?xf32>654 // CHECK-SAME: %[[OUT_1:[0-9a-z]+]]: tensor<?xf32>655 // CHECK-SAME: %[[OUT_2:[0-9a-z]+]]: tensor<?xf32>656 // CHECK-SAME: %[[OUT_3:[0-9a-z]+]]: tensor<?xf32>657 func.func @fuse_tileable_using_new_handle(%idx: index, %in: tensor<?xf32>, %out_1: tensor<?xf32>, %out_2: tensor<?xf32>, %out_3: tensor<?xf32>)658 -> (tensor<?xf32>, tensor<?xf32>) {659 %cst = arith.constant 4.200000e+01 : f32660 %c0 = arith.constant 0 : index661 662 %0 = linalg.generic {663 indexing_maps = [#map3, #map3], iterator_types = ["parallel"]664 } ins(%in : tensor<?xf32>) outs(%out_1 : tensor<?xf32>) {665 ^bb0(%a: f32, %b: f32):666 %d = arith.addf %a, %b : f32667 linalg.yield %d : f32668 } -> tensor<?xf32>669 670 %1 = linalg.generic {671 indexing_maps = [#map3, #map3], iterator_types = ["parallel"]672 } ins(%0 : tensor<?xf32>) outs(%out_1 : tensor<?xf32>) {673 ^bb0(%a: f32, %b: f32):674 %d = arith.mulf %a, %b : f32675 linalg.yield %d : f32676 } -> tensor<?xf32>677 %d0 = tensor.dim %out_1, %c0 : tensor<?xf32>678 679 %2 = affine.apply #map0()[%d0, %idx]680 681 // CHECK: %[[R0:.*]]:2 = scf.forall (%[[ARG5:.*]]) in (%{{.*}}) shared_outs(%[[ARG6:.*]] = %[[OUT_2]], %[[ARG7:.*]] = %[[OUT_1]])682 // CHECK-SAME: -> (tensor<?xf32>, tensor<?xf32>) {683 %3 = scf.forall (%i) in (%2) shared_outs(%o = %out_2) -> (tensor<?xf32>) {684 // CHECK: %[[I0:.*]] = affine.apply {{.*}}685 %4 = affine.apply #map1(%i)[%idx]686 // CHECK: %[[I1:.*]] = affine.min {{.*}}687 %5 = affine.min #map2(%i)[%d0, %idx]688 %6 = tensor.extract_slice %o[%4] [%5] [1] : tensor<?xf32> to tensor<?xf32>689 690 // CHECK: linalg.generic691 // CHECK: %[[T1:.*]] = linalg.generic {{.*}}692 // CHECK: %[[T2:.*]] = linalg.generic {{.*}}693 %7 = tensor.extract_slice %1[%4] [%5] [1] : tensor<?xf32> to tensor<?xf32>694 695 %8 = linalg.exp ins(%7 : tensor<?xf32>) outs(%6 : tensor<?xf32>) -> tensor<?xf32>696 scf.forall.in_parallel {697 // CHECK: tensor.parallel_insert_slice %[[T2]] into %[[ARG7]][%[[I0]]] [%[[I1]]] [1] : tensor<?xf32> into tensor<?xf32>698 tensor.parallel_insert_slice %8 into %o[%2] [%5] [1] : tensor<?xf32> into tensor<?xf32>699 }700 }701 // CHECK: return %[[R0]]#0, %[[R0]]#1702 func.return %3, %1 : tensor<?xf32>, tensor<?xf32>703 // CHECK: }704 }705 706 module attributes {transform.with_named_sequence} {707 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {708 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.op<"linalg.generic">709 %add, %reduce = transform.split_handle %0 : (!transform.op<"linalg.generic">) -> (!transform.op<"linalg.generic">, !transform.op<"linalg.generic">)710 %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.op<"scf.forall">711 712 %fused_ops, %new_forall = transform.structured.fuse_into_containing_op %reduce into %1713 : (!transform.op<"linalg.generic">, !transform.op<"scf.forall">) -> (!transform.any_op, !transform.op<"scf.forall">)714 %fused_ops_2, %new_forall_2 = transform.structured.fuse_into_containing_op %add into %new_forall715 : (!transform.op<"linalg.generic">, !transform.op<"scf.forall">) -> (!transform.any_op, !transform.op<"scf.forall">)716 transform.yield717 }718 }719}720 721// -----722 723// This is a regression test. Make sure that the transform succeeds and valid724// IR is generated.725 726module {727 // CHECK-LABEL: func.func @softmax_dispatch_0_generic_16x128x128_f32728 func.func @softmax_dispatch_0_generic_16x128x128_f32() -> tensor<16x128x128xf32> {729 %c0 = arith.constant 0 : index730 %cst = arith.constant dense<5.000000e+00> : tensor<16x128x128xf32>731 %cst_1 = arith.constant 5.000000e+00 : f32732 %1 = tensor.empty() : tensor<16x128xf32>733 %2 = tensor.empty() : tensor<16x128x128xf32>734 %3 = linalg.fill ins(%cst_1 : f32) outs(%1 : tensor<16x128xf32>) -> tensor<16x128xf32>735 %4 = linalg.fill ins(%cst_1 : f32) outs(%1 : tensor<16x128xf32>) -> tensor<16x128xf32>736 %5 = linalg.generic {producer, indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>], iterator_types = ["parallel", "parallel", "reduction"]} ins(%cst : tensor<16x128x128xf32>) outs(%4 : tensor<16x128xf32>) {737 ^bb0(%in: f32, %out: f32):738 %8 = arith.maximumf %in, %out : f32739 linalg.yield %8 : f32740 } -> tensor<16x128xf32>741 %c16 = arith.constant 16 : index742 %c32 = arith.constant 32 : index743 %7 = scf.forall (%arg0, %arg1) in (16, 32) shared_outs(%arg2 = %2) -> (tensor<16x128x128xf32>) {744 %11 = affine.apply affine_map<(d0) -> (d0 * 4)>(%arg1)745 %extracted_slice = tensor.extract_slice %5[%arg0, %11] [1, 4] [1, 1] : tensor<16x128xf32> to tensor<1x4xf32>746 %extracted_slice_3 = tensor.extract_slice %2[%arg0, %11, 0] [1, 4, 128] [1, 1, 1] : tensor<16x128x128xf32> to tensor<1x4x128xf32>747 %extracted_slice_4 = tensor.extract_slice %3[%arg0, %11] [1, 4] [1, 1] : tensor<16x128xf32> to tensor<1x4xf32>748 %15:2 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>], iterator_types = ["parallel", "parallel", "reduction"]} ins(%extracted_slice : tensor<1x4xf32>) outs(%extracted_slice_3, %extracted_slice_4 : tensor<1x4x128xf32>, tensor<1x4xf32>) {749 ^bb0(%in: f32, %out: f32, %out_9: f32):750 %22 = arith.subf %cst_1, %in : f32751 %23 = math.exp %22 : f32752 %24 = arith.addf %23, %out_9 : f32753 linalg.yield %23, %24 : f32, f32754 } -> (tensor<1x4x128xf32>, tensor<1x4xf32>)755 %extracted_slice_5 = tensor.extract_slice %5[%arg0, %11] [1, 4] [1, 1] : tensor<16x128xf32> to tensor<1x4xf32>756 %extracted_slice_6 = tensor.extract_slice %2[%arg0, %11, 0] [1, 4, 128] [1, 1, 1] : tensor<16x128x128xf32> to tensor<1x4x128xf32>757 %extracted_slice_7 = tensor.extract_slice %3[%arg0, %11] [1, 4] [1, 1] : tensor<16x128xf32> to tensor<1x4xf32>758 %19:2 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>], iterator_types = ["parallel", "parallel", "reduction"]} ins(%extracted_slice_5 : tensor<1x4xf32>) outs(%extracted_slice_6, %extracted_slice_7 : tensor<1x4x128xf32>, tensor<1x4xf32>) {759 ^bb0(%in: f32, %out: f32, %out_9: f32):760 %22 = arith.subf %cst_1, %in : f32761 %23 = math.exp %22 : f32762 %24 = arith.addf %23, %out_9 : f32763 linalg.yield %23, %24 : f32, f32764 } -> (tensor<1x4x128xf32>, tensor<1x4xf32>)765 %extracted_slice_8 = tensor.extract_slice %arg2[%arg0, %11, 0] [1, 4, 128] [1, 1, 1] : tensor<16x128x128xf32> to tensor<1x4x128xf32>766 %20 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%15#0, %19#1 : tensor<1x4x128xf32>, tensor<1x4xf32>) outs(%extracted_slice_8 : tensor<1x4x128xf32>) {767 ^bb0(%in: f32, %in_9: f32, %out: f32):768 %22 = arith.divf %in, %in_9 : f32769 linalg.yield %22 : f32770 } -> tensor<1x4x128xf32>771 scf.forall.in_parallel {772 tensor.parallel_insert_slice %20 into %arg2[%arg0, %11, 0] [1, 4, 128] [1, 1, 1] : tensor<1x4x128xf32> into tensor<16x128x128xf32>773 }774 }775 return %7 : tensor<16x128x128xf32>776 }777 778 module attributes {transform.with_named_sequence} {779 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {780 %0 = transform.structured.match attributes{producer} in %arg1 : (!transform.any_op) -> !transform.op<"linalg.generic">781 %1 = transform.structured.match ops{["scf.forall"]} in %arg1 : (!transform.any_op) -> !transform.op<"scf.forall">782 transform.structured.fuse_into_containing_op %0 into %1783 : (!transform.op<"linalg.generic">, !transform.op<"scf.forall">) -> (!transform.any_op, !transform.any_op)784 transform.yield785 }786 }787}788 789 790////////////////////////////////////////////////////////////////////////////////791// Tests below are expected to fail.792////////////////////////////////////////////////////////////////////////////////793 794// -----795 796// NO-CHECK-LABEL-ON-EXPECTED-ERROR797func.func @copy_1d_1024xf16(%arg0: tensor<123x456xf32>, %arg1: tensor<456x789xf32>, %arg2 : tensor<123x789xf32>) -> tensor<123x789xf32> {798 %0 = arith.constant 0.000000e+00 : f32799 %1 = linalg.fill ins(%0 : f32) outs(%arg2 : tensor<123x789xf32>) -> tensor<123x789xf32>800 // expected-note @below {{containing op}}801 %2 = linalg.matmul ins(%arg0, %arg1 : tensor<123x456xf32>, tensor<456x789xf32>) outs(%1 : tensor<123x789xf32>) -> tensor<123x789xf32>802 return %2 : tensor<123x789xf32>803}804 805module attributes {transform.with_named_sequence} {806 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {807 %0 = transform.structured.match ops{["linalg.fill"]} in %arg1808 : (!transform.any_op) -> !transform.any_op809 %1 = transform.structured.match ops{["linalg.matmul"]} in %arg1810 : (!transform.any_op) -> !transform.any_op811 %tiled_op, %forall_op = transform.structured.tile_using_forall %1812 num_threads [] tile_sizes [50, 16]813 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)814 // Note that we pass in %tiled_op, which isn't a container op.815 // expected-error @+2 {{could not find next producer to fuse into container}}816 %fused_op, %new_containing_op =817 transform.structured.fuse_into_containing_op %0 into %tiled_op818 : (!transform.any_op, !transform.any_op)819 -> (!transform.any_op, !transform.any_op)820 transform.yield821 }822}823