1157 lines · plain
1// RUN: mlir-opt --transform-interpreter --cse --split-input-file --verify-diagnostics %s | FileCheck %s2 3#map = affine_map<(d0) -> (d0)>4module {5 func.func @fuse_tileable_consumer_scf_for(%arg0: tensor<32xf32>, %arg1: tensor<32xf32>, %arg2: tensor<64xf32>) -> tensor<64xf32> {6 %c4 = arith.constant 4 : index7 %c64 = arith.constant 64 : index8 %c0 = arith.constant 0 : index9 %1:2 = scf.for %arg3 = %c0 to %c64 step %c4 iter_args(%arg4 = %arg2, %arg5 = %arg2) -> (tensor<64xf32>, tensor<64xf32>) {10 %extracted_slice = tensor.extract_slice %arg4[%arg3] [32] [1] : tensor<64xf32> to tensor<32xf32>11 %3 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} ins(%arg0, %arg1 : tensor<32xf32>, tensor<32xf32>) outs(%extracted_slice : tensor<32xf32>) {12 ^bb0(%in: f32, %in_16: f32, %out: f32):13 %13 = arith.mulf %in, %in_16 : f3214 %14 = arith.addf %out, %13 : f3215 linalg.yield %14 : f3216 } -> tensor<32xf32>17 %4 = tensor.insert_slice %3 into %arg4[%arg3] [32] [1] : tensor<32xf32> into tensor<64xf32>18 scf.yield %arg5, %4 : tensor<64xf32>, tensor<64xf32>19 }20 %in_operand_2 = tensor.empty() : tensor<64xf32>21 %out_operand_3 = tensor.empty() : tensor<64xf32>22 %2 = linalg.add ins(%1#1, %in_operand_2 : tensor<64xf32>, tensor<64xf32>) outs(%out_operand_3 : tensor<64xf32>) -> tensor<64xf32>23 return %2 : tensor<64xf32>24 }25}26 27module attributes {transform.with_named_sequence} {28 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {29 %loop = transform.structured.match ops{["scf.for"]} in %arg130 : (!transform.any_op) -> !transform.any_op31 %yield = transform.structured.match ops{["tensor.insert_slice"]} in %arg132 : (!transform.any_op) -> !transform.any_op33 %a, %b = transform.test.fuse_consumer_using_slice %yield in (%loop)34 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)35 transform.yield36 }37}38// CHECK: func.func @fuse_tileable_consumer_scf_for(39// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<32xf32>40// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<32xf32>41// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<64xf32>)42// CHECK: %[[C0:.*]] = arith.constant 0 : index43// CHECK: %0 = tensor.empty() : tensor<64xf32>44// CHECK: %[[FINAL_RESULT:.*]]:3 = scf.for %[[IV:.*]] = %[[C0]]45// CHECK-SAME: iter_args(%[[FIRST_OUT_ARG:.*]] = %[[ARG2]], %[[SECOND_OUT_ARG:.*]] = %[[ARG2]], %[[ELEM_OUT_ARG:.*]] = %0)46// CHECK-SAME: {47// CHECK: %[[MAT_OUT_SLICE:.*]] = tensor.extract_slice %[[FIRST_OUT_ARG]][%[[IV]]] [32] [1]48// CHECK: %[[MAT_OUT:.*]] = linalg.generic49// CHECK-SAME: outs(%[[MAT_OUT_SLICE]] : tensor<32xf32>)50// CHECK: %[[INSERT_MAT:.*]] = tensor.insert_slice %[[MAT_OUT]] into %[[FIRST_OUT_ARG]][%[[IV]]] [32] [1]51// CHECK: %[[SLICE_OPERAND2:.*]] = tensor.extract_slice %0[%[[IV]]] [32] [1]52// CHECK: %[[SLICE_OUT:.*]] = tensor.extract_slice %[[ELEM_OUT_ARG]][%[[IV]]] [32] [1]53// CHECK: %[[ELEM_OUT:.*]] = linalg.add54// CHECK-SAME: ins(%[[MAT_OUT]], %[[SLICE_OPERAND2]] :55// CHECK-SAME: outs(%[[SLICE_OUT]] :56// CHECK: %[[INSERT_ELEM:.*]] = tensor.insert_slice %[[ELEM_OUT]] into %[[ELEM_OUT_ARG]][%[[IV]]] [32] [1]57// CHECK: scf.yield %[[SECOND_OUT_ARG]], %[[INSERT_MAT]], %[[INSERT_ELEM]] :58// CHECK: }59// CHECK: return %[[FINAL_RESULT]]#2 :60 61// -----62 63module {64 func.func @fuse_tileable_consumer_scf_forall(%arg0: tensor<32x32xf32>, %arg1: tensor<32x32xf32>, %arg2: tensor<64x64xf32>) -> tensor<64x64xf32> {65 %c4 = arith.constant 4 : index66 %c64 = arith.constant 64 : index67 %c0 = arith.constant 0 : index68 %1:2 = scf.forall (%arg3, %arg4) in (2, 2) shared_outs(%arg5 = %arg2, %arg6 = %arg2) -> (tensor<64x64xf32>, tensor<64x64xf32>) {69 %extracted_slice = tensor.extract_slice %arg5[%arg3, %arg4] [32, 32] [1, 1] : tensor<64x64xf32> to tensor<32x32xf32>70 %extracted_slice_1 = tensor.extract_slice %arg6[%arg3, %arg4] [32, 32] [1, 1] : tensor<64x64xf32> to tensor<32x32xf32>71 %3 = linalg.matmul ins(%arg0, %arg1 : tensor<32x32xf32>, tensor<32x32xf32>) outs(%extracted_slice : tensor<32x32xf32>) -> tensor<32x32xf32>72 scf.forall.in_parallel {73 tensor.parallel_insert_slice %3 into %arg6[%arg3, %arg4] [32, 32] [1, 1] : tensor<32x32xf32> into tensor<64x64xf32>74 tensor.parallel_insert_slice %extracted_slice_1 into %arg5[%arg3, %arg4] [32, 32] [1, 1] : tensor<32x32xf32> into tensor<64x64xf32>75 }76 }77 %in_operand_2 = tensor.empty() : tensor<64x64xf32>78 %out_operand_3 = tensor.empty() : tensor<64x64xf32>79 %2 = linalg.add ins(%1#1, %in_operand_2 : tensor<64x64xf32>, tensor<64x64xf32>) outs(%out_operand_3 : tensor<64x64xf32>) -> tensor<64x64xf32>80 return %2 : tensor<64x64xf32>81 }82}83 84module attributes {transform.with_named_sequence} {85 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {86 %slice_ops = transform.structured.match ops{["tensor.parallel_insert_slice"]} in %arg187 : (!transform.any_op) -> !transform.any_op88 %loop = transform.structured.match ops{["scf.forall"]} in %arg189 : (!transform.any_op) -> !transform.any_op90 %first_slice_op, %second_slice_op = transform.split_handle %slice_ops91 : (!transform.any_op)92 -> (!transform.any_op, !transform.any_op)93 %a, %b = transform.test.fuse_consumer_using_slice %first_slice_op in (%loop)94 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)95 transform.yield96 }97}98// CHECK: func.func @fuse_tileable_consumer_scf_forall(99// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<32x32xf32>100// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<32x32xf32>101// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<64x64xf32>)102// CHECK: %[[OUT_INIT:.*]] = tensor.empty() : tensor<64x64xf32>103// CHECK: %[[FINAL_RESULT:.*]]:3 = scf.forall (%[[IV1:.*]], %[[IV2:.*]]) in (2, 2)104// CHECK-SAME: shared_outs(%[[FIRST_OUT_ARG:.*]] = %[[ARG2]], %[[SECOND_OUT_ARG:.*]] = %[[ARG2]], %[[ELEM_OUT_ARG:.*]] = %[[OUT_INIT]])105// CHECK-SAME: {106// CHECK: %[[MAT_OUT_SLICE:.*]] = tensor.extract_slice %[[FIRST_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]107// CHECK: %[[SECOND_ARG_SLICE:.*]] = tensor.extract_slice %[[SECOND_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]108// CHECK: %[[MAT_OUT:.*]] = linalg.matmul109// CHECK-SAME: outs(%[[MAT_OUT_SLICE]] :110// CHECK: %[[SLICE_OPERAND2:.*]] = tensor.extract_slice %[[OUT_INIT]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]111// CHECK: %[[SLICE_OUT:.*]] = tensor.extract_slice %[[ELEM_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]112// CHECK: %[[ELEM_OUT:.*]] = linalg.add113// CHECK-SAME: ins(%[[MAT_OUT]], %[[SLICE_OPERAND2]] :114// CHECK-SAME: outs(%[[SLICE_OUT]] :115// CHECK: scf.forall.in_parallel {116// CHECK: tensor.parallel_insert_slice %[[MAT_OUT]] into %[[SECOND_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]117// CHECK: tensor.parallel_insert_slice %[[SECOND_ARG_SLICE]] into %[[FIRST_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]118// CHECK: tensor.parallel_insert_slice %[[ELEM_OUT]] into %[[ELEM_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]119// CHECK: }120// CHECK: }121// CHECK: return %[[FINAL_RESULT]]#2 :122 123// -----124 125#map = affine_map<(d0) -> (d0)>126module {127 func.func @fuse_tileable_consumer_scf_for_multi_yielding_consumer(%arg0: tensor<32xf32>, %arg1: tensor<32xf32>, %arg2: tensor<64xf32>) -> tensor<64xf32> {128 %c4 = arith.constant 4 : index129 %c64 = arith.constant 64 : index130 %c0 = arith.constant 0 : index131 %1:2 = scf.for %arg3 = %c0 to %c64 step %c4 iter_args(%arg4 = %arg2, %arg5 = %arg2) -> (tensor<64xf32>, tensor<64xf32>) {132 %extracted_slice = tensor.extract_slice %arg4[%arg3] [32] [1] : tensor<64xf32> to tensor<32xf32>133 %3 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} ins(%arg0, %arg1 : tensor<32xf32>, tensor<32xf32>) outs(%extracted_slice : tensor<32xf32>) {134 ^bb0(%in: f32, %in_16: f32, %out: f32):135 %13 = arith.mulf %in, %in_16 : f32136 %14 = arith.addf %out, %13 : f32137 linalg.yield %14 : f32138 } -> tensor<32xf32>139 %4 = tensor.insert_slice %3 into %arg4[%arg3] [32] [1] : tensor<32xf32> into tensor<64xf32>140 scf.yield %arg5, %4 : tensor<64xf32>, tensor<64xf32>141 }142 %in_operand_2 = tensor.empty() : tensor<64xf32>143 %out_operand_3 = tensor.empty() : tensor<64xf32>144 %out_operand_4 = tensor.empty() : tensor<64xf32>145 %2:2 = linalg.generic {indexing_maps = [#map, #map, #map, #map], iterator_types = ["parallel"]} ins(%1#1, %in_operand_2 : tensor<64xf32>, tensor<64xf32>) outs(%out_operand_3, %out_operand_4 : tensor<64xf32>, tensor<64xf32>) {146 ^bb0(%in: f32, %in_16: f32, %out_0: f32, %out_1: f32):147 %13 = arith.mulf %in, %in_16 : f32148 %14 = arith.subf %out_0, %13 : f32149 %15 = arith.addf %out_1, %in : f32150 linalg.yield %14, %15 : f32, f32151 } -> (tensor<64xf32>, tensor<64xf32>)152 return %2#1 : tensor<64xf32>153 }154}155 156module attributes {transform.with_named_sequence} {157 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {158 %yield = transform.structured.match ops{["tensor.insert_slice"]} in %arg1159 : (!transform.any_op) -> !transform.any_op160 %loop = transform.structured.match ops{["scf.for"]} in %arg1161 : (!transform.any_op) -> !transform.any_op162 %a, %b = transform.test.fuse_consumer_using_slice %yield in (%loop)163 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)164 transform.yield165 }166}167// CHECK: func.func @fuse_tileable_consumer_scf_for_multi_yielding_consumer(168// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<32xf32>169// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<32xf32>170// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<64xf32>)171// CHECK: %[[C0:.*]] = arith.constant 0 : index172// CHECK: %0 = tensor.empty() : tensor<64xf32>173// CHECK: %[[FINAL_RESULT:.*]]:4 = scf.for %[[IV:.*]] = %[[C0]]174// CHECK-SAME: iter_args(%[[FIRST_OUT_ARG:.*]] = %[[ARG2]], %[[SECOND_OUT_ARG:.*]] = %[[ARG2]], %[[ELEM_OUT_ARG_0:.*]] = %0, %[[ELEM_OUT_ARG_1:.*]] = %0)175// CHECK-SAME: {176// CHECK: %[[MAT_OUT_SLICE:.*]] = tensor.extract_slice %[[FIRST_OUT_ARG]][%[[IV]]] [32] [1]177// CHECK: %[[MAT_OUT:.*]] = linalg.generic178// CHECK-SAME: outs(%[[MAT_OUT_SLICE]] : tensor<32xf32>)179// CHECK: %[[INSERT_MAT:.*]] = tensor.insert_slice %[[MAT_OUT]] into %[[FIRST_OUT_ARG]][%[[IV]]] [32] [1]180// CHECK: %[[SLICE_OPERAND2:.*]] = tensor.extract_slice %0[%[[IV]]] [32] [1]181// CHECK: %[[SLICE_OUT_0:.*]] = tensor.extract_slice %[[ELEM_OUT_ARG_0]][%[[IV]]] [32] [1]182// CHECK: %[[SLICE_OUT_1:.*]] = tensor.extract_slice %[[ELEM_OUT_ARG_1]][%[[IV]]] [32] [1]183// CHECK: %[[ELEM_OUT:.*]]:2 = linalg.generic184// CHECK-SAME: ins(%[[MAT_OUT]], %[[SLICE_OPERAND2]] :185// CHECK-SAME: outs(%[[SLICE_OUT_0]], %[[SLICE_OUT_1]] :186// CHECK: %[[INSERT_ELEM_0:.*]] = tensor.insert_slice %[[ELEM_OUT]]#0 into %[[ELEM_OUT_ARG_0]][%[[IV]]] [32] [1]187// CHECK: %[[INSERT_ELEM_1:.*]] = tensor.insert_slice %[[ELEM_OUT]]#1 into %[[ELEM_OUT_ARG_1]][%[[IV]]] [32] [1]188// CHECK: scf.yield %[[SECOND_OUT_ARG]], %[[INSERT_MAT]], %[[INSERT_ELEM_0]], %[[INSERT_ELEM_1]] :189// CHECK: }190// CHECK: return %[[FINAL_RESULT]]#3 :191 192// -----193 194#map = affine_map<(d0, d1) -> (d0, d1)>195module {196 func.func @fuse_tileable_consumer_scf_forall_multi_yielding_consumer(%arg0: tensor<32x32xf32>, %arg1: tensor<32x32xf32>, %arg2: tensor<64x64xf32>, %arg3: tensor<64x32xf32>) -> (tensor<64x64xf32>, tensor<2048xf32>) {197 %c4 = arith.constant 4 : index198 %c64 = arith.constant 64 : index199 %c0 = arith.constant 0 : index200 %0:2 = scf.forall (%arg4, %arg5) in (2, 2) shared_outs(%arg6 = %arg3, %arg7 = %arg2) -> (tensor<64x32xf32>, tensor<64x64xf32>) {201 %extracted_slice = tensor.extract_slice %arg6[%arg4, %arg5] [32, 32] [1, 1] : tensor<64x32xf32> to tensor<32x32xf32>202 %extracted_slice_0 = tensor.extract_slice %arg7[%arg4, %arg5] [32, 32] [1, 1] : tensor<64x64xf32> to tensor<32x32xf32>203 %6 = linalg.matmul ins(%arg0, %arg1 : tensor<32x32xf32>, tensor<32x32xf32>) outs(%extracted_slice : tensor<32x32xf32>) -> tensor<32x32xf32>204 scf.forall.in_parallel {205 tensor.parallel_insert_slice %6 into %arg7[%arg4, %arg5] [32, 32] [1, 1] : tensor<32x32xf32> into tensor<64x64xf32>206 tensor.parallel_insert_slice %extracted_slice_0 into %arg6[%arg4, %arg5] [32, 32] [1, 1] : tensor<32x32xf32> into tensor<64x32xf32>207 }208 }209 %1 = tensor.empty() : tensor<64x64xf32>210 %2 = tensor.empty() : tensor<64x64xf32>211 %3 = tensor.empty() : tensor<64x64xf32>212 %4:2 = linalg.generic {indexing_maps = [#map, #map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%0#1, %1 : tensor<64x64xf32>, tensor<64x64xf32>) outs(%2, %3 : tensor<64x64xf32>, tensor<64x64xf32>) {213 ^bb0(%in: f32, %in_0: f32, %out: f32, %out_1: f32):214 %6 = arith.mulf %in, %in_0 : f32215 %7 = arith.subf %out, %6 : f32216 %8 = arith.addf %out_1, %in : f32217 linalg.yield %7, %8 : f32, f32218 } -> (tensor<64x64xf32>, tensor<64x64xf32>)219 %5 = tensor.empty() : tensor<2048xf32>220 %unpack = linalg.unpack %0#0 outer_dims_perm = [0] inner_dims_pos = [0] inner_tiles = [32] into %5 : tensor<64x32xf32> -> tensor<2048xf32>221 return %4#1, %unpack : tensor<64x64xf32>, tensor<2048xf32>222 }223}224 225module attributes {transform.with_named_sequence} {226 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {227 %slice_ops = transform.structured.match ops{["tensor.parallel_insert_slice"]} in %arg1228 : (!transform.any_op) -> !transform.any_op229 %loop = transform.structured.match ops{["scf.forall"]} in %arg1230 : (!transform.any_op) -> !transform.any_op231 %first_slice_op, %second_slice_op = transform.split_handle %slice_ops232 : (!transform.any_op)233 -> (!transform.any_op, !transform.any_op)234 %a, %b = transform.test.fuse_consumer_using_slice %first_slice_op in (%loop)235 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)236 transform.yield237 }238}239// CHECK: func.func @fuse_tileable_consumer_scf_forall_multi_yielding_consumer(240// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<32x32xf32>241// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<32x32xf32>242// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<64x64xf32>243// CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]: tensor<64x32xf32>)244// CHECK: %[[OUT_INIT:.*]] = tensor.empty() : tensor<64x64xf32>245// CHECK: %[[FINAL_RESULT:.*]]:4 = scf.forall (%[[IV1:.*]], %[[IV2:.*]]) in (2, 2)246// CHECK-SAME: shared_outs(%[[FIRST_OUT_ARG:.*]] = %[[ARG3]], %[[SECOND_OUT_ARG:.*]] = %[[ARG2]], %[[ELEM_OUT_ARG_0:.*]] = %[[OUT_INIT]], %[[ELEM_OUT_ARG_1:.*]] = %[[OUT_INIT]])247// CHECK-SAME: {248// CHECK: %[[MAT_OUT_SLICE:.*]] = tensor.extract_slice %[[FIRST_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]249// CHECK: %[[SECOND_ARG_SLICE:.*]] = tensor.extract_slice %[[SECOND_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]250// CHECK: %[[MAT_OUT:.*]] = linalg.matmul251// CHECK-SAME: outs(%[[MAT_OUT_SLICE]] :252// CHECK: %[[SLICE_OPERAND2:.*]] = tensor.extract_slice %[[OUT_INIT]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]253// CHECK: %[[SLICE_OUT_0:.*]] = tensor.extract_slice %[[ELEM_OUT_ARG_0]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]254// CHECK: %[[SLICE_OUT_1:.*]] = tensor.extract_slice %[[ELEM_OUT_ARG_1]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]255// CHECK: %[[ELEM_OUT:.*]]:2 = linalg.generic256// CHECK-SAME: ins(%[[MAT_OUT]], %[[SLICE_OPERAND2]] :257// CHECK-SAME: outs(%[[SLICE_OUT_0]], %[[SLICE_OUT_1]] :258// CHECK: scf.forall.in_parallel {259// CHECK: tensor.parallel_insert_slice %[[MAT_OUT]] into %[[SECOND_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]260// CHECK: tensor.parallel_insert_slice %[[SECOND_ARG_SLICE]] into %[[FIRST_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]261// CHECK: tensor.parallel_insert_slice %[[ELEM_OUT]]#0 into %[[ELEM_OUT_ARG_0]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]262// CHECK: tensor.parallel_insert_slice %[[ELEM_OUT]]#1 into %[[ELEM_OUT_ARG_1]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]263// CHECK: }264// CHECK: }265// CHECK: %[[UNPACK:.*]] = linalg.unpack %[[FINAL_RESULT]]#0 outer_dims_perm = [0] inner_dims_pos = [0] inner_tiles = [32] into %{{.*}} : tensor<64x32xf32> -> tensor<2048xf32>266// CHECK: return %[[FINAL_RESULT]]#3, %[[UNPACK]] :267 268// -----269 270#map = affine_map<(d0, d1) -> (d0, d1)>271module {272 func.func @fuse_unpack_consumer_into_scf_forall(%arg0: tensor<32x32xf32>, %arg1: tensor<32x32xf32>, %arg2: tensor<64x32xf32>) -> tensor<2048xf32> {273 %c4 = arith.constant 4 : index274 %c64 = arith.constant 64 : index275 %c0 = arith.constant 0 : index276 %1 = scf.forall (%arg3, %arg4) = (0, 0) to (64, 32) step (32, 32) shared_outs(%arg5 = %arg2) -> (tensor<64x32xf32>) {277 %extracted_slice = tensor.extract_slice %arg5[%arg3, %arg4] [32, 32] [1, 1] : tensor<64x32xf32> to tensor<32x32xf32>278 %3 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0, %arg1 : tensor<32x32xf32>, tensor<32x32xf32>) outs(%extracted_slice : tensor<32x32xf32>) {279 ^bb0(%in: f32, %in_16: f32, %out: f32):280 %13 = arith.mulf %in, %in_16 : f32281 %14 = arith.addf %out, %13 : f32282 linalg.yield %14 : f32283 } -> tensor<32x32xf32>284 scf.forall.in_parallel {285 tensor.parallel_insert_slice %3 into %arg5[%arg3, %arg4] [32, 32] [1, 1] : tensor<32x32xf32> into tensor<64x32xf32>286 }287 }288 %output = tensor.empty() : tensor<2048xf32>289 %unpack = linalg.unpack %1 outer_dims_perm = [0] inner_dims_pos = [0] inner_tiles = [32] into %output : tensor<64x32xf32> -> tensor<2048xf32>290 return %unpack : tensor<2048xf32>291 }292}293 294module attributes {transform.with_named_sequence} {295 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {296 %slice_op = transform.structured.match ops{["tensor.parallel_insert_slice"]} in %arg1297 : (!transform.any_op) -> !transform.any_op298 %loop = transform.structured.match ops{["scf.forall"]} in %arg1299 : (!transform.any_op) -> !transform.any_op300 %a, %b = transform.test.fuse_consumer_using_slice %slice_op in (%loop)301 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)302 transform.yield303 }304}305// CHECK-DAG: #[[UNPACK_RESULT_OFFSET_MAP:.*]] = affine_map<(d0) -> (d0 * 32)>306// CHECK-DAG: #[[UNPACK_RESULT_SIZE_MAP:.*]] = affine_map<(d0) -> (1024, d0 * -32 + 2048)>307// CHECK: func.func @fuse_unpack_consumer_into_scf_forall(308// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<32x32xf32>309// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<32x32xf32>310// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<64x32xf32>)311// CHECK: %[[OUT_INIT:.*]] = tensor.empty() : tensor<2048xf32>312// CHECK: %[[FINAL_RESULT:.*]]:2 = scf.forall (%[[IV1:.*]], %[[IV2:.*]]) = (0, 0) to (64, 32) step (32, 32)313// CHECK-SAME: shared_outs(%[[FIRST_OUT_ARG:.*]] = %[[ARG2]], %[[UNPACK_OUT_ARG:.*]] = %[[OUT_INIT]])314// CHECK-SAME: {315// CHECK: %[[GENERIC_OUT_SLICE:.*]] = tensor.extract_slice %[[FIRST_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]316// CHECK: %[[GENERIC_OUT:.*]] = linalg.generic317// CHECK-SAME: outs(%[[GENERIC_OUT_SLICE]] :318// CHECK-DAG: %[[UNPACK_RESULT_OFFSET:.*]] = affine.apply #[[UNPACK_RESULT_OFFSET_MAP]](%[[IV1]])319// CHECK-DAG: %[[UNPACK_RESULT_SIZE:.*]] = affine.min #[[UNPACK_RESULT_SIZE_MAP]](%[[IV1]])320// CHECK: %[[TILED_UNPACK_DEST:.*]] = tensor.extract_slice %[[UNPACK_OUT_ARG]][%[[UNPACK_RESULT_OFFSET]]] [%[[UNPACK_RESULT_SIZE]]] [1]321// CHECK: %[[TILED_UNPACK_OUT:.*]] = linalg.unpack %[[GENERIC_OUT]]322// CHECK-SAME: outer_dims_perm = [0] inner_dims_pos = [0] inner_tiles = [32]323// CHECK-SAME: into %[[TILED_UNPACK_DEST]]324// CHECK: scf.forall.in_parallel {325// CHECK: tensor.parallel_insert_slice %[[GENERIC_OUT]] into %[[FIRST_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]326// CHECK: tensor.parallel_insert_slice %[[TILED_UNPACK_OUT]] into %[[UNPACK_OUT_ARG]][%[[UNPACK_RESULT_OFFSET]]] [%[[UNPACK_RESULT_SIZE]]] [1]327// CHECK: }328// CHECK: }329// CHECK: return %[[FINAL_RESULT]]#1 :330 331// -----332 333#map = affine_map<(d0, d1) -> (d0, d1)>334module {335 func.func @fuse_unaligned_unpack_consumer_into_scf_forall(%arg0: tensor<32x32xf32>, %arg1: tensor<32x32xf32>, %arg2: tensor<64x32xf32>) -> tensor<2047xf32> {336 %c4 = arith.constant 4 : index337 %c64 = arith.constant 64 : index338 %c0 = arith.constant 0 : index339 %1 = scf.forall (%arg3, %arg4) = (0, 0) to (64, 32) step (32, 32) shared_outs(%arg5 = %arg2) -> (tensor<64x32xf32>) {340 %extracted_slice = tensor.extract_slice %arg5[%arg3, %arg4] [32, 32] [1, 1] : tensor<64x32xf32> to tensor<32x32xf32>341 %3 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0, %arg1 : tensor<32x32xf32>, tensor<32x32xf32>) outs(%extracted_slice : tensor<32x32xf32>) {342 ^bb0(%in: f32, %in_16: f32, %out: f32):343 %13 = arith.mulf %in, %in_16 : f32344 %14 = arith.addf %out, %13 : f32345 linalg.yield %14 : f32346 } -> tensor<32x32xf32>347 scf.forall.in_parallel {348 tensor.parallel_insert_slice %3 into %arg5[%arg3, %arg4] [32, 32] [1, 1] : tensor<32x32xf32> into tensor<64x32xf32>349 }350 }351 %output = tensor.empty() : tensor<2047xf32>352 %unpack = linalg.unpack %1 outer_dims_perm = [0] inner_dims_pos = [0] inner_tiles = [32] into %output : tensor<64x32xf32> -> tensor<2047xf32>353 return %unpack : tensor<2047xf32>354 }355}356 357module attributes {transform.with_named_sequence} {358 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {359 %slice_op = transform.structured.match ops{["tensor.parallel_insert_slice"]} in %arg1360 : (!transform.any_op) -> !transform.any_op361 %loop = transform.structured.match ops{["scf.forall"]} in %arg1362 : (!transform.any_op) -> !transform.any_op363 %a, %b = transform.test.fuse_consumer_using_slice %slice_op in (%loop)364 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)365 transform.yield366 }367}368// CHECK-DAG: #[[UNPACK_RESULT_OFFSET_MAP:.*]] = affine_map<(d0) -> (d0 * 32)>369// CHECK-DAG: #[[UNPACK_RESULT_SIZE_MAP:.*]] = affine_map<(d0) -> (1024, d0 * -32 + 2047)>370// CHECK: func.func @fuse_unaligned_unpack_consumer_into_scf_forall(371// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<32x32xf32>372// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<32x32xf32>373// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<64x32xf32>)374// CHECK: %[[OUT_INIT:.*]] = tensor.empty() : tensor<2047xf32>375// CHECK: %[[FINAL_RESULT:.*]]:2 = scf.forall (%[[IV1:.*]], %[[IV2:.*]]) = (0, 0) to (64, 32) step (32, 32)376// CHECK-SAME: shared_outs(%[[FIRST_OUT_ARG:.*]] = %[[ARG2]], %[[UNPACK_OUT_ARG:.*]] = %[[OUT_INIT]])377// CHECK-SAME: {378// CHECK: %[[GENERIC_OUT_SLICE:.*]] = tensor.extract_slice %[[FIRST_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]379// CHECK: %[[GENERIC_OUT:.*]] = linalg.generic380// CHECK-SAME: outs(%[[GENERIC_OUT_SLICE]] :381// CHECK-DAG: %[[UNPACK_RESULT_OFFSET:.*]] = affine.apply #[[UNPACK_RESULT_OFFSET_MAP]](%[[IV1]])382// CHECK-DAG: %[[UNPACK_RESULT_SIZE:.*]] = affine.min #[[UNPACK_RESULT_SIZE_MAP]](%[[IV1]])383// CHECK: %[[TILED_UNPACK_DEST:.*]] = tensor.extract_slice %[[UNPACK_OUT_ARG]][%[[UNPACK_RESULT_OFFSET]]] [%[[UNPACK_RESULT_SIZE]]] [1]384// CHECK: %[[TILED_UNPACK_OUT:.*]] = linalg.unpack %[[GENERIC_OUT]]385// CHECK-SAME: outer_dims_perm = [0] inner_dims_pos = [0] inner_tiles = [32]386// CHECK-SAME: into %[[TILED_UNPACK_DEST]]387// CHECK: scf.forall.in_parallel {388// CHECK: tensor.parallel_insert_slice %[[GENERIC_OUT]] into %[[FIRST_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]389// CHECK: tensor.parallel_insert_slice %[[TILED_UNPACK_OUT]] into %[[UNPACK_OUT_ARG]][%[[UNPACK_RESULT_OFFSET]]] [%[[UNPACK_RESULT_SIZE]]] [1]390// CHECK: }391// CHECK: }392// CHECK: return %[[FINAL_RESULT]]#1 :393 394// -----395 396#map = affine_map<(d0, d1) -> (d0, d1)>397module {398 func.func @fuse_perfect_tiling_pack_consumer(%arg0: tensor<32x32xf32>, %arg1: tensor<32x32xf32>, %arg2: tensor<64x32xf32>) -> tensor<4x32x16xf32> {399 %c4 = arith.constant 4 : index400 %c64 = arith.constant 64 : index401 %c0 = arith.constant 0 : index402 %1 = scf.forall (%arg3, %arg4) in (2, 1) shared_outs(%arg5 = %arg2) -> (tensor<64x32xf32>) {403 %extracted_slice = tensor.extract_slice %arg5[%arg3, %arg4] [32, 32] [1, 1] : tensor<64x32xf32> to tensor<32x32xf32>404 %3 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0, %arg1 : tensor<32x32xf32>, tensor<32x32xf32>) outs(%extracted_slice : tensor<32x32xf32>) {405 ^bb0(%in: f32, %in_16: f32, %out: f32):406 %13 = arith.mulf %in, %in_16 : f32407 %14 = arith.addf %out, %13 : f32408 linalg.yield %14 : f32409 } -> tensor<32x32xf32>410 scf.forall.in_parallel {411 tensor.parallel_insert_slice %3 into %arg5[%arg3, %arg4] [32, 32] [1, 1] : tensor<32x32xf32> into tensor<64x32xf32>412 }413 }414 %output = tensor.empty() : tensor<4x32x16xf32>415 %pack = linalg.pack %1 inner_dims_pos = [0] inner_tiles = [16] into %output : tensor<64x32xf32> -> tensor<4x32x16xf32>416 return %pack : tensor<4x32x16xf32>417 }418}419 420module attributes {transform.with_named_sequence} {421 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {422 %slice_op = transform.structured.match ops{["tensor.parallel_insert_slice"]} in %arg1423 : (!transform.any_op) -> !transform.any_op424 %loop = transform.structured.match ops{["scf.forall"]} in %arg1425 : (!transform.any_op) -> !transform.any_op426 %a, %b = transform.test.fuse_consumer_using_slice %slice_op in (%loop)427 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)428 transform.yield429 }430}431// CHECK: #[[PACK_RESULT_MAP:.*]] = affine_map<(d0) -> (d0 floordiv 16)>432// CHECK: func.func @fuse_perfect_tiling_pack_consumer(433// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<32x32xf32>434// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<32x32xf32>435// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<64x32xf32>)436// CHECK: %[[OUT_INIT:.*]] = tensor.empty() : tensor<4x32x16xf32>437// CHECK: %[[FINAL_RESULT:.*]]:2 = scf.forall (%[[IV1:.*]], %[[IV2:.*]]) in (2, 1)438// CHECK-SAME: shared_outs(%[[FIRST_OUT_ARG:.*]] = %[[ARG2]], %[[PACK_OUT_ARG:.*]] = %[[OUT_INIT]])439// CHECK-SAME: {440// CHECK: %[[GENERIC_OUT_SLICE:.*]] = tensor.extract_slice %[[FIRST_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]441// CHECK: %[[GENERIC_OUT:.*]] = linalg.generic442// CHECK-SAME: outs(%[[GENERIC_OUT_SLICE]] :443// CHECK: %[[PACK_RESULT_OFFSET:.*]] = affine.apply #[[PACK_RESULT_MAP]](%[[IV1]])444// CHECK: %[[TILED_PACK_DEST:.*]] = tensor.extract_slice %[[PACK_OUT_ARG]][%[[PACK_RESULT_OFFSET]], %[[IV2]], 0] [2, 32, 16] [1, 1, 1]445// CHECK: %[[TILED_PACK_OUT:.*]] = linalg.pack %[[GENERIC_OUT]]446// CHECK-SAME: inner_dims_pos = [0] inner_tiles = [16]447// CHECK-SAME: into %[[TILED_PACK_DEST]]448// CHECK: scf.forall.in_parallel {449// CHECK: tensor.parallel_insert_slice %[[GENERIC_OUT]] into %[[FIRST_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]450// CHECK: tensor.parallel_insert_slice %[[TILED_PACK_OUT]] into %[[PACK_OUT_ARG]][%[[PACK_RESULT_OFFSET]], %[[IV2]], 0] [2, 32, 16] [1, 1, 1]451 452// -----453 454#map = affine_map<(d0) -> (-d0 + 4, 16)>455func.func @fuse_pack_consumer_if_single_iteration(%arg0: tensor<4x4xf32>) -> tensor<1x4x16x1xf32> {456 %0 = tensor.empty() : tensor<1x4x16x1xf32>457 %1 = tensor.empty() : tensor<4x4xf32>458 %2 = scf.forall (%arg1) = (0) to (4) step (16) shared_outs(%arg2 = %1) -> (tensor<4x4xf32>) {459 %3 = affine.min #map(%arg1)460 %extracted_slice = tensor.extract_slice %arg0[%arg1, 0] [%3, 4] [1, 1] : tensor<4x4xf32> to tensor<?x4xf32>461 %extracted_slice_0 = tensor.extract_slice %arg2[%arg1, 0] [%3, 4] [1, 1] : tensor<4x4xf32> to tensor<?x4xf32>462 %4 = linalg.exp ins(%extracted_slice : tensor<?x4xf32>) outs(%extracted_slice_0 : tensor<?x4xf32>) -> tensor<?x4xf32>463 scf.forall.in_parallel {464 tensor.parallel_insert_slice %4 into %arg2[%arg1, 0] [%3, 4] [1, 1] : tensor<?x4xf32> into tensor<4x4xf32>465 }466 }467 %cst = arith.constant 0.000000e+00 : f32468 %pack = linalg.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [16, 1] into %0 : tensor<4x4xf32> -> tensor<1x4x16x1xf32>469 return %pack : tensor<1x4x16x1xf32>470}471 472module attributes {transform.with_named_sequence} {473 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {474 %0 = transform.structured.match ops{["tensor.parallel_insert_slice"]} in %arg0 : (!transform.any_op) -> !transform.any_op475 %1 = transform.structured.match ops{["scf.forall"]} in %arg0 : (!transform.any_op) -> !transform.any_op476 %consumer, %fused_consumer = transform.test.fuse_consumer_using_slice %0 in(%1) : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)477 transform.yield478 }479}480// CHECK: #[[MAP:.*]] = affine_map<(d0) -> (-d0 + 4, 16)>481// CHECK: func.func @fuse_pack_consumer_if_single_iteration(482// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]483// CHECK-DAG: %[[PACK_INIT:.*]] = tensor.empty() : tensor<1x4x16x1xf32>484// CHECK-DAG: %[[ELEM_INIT:.*]] = tensor.empty() : tensor<4x4xf32>485// CHECK-DAG: %[[PAD_VAL:.*]] = arith.constant 0.000000e+00 : f32486// CHECK: %{{.*}}:2 = scf.forall (%[[IV:.*]]) = (0) to (4) step (16)487// CHECK-SAME: shared_outs(%[[ELEM_OUT_ARG:.*]] = %[[ELEM_INIT]], %[[PACK_OUT_ARG:.*]] = %[[PACK_INIT]])488// CHECK-DAG: %[[SIZE:.+]] = affine.min #[[MAP]](%[[IV]])489// CHECK-DAG: %[[ELEM_SRC:.*]] = tensor.extract_slice %[[ARG0]][%[[IV]], 0] [%[[SIZE]], 4] [1, 1]490// CHECK-DAG: %[[ELEM_DEST:.*]] = tensor.extract_slice %[[ELEM_OUT_ARG]][%[[IV]], 0] [%[[SIZE]], 4] [1, 1]491// CHECK: %[[ELEM:.*]] = linalg.exp492// CHECK-SAME: ins(%[[ELEM_SRC]]493// CHECK-SAME: outs(%[[ELEM_DEST]]494// CHECK-DAG: %[[TILED_PACK_DEST:.*]] = tensor.extract_slice %[[PACK_OUT_ARG]][%[[IV]], 0, 0, 0] [1, 4, 16, 1] [1, 1, 1, 1]495// CHECK: %[[PACK:.*]] = linalg.pack %[[ELEM]]496// CHECK-SAME: padding_value(%[[PAD_VAL]] : f32)497// CHECK-SAME: outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [16, 1]498// CHECK-SAME: into %[[TILED_PACK_DEST]]499// CHECK: scf.forall.in_parallel {500// CHECK: tensor.parallel_insert_slice %[[ELEM]] into %[[ELEM_OUT_ARG]][%[[IV]], 0] [%[[SIZE]], 4] [1, 1]501// CHECK: tensor.parallel_insert_slice %[[PACK]] into %[[PACK_OUT_ARG]][%[[IV]], 0, 0, 0] [1, 4, 16, 1] [1, 1, 1, 1]502 503// -----504 505func.func @fuse_perfect_tiling_pack_consumer_with_outer_dims_perm(%arg0: tensor<64x32xf32>, %arg1: tensor<64x32xf32>, %arg2: tensor<2x64x16x1xf32>) -> tensor<2x64x16x1xf32> {506 %0 = scf.forall (%arg3) = (0) to (32) step (16) shared_outs(%arg4 = %arg1) -> (tensor<64x32xf32>) {507 %src = tensor.extract_slice %arg0[0, %arg3] [64, 16] [1, 1] : tensor<64x32xf32> to tensor<64x16xf32>508 %dest = tensor.extract_slice %arg4[0, %arg3] [64, 16] [1, 1] : tensor<64x32xf32> to tensor<64x16xf32>509 %1 = linalg.exp ins(%src : tensor<64x16xf32>) outs(%dest : tensor<64x16xf32>) -> tensor<64x16xf32>510 scf.forall.in_parallel {511 tensor.parallel_insert_slice %1 into %arg4[0, %arg3] [64, 16] [1, 1] : tensor<64x16xf32> into tensor<64x32xf32>512 }513 }514 %pack = linalg.pack %0 outer_dims_perm = [1, 0] inner_dims_pos = [1, 0] inner_tiles = [16, 1] into %arg2 : tensor<64x32xf32> -> tensor<2x64x16x1xf32>515 return %pack : tensor<2x64x16x1xf32>516}517 518module attributes {transform.with_named_sequence} {519 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {520 %0 = transform.structured.match ops{["tensor.parallel_insert_slice"]} in %arg0 : (!transform.any_op) -> !transform.any_op521 %1 = transform.structured.match ops{["scf.forall"]} in %arg0 : (!transform.any_op) -> !transform.any_op522 %consumer, %fused_consumer = transform.test.fuse_consumer_using_slice %0 in(%1) : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)523 transform.yield524 }525}526// CHECK: #[[PACK_RESULT_MAP:.*]] = affine_map<(d0) -> (d0 floordiv 16)>527// CHECK: func.func @fuse_perfect_tiling_pack_consumer_with_outer_dims_perm(528// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]529// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]530// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]531// CHECK: %{{.*}}:2 = scf.forall (%[[IV:.*]]) = (0) to (32) step (16)532// CHECK-SAME: shared_outs(%[[FIRST_OUT_ARG:.*]] = %[[ARG1]], %[[PACK_OUT_ARG:.*]] = %[[ARG2]])533// CHECK: %[[ELEM_SRC:.*]] = tensor.extract_slice %[[ARG0]][0, %[[IV]]] [64, 16] [1, 1]534// CHECK: %[[ELEM_DEST:.*]] = tensor.extract_slice %[[FIRST_OUT_ARG]][0, %[[IV]]] [64, 16] [1, 1]535// CHECK: %[[ELEM:.*]] = linalg.exp536// CHECK-SAME: ins(%[[ELEM_SRC]]537// CHECK-SAME: outs(%[[ELEM_DEST]]538// CHECK-DAG: %[[PACK_RESULT_OFFSET:.*]] = affine.apply #[[PACK_RESULT_MAP]](%[[IV]])539// CHECK-DAG: %[[TILED_PACK_DEST:.*]] = tensor.extract_slice %[[PACK_OUT_ARG]][%[[PACK_RESULT_OFFSET]], 0, 0, 0] [1, 64, 16, 1] [1, 1, 1, 1]540// CHECK: %[[PACK:.*]] = linalg.pack %[[ELEM]]541// CHECK-SAME: outer_dims_perm = [1, 0] inner_dims_pos = [1, 0] inner_tiles = [16, 1]542// CHECK-SAME: into %[[TILED_PACK_DEST]]543// CHECK: scf.forall.in_parallel {544// CHECK: tensor.parallel_insert_slice %[[ELEM]] into %[[FIRST_OUT_ARG]][0, %[[IV]]] [64, 16] [1, 1]545// CHECK: tensor.parallel_insert_slice %[[PACK]] into %[[PACK_OUT_ARG]][%[[PACK_RESULT_OFFSET]], 0, 0, 0] [1, 64, 16, 1] [1, 1, 1, 1]546 547// -----548 549// It is valid to fuse the pack op in perfect tiling scenario when the dimension550// is dynamic and padding is not needed.551 552func.func @fuse_pack_consumer_with_no_pad_dynamic_dim(%arg0: tensor<64x?xf32>, %arg1: tensor<64x?xf32>, %1: tensor<64x?x16xf32>) -> tensor<64x?x16xf32> {553 %c1 = arith.constant 1 : index554 %d1 = tensor.dim %arg0, %c1 : tensor<64x?xf32>555 %0 = scf.forall (%arg2) = (0) to (%d1) step (16) shared_outs(%arg3 = %arg1) -> (tensor<64x?xf32>) {556 %src = tensor.extract_slice %arg0[0, %arg2] [64, 16] [1, 1] : tensor<64x?xf32> to tensor<64x16xf32>557 %dest = tensor.extract_slice %arg3[0, %arg2] [64, 16] [1, 1] : tensor<64x?xf32> to tensor<64x16xf32>558 %2 = linalg.exp ins(%src : tensor<64x16xf32>) outs(%dest : tensor<64x16xf32>) -> tensor<64x16xf32>559 scf.forall.in_parallel {560 tensor.parallel_insert_slice %2 into %arg3[0, %arg2] [64, 16] [1, 1] : tensor<64x16xf32> into tensor<64x?xf32>561 }562 }563 %pack = linalg.pack %0 inner_dims_pos = [1] inner_tiles = [16] into %1 : tensor<64x?xf32> -> tensor<64x?x16xf32>564 return %pack : tensor<64x?x16xf32>565}566 567module attributes {transform.with_named_sequence} {568 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {569 %0 = transform.structured.match ops{["tensor.parallel_insert_slice"]} in %arg0 : (!transform.any_op) -> !transform.any_op570 %1 = transform.structured.match ops{["scf.forall"]} in %arg0 : (!transform.any_op) -> !transform.any_op571 %consumer, %fused_consumer = transform.test.fuse_consumer_using_slice %0 in(%1) : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)572 transform.yield573 }574}575// CHECK: #[[PACK_RESULT_MAP:.*]] = affine_map<(d0) -> (d0 floordiv 16)>576// CHECK: func.func @fuse_pack_consumer_with_no_pad_dynamic_dim(577// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]578// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]579// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]580// CHECK: %{{.*}}:2 = scf.forall (%[[IV:.*]]) = (0) to (%{{.+}}) step (16)581// CHECK-SAME: shared_outs(%[[FIRST_OUT_ARG:.*]] = %[[ARG1]], %[[PACK_OUT_ARG:.*]] = %[[ARG2]])582// CHECK: %[[ELEM_SRC:.*]] = tensor.extract_slice %[[ARG0]][0, %[[IV]]] [64, 16] [1, 1]583// CHECK: %[[ELEM_DEST:.*]] = tensor.extract_slice %[[FIRST_OUT_ARG]][0, %[[IV]]] [64, 16] [1, 1]584// CHECK: %[[ELEM:.*]] = linalg.exp585// CHECK-SAME: ins(%[[ELEM_SRC]]586// CHECK-SAME: outs(%[[ELEM_DEST]]587// CHECK-DAG: %[[PACK_RESULT_OFFSET:.*]] = affine.apply #[[PACK_RESULT_MAP]](%[[IV]])588// CHECK-DAG: %[[TILED_PACK_DEST:.*]] = tensor.extract_slice %[[PACK_OUT_ARG]][0, %[[PACK_RESULT_OFFSET]], 0] [64, 1, 16] [1, 1, 1]589// CHECK: %[[PACK:.*]] = linalg.pack %[[ELEM]]590// CHECK-SAME: inner_dims_pos = [1] inner_tiles = [16]591// CHECK-SAME: into %[[TILED_PACK_DEST]]592// CHECK: scf.forall.in_parallel {593// CHECK: tensor.parallel_insert_slice %[[ELEM]] into %[[FIRST_OUT_ARG]][0, %[[IV]]] [64, 16] [1, 1]594// CHECK: tensor.parallel_insert_slice %[[PACK]] into %[[PACK_OUT_ARG]][0, %[[PACK_RESULT_OFFSET]], 0] [64, 1, 16] [1, 1, 1]595 596// -----597 598// It is valid to fuse the pack op with padding semantics if it is a perfect599// tiling case.600 601func.func @fuse_pack_consumer_with_padding_semantics(%arg0: tensor<64x32xf32>, %arg1: tensor<64x32xf32>) -> tensor<22x2x3x16xf32> {602 %0 = scf.forall (%arg2, %arg3) = (0, 0) to (64, 32) step (15, 16) shared_outs(%arg4 = %arg1) -> (tensor<64x32xf32>) {603 %size = affine.min affine_map<(d0) -> (-d0 + 64, 15)>(%arg2)604 %src = tensor.extract_slice %arg0[%arg2, %arg3] [%size, 16] [1, 1] : tensor<64x32xf32> to tensor<?x16xf32>605 %dest = tensor.extract_slice %arg4[%arg2, %arg3] [%size, 16] [1, 1] : tensor<64x32xf32> to tensor<?x16xf32>606 %2 = linalg.exp ins(%src : tensor<?x16xf32>) outs(%dest : tensor<?x16xf32>) -> tensor<?x16xf32>607 scf.forall.in_parallel {608 tensor.parallel_insert_slice %2 into %arg4[%arg2, %arg3] [%size, 16] [1, 1] : tensor<?x16xf32> into tensor<64x32xf32>609 }610 }611 %1 = tensor.empty() : tensor<22x2x3x16xf32>612 %cst = arith.constant 0.000000e+00 : f32613 %pack = linalg.pack %0 padding_value(%cst : f32) inner_dims_pos = [0, 1] inner_tiles = [3, 16] into %1 : tensor<64x32xf32> -> tensor<22x2x3x16xf32>614 return %pack : tensor<22x2x3x16xf32>615}616 617module attributes {transform.with_named_sequence} {618 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {619 %0 = transform.structured.match ops{["tensor.parallel_insert_slice"]} in %arg0 : (!transform.any_op) -> !transform.any_op620 %1 = transform.structured.match ops{["scf.forall"]} in %arg0 : (!transform.any_op) -> !transform.any_op621 %consumer, %fused_consumer = transform.test.fuse_consumer_using_slice %0 in(%1) : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)622 transform.yield623 }624}625// CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0) -> (-d0 + 64, 15)>626// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0) -> (d0 floordiv 3)>627// CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0) -> (d0 ceildiv 3)>628// CHECK-DAG: #[[MAP3:.*]] = affine_map<(d0) -> (d0 floordiv 16)>629// CHECK: func.func @fuse_pack_consumer_with_padding_semantics(630// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]631// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]632// CHECK-DAG: %[[OUT_INIT:.*]] = tensor.empty() : tensor<22x2x3x16xf32>633// CHECK-DAG: %[[PAD_VAL:.*]] = arith.constant 0.000000e+00 : f32634// CHECK: %{{.*}}:2 = scf.forall (%[[I:.*]], %[[J:.*]]) = (0, 0) to (64, 32) step (15, 16)635// CHECK-SAME: shared_outs(%[[ELEM_OUT:.*]] = %[[ARG1]], %[[PACK_OUT:.*]] = %[[OUT_INIT]])636// CHECK: %[[SIZE:.+]] = affine.min #[[MAP0]](%[[I]])637// CHECK: %[[ELEM_SRC:.*]] = tensor.extract_slice %[[ARG0]]638// CHECK-SAME: [%[[I]], %[[J]]] [%[[SIZE]], 16] [1, 1]639// CHECK: %[[ELEM_DEST:.*]] = tensor.extract_slice %[[ELEM_OUT]]640// CHECK-SAME: [%[[I]], %[[J]]] [%[[SIZE]], 16] [1, 1]641// CHECK: %[[ELEM:.*]] = linalg.exp642// CHECK-SAME: ins(%[[ELEM_SRC]]643// CHECK-SAME: outs(%[[ELEM_DEST]]644// CHECK-DAG: %[[D0_OFFSET:.*]] = affine.apply #[[MAP1]](%[[I]])645// CHECK-DAG: %[[D0_SIZE:.*]] = affine.apply #[[MAP2]](%[[SIZE]])646// CHECK-DAG: %[[D1_OFFSET:.*]] = affine.apply #[[MAP3]](%[[J]])647// CHECK-DAG: %[[PACK_INIT:.*]] = tensor.extract_slice %[[PACK_OUT]]648// CHECK-SAME: [%[[D0_OFFSET]], %[[D1_OFFSET]], 0, 0] [%[[D0_SIZE]], 1, 3, 16] [1, 1, 1, 1]649// CHECK: %[[PACK:.*]] = linalg.pack %[[ELEM]]650// CHECK-SAME: padding_value(%[[PAD_VAL]] : f32)651// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [3, 16]652// CHECK-SAME: into %[[TILED_PACK_DEST]]653// CHECK: scf.forall.in_parallel {654// CHECK: tensor.parallel_insert_slice %[[ELEM]] into %[[ELEM_OUT]]655// CHECK-SAME: [%[[I]], %[[J]]] [%[[SIZE]], 16] [1, 1]656// CHECK: tensor.parallel_insert_slice %[[PACK]] into %[[PACK_OUT]]657// CHECK-SAME: [%[[D0_OFFSET]], %[[D1_OFFSET]], 0, 0] [%[[D0_SIZE]], 1, 3, 16] [1, 1, 1, 1]658 659// -----660 661// Imperfect tiling is not supported in pack op consumer fusion.662 663#map = affine_map<(d0) -> (d0 * 5)>664#map1 = affine_map<(d0) -> (d0)>665func.func @nofuse_pack_with_imperfect_tiling(%arg0: tensor<30xf32>) -> tensor<5x6xf32> {666 %0 = tensor.empty() : tensor<30xf32>667 %1 = scf.forall (%arg1) in (6) shared_outs(%arg2 = %0) -> (tensor<30xf32>) {668 %3 = affine.apply #map(%arg1)669 %extracted_slice = tensor.extract_slice %arg0[%3] [5] [1] : tensor<30xf32> to tensor<5xf32>670 %extracted_slice_0 = tensor.extract_slice %arg2[%3] [5] [1] : tensor<30xf32> to tensor<5xf32>671 %4 = linalg.generic {indexing_maps = [#map1, #map1], iterator_types = ["parallel"]} ins(%extracted_slice : tensor<5xf32>) outs(%extracted_slice_0 : tensor<5xf32>) {672 ^bb0(%in: f32, %out: f32):673 %5 = arith.addf %in, %in : f32674 linalg.yield %5 : f32675 } -> tensor<5xf32>676 scf.forall.in_parallel {677 // expected-error @below {{failed to fuse consumer of slice}}678 tensor.parallel_insert_slice %4 into %arg2[%3] [5] [1] : tensor<5xf32> into tensor<30xf32>679 }680 }681 %2 = tensor.empty() : tensor<5x6xf32>682 %pack = linalg.pack %1 outer_dims_perm = [0] inner_dims_pos = [0] inner_tiles = [6] into %2 : tensor<30xf32> -> tensor<5x6xf32>683 return %pack : tensor<5x6xf32>684}685 686module attributes {transform.with_named_sequence} {687 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {688 %0 = transform.structured.match ops{["tensor.parallel_insert_slice"]} in %arg0 : (!transform.any_op) -> !transform.any_op689 %1 = transform.structured.match ops{["scf.forall"]} in %arg0 : (!transform.any_op) -> !transform.any_op690 %consumer, %fused_consumer = transform.test.fuse_consumer_using_slice %0 in(%1) : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)691 transform.yield692 }693}694 695// -----696 697module {698 func.func @fuse_add_multiple_tilable_consumers(%arg0: tensor<256x256xf32>, %arg1: tensor<256x256xf32>, %arg2: tensor<256x256xf32>) -> (tensor<256x256xf32>, tensor<256x256xf32>) {699 %c0 = arith.constant 0 : index700 %c64 = arith.constant 64 : index701 %c256 = arith.constant 256 : index702 %cst = arith.constant 0.000000e+00 : f32703 %dest0 = tensor.empty() : tensor<256x256xf32>704 %1 = scf.for %arg3 = %c0 to %c256 step %c64 iter_args(%arg4 = %dest0) -> (tensor<256x256xf32>) {705 %extracted_slice_1 = tensor.extract_slice %arg4[%arg3, 0] [64, 256] [1, 1] : tensor<256x256xf32> to tensor<64x256xf32>706 %extracted_slice_2 = tensor.extract_slice %arg0[%arg3, 0] [64, 256] [1, 1] : tensor<256x256xf32> to tensor<64x256xf32>707 %extracted_slice_3 = tensor.extract_slice %arg1[%arg3, 0] [64, 256] [1, 1] : tensor<256x256xf32> to tensor<64x256xf32>708 %3 = linalg.add ins(%extracted_slice_2, %extracted_slice_3 : tensor<64x256xf32>, tensor<64x256xf32>) outs(%extracted_slice_1 : tensor<64x256xf32>) -> tensor<64x256xf32>709 %insert_slice = tensor.insert_slice %3 into %arg4[%arg3, 0] [64, 256] [1, 1] : tensor<64x256xf32> into tensor<256x256xf32>710 scf.yield %insert_slice : tensor<256x256xf32>711 }712 %4 = linalg.mul ins(%1, %arg2 : tensor<256x256xf32>, tensor<256x256xf32>) outs(%dest0 : tensor<256x256xf32>) -> tensor<256x256xf32>713 %5 = linalg.exp ins(%1 : tensor<256x256xf32>) outs(%dest0 : tensor<256x256xf32>) -> tensor<256x256xf32>714 return %4, %5 : tensor<256x256xf32>, tensor<256x256xf32>715 }716}717 718module attributes {transform.with_named_sequence} {719 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {720 %slice_op = transform.structured.match ops{["tensor.insert_slice"]} in %arg1721 : (!transform.any_op) -> !transform.any_op722 %loop = transform.structured.match ops{["scf.for"]} in %arg1723 : (!transform.any_op) -> !transform.any_op724 %a, %b = transform.test.fuse_consumer_using_slice %slice_op in (%loop) num_consumer_to_fuse = 2725 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)726 transform.yield727 }728}729// CHECK: func.func @fuse_add_multiple_tilable_consumers(730// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<256x256xf32>731// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<256x256xf32>732// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<256x256xf32>733// CHECK: %[[dest0:.*]] = tensor.empty() : tensor<256x256xf32>734// CHECK: %[[LOOP_RESULT:.*]]:3 = scf.for %[[IV1:.*]] = %[[C0]]735// CHECK-SAME: iter_args(%[[FIRST_OUT_ARG:.*]] = %[[dest0]], %[[SECOND_OUT_ARG:.*]] = %[[dest0]], %[[THIRD_OUT_ARG:.*]] = %[[dest0]])736// CHECK-SAME: {737// CHECK: %[[ADD_OUT_SLICE:.*]] = tensor.extract_slice %[[FIRST_OUT_ARG]][%[[IV1]], 0] [64, 256] [1, 1]738// CHECK: %[[ADD_INS0_SLICE:.*]] = tensor.extract_slice %[[ARG0]][%[[IV1]], 0] [64, 256] [1, 1]739// CHECK: %[[ADD_INS1_SLICE:.*]] = tensor.extract_slice %[[ARG1]][%[[IV1]], 0] [64, 256] [1, 1]740// CHECK: %[[TILED_ADD_OUT:.*]] = linalg.add741// CHECK-SAME: ins(%[[ADD_INS0_SLICE]], %[[ADD_INS1_SLICE]] :742// CHECK-SAME: outs(%[[ADD_OUT_SLICE]] :743// CHECK: %[[INSERT_ADD:.*]] = tensor.insert_slice %[[TILED_ADD_OUT]] into %[[FIRST_OUT_ARG]][%[[IV1]], 0] [64, 256] [1, 1]744// CHECK: %[[EXP_OUT_SLICE:.*]] = tensor.extract_slice %[[SECOND_OUT_ARG]][%[[IV1]], 0] [64, 256] [1, 1]745// CHECK: %[[TILED_EXP_OUT:.*]] = linalg.exp746// CHECK-SAME: ins(%[[TILED_ADD_OUT]] :747// CHECK-SAME: outs(%[[EXP_OUT_SLICE]] :748// CHECK: %[[MUL_INS2_SLICE:.*]] = tensor.extract_slice %[[ARG2]][%[[IV1]], 0] [64, 256] [1, 1]749// CHECK: %[[MUL_OUT_SLICE:.*]] = tensor.extract_slice %[[THIRD_OUT_ARG]][%[[IV1]], 0] [64, 256] [1, 1]750// CHECK: %[[TILED_MUL_OUT:.*]] = linalg.mul751// CHECK-SAME: ins(%[[TILED_ADD_OUT]], %[[MUL_INS2_SLICE]] :752// CHECK-SAME: outs(%[[MUL_OUT_SLICE]] :753// CHECK: %[[INSERT_EXP:.*]] = tensor.insert_slice %[[TILED_EXP_OUT]] into %[[SECOND_OUT_ARG]][%[[IV1]], 0] [64, 256] [1, 1]754// CHECK: %[[INSERT_MUL:.*]] = tensor.insert_slice %[[TILED_MUL_OUT]] into %[[THIRD_OUT_ARG]][%[[IV1]], 0] [64, 256] [1, 1]755// CHECK: scf.yield %[[INSERT_ADD]], %[[INSERT_EXP]], %[[INSERT_MUL]] :756// CHECK: }757// CHECK: return %[[LOOP_RESULT]]#2, %[[LOOP_RESULT]]#1 :758 759// -----760 761module {762 func.func @no_fuse_only_dps_consumer(%arg0: tensor<256x256xf32>, %arg1: tensor<256x256xf32>, %arg2: tensor<256x256xf32>) -> (tensor<256x256xf32>, tensor<258x258xf32>) {763 %c0 = arith.constant 0 : index764 %c64 = arith.constant 64 : index765 %c256 = arith.constant 256 : index766 %cst = arith.constant 0.000000e+00 : f32767 %dest0 = tensor.empty() : tensor<256x256xf32>768 %1 = scf.for %arg3 = %c0 to %c256 step %c64 iter_args(%arg4 = %dest0) -> (tensor<256x256xf32>) {769 %extracted_slice_1 = tensor.extract_slice %arg4[%arg3, 0] [64, 256] [1, 1] : tensor<256x256xf32> to tensor<64x256xf32>770 %extracted_slice_2 = tensor.extract_slice %arg0[%arg3, 0] [64, 256] [1, 1] : tensor<256x256xf32> to tensor<64x256xf32>771 %extracted_slice_3 = tensor.extract_slice %arg1[%arg3, 0] [64, 256] [1, 1] : tensor<256x256xf32> to tensor<64x256xf32>772 %3 = linalg.add ins(%extracted_slice_2, %extracted_slice_3 : tensor<64x256xf32>, tensor<64x256xf32>) outs(%extracted_slice_1 : tensor<64x256xf32>) -> tensor<64x256xf32>773 %insert_slice = tensor.insert_slice %3 into %arg4[%arg3, 0] [64, 256] [1, 1] : tensor<64x256xf32> into tensor<256x256xf32>774 scf.yield %insert_slice : tensor<256x256xf32>775 }776 %dest1 = tensor.empty() : tensor<258x258xf32>777 %4 = tensor.insert_slice %1 into %dest1[0, 0] [256, 256] [1, 1] : tensor<256x256xf32> into tensor<258x258xf32>778 %5 = linalg.mul ins(%1, %arg2 : tensor<256x256xf32>, tensor<256x256xf32>) outs(%dest0 : tensor<256x256xf32>) -> tensor<256x256xf32>779 return %5, %4 : tensor<256x256xf32>, tensor<258x258xf32>780 }781}782 783module attributes {transform.with_named_sequence} {784 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {785 %slice_ops = transform.structured.match ops{["tensor.insert_slice"]} in %arg1786 : (!transform.any_op) -> !transform.any_op787 %loop = transform.structured.match ops{["scf.for"]} in %arg1 : (!transform.any_op) -> !transform.any_op788 %slice_op, %other_slice = transform.split_handle %slice_ops : (!transform.any_op) -> (!transform.any_op, !transform.any_op)789 %a, %b = transform.test.fuse_consumer_using_slice %slice_op in (%loop) num_consumer_to_fuse = 1790 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)791 transform.yield792 }793}794// CHECK: func.func @no_fuse_only_dps_consumer(795// CHECK: %[[LOOP_RESULT:.*]]:2 = scf.for {{.*}} {796// CHECK: linalg.add797// CHECK: linalg.mul798// CHECK: scf.yield799// CHECK: }800// CHECK: %[[RES_SLICE:.+]] = tensor.insert_slice801// CHECK: return %[[LOOP_RESULT]]#1, %[[RES_SLICE]]802 803// -----804 805#map = affine_map<(d0, d1, d2) -> (d0, d1)>806#map1 = affine_map<(d0, d1, d2) -> (d2)>807#map2 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>808module {809 func.func @fuse_with_tilable_consumer_with_projected_permutations(%arg0: tensor<256x256xf32>, %arg1: tensor<256x256xf32>, %arg2: tensor<24xf32>) -> tensor<256x256x24xf32> {810 %c0 = arith.constant 0 : index811 %c64 = arith.constant 64 : index812 %c256 = arith.constant 256 : index813 %0 = tensor.empty() : tensor<256x256xf32>814 %1 = scf.for %arg3 = %c0 to %c256 step %c64 iter_args(%arg4 = %0) -> (tensor<256x256xf32>) {815 %extracted_slice = tensor.extract_slice %arg4[%arg3, 0] [64, 256] [1, 1] : tensor<256x256xf32> to tensor<64x256xf32>816 %extracted_slice_0 = tensor.extract_slice %arg0[%arg3, 0] [64, 256] [1, 1] : tensor<256x256xf32> to tensor<64x256xf32>817 %extracted_slice_1 = tensor.extract_slice %arg1[%arg3, 0] [64, 256] [1, 1] : tensor<256x256xf32> to tensor<64x256xf32>818 %4 = linalg.add ins(%extracted_slice_0, %extracted_slice_1 : tensor<64x256xf32>, tensor<64x256xf32>) outs(%extracted_slice : tensor<64x256xf32>) -> tensor<64x256xf32>819 %inserted_slice = tensor.insert_slice %4 into %arg4[%arg3, 0] [64, 256] [1, 1] : tensor<64x256xf32> into tensor<256x256xf32>820 scf.yield %inserted_slice : tensor<256x256xf32>821 }822 %2 = tensor.empty() : tensor<256x256x24xf32>823 %3 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel"]} ins(%1, %arg2 : tensor<256x256xf32>, tensor<24xf32>) outs(%2 : tensor<256x256x24xf32>) {824 ^bb0(%in: f32, %in_0: f32, %out: f32):825 %4 = arith.addf %in, %in_0 : f32826 linalg.yield %4 : f32827 } -> tensor<256x256x24xf32>828 return %3 : tensor<256x256x24xf32>829 }830}831 832// CHECK: func.func @fuse_with_tilable_consumer_with_projected_permutations(%[[VAL_0:.*]]: tensor<256x256xf32>, %[[VAL_1:.*]]: tensor<256x256xf32>, %[[VAL_2:.*]]: tensor<24xf32>) -> tensor<256x256x24xf32> {833// CHECK: %[[VAL_3:.*]] = arith.constant 0 : index834// CHECK: %[[VAL_4:.*]] = arith.constant 64 : index835// CHECK: %[[VAL_5:.*]] = arith.constant 256 : index836// CHECK: %[[VAL_6:.*]] = tensor.empty() : tensor<256x256xf32>837// CHECK: %[[VAL_7:.*]] = tensor.empty() : tensor<256x256x24xf32>838// CHECK: %[[VAL_8:.*]]:2 = scf.for %[[VAL_9:.*]] = %[[VAL_3]] to %[[VAL_5]] step %[[VAL_4]] iter_args(%[[VAL_10:.*]] = %[[VAL_6]], %[[VAL_11:.*]] = %[[VAL_7]]) -> (tensor<256x256xf32>, tensor<256x256x24xf32>) {839// CHECK: %[[VAL_12:.*]] = tensor.extract_slice %[[VAL_10]]{{\[}}%[[VAL_9]], 0] [64, 256] [1, 1]840// CHECK: %[[VAL_13:.*]] = tensor.extract_slice %[[VAL_0]]{{\[}}%[[VAL_9]], 0] [64, 256] [1, 1]841// CHECK: %[[VAL_14:.*]] = tensor.extract_slice %[[VAL_1]]{{\[}}%[[VAL_9]], 0] [64, 256] [1, 1]842// CHECK: %[[VAL_15:.*]] = linalg.add ins(%[[VAL_13]], %[[VAL_14]] : tensor<64x256xf32>, tensor<64x256xf32>) outs(%[[VAL_12]] : tensor<64x256xf32>) -> tensor<64x256xf32>843// CHECK: %[[VAL_16:.*]] = tensor.insert_slice %[[VAL_15]] into %[[VAL_10]]{{\[}}%[[VAL_9]], 0] [64, 256] [1, 1]844// CHECK: %[[VAL_17:.*]] = tensor.extract_slice %[[VAL_2]][0] [24] [1] : tensor<24xf32> to tensor<24xf32>845// CHECK: %[[VAL_18:.*]] = tensor.extract_slice %[[VAL_11]]{{\[}}%[[VAL_9]], 0, 0] [64, 256, 24] [1, 1, 1]846// CHECK: %[[VAL_19:.*]] = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel"]} ins(%[[VAL_15]], %[[VAL_17]] : tensor<64x256xf32>, tensor<24xf32>) outs(%[[VAL_18]] : tensor<64x256x24xf32>) {847// CHECK: ^bb0(%[[VAL_20:.*]]: f32, %[[VAL_21:.*]]: f32, %[[VAL_22:.*]]: f32):848// CHECK: %[[VAL_23:.*]] = arith.addf %[[VAL_20]], %[[VAL_21]] : f32849// CHECK: linalg.yield %[[VAL_23]] : f32850// CHECK: } -> tensor<64x256x24xf32>851// CHECK: %[[VAL_24:.*]] = tensor.insert_slice %[[VAL_25:.*]] into %[[VAL_11]]{{\[}}%[[VAL_9]], 0, 0] [64, 256, 24] [1, 1, 1]852// CHECK: scf.yield %[[VAL_16]], %[[VAL_24]] : tensor<256x256xf32>, tensor<256x256x24xf32>853// CHECK: }854 855module attributes {transform.with_named_sequence} {856 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {857 %slice_op = transform.structured.match ops{["tensor.insert_slice"]} in %arg1858 : (!transform.any_op) -> !transform.any_op859 %loop = transform.structured.match ops{["scf.for"]} in %arg1860 : (!transform.any_op) -> !transform.any_op861 %a, %b = transform.test.fuse_consumer_using_slice %slice_op in (%loop) num_consumer_to_fuse = 1862 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)863 transform.yield864 }865}866 867// -----868 869func.func @multi_slice_fusion1(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>, %arg3 : index) -> tensor<?xf32> {870 %c0 = arith.constant 0 : index871 %c1 = arith.constant 1 : index872 %dim0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>873 %dim1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>874 %loop:2 = scf.forall (%iv0) = (%c0) to (%dim0) step (%arg3) shared_outs(%init0 = %arg1, %init1 = %arg2) -> (tensor<?xf32>, tensor<?xf32>) {875 %tilesize = affine.min affine_map<(d0)[s0, s1] -> (s1, s0 - d0)>(%iv0)[%dim0, %arg3]876 %arg0_slice = tensor.extract_slice %arg0[%iv0, 0] [%tilesize, %dim1] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>877 %init0_slice = tensor.extract_slice %init0[%iv0] [%tilesize] [1] : tensor<?xf32> to tensor<?xf32>878 %init1_slice = tensor.extract_slice %init1[%iv0] [%tilesize] [1] : tensor<?xf32> to tensor<?xf32>879 %generic:2 = linalg.generic {880 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0)>],881 iterator_types = ["parallel", "reduction"]}882 ins(%arg0_slice : tensor<?x?xf32>) outs(%init0_slice, %init1_slice : tensor<?xf32>, tensor<?xf32>) {883 ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):884 %0 = arith.mulf %b0, %b1 : f32885 %1 = arith.addf %b0, %b2 : f32886 linalg.yield %0, %1 : f32, f32887 } -> (tensor<?xf32>, tensor<?xf32>)888 scf.forall.in_parallel {889 tensor.parallel_insert_slice %generic#0 into %init0[%iv0] [%tilesize] [1] : tensor<?xf32> into tensor<?xf32>890 tensor.parallel_insert_slice %generic#1 into %init1[%iv0] [%tilesize] [1] : tensor<?xf32> into tensor<?xf32>891 }892 }893 %empty = tensor.empty(%dim0) : tensor<?xf32>894 %result = linalg.generic {895 indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],896 iterator_types = ["parallel"]}897 ins(%loop#0, %loop#1 : tensor<?xf32>, tensor<?xf32>) outs(%empty : tensor<?xf32>) {898 ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):899 %0 = arith.addf %b0, %b1 : f32900 linalg.yield %0 : f32901 } -> tensor<?xf32>902 return %result : tensor<?xf32>903}904// CHECK-LABEL: func @multi_slice_fusion1(905// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32>906// CHECK: %[[C0:.+]] = arith.constant 0907// CHECK: %[[DIM0:.+]] = tensor.dim %[[ARG0]], %[[C0]]908// CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM0]])909// CHECK: %[[RESULT:.+]]:3 = scf.forall (%[[IV:.+]]) =910// CHECK-SAME: , %[[INIT:[a-zA-Z0-9]+]] = %[[EMPTY]])911// CHECK: %[[TILESIZE:.+]] = affine.min912// CHECK-DAG: %[[GENERIC:.+]]:2 = linalg.generic913// CHECK-DAG: %[[INIT_SLICE:.+]] = tensor.extract_slice %[[INIT]][%[[IV]]] [%[[TILESIZE]]]914// CHECK: %[[FUSED:.+]] = linalg.generic915// CHECK-SAME: ins(%[[GENERIC]]#0, %[[GENERIC]]#1 :916// CHECK: tensor.parallel_insert_slice %[[FUSED]] into %[[INIT]][%[[IV]]] [%[[TILESIZE]]]917// CHECK: return %[[RESULT]]#2918 919module attributes {transform.with_named_sequence} {920 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {921 %loop = transform.structured.match ops{["scf.forall"]} in %arg1922 : (!transform.any_op) -> !transform.any_op923 %yield = transform.structured.match ops{["tensor.parallel_insert_slice"]} in %arg1924 : (!transform.any_op) -> !transform.any_op925 %yield0, %yield1 = transform.split_handle %yield : (!transform.any_op) -> (!transform.any_op, !transform.any_op)926 %a, %b = transform.test.fuse_consumer_using_slice %yield0, %yield1 in (%loop)927 : (!transform.any_op, !transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)928 transform.yield929 }930}931 932// -----933 934// Check that when the given operand tiles are inconsistent, tiling fails.935 936func.func @multi_slice_fusion2(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>, %arg3 : index) -> tensor<?xf32> {937 %c0 = arith.constant 0 : index938 %c1 = arith.constant 1 : index939 %dim0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>940 %dim1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>941 %loop:2 = scf.forall (%iv0) = (%c0) to (%dim0) step (%arg3) shared_outs(%init0 = %arg1, %init1 = %arg2) -> (tensor<?xf32>, tensor<?xf32>) {942 %tilesize = affine.min affine_map<(d0)[s0, s1] -> (s1, s0 - d0)>(%iv0)[%dim0, %arg3]943 %arg0_slice = tensor.extract_slice %arg0[%iv0, 0] [%tilesize, %dim1] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>944 %init0_slice = tensor.extract_slice %init0[%iv0] [%tilesize] [1] : tensor<?xf32> to tensor<?xf32>945 %generic0 = linalg.generic {946 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>],947 iterator_types = ["parallel", "reduction"]}948 ins(%arg0_slice : tensor<?x?xf32>) outs(%init0_slice : tensor<?xf32>) {949 ^bb0(%b0 : f32, %b1 : f32):950 %0 = arith.mulf %b0, %b1 : f32951 linalg.yield %0 : f32952 } -> tensor<?xf32>953 %init1_slice = tensor.extract_slice %init1[%iv0] [%tilesize] [1] : tensor<?xf32> to tensor<?xf32>954 %generic1 = linalg.generic {955 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>],956 iterator_types = ["parallel", "reduction"]}957 ins(%arg0_slice : tensor<?x?xf32>) outs(%init1_slice: tensor<?xf32>) {958 ^bb0(%b0 : f32, %b1 : f32):959 %0 = arith.addf %b0, %b1 : f32960 linalg.yield %0: f32961 } -> tensor<?xf32>962 scf.forall.in_parallel {963 tensor.parallel_insert_slice %generic0 into %init0[%iv0] [%tilesize] [1] : tensor<?xf32> into tensor<?xf32>964 tensor.parallel_insert_slice %generic1 into %init1[%iv0] [%tilesize] [1] : tensor<?xf32> into tensor<?xf32>965 }966 }967 %empty = tensor.empty(%dim0) : tensor<?xf32>968 %result = linalg.generic {969 indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],970 iterator_types = ["parallel"]}971 ins(%loop#0, %loop#1 : tensor<?xf32>, tensor<?xf32>) outs(%empty : tensor<?xf32>) {972 ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):973 %0 = arith.addf %b0, %b1 : f32974 linalg.yield %0 : f32975 } -> tensor<?xf32>976 return %result : tensor<?xf32>977}978// CHECK-LABEL: func @multi_slice_fusion2(979// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32>980// CHECK: %[[C0:.+]] = arith.constant 0981// CHECK: %[[DIM0:.+]] = tensor.dim %[[ARG0]], %[[C0]]982// CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM0]])983// CHECK: %[[RESULT:.+]]:3 = scf.forall (%[[IV:.+]]) =984// CHECK-SAME: , %[[INIT:[a-zA-Z0-9]+]] = %[[EMPTY]])985// CHECK: %[[TILESIZE:.+]] = affine.min986// CHECK: %[[GENERIC0:.+]] = linalg.generic987// CHECK: %[[GENERIC1:.+]] = linalg.generic988// CHECK-DAG: %[[INIT_SLICE:.+]] = tensor.extract_slice %[[INIT]][%[[IV]]] [%[[TILESIZE]]]989// CHECK: %[[FUSED:.+]] = linalg.generic990// CHECK-SAME: ins(%[[GENERIC0]], %[[GENERIC1]] :991// CHECK: tensor.parallel_insert_slice %[[FUSED]] into %[[INIT]][%[[IV]]] [%[[TILESIZE]]]992// CHECK: return %[[RESULT]]#2993 994module attributes {transform.with_named_sequence} {995 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {996 %loop = transform.structured.match ops{["scf.forall"]} in %arg1997 : (!transform.any_op) -> !transform.any_op998 %yield = transform.structured.match ops{["tensor.parallel_insert_slice"]} in %arg1999 : (!transform.any_op) -> !transform.any_op1000 %yield0, %yield1 = transform.split_handle %yield : (!transform.any_op) -> (!transform.any_op, !transform.any_op)1001 %a, %b = transform.test.fuse_consumer_using_slice %yield0, %yield1 in (%loop)1002 : (!transform.any_op, !transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)1003 transform.yield1004 }1005}1006 1007// -----1008 1009func.func @multi_slice_fusion_with_broadcast(%arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?xf32>,1010 %arg3 : index, %arg4 : index) -> tensor<?x?xf32> {1011 %c0 = arith.constant 0 : index1012 %c1 = arith.constant 1 : index1013 %c2 = arith.constant 2 : index1014 %dim0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>1015 %dim1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>1016 %dim2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32>1017 %loop:2 = scf.forall (%iv0, %iv1) = (%c0, %c0) to (%dim0, %dim1) step (%arg3, %arg4)1018 shared_outs(%init0 = %arg1, %init1 = %arg2) -> (tensor<?x?xf32>, tensor<?xf32>) {1019 %tilesize0 = affine.min affine_map<(d0)[s0, s1] -> (s1, s0 - d0)>(%iv0)[%dim0, %arg3]1020 %tilesize1 = affine.min affine_map<(d0)[s0, s1] -> (s1, s0 - d0)>(%iv1)[%dim1, %arg4]1021 %arg0_slice = tensor.extract_slice %arg0[%iv0, %iv1, 0] [%tilesize0, %tilesize1, %dim2] [1, 1, 1]1022 : tensor<?x?x?xf32> to tensor<?x?x?xf32>1023 %init0_slice = tensor.extract_slice %init0[%iv0, %iv1] [%tilesize0, %tilesize1] [1, 1]1024 : tensor<?x?xf32> to tensor<?x?xf32>1025 %generic0 = linalg.generic {1026 indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>],1027 iterator_types = ["parallel", "parallel", "reduction"]}1028 ins(%arg0_slice : tensor<?x?x?xf32>) outs(%init0_slice : tensor<?x?xf32>) {1029 ^bb0(%b0 : f32, %b1 : f32):1030 %0 = arith.mulf %b0, %b1 : f321031 linalg.yield %0 : f321032 } -> tensor<?x?xf32>1033 %init1_slice = tensor.extract_slice %init1[%iv0] [%tilesize0] [1] : tensor<?xf32> to tensor<?xf32>1034 %generic1 = linalg.generic {1035 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>],1036 iterator_types = ["parallel", "reduction"]}1037 ins(%generic0 : tensor<?x?xf32>) outs(%init1_slice: tensor<?xf32>) {1038 ^bb0(%b0 : f32, %b1 : f32):1039 %0 = arith.addf %b0, %b1 : f321040 linalg.yield %0: f321041 } -> tensor<?xf32>1042 scf.forall.in_parallel {1043 tensor.parallel_insert_slice %generic0 into %init0[%iv0, %iv1] [%tilesize0, %tilesize1] [1, 1]1044 : tensor<?x?xf32> into tensor<?x?xf32>1045 tensor.parallel_insert_slice %generic1 into %init1[%iv0] [%tilesize0] [1] : tensor<?xf32> into tensor<?xf32>1046 }1047 }1048 %empty = tensor.empty(%dim0, %dim1) : tensor<?x?xf32>1049 %result = linalg.generic {1050 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0, d1)>],1051 iterator_types = ["parallel", "parallel"]}1052 ins(%loop#0, %loop#1 : tensor<?x?xf32>, tensor<?xf32>) outs(%empty : tensor<?x?xf32>) {1053 ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):1054 %0 = arith.addf %b0, %b1 : f321055 linalg.yield %0 : f321056 } -> tensor<?x?xf32>1057 return %result : tensor<?x?xf32>1058}1059module attributes {transform.with_named_sequence} {1060 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {1061 %loop = transform.structured.match ops{["scf.forall"]} in %arg11062 : (!transform.any_op) -> !transform.any_op1063 %yield = transform.structured.match ops{["tensor.parallel_insert_slice"]} in %arg11064 : (!transform.any_op) -> !transform.any_op1065 %yield0, %yield1 = transform.split_handle %yield : (!transform.any_op) -> (!transform.any_op, !transform.any_op)1066 %a, %b = transform.test.fuse_consumer_using_slice %yield0, %yield1 in (%loop)1067 : (!transform.any_op, !transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)1068 transform.yield1069 }1070}1071// CHECK-LABEL: func @multi_slice_fusion_with_broadcast(1072// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32>1073// CHECK-DAG: %[[C0:.+]] = arith.constant 01074// CHECK-DAG: %[[C1:.+]] = arith.constant 11075// CHECK-DAG: %[[DIM0:.+]] = tensor.dim %[[ARG0]], %[[C0]]1076// CHECK-DAG: %[[DIM1:.+]] = tensor.dim %[[ARG0]], %[[C1]]1077// CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM0]], %[[DIM1]])1078// CHECK: %[[RESULT:.+]]:3 = scf.forall (%[[IV0:[a-zA-Z0-9]+]], %[[IV1:[a-zA-Z0-9]+]]) =1079// CHECK-SAME: , %[[INIT:[a-zA-Z0-9]+]] = %[[EMPTY]])1080// CHECK-DAG: %[[TILESIZE0:.+]] = affine.min {{.+}}(%[[IV0]])1081// CHECK-DAG: %[[TILESIZE1:.+]] = affine.min {{.+}}(%[[IV1]])1082// CHECK: %[[GENERIC0:.+]] = linalg.generic1083// CHECK: %[[GENERIC1:.+]] = linalg.generic1084// CHECK-DAG: %[[INIT_SLICE:.+]] = tensor.extract_slice %[[INIT]][%[[IV0]], %[[IV1]]] [%[[TILESIZE0]], %[[TILESIZE1]]]1085// CHECK: %[[FUSED:.+]] = linalg.generic1086// CHECK-SAME: ins(%[[GENERIC0]], %[[GENERIC1]] :1087// CHECK: tensor.parallel_insert_slice %[[FUSED]] into %[[INIT]][%[[IV0]], %[[IV1]]] [%[[TILESIZE0]], %[[TILESIZE1]]]1088// CHECK: return %[[RESULT]]#21089 1090// -----1091 1092func.func @multi_slice_fusion_invalid(%arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>,1093 %arg3 : index, %arg4 : index) -> tensor<?x?xf32> {1094 %c0 = arith.constant 0 : index1095 %c1 = arith.constant 1 : index1096 %c2 = arith.constant 2 : index1097 %dim0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>1098 %dim1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>1099 %dim2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32>1100 %loop:2 = scf.forall (%iv0, %iv1) = (%c0, %c0) to (%dim0, %dim1) step (%arg3, %arg4)1101 shared_outs(%init0 = %arg1, %init1 = %arg2) -> (tensor<?x?xf32>, tensor<?x?xf32>) {1102 %tilesize0 = affine.min affine_map<(d0)[s0, s1] -> (s1, s0 - d0)>(%iv0)[%dim0, %arg3]1103 %tilesize1 = affine.min affine_map<(d0)[s0, s1] -> (s1, s0 - d0)>(%iv1)[%dim1, %arg4]1104 %arg0_slice = tensor.extract_slice %arg0[%iv0, %iv1, 0] [%tilesize0, %tilesize1, %dim2] [1, 1, 1]1105 : tensor<?x?x?xf32> to tensor<?x?x?xf32>1106 %init0_slice = tensor.extract_slice %init0[%iv0, %iv1] [%tilesize0, %tilesize1] [1, 1]1107 : tensor<?x?xf32> to tensor<?x?xf32>1108 %generic0 = linalg.generic {1109 indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>],1110 iterator_types = ["parallel", "parallel", "reduction"]}1111 ins(%arg0_slice : tensor<?x?x?xf32>) outs(%init0_slice : tensor<?x?xf32>) {1112 ^bb0(%b0 : f32, %b1 : f32):1113 %0 = arith.mulf %b0, %b1 : f321114 linalg.yield %0 : f321115 } -> tensor<?x?xf32>1116 %init1_slice = tensor.extract_slice %init1[%iv0, %iv1] [%tilesize0, %tilesize1] [1, 1]1117 : tensor<?x?xf32> to tensor<?x?xf32>1118 %generic1 = linalg.generic {1119 indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>],1120 iterator_types = ["parallel", "parallel", "reduction"]}1121 ins(%arg0_slice : tensor<?x?x?xf32>) outs(%init1_slice: tensor<?x?xf32>) {1122 ^bb0(%b0 : f32, %b1 : f32):1123 %0 = arith.addf %b0, %b1 : f321124 linalg.yield %0: f321125 } -> tensor<?x?xf32>1126 scf.forall.in_parallel {1127 // expected-error @below {{failed to fuse consumer of slice}}1128 tensor.parallel_insert_slice %generic0 into %init0[%iv0, %iv1] [%tilesize0, %tilesize1] [1, 1]1129 : tensor<?x?xf32> into tensor<?x?xf32>1130 tensor.parallel_insert_slice %generic1 into %init1[%iv0, %iv1] [%tilesize0, %tilesize1] [1, 1]1131 : tensor<?x?xf32> into tensor<?x?xf32>1132 }1133 }1134 %empty = tensor.empty(%dim0, %dim1) : tensor<?x?xf32>1135 %result = linalg.generic {1136 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d0, d1)>],1137 iterator_types = ["parallel", "parallel"]}1138 ins(%loop#0, %loop#1 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%empty : tensor<?x?xf32>) {1139 ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):1140 %0 = arith.addf %b0, %b1 : f321141 linalg.yield %0 : f321142 } -> tensor<?x?xf32>1143 return %result : tensor<?x?xf32>1144}1145module attributes {transform.with_named_sequence} {1146 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {1147 %loop = transform.structured.match ops{["scf.forall"]} in %arg11148 : (!transform.any_op) -> !transform.any_op1149 %yield = transform.structured.match ops{["tensor.parallel_insert_slice"]} in %arg11150 : (!transform.any_op) -> !transform.any_op1151 %yield0, %yield1 = transform.split_handle %yield : (!transform.any_op) -> (!transform.any_op, !transform.any_op)1152 %a, %b = transform.test.fuse_consumer_using_slice %yield0, %yield1 in (%loop)1153 : (!transform.any_op, !transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)1154 transform.yield1155 }1156}1157