1155 lines · plain
1// RUN: mlir-opt --transform-interpreter --cse --split-input-file --verify-diagnostics --mlir-print-local-scope %s | FileCheck %s2 3#map = affine_map<(d0) -> (d0)>4module {5 func.func @fuse_tilable_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 %add = transform.structured.match ops{["linalg.add"]} in %arg132 : (!transform.any_op) -> !transform.any_op33 %a, %new_loop = transform.test.fuse_consumer %add into (%loop)34 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)35 transform.yield36 }37}38// CHECK: func.func @fuse_tilable_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 63#map = affine_map<(d0) -> (d0)>64module {65 func.func @fuse_tilable_consumer_nested_scf_for(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2 : tensor<?x?xf32>,66 %lb0 : index, %ub0 : index, %step0 : index,67 %lb1 : index, %ub1 : index, %step1 : index) -> tensor<?x?xf32> {68 %0 = scf.for %arg3 = %lb0 to %ub0 step %step0 iter_args(%init0 = %arg0) -> tensor<?x?xf32> {69 %1 = scf.for %arg4 = %lb1 to %ub1 step %step1 iter_args(%init1 = %init0) -> tensor<?x?xf32> {70 %extracted_slice = tensor.extract_slice %init1[%arg3, %arg4] [%step0, %step1] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>71 %2 = tensor.insert_slice %extracted_slice into %init1[%arg3, %arg4] [%step0, %step1] [1, 1] : tensor<?x?xf32> into tensor<?x?xf32>72 scf.yield %2 : tensor<?x?xf32>73 }74 scf.yield %1 : tensor<?x?xf32>75 }76 %2 = linalg.add ins(%0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>77 return %2 : tensor<?x?xf32>78 }79}80 81module attributes {transform.with_named_sequence} {82 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {83 %loops = transform.structured.match ops{["scf.for"]} in %arg184 : (!transform.any_op) -> !transform.any_op85 %loop0, %loop1 = transform.split_handle %loops86 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)87 %add = transform.structured.match ops{["linalg.add"]} in %arg188 : (!transform.any_op) -> !transform.any_op89 %a, %new_loop0, %new_loop1 = transform.test.fuse_consumer %add into (%loop0, %loop1)90 : (!transform.any_op, !transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)91 transform.yield92 }93}94// CHECK: func @fuse_tilable_consumer_nested_scf_for(95// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?xf32>96// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>97// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: tensor<?x?xf32>98// CHECK: %[[OUTER_RESULT:.+]]:2 = scf.for99// CHECK-SAME: iter_args(%[[INIT00:[a-zA-Z0-9_]+]] = %[[ARG0]], %[[INIT01:[a-zA-Z0-9_]+]] = %[[ARG2]])100// CHECK: %[[INNER_RESULT:.+]]:2 = scf.for101// CHECK-SAME: iter_args(%[[INIT10:[a-zA-Z0-9_]+]] = %[[INIT00]], %[[INIT11:[a-zA-Z0-9_]+]] = %[[INIT01]])102// CHECK-DAG: %[[OPERAND1:.+]] = tensor.extract_slice %[[INIT10]]103// CHECK-DAG: %[[OLD_INSERT_SLICE:.+]] = tensor.insert_slice %[[OPERAND1]] into %[[INIT10]]104// CHECK-DAG: %[[OPERAND2:.+]] = tensor.extract_slice %[[ARG1]]105// CHECK-DAG: %[[INIT:.+]] = tensor.extract_slice %[[INIT11]]106// CHECK: %[[ADD:.+]] = linalg.add107// CHECK-SAME: ins(%[[OPERAND1]], %[[OPERAND2]] :108// CHECK-SAME: outs(%[[INIT]] :109// CHECK: %[[INSERT_SLICE:.+]] = tensor.insert_slice %[[ADD]] into %[[INIT11]]110// CHECK: scf.yield %[[OLD_INSERT_SLICE]], %[[INSERT_SLICE]]111// CHECK: scf.yield %[[INNER_RESULT]]#0, %[[INNER_RESULT]]#1112// CHECK: return %[[OUTER_RESULT]]#1113 114// -----115 116module {117 func.func @fuse_tilable_consumer_scf_forall(%arg0: tensor<32x32xf32>, %arg1: tensor<32x32xf32>, %arg2: tensor<64x64xf32>) -> tensor<64x64xf32> {118 %c4 = arith.constant 4 : index119 %c64 = arith.constant 64 : index120 %c0 = arith.constant 0 : index121 %1:2 = scf.forall (%arg3, %arg4) in (2, 2) shared_outs(%arg5 = %arg2, %arg6 = %arg2) -> (tensor<64x64xf32>, tensor<64x64xf32>) {122 %extracted_slice = tensor.extract_slice %arg5[%arg3, %arg4] [32, 32] [1, 1] : tensor<64x64xf32> to tensor<32x32xf32>123 %extracted_slice_1 = tensor.extract_slice %arg6[%arg3, %arg4] [32, 32] [1, 1] : tensor<64x64xf32> to tensor<32x32xf32>124 %3 = linalg.matmul ins(%arg0, %arg1 : tensor<32x32xf32>, tensor<32x32xf32>) outs(%extracted_slice : tensor<32x32xf32>) -> tensor<32x32xf32>125 scf.forall.in_parallel {126 tensor.parallel_insert_slice %3 into %arg6[%arg3, %arg4] [32, 32] [1, 1] : tensor<32x32xf32> into tensor<64x64xf32>127 tensor.parallel_insert_slice %extracted_slice_1 into %arg5[%arg3, %arg4] [32, 32] [1, 1] : tensor<32x32xf32> into tensor<64x64xf32>128 }129 }130 %in_operand_2 = tensor.empty() : tensor<64x64xf32>131 %out_operand_3 = tensor.empty() : tensor<64x64xf32>132 %2 = linalg.add ins(%1#1, %in_operand_2 : tensor<64x64xf32>, tensor<64x64xf32>) outs(%out_operand_3 : tensor<64x64xf32>) -> tensor<64x64xf32>133 return %2 : tensor<64x64xf32>134 }135}136 137module attributes {transform.with_named_sequence} {138 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {139 %add = transform.structured.match ops{["linalg.add"]} in %arg1140 : (!transform.any_op) -> !transform.any_op141 %loop = transform.structured.match ops{["scf.forall"]} in %arg1142 : (!transform.any_op) -> !transform.any_op143 %a, %new_loop = transform.test.fuse_consumer %add into (%loop)144 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)145 transform.yield146 }147}148// CHECK: func.func @fuse_tilable_consumer_scf_forall(149// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<32x32xf32>150// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<32x32xf32>151// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<64x64xf32>)152// CHECK: %[[OUT_INIT:.*]] = tensor.empty() : tensor<64x64xf32>153// CHECK: %[[FINAL_RESULT:.*]]:3 = scf.forall (%[[IV1:.*]], %[[IV2:.*]]) in (2, 2)154// CHECK-SAME: shared_outs(%[[FIRST_OUT_ARG:.*]] = %[[ARG2]], %[[SECOND_OUT_ARG:.*]] = %[[ARG2]], %[[ELEM_OUT_ARG:.*]] = %[[OUT_INIT]])155// CHECK-SAME: {156// CHECK: %[[MAT_OUT_SLICE:.*]] = tensor.extract_slice %[[FIRST_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]157// CHECK: %[[SECOND_ARG_SLICE:.*]] = tensor.extract_slice %[[SECOND_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]158// CHECK: %[[MAT_OUT:.*]] = linalg.matmul159// CHECK-SAME: outs(%[[MAT_OUT_SLICE]] :160// CHECK: %[[SLICE_OPERAND2:.*]] = tensor.extract_slice %[[OUT_INIT]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]161// CHECK: %[[SLICE_OUT:.*]] = tensor.extract_slice %[[ELEM_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]162// CHECK: %[[ELEM_OUT:.*]] = linalg.add163// CHECK-SAME: ins(%[[MAT_OUT]], %[[SLICE_OPERAND2]] :164// CHECK-SAME: outs(%[[SLICE_OUT]] :165// CHECK: scf.forall.in_parallel {166// CHECK: tensor.parallel_insert_slice %[[MAT_OUT]] into %[[SECOND_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]167// CHECK: tensor.parallel_insert_slice %[[SECOND_ARG_SLICE]] into %[[FIRST_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]168// CHECK: tensor.parallel_insert_slice %[[ELEM_OUT]] into %[[ELEM_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]169// CHECK: }170// CHECK: }171// CHECK: return %[[FINAL_RESULT]]#2 :172 173// -----174 175#map = affine_map<(d0) -> (d0)>176module {177 func.func @fuse_tilable_consumer_scf_for_multi_yielding_consumer(%arg0: tensor<32xf32>, %arg1: tensor<32xf32>, %arg2: tensor<64xf32>) -> tensor<64xf32> {178 %c4 = arith.constant 4 : index179 %c64 = arith.constant 64 : index180 %c0 = arith.constant 0 : index181 %1:2 = scf.for %arg3 = %c0 to %c64 step %c4 iter_args(%arg4 = %arg2, %arg5 = %arg2) -> (tensor<64xf32>, tensor<64xf32>) {182 %extracted_slice = tensor.extract_slice %arg4[%arg3] [32] [1] : tensor<64xf32> to tensor<32xf32>183 %3 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} ins(%arg0, %arg1 : tensor<32xf32>, tensor<32xf32>) outs(%extracted_slice : tensor<32xf32>) {184 ^bb0(%in: f32, %in_16: f32, %out: f32):185 %13 = arith.mulf %in, %in_16 : f32186 %14 = arith.addf %out, %13 : f32187 linalg.yield %14 : f32188 } -> tensor<32xf32>189 %4 = tensor.insert_slice %3 into %arg4[%arg3] [32] [1] : tensor<32xf32> into tensor<64xf32>190 scf.yield %arg5, %4 : tensor<64xf32>, tensor<64xf32>191 }192 %in_operand_2 = tensor.empty() : tensor<64xf32>193 %out_operand_3 = tensor.empty() : tensor<64xf32>194 %out_operand_4 = tensor.empty() : tensor<64xf32>195 %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>) {196 ^bb0(%in: f32, %in_16: f32, %out_0: f32, %out_1: f32):197 %13 = arith.mulf %in, %in_16 : f32198 %14 = arith.subf %out_0, %13 : f32199 %15 = arith.addf %out_1, %in : f32200 linalg.yield %14, %15 : f32, f32201 } -> (tensor<64xf32>, tensor<64xf32>)202 return %2#1 : tensor<64xf32>203 }204}205 206module attributes {transform.with_named_sequence} {207 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {208 %generics = transform.structured.match ops{["linalg.generic"]} in %arg1209 : (!transform.any_op) -> !transform.any_op210 %producer, %consumer = transform.split_handle %generics211 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)212 %loop = transform.structured.match ops{["scf.for"]} in %arg1213 : (!transform.any_op) -> !transform.any_op214 %a, %new_loop = transform.test.fuse_consumer %consumer into (%loop)215 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)216 transform.yield217 }218}219// CHECK: func.func @fuse_tilable_consumer_scf_for_multi_yielding_consumer(220// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<32xf32>221// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<32xf32>222// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<64xf32>)223// CHECK: %[[C0:.*]] = arith.constant 0 : index224// CHECK: %0 = tensor.empty() : tensor<64xf32>225// CHECK: %[[FINAL_RESULT:.*]]:4 = scf.for %[[IV:.*]] = %[[C0]]226// CHECK-SAME: iter_args(%[[FIRST_OUT_ARG:.*]] = %[[ARG2]], %[[SECOND_OUT_ARG:.*]] = %[[ARG2]], %[[ELEM_OUT_ARG_0:.*]] = %0, %[[ELEM_OUT_ARG_1:.*]] = %0)227// CHECK-SAME: {228// CHECK: %[[MAT_OUT_SLICE:.*]] = tensor.extract_slice %[[FIRST_OUT_ARG]][%[[IV]]] [32] [1]229// CHECK: %[[MAT_OUT:.*]] = linalg.generic230// CHECK-SAME: outs(%[[MAT_OUT_SLICE]] : tensor<32xf32>)231// CHECK: %[[INSERT_MAT:.*]] = tensor.insert_slice %[[MAT_OUT]] into %[[FIRST_OUT_ARG]][%[[IV]]] [32] [1]232// CHECK: %[[SLICE_OPERAND2:.*]] = tensor.extract_slice %0[%[[IV]]] [32] [1]233// CHECK: %[[SLICE_OUT_0:.*]] = tensor.extract_slice %[[ELEM_OUT_ARG_0]][%[[IV]]] [32] [1]234// CHECK: %[[SLICE_OUT_1:.*]] = tensor.extract_slice %[[ELEM_OUT_ARG_1]][%[[IV]]] [32] [1]235// CHECK: %[[ELEM_OUT:.*]]:2 = linalg.generic236// CHECK-SAME: ins(%[[MAT_OUT]], %[[SLICE_OPERAND2]] :237// CHECK-SAME: outs(%[[SLICE_OUT_0]], %[[SLICE_OUT_1]] :238// CHECK: %[[INSERT_ELEM_0:.*]] = tensor.insert_slice %[[ELEM_OUT]]#0 into %[[ELEM_OUT_ARG_0]][%[[IV]]] [32] [1]239// CHECK: %[[INSERT_ELEM_1:.*]] = tensor.insert_slice %[[ELEM_OUT]]#1 into %[[ELEM_OUT_ARG_1]][%[[IV]]] [32] [1]240// CHECK: scf.yield %[[SECOND_OUT_ARG]], %[[INSERT_MAT]], %[[INSERT_ELEM_0]], %[[INSERT_ELEM_1]] :241// CHECK: }242// CHECK: return %[[FINAL_RESULT]]#3 :243 244// -----245 246#map = affine_map<(d0, d1) -> (d0, d1)>247module {248 func.func @fuse_tilable_consumer_scf_forall_multi_yielding_consumer(%arg0: tensor<32x32xf32>, %arg1: tensor<32x32xf32>, %arg2: tensor<64x64xf32>, %arg3: tensor<64x32xf32>) -> (tensor<64x64xf32>, tensor<2048xf32>) {249 %c4 = arith.constant 4 : index250 %c64 = arith.constant 64 : index251 %c0 = arith.constant 0 : index252 %0:2 = scf.forall (%arg4, %arg5) in (2, 2) shared_outs(%arg6 = %arg3, %arg7 = %arg2) -> (tensor<64x32xf32>, tensor<64x64xf32>) {253 %extracted_slice = tensor.extract_slice %arg6[%arg4, %arg5] [32, 32] [1, 1] : tensor<64x32xf32> to tensor<32x32xf32>254 %extracted_slice_0 = tensor.extract_slice %arg7[%arg4, %arg5] [32, 32] [1, 1] : tensor<64x64xf32> to tensor<32x32xf32>255 %6 = linalg.matmul ins(%arg0, %arg1 : tensor<32x32xf32>, tensor<32x32xf32>) outs(%extracted_slice : tensor<32x32xf32>) -> tensor<32x32xf32>256 scf.forall.in_parallel {257 tensor.parallel_insert_slice %6 into %arg7[%arg4, %arg5] [32, 32] [1, 1] : tensor<32x32xf32> into tensor<64x64xf32>258 tensor.parallel_insert_slice %extracted_slice_0 into %arg6[%arg4, %arg5] [32, 32] [1, 1] : tensor<32x32xf32> into tensor<64x32xf32>259 }260 }261 %1 = tensor.empty() : tensor<64x64xf32>262 %2 = tensor.empty() : tensor<64x64xf32>263 %3 = tensor.empty() : tensor<64x64xf32>264 %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>) {265 ^bb0(%in: f32, %in_0: f32, %out: f32, %out_1: f32):266 %6 = arith.mulf %in, %in_0 : f32267 %7 = arith.subf %out, %6 : f32268 %8 = arith.addf %out_1, %in : f32269 linalg.yield %7, %8 : f32, f32270 } -> (tensor<64x64xf32>, tensor<64x64xf32>)271 %5 = tensor.empty() : tensor<2048xf32>272 %unpack = linalg.unpack %0#0 outer_dims_perm = [0] inner_dims_pos = [0] inner_tiles = [32] into %5 : tensor<64x32xf32> -> tensor<2048xf32>273 return %4#1, %unpack : tensor<64x64xf32>, tensor<2048xf32>274 }275}276 277module attributes {transform.with_named_sequence} {278 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {279 %generic = transform.structured.match ops{["linalg.generic"]} in %arg1280 : (!transform.any_op) -> !transform.any_op281 %loop = transform.structured.match ops{["scf.forall"]} in %arg1282 : (!transform.any_op) -> !transform.any_op283 %a, %new_loops = transform.test.fuse_consumer %generic into (%loop)284 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)285 transform.yield286 }287}288// CHECK: func.func @fuse_tilable_consumer_scf_forall_multi_yielding_consumer(289// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<32x32xf32>290// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<32x32xf32>291// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<64x64xf32>292// CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]: tensor<64x32xf32>)293// CHECK: %[[OUT_INIT:.*]] = tensor.empty() : tensor<64x64xf32>294// CHECK: %[[FINAL_RESULT:.*]]:4 = scf.forall (%[[IV1:.*]], %[[IV2:.*]]) in (2, 2)295// CHECK-SAME: shared_outs(%[[FIRST_OUT_ARG:.*]] = %[[ARG3]], %[[SECOND_OUT_ARG:.*]] = %[[ARG2]], %[[ELEM_OUT_ARG_0:.*]] = %[[OUT_INIT]], %[[ELEM_OUT_ARG_1:.*]] = %[[OUT_INIT]])296// CHECK-SAME: {297// CHECK: %[[MAT_OUT_SLICE:.*]] = tensor.extract_slice %[[FIRST_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]298// CHECK: %[[SECOND_ARG_SLICE:.*]] = tensor.extract_slice %[[SECOND_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]299// CHECK: %[[MAT_OUT:.*]] = linalg.matmul300// CHECK-SAME: outs(%[[MAT_OUT_SLICE]] :301// CHECK: %[[SLICE_OPERAND2:.*]] = tensor.extract_slice %[[OUT_INIT]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]302// CHECK: %[[SLICE_OUT_0:.*]] = tensor.extract_slice %[[ELEM_OUT_ARG_0]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]303// CHECK: %[[SLICE_OUT_1:.*]] = tensor.extract_slice %[[ELEM_OUT_ARG_1]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]304// CHECK: %[[ELEM_OUT:.*]]:2 = linalg.generic305// CHECK-SAME: ins(%[[MAT_OUT]], %[[SLICE_OPERAND2]] :306// CHECK-SAME: outs(%[[SLICE_OUT_0]], %[[SLICE_OUT_1]] :307// CHECK: scf.forall.in_parallel {308// CHECK: tensor.parallel_insert_slice %[[MAT_OUT]] into %[[SECOND_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]309// CHECK: tensor.parallel_insert_slice %[[SECOND_ARG_SLICE]] into %[[FIRST_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]310// CHECK: tensor.parallel_insert_slice %[[ELEM_OUT]]#0 into %[[ELEM_OUT_ARG_0]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]311// CHECK: tensor.parallel_insert_slice %[[ELEM_OUT]]#1 into %[[ELEM_OUT_ARG_1]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]312// CHECK: }313// CHECK: }314// CHECK: %[[UNPACK:.*]] = linalg.unpack %[[FINAL_RESULT]]#0 outer_dims_perm = [0] inner_dims_pos = [0] inner_tiles = [32] into %{{.*}} : tensor<64x32xf32> -> tensor<2048xf32>315// CHECK: return %[[FINAL_RESULT]]#3, %[[UNPACK]] :316 317// -----318 319#map = affine_map<(d0, d1) -> (d0, d1)>320module {321 func.func @fuse_unpack_consumer_into_scf_forall(%arg0: tensor<32x32xf32>, %arg1: tensor<32x32xf32>, %arg2: tensor<64x32xf32>) -> tensor<2048xf32> {322 %c4 = arith.constant 4 : index323 %c64 = arith.constant 64 : index324 %c0 = arith.constant 0 : index325 %1 = scf.forall (%arg3, %arg4) = (0, 0) to (64, 32) step (32, 32) shared_outs(%arg5 = %arg2) -> (tensor<64x32xf32>) {326 %extracted_slice = tensor.extract_slice %arg5[%arg3, %arg4] [32, 32] [1, 1] : tensor<64x32xf32> to tensor<32x32xf32>327 %3 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0, %arg1 : tensor<32x32xf32>, tensor<32x32xf32>) outs(%extracted_slice : tensor<32x32xf32>) {328 ^bb0(%in: f32, %in_16: f32, %out: f32):329 %13 = arith.mulf %in, %in_16 : f32330 %14 = arith.addf %out, %13 : f32331 linalg.yield %14 : f32332 } -> tensor<32x32xf32>333 scf.forall.in_parallel {334 tensor.parallel_insert_slice %3 into %arg5[%arg3, %arg4] [32, 32] [1, 1] : tensor<32x32xf32> into tensor<64x32xf32>335 }336 }337 %output = tensor.empty() : tensor<2048xf32>338 %unpack = linalg.unpack %1 outer_dims_perm = [0] inner_dims_pos = [0] inner_tiles = [32] into %output : tensor<64x32xf32> -> tensor<2048xf32>339 return %unpack : tensor<2048xf32>340 }341}342 343module attributes {transform.with_named_sequence} {344 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {345 %consumer = transform.structured.match ops{["linalg.unpack"]} in %arg1346 : (!transform.any_op) -> !transform.any_op347 %loop = transform.structured.match ops{["scf.forall"]} in %arg1348 : (!transform.any_op) -> !transform.any_op349 %a, %new_loop = transform.test.fuse_consumer %consumer into (%loop)350 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)351 transform.yield352 }353}354// CHECK: func.func @fuse_unpack_consumer_into_scf_forall(355// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<32x32xf32>356// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<32x32xf32>357// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<64x32xf32>)358// CHECK: %[[OUT_INIT:.*]] = tensor.empty() : tensor<2048xf32>359// CHECK: %[[FINAL_RESULT:.*]]:2 = scf.forall (%[[IV1:.*]], %[[IV2:.*]]) = (0, 0) to (64, 32) step (32, 32)360// CHECK-SAME: shared_outs(%[[FIRST_OUT_ARG:.*]] = %[[ARG2]], %[[UNPACK_OUT_ARG:.*]] = %[[OUT_INIT]])361// CHECK-SAME: {362// CHECK: %[[GENERIC_OUT_SLICE:.*]] = tensor.extract_slice %[[FIRST_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]363// CHECK: %[[GENERIC_OUT:.*]] = linalg.generic364// CHECK-SAME: outs(%[[GENERIC_OUT_SLICE]] :365// CHECK-DAG: %[[UNPACK_RESULT_OFFSET:.*]] = affine.apply affine_map<(d0) -> (d0 * 32)>(%[[IV1]])366// CHECK-DAG: %[[UNPACK_RESULT_SIZE:.*]] = affine.min affine_map<(d0) -> (1024, d0 * -32 + 2048)>(%[[IV1]])367// CHECK: %[[TILED_UNPACK_DEST:.*]] = tensor.extract_slice %[[UNPACK_OUT_ARG]][%[[UNPACK_RESULT_OFFSET]]] [%[[UNPACK_RESULT_SIZE]]] [1]368// CHECK: %[[TILED_UNPACK_OUT:.*]] = linalg.unpack %[[GENERIC_OUT]]369// CHECK-SAME: outer_dims_perm = [0] inner_dims_pos = [0] inner_tiles = [32]370// CHECK-SAME: into %[[TILED_UNPACK_DEST]]371// CHECK: scf.forall.in_parallel {372// CHECK: tensor.parallel_insert_slice %[[GENERIC_OUT]] into %[[FIRST_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]373// CHECK: tensor.parallel_insert_slice %[[TILED_UNPACK_OUT]] into %[[UNPACK_OUT_ARG]][%[[UNPACK_RESULT_OFFSET]]] [%[[UNPACK_RESULT_SIZE]]] [1]374// CHECK: }375// CHECK: }376// CHECK: return %[[FINAL_RESULT]]#1 :377 378// -----379 380#map = affine_map<(d0, d1) -> (d0, d1)>381module {382 func.func @fuse_unaligned_unpack_consumer_into_scf_forall(%arg0: tensor<32x32xf32>, %arg1: tensor<32x32xf32>, %arg2: tensor<64x32xf32>) -> tensor<2047xf32> {383 %c4 = arith.constant 4 : index384 %c64 = arith.constant 64 : index385 %c0 = arith.constant 0 : index386 %1 = scf.forall (%arg3, %arg4) = (0, 0) to (64, 32) step (32, 32) shared_outs(%arg5 = %arg2) -> (tensor<64x32xf32>) {387 %extracted_slice = tensor.extract_slice %arg5[%arg3, %arg4] [32, 32] [1, 1] : tensor<64x32xf32> to tensor<32x32xf32>388 %3 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0, %arg1 : tensor<32x32xf32>, tensor<32x32xf32>) outs(%extracted_slice : tensor<32x32xf32>) {389 ^bb0(%in: f32, %in_16: f32, %out: f32):390 %13 = arith.mulf %in, %in_16 : f32391 %14 = arith.addf %out, %13 : f32392 linalg.yield %14 : f32393 } -> tensor<32x32xf32>394 scf.forall.in_parallel {395 tensor.parallel_insert_slice %3 into %arg5[%arg3, %arg4] [32, 32] [1, 1] : tensor<32x32xf32> into tensor<64x32xf32>396 }397 }398 %output = tensor.empty() : tensor<2047xf32>399 %unpack = linalg.unpack %1 outer_dims_perm = [0] inner_dims_pos = [0] inner_tiles = [32] into %output : tensor<64x32xf32> -> tensor<2047xf32>400 return %unpack : tensor<2047xf32>401 }402}403 404module attributes {transform.with_named_sequence} {405 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {406 %consumer = transform.structured.match ops{["linalg.unpack"]} in %arg1407 : (!transform.any_op) -> !transform.any_op408 %loop = transform.structured.match ops{["scf.forall"]} in %arg1409 : (!transform.any_op) -> !transform.any_op410 %a, %new_loop = transform.test.fuse_consumer %consumer into (%loop)411 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)412 transform.yield413 }414}415// CHECK: func.func @fuse_unaligned_unpack_consumer_into_scf_forall(416// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<32x32xf32>417// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<32x32xf32>418// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<64x32xf32>)419// CHECK: %[[OUT_INIT:.*]] = tensor.empty() : tensor<2047xf32>420// CHECK: %[[FINAL_RESULT:.*]]:2 = scf.forall (%[[IV1:.*]], %[[IV2:.*]]) = (0, 0) to (64, 32) step (32, 32)421// CHECK-SAME: shared_outs(%[[FIRST_OUT_ARG:.*]] = %[[ARG2]], %[[UNPACK_OUT_ARG:.*]] = %[[OUT_INIT]])422// CHECK-SAME: {423// CHECK: %[[GENERIC_OUT_SLICE:.*]] = tensor.extract_slice %[[FIRST_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]424// CHECK: %[[GENERIC_OUT:.*]] = linalg.generic425// CHECK-SAME: outs(%[[GENERIC_OUT_SLICE]] :426// CHECK-DAG: %[[UNPACK_RESULT_OFFSET:.*]] = affine.apply affine_map<(d0) -> (d0 * 32)>(%[[IV1]])427// CHECK-DAG: %[[UNPACK_RESULT_SIZE:.*]] = affine.min affine_map<(d0) -> (1024, d0 * -32 + 2047)>(%[[IV1]])428// CHECK: %[[TILED_UNPACK_DEST:.*]] = tensor.extract_slice %[[UNPACK_OUT_ARG]][%[[UNPACK_RESULT_OFFSET]]] [%[[UNPACK_RESULT_SIZE]]] [1]429// CHECK: %[[TILED_UNPACK_OUT:.*]] = linalg.unpack %[[GENERIC_OUT]]430// CHECK-SAME: outer_dims_perm = [0] inner_dims_pos = [0] inner_tiles = [32]431// CHECK-SAME: into %[[TILED_UNPACK_DEST]]432// CHECK: scf.forall.in_parallel {433// CHECK: tensor.parallel_insert_slice %[[GENERIC_OUT]] into %[[FIRST_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]434// CHECK: tensor.parallel_insert_slice %[[TILED_UNPACK_OUT]] into %[[UNPACK_OUT_ARG]][%[[UNPACK_RESULT_OFFSET]]] [%[[UNPACK_RESULT_SIZE]]] [1]435// CHECK: }436// CHECK: }437// CHECK: return %[[FINAL_RESULT]]#1 :438 439// -----440 441#map = affine_map<(d0, d1) -> (d0, d1)>442module {443 func.func @fuse_perfect_tiling_pack_consumer(%arg0: tensor<32x32xf32>, %arg1: tensor<32x32xf32>, %arg2: tensor<64x32xf32>) -> tensor<4x32x16xf32> {444 %c4 = arith.constant 4 : index445 %c64 = arith.constant 64 : index446 %c0 = arith.constant 0 : index447 %1 = scf.forall (%arg3, %arg4) in (2, 1) shared_outs(%arg5 = %arg2) -> (tensor<64x32xf32>) {448 %extracted_slice = tensor.extract_slice %arg5[%arg3, %arg4] [32, 32] [1, 1] : tensor<64x32xf32> to tensor<32x32xf32>449 %3 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0, %arg1 : tensor<32x32xf32>, tensor<32x32xf32>) outs(%extracted_slice : tensor<32x32xf32>) {450 ^bb0(%in: f32, %in_16: f32, %out: f32):451 %13 = arith.mulf %in, %in_16 : f32452 %14 = arith.addf %out, %13 : f32453 linalg.yield %14 : f32454 } -> tensor<32x32xf32>455 scf.forall.in_parallel {456 tensor.parallel_insert_slice %3 into %arg5[%arg3, %arg4] [32, 32] [1, 1] : tensor<32x32xf32> into tensor<64x32xf32>457 }458 }459 %output = tensor.empty() : tensor<4x32x16xf32>460 %pack = linalg.pack %1 inner_dims_pos = [0] inner_tiles = [16] into %output : tensor<64x32xf32> -> tensor<4x32x16xf32>461 return %pack : tensor<4x32x16xf32>462 }463}464 465module attributes {transform.with_named_sequence} {466 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {467 %consumer = transform.structured.match ops{["linalg.pack"]} in %arg1468 : (!transform.any_op) -> !transform.any_op469 %loop = transform.structured.match ops{["scf.forall"]} in %arg1470 : (!transform.any_op) -> !transform.any_op471 %a, %new_loop = transform.test.fuse_consumer %consumer into (%loop)472 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)473 transform.yield474 }475}476// CHECK: func.func @fuse_perfect_tiling_pack_consumer(477// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<32x32xf32>478// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<32x32xf32>479// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<64x32xf32>)480// CHECK: %[[OUT_INIT:.*]] = tensor.empty() : tensor<4x32x16xf32>481// CHECK: %[[FINAL_RESULT:.*]]:2 = scf.forall (%[[IV1:.*]], %[[IV2:.*]]) in (2, 1)482// CHECK-SAME: shared_outs(%[[FIRST_OUT_ARG:.*]] = %[[ARG2]], %[[PACK_OUT_ARG:.*]] = %[[OUT_INIT]])483// CHECK-SAME: {484// CHECK: %[[GENERIC_OUT_SLICE:.*]] = tensor.extract_slice %[[FIRST_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]485// CHECK: %[[GENERIC_OUT:.*]] = linalg.generic486// CHECK-SAME: outs(%[[GENERIC_OUT_SLICE]] :487// CHECK: %[[PACK_RESULT_OFFSET:.*]] = affine.apply affine_map<(d0) -> (d0 floordiv 16)>(%[[IV1]])488// CHECK: %[[TILED_PACK_DEST:.*]] = tensor.extract_slice %[[PACK_OUT_ARG]][%[[PACK_RESULT_OFFSET]], %[[IV2]], 0] [2, 32, 16] [1, 1, 1]489// CHECK: %[[TILED_PACK_OUT:.*]] = linalg.pack %[[GENERIC_OUT]]490// CHECK-SAME: inner_dims_pos = [0] inner_tiles = [16]491// CHECK-SAME: into %[[TILED_PACK_DEST]]492// CHECK: scf.forall.in_parallel {493// CHECK: tensor.parallel_insert_slice %[[GENERIC_OUT]] into %[[FIRST_OUT_ARG]][%[[IV1]], %[[IV2]]] [32, 32] [1, 1]494// CHECK: tensor.parallel_insert_slice %[[TILED_PACK_OUT]] into %[[PACK_OUT_ARG]][%[[PACK_RESULT_OFFSET]], %[[IV2]], 0] [2, 32, 16] [1, 1, 1]495 496// -----497 498#map = affine_map<(d0) -> (-d0 + 4, 16)>499func.func @fuse_pack_consumer_if_single_iteration(%arg0: tensor<4x4xf32>) -> tensor<1x4x16x1xf32> {500 %0 = tensor.empty() : tensor<1x4x16x1xf32>501 %1 = tensor.empty() : tensor<4x4xf32>502 %2 = scf.forall (%arg1) = (0) to (4) step (16) shared_outs(%arg2 = %1) -> (tensor<4x4xf32>) {503 %3 = affine.min #map(%arg1)504 %extracted_slice = tensor.extract_slice %arg0[%arg1, 0] [%3, 4] [1, 1] : tensor<4x4xf32> to tensor<?x4xf32>505 %extracted_slice_0 = tensor.extract_slice %arg2[%arg1, 0] [%3, 4] [1, 1] : tensor<4x4xf32> to tensor<?x4xf32>506 %4 = linalg.exp ins(%extracted_slice : tensor<?x4xf32>) outs(%extracted_slice_0 : tensor<?x4xf32>) -> tensor<?x4xf32>507 scf.forall.in_parallel {508 tensor.parallel_insert_slice %4 into %arg2[%arg1, 0] [%3, 4] [1, 1] : tensor<?x4xf32> into tensor<4x4xf32>509 }510 }511 %cst = arith.constant 0.000000e+00 : f32512 %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>513 return %pack : tensor<1x4x16x1xf32>514}515 516module attributes {transform.with_named_sequence} {517 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {518 %consumer = transform.structured.match ops{["linalg.pack"]} in %arg0 : (!transform.any_op) -> !transform.any_op519 %1 = transform.structured.match ops{["scf.forall"]} in %arg0 : (!transform.any_op) -> !transform.any_op520 %fused_consumer, %new_loop = transform.test.fuse_consumer %consumer into(%1) : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)521 transform.yield522 }523}524// CHECK: func.func @fuse_pack_consumer_if_single_iteration(525// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]526// CHECK-DAG: %[[PACK_INIT:.*]] = tensor.empty() : tensor<1x4x16x1xf32>527// CHECK-DAG: %[[ELEM_INIT:.*]] = tensor.empty() : tensor<4x4xf32>528// CHECK-DAG: %[[PAD_VAL:.*]] = arith.constant 0.000000e+00 : f32529// CHECK: %{{.*}}:2 = scf.forall (%[[IV:.*]]) = (0) to (4) step (16)530// CHECK-SAME: shared_outs(%[[ELEM_OUT_ARG:.*]] = %[[ELEM_INIT]], %[[PACK_OUT_ARG:.*]] = %[[PACK_INIT]])531// CHECK-DAG: %[[SIZE:.+]] = affine.min affine_map<(d0) -> (-d0 + 4, 16)>(%[[IV]])532// CHECK-DAG: %[[ELEM_SRC:.*]] = tensor.extract_slice %[[ARG0]][%[[IV]], 0] [%[[SIZE]], 4] [1, 1]533// CHECK-DAG: %[[ELEM_DEST:.*]] = tensor.extract_slice %[[ELEM_OUT_ARG]][%[[IV]], 0] [%[[SIZE]], 4] [1, 1]534// CHECK: %[[ELEM:.*]] = linalg.exp535// CHECK-SAME: ins(%[[ELEM_SRC]]536// CHECK-SAME: outs(%[[ELEM_DEST]]537// CHECK-DAG: %[[TILED_PACK_DEST:.*]] = tensor.extract_slice %[[PACK_OUT_ARG]][%[[IV]], 0, 0, 0] [1, 4, 16, 1] [1, 1, 1, 1]538// CHECK: %[[PACK:.*]] = linalg.pack %[[ELEM]]539// CHECK-SAME: padding_value(%[[PAD_VAL]] : f32)540// CHECK-SAME: outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [16, 1]541// CHECK-SAME: into %[[TILED_PACK_DEST]]542// CHECK: scf.forall.in_parallel {543// CHECK: tensor.parallel_insert_slice %[[ELEM]] into %[[ELEM_OUT_ARG]][%[[IV]], 0] [%[[SIZE]], 4] [1, 1]544// CHECK: tensor.parallel_insert_slice %[[PACK]] into %[[PACK_OUT_ARG]][%[[IV]], 0, 0, 0] [1, 4, 16, 1] [1, 1, 1, 1]545 546// -----547 548func.func @fuse_perfect_tiling_pack_consumer_with_outer_dims_perm(%arg0: tensor<64x32xf32>, %arg1: tensor<64x32xf32>, %arg2: tensor<2x64x16x1xf32>) -> tensor<2x64x16x1xf32> {549 %0 = scf.forall (%arg3) = (0) to (32) step (16) shared_outs(%arg4 = %arg1) -> (tensor<64x32xf32>) {550 %src = tensor.extract_slice %arg0[0, %arg3] [64, 16] [1, 1] : tensor<64x32xf32> to tensor<64x16xf32>551 %dest = tensor.extract_slice %arg4[0, %arg3] [64, 16] [1, 1] : tensor<64x32xf32> to tensor<64x16xf32>552 %1 = linalg.exp ins(%src : tensor<64x16xf32>) outs(%dest : tensor<64x16xf32>) -> tensor<64x16xf32>553 scf.forall.in_parallel {554 tensor.parallel_insert_slice %1 into %arg4[0, %arg3] [64, 16] [1, 1] : tensor<64x16xf32> into tensor<64x32xf32>555 }556 }557 %pack = linalg.pack %0 outer_dims_perm = [1, 0] inner_dims_pos = [1, 0] inner_tiles = [16, 1] into %arg2 : tensor<64x32xf32> -> tensor<2x64x16x1xf32>558 return %pack : tensor<2x64x16x1xf32>559}560 561module attributes {transform.with_named_sequence} {562 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {563 %0 = transform.structured.match ops{["linalg.pack"]} in %arg0 : (!transform.any_op) -> !transform.any_op564 %1 = transform.structured.match ops{["scf.forall"]} in %arg0 : (!transform.any_op) -> !transform.any_op565 %fused_consumer, %new_loop = transform.test.fuse_consumer %0 into(%1) : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)566 transform.yield567 }568}569// CHECK: func.func @fuse_perfect_tiling_pack_consumer_with_outer_dims_perm(570// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]571// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]572// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]573// CHECK: %{{.*}}:2 = scf.forall (%[[IV:.*]]) = (0) to (32) step (16)574// CHECK-SAME: shared_outs(%[[FIRST_OUT_ARG:.*]] = %[[ARG1]], %[[PACK_OUT_ARG:.*]] = %[[ARG2]])575// CHECK: %[[ELEM_SRC:.*]] = tensor.extract_slice %[[ARG0]][0, %[[IV]]] [64, 16] [1, 1]576// CHECK: %[[ELEM_DEST:.*]] = tensor.extract_slice %[[FIRST_OUT_ARG]][0, %[[IV]]] [64, 16] [1, 1]577// CHECK: %[[ELEM:.*]] = linalg.exp578// CHECK-SAME: ins(%[[ELEM_SRC]]579// CHECK-SAME: outs(%[[ELEM_DEST]]580// CHECK-DAG: %[[PACK_RESULT_OFFSET:.*]] = affine.apply affine_map<(d0) -> (d0 floordiv 16)>(%[[IV]])581// 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]582// CHECK: %[[PACK:.*]] = linalg.pack %[[ELEM]]583// CHECK-SAME: outer_dims_perm = [1, 0] inner_dims_pos = [1, 0] inner_tiles = [16, 1]584// CHECK-SAME: into %[[TILED_PACK_DEST]]585// CHECK: scf.forall.in_parallel {586// CHECK: tensor.parallel_insert_slice %[[ELEM]] into %[[FIRST_OUT_ARG]][0, %[[IV]]] [64, 16] [1, 1]587// CHECK: tensor.parallel_insert_slice %[[PACK]] into %[[PACK_OUT_ARG]][%[[PACK_RESULT_OFFSET]], 0, 0, 0] [1, 64, 16, 1] [1, 1, 1, 1]588 589// -----590 591// It is valid to fuse the pack op in perfect tiling scenario when the dimension592// is dynamic and padding is not needed.593 594func.func @fuse_pack_consumer_with_no_pad_dynamic_dim(%arg0: tensor<64x?xf32>, %arg1: tensor<64x?xf32>, %1: tensor<64x?x16xf32>) -> tensor<64x?x16xf32> {595 %c1 = arith.constant 1 : index596 %d1 = tensor.dim %arg0, %c1 : tensor<64x?xf32>597 %0 = scf.forall (%arg2) = (0) to (%d1) step (16) shared_outs(%arg3 = %arg1) -> (tensor<64x?xf32>) {598 %src = tensor.extract_slice %arg0[0, %arg2] [64, 16] [1, 1] : tensor<64x?xf32> to tensor<64x16xf32>599 %dest = tensor.extract_slice %arg3[0, %arg2] [64, 16] [1, 1] : tensor<64x?xf32> to tensor<64x16xf32>600 %2 = linalg.exp ins(%src : tensor<64x16xf32>) outs(%dest : tensor<64x16xf32>) -> tensor<64x16xf32>601 scf.forall.in_parallel {602 tensor.parallel_insert_slice %2 into %arg3[0, %arg2] [64, 16] [1, 1] : tensor<64x16xf32> into tensor<64x?xf32>603 }604 }605 %pack = linalg.pack %0 inner_dims_pos = [1] inner_tiles = [16] into %1 : tensor<64x?xf32> -> tensor<64x?x16xf32>606 return %pack : tensor<64x?x16xf32>607}608 609module attributes {transform.with_named_sequence} {610 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {611 %0 = transform.structured.match ops{["linalg.pack"]} in %arg0 : (!transform.any_op) -> !transform.any_op612 %1 = transform.structured.match ops{["scf.forall"]} in %arg0 : (!transform.any_op) -> !transform.any_op613 %fused_consumer, %new_loop = transform.test.fuse_consumer %0 into(%1) : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)614 transform.yield615 }616}617// CHECK: func.func @fuse_pack_consumer_with_no_pad_dynamic_dim(618// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]619// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]620// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]621// CHECK: %{{.*}}:2 = scf.forall (%[[IV:.*]]) = (0) to (%{{.+}}) step (16)622// CHECK-SAME: shared_outs(%[[FIRST_OUT_ARG:.*]] = %[[ARG1]], %[[PACK_OUT_ARG:.*]] = %[[ARG2]])623// CHECK: %[[ELEM_SRC:.*]] = tensor.extract_slice %[[ARG0]][0, %[[IV]]] [64, 16] [1, 1]624// CHECK: %[[ELEM_DEST:.*]] = tensor.extract_slice %[[FIRST_OUT_ARG]][0, %[[IV]]] [64, 16] [1, 1]625// CHECK: %[[ELEM:.*]] = linalg.exp626// CHECK-SAME: ins(%[[ELEM_SRC]]627// CHECK-SAME: outs(%[[ELEM_DEST]]628// CHECK-DAG: %[[PACK_RESULT_OFFSET:.*]] = affine.apply affine_map<(d0) -> (d0 floordiv 16)>(%[[IV]])629// CHECK-DAG: %[[TILED_PACK_DEST:.*]] = tensor.extract_slice %[[PACK_OUT_ARG]][0, %[[PACK_RESULT_OFFSET]], 0] [64, 1, 16] [1, 1, 1]630// CHECK: %[[PACK:.*]] = linalg.pack %[[ELEM]]631// CHECK-SAME: inner_dims_pos = [1] inner_tiles = [16]632// CHECK-SAME: into %[[TILED_PACK_DEST]]633// CHECK: scf.forall.in_parallel {634// CHECK: tensor.parallel_insert_slice %[[ELEM]] into %[[FIRST_OUT_ARG]][0, %[[IV]]] [64, 16] [1, 1]635// CHECK: tensor.parallel_insert_slice %[[PACK]] into %[[PACK_OUT_ARG]][0, %[[PACK_RESULT_OFFSET]], 0] [64, 1, 16] [1, 1, 1]636 637// -----638 639// It is valid to fuse the pack op with padding semantics if it is a perfect640// tiling case.641 642func.func @fuse_pack_consumer_with_padding_semantics(%arg0: tensor<64x32xf32>, %arg1: tensor<64x32xf32>) -> tensor<22x2x3x16xf32> {643 %0 = scf.forall (%arg2, %arg3) = (0, 0) to (64, 32) step (15, 16) shared_outs(%arg4 = %arg1) -> (tensor<64x32xf32>) {644 %size = affine.min affine_map<(d0) -> (-d0 + 64, 15)>(%arg2)645 %src = tensor.extract_slice %arg0[%arg2, %arg3] [%size, 16] [1, 1] : tensor<64x32xf32> to tensor<?x16xf32>646 %dest = tensor.extract_slice %arg4[%arg2, %arg3] [%size, 16] [1, 1] : tensor<64x32xf32> to tensor<?x16xf32>647 %2 = linalg.exp ins(%src : tensor<?x16xf32>) outs(%dest : tensor<?x16xf32>) -> tensor<?x16xf32>648 scf.forall.in_parallel {649 tensor.parallel_insert_slice %2 into %arg4[%arg2, %arg3] [%size, 16] [1, 1] : tensor<?x16xf32> into tensor<64x32xf32>650 }651 }652 %1 = tensor.empty() : tensor<22x2x3x16xf32>653 %cst = arith.constant 0.000000e+00 : f32654 %pack = linalg.pack %0 padding_value(%cst : f32) inner_dims_pos = [0, 1] inner_tiles = [3, 16] into %1 : tensor<64x32xf32> -> tensor<22x2x3x16xf32>655 return %pack : tensor<22x2x3x16xf32>656}657 658module attributes {transform.with_named_sequence} {659 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {660 %0 = transform.structured.match ops{["linalg.pack"]} in %arg0 : (!transform.any_op) -> !transform.any_op661 %1 = transform.structured.match ops{["scf.forall"]} in %arg0 : (!transform.any_op) -> !transform.any_op662 %fused_consumer, %new_loop = transform.test.fuse_consumer %0 into(%1) : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)663 transform.yield664 }665}666// CHECK: func.func @fuse_pack_consumer_with_padding_semantics(667// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]668// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]669// CHECK-DAG: %[[OUT_INIT:.*]] = tensor.empty() : tensor<22x2x3x16xf32>670// CHECK-DAG: %[[PAD_VAL:.*]] = arith.constant 0.000000e+00 : f32671// CHECK: %{{.*}}:2 = scf.forall (%[[I:.*]], %[[J:.*]]) = (0, 0) to (64, 32) step (15, 16)672// CHECK-SAME: shared_outs(%[[ELEM_OUT:.*]] = %[[ARG1]], %[[PACK_OUT:.*]] = %[[OUT_INIT]])673// CHECK: %[[SIZE:.+]] = affine.min affine_map<(d0) -> (-d0 + 64, 15)>(%[[I]])674// CHECK: %[[ELEM_SRC:.*]] = tensor.extract_slice %[[ARG0]]675// CHECK-SAME: [%[[I]], %[[J]]] [%[[SIZE]], 16] [1, 1]676// CHECK: %[[ELEM_DEST:.*]] = tensor.extract_slice %[[ELEM_OUT]]677// CHECK-SAME: [%[[I]], %[[J]]] [%[[SIZE]], 16] [1, 1]678// CHECK: %[[ELEM:.*]] = linalg.exp679// CHECK-SAME: ins(%[[ELEM_SRC]]680// CHECK-SAME: outs(%[[ELEM_DEST]]681// CHECK-DAG: %[[D0_OFFSET:.*]] = affine.apply affine_map<(d0) -> (d0 floordiv 3)>(%[[I]])682// CHECK-DAG: %[[D0_SIZE:.*]] = affine.apply affine_map<(d0) -> (d0 ceildiv 3)>(%[[SIZE]])683// CHECK-DAG: %[[D1_OFFSET:.*]] = affine.apply affine_map<(d0) -> (d0 floordiv 16)>(%[[J]])684// CHECK-DAG: %[[PACK_INIT:.*]] = tensor.extract_slice %[[PACK_OUT]]685// CHECK-SAME: [%[[D0_OFFSET]], %[[D1_OFFSET]], 0, 0] [%[[D0_SIZE]], 1, 3, 16] [1, 1, 1, 1]686// CHECK: %[[PACK:.*]] = linalg.pack %[[ELEM]]687// CHECK-SAME: padding_value(%[[PAD_VAL]] : f32)688// CHECK-SAME: inner_dims_pos = [0, 1] inner_tiles = [3, 16]689// CHECK-SAME: into %[[TILED_PACK_DEST]]690// CHECK: scf.forall.in_parallel {691// CHECK: tensor.parallel_insert_slice %[[ELEM]] into %[[ELEM_OUT]]692// CHECK-SAME: [%[[I]], %[[J]]] [%[[SIZE]], 16] [1, 1]693// CHECK: tensor.parallel_insert_slice %[[PACK]] into %[[PACK_OUT]]694// CHECK-SAME: [%[[D0_OFFSET]], %[[D1_OFFSET]], 0, 0] [%[[D0_SIZE]], 1, 3, 16] [1, 1, 1, 1]695 696// -----697 698// Imperfect tiling is not supported in pack op consumer fusion.699 700#map = affine_map<(d0) -> (d0 * 5)>701#map1 = affine_map<(d0) -> (d0)>702func.func @nofuse_pack_with_imperfect_tiling(%arg0: tensor<30xf32>) -> tensor<5x6xf32> {703 %0 = tensor.empty() : tensor<30xf32>704 %1 = scf.forall (%arg1) in (6) shared_outs(%arg2 = %0) -> (tensor<30xf32>) {705 %3 = affine.apply #map(%arg1)706 %extracted_slice = tensor.extract_slice %arg0[%3] [5] [1] : tensor<30xf32> to tensor<5xf32>707 %extracted_slice_0 = tensor.extract_slice %arg2[%3] [5] [1] : tensor<30xf32> to tensor<5xf32>708 %4 = linalg.generic {indexing_maps = [#map1, #map1], iterator_types = ["parallel"]} ins(%extracted_slice : tensor<5xf32>) outs(%extracted_slice_0 : tensor<5xf32>) {709 ^bb0(%in: f32, %out: f32):710 %5 = arith.addf %in, %in : f32711 linalg.yield %5 : f32712 } -> tensor<5xf32>713 scf.forall.in_parallel {714 715 tensor.parallel_insert_slice %4 into %arg2[%3] [5] [1] : tensor<5xf32> into tensor<30xf32>716 }717 }718 %2 = tensor.empty() : tensor<5x6xf32>719 // expected-error @below {{failed to fuse consumer of slice}}720 %pack = linalg.pack %1 outer_dims_perm = [0] inner_dims_pos = [0] inner_tiles = [6] into %2 : tensor<30xf32> -> tensor<5x6xf32>721 return %pack : tensor<5x6xf32>722}723 724module attributes {transform.with_named_sequence} {725 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {726 %0 = transform.structured.match ops{["linalg.pack"]} in %arg0 : (!transform.any_op) -> !transform.any_op727 %1 = transform.structured.match ops{["scf.forall"]} in %arg0 : (!transform.any_op) -> !transform.any_op728 %fused_consumer, %new_loop = transform.test.fuse_consumer %0 into(%1) : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)729 transform.yield730 }731}732 733// -----734 735module {736 func.func @fuse_add_multiple_tilable_consumers(%arg0: tensor<256x256xf32>, %arg1: tensor<256x256xf32>, %arg2: tensor<256x256xf32>) -> (tensor<256x256xf32>, tensor<256x256xf32>) {737 %c0 = arith.constant 0 : index738 %c64 = arith.constant 64 : index739 %c256 = arith.constant 256 : index740 %cst = arith.constant 0.000000e+00 : f32741 %dest0 = tensor.empty() : tensor<256x256xf32>742 %1 = scf.for %arg3 = %c0 to %c256 step %c64 iter_args(%arg4 = %dest0) -> (tensor<256x256xf32>) {743 %extracted_slice_1 = tensor.extract_slice %arg4[%arg3, 0] [64, 256] [1, 1] : tensor<256x256xf32> to tensor<64x256xf32>744 %extracted_slice_2 = tensor.extract_slice %arg0[%arg3, 0] [64, 256] [1, 1] : tensor<256x256xf32> to tensor<64x256xf32>745 %extracted_slice_3 = tensor.extract_slice %arg1[%arg3, 0] [64, 256] [1, 1] : tensor<256x256xf32> to tensor<64x256xf32>746 %3 = linalg.add ins(%extracted_slice_2, %extracted_slice_3 : tensor<64x256xf32>, tensor<64x256xf32>) outs(%extracted_slice_1 : tensor<64x256xf32>) -> tensor<64x256xf32>747 %insert_slice = tensor.insert_slice %3 into %arg4[%arg3, 0] [64, 256] [1, 1] : tensor<64x256xf32> into tensor<256x256xf32>748 scf.yield %insert_slice : tensor<256x256xf32>749 }750 %4 = linalg.mul ins(%1, %arg2 : tensor<256x256xf32>, tensor<256x256xf32>) outs(%dest0 : tensor<256x256xf32>) -> tensor<256x256xf32>751 %5 = linalg.exp ins(%1 : tensor<256x256xf32>) outs(%dest0 : tensor<256x256xf32>) -> tensor<256x256xf32>752 return %4, %5 : tensor<256x256xf32>, tensor<256x256xf32>753 }754}755 756module attributes {transform.with_named_sequence} {757 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {758 %mulop = transform.structured.match ops{["linalg.mul"]} in %arg1759 : (!transform.any_op) -> !transform.any_op760 %loop = transform.structured.match ops{["scf.for"]} in %arg1761 : (!transform.any_op) -> !transform.any_op762 %fused_consumer, %new_loop = transform.test.fuse_consumer %mulop into (%loop)763 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)764 %expop = transform.structured.match ops{["linalg.exp"]} in %arg1765 : (!transform.any_op) -> !transform.any_op766 %fused_consumer_2, %new_loop_2 = transform.test.fuse_consumer %expop into (%new_loop)767 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)768 transform.yield769 }770}771// CHECK: func.func @fuse_add_multiple_tilable_consumers(772// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<256x256xf32>773// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<256x256xf32>774// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<256x256xf32>775// CHECK: %[[dest0:.*]] = tensor.empty() : tensor<256x256xf32>776// CHECK: %[[LOOP_RESULT:.*]]:3 = scf.for %[[IV1:.*]] = %[[C0]]777// CHECK-SAME: iter_args(%[[FIRST_OUT_ARG:.*]] = %[[dest0]], %[[SECOND_OUT_ARG:.*]] = %[[dest0]], %[[THIRD_OUT_ARG:.*]] = %[[dest0]])778// CHECK-SAME: {779// CHECK: %[[ADD_OUT_SLICE:.*]] = tensor.extract_slice %[[FIRST_OUT_ARG]][%[[IV1]], 0] [64, 256] [1, 1]780// CHECK: %[[ADD_INS0_SLICE:.*]] = tensor.extract_slice %[[ARG0]][%[[IV1]], 0] [64, 256] [1, 1]781// CHECK: %[[ADD_INS1_SLICE:.*]] = tensor.extract_slice %[[ARG1]][%[[IV1]], 0] [64, 256] [1, 1]782// CHECK: %[[TILED_ADD_OUT:.*]] = linalg.add783// CHECK-SAME: ins(%[[ADD_INS0_SLICE]], %[[ADD_INS1_SLICE]] :784// CHECK-SAME: outs(%[[ADD_OUT_SLICE]] :785// CHECK: %[[INSERT_ADD:.*]] = tensor.insert_slice %[[TILED_ADD_OUT]] into %[[FIRST_OUT_ARG]][%[[IV1]], 0] [64, 256] [1, 1]786// CHECK: %[[MUL_INS2_SLICE:.*]] = tensor.extract_slice %[[ARG2]][%[[IV1]], 0] [64, 256] [1, 1]787// CHECK: %[[MUL_OUT_SLICE:.*]] = tensor.extract_slice %[[SECOND_OUT_ARG]][%[[IV1]], 0] [64, 256] [1, 1]788// CHECK: %[[TILED_MUL_OUT:.*]] = linalg.mul789// CHECK-SAME: ins(%[[TILED_ADD_OUT]], %[[MUL_INS2_SLICE]] :790// CHECK-SAME: outs(%[[MUL_OUT_SLICE]] :791// CHECK: %[[EXP_OUT_SLICE:.*]] = tensor.extract_slice %[[THIRD_OUT_ARG]][%[[IV1]], 0] [64, 256] [1, 1]792// CHECK: %[[TILED_EXP_OUT:.*]] = linalg.exp793// CHECK-SAME: ins(%[[TILED_ADD_OUT]] :794// CHECK-SAME: outs(%[[EXP_OUT_SLICE]] :795// CHECK: %[[INSERT_MUL:.*]] = tensor.insert_slice %[[TILED_MUL_OUT]] into %[[SECOND_OUT_ARG]][%[[IV1]], 0] [64, 256] [1, 1]796// CHECK: %[[INSERT_EXP:.*]] = tensor.insert_slice %[[TILED_EXP_OUT]] into %[[THIRD_OUT_ARG]][%[[IV1]], 0] [64, 256] [1, 1]797// CHECK: scf.yield %[[INSERT_ADD]], %[[INSERT_MUL]], %[[INSERT_EXP]] :798// CHECK: }799// CHECK: return %[[LOOP_RESULT]]#1, %[[LOOP_RESULT]]#2 :800 801// -----802 803#map = affine_map<(d0, d1, d2) -> (d0, d1)>804#map1 = affine_map<(d0, d1, d2) -> (d2)>805#map2 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>806module {807 func.func @fuse_with_tilable_consumer_with_projected_permutations(%arg0: tensor<256x256xf32>, %arg1: tensor<256x256xf32>, %arg2: tensor<24xf32>) -> tensor<256x256x24xf32> {808 %c0 = arith.constant 0 : index809 %c64 = arith.constant 64 : index810 %c256 = arith.constant 256 : index811 %0 = tensor.empty() : tensor<256x256xf32>812 %1 = scf.for %arg3 = %c0 to %c256 step %c64 iter_args(%arg4 = %0) -> (tensor<256x256xf32>) {813 %extracted_slice = tensor.extract_slice %arg4[%arg3, 0] [64, 256] [1, 1] : tensor<256x256xf32> to tensor<64x256xf32>814 %extracted_slice_0 = tensor.extract_slice %arg0[%arg3, 0] [64, 256] [1, 1] : tensor<256x256xf32> to tensor<64x256xf32>815 %extracted_slice_1 = tensor.extract_slice %arg1[%arg3, 0] [64, 256] [1, 1] : tensor<256x256xf32> to tensor<64x256xf32>816 %4 = linalg.add ins(%extracted_slice_0, %extracted_slice_1 : tensor<64x256xf32>, tensor<64x256xf32>) outs(%extracted_slice : tensor<64x256xf32>) -> tensor<64x256xf32>817 %inserted_slice = tensor.insert_slice %4 into %arg4[%arg3, 0] [64, 256] [1, 1] : tensor<64x256xf32> into tensor<256x256xf32>818 scf.yield %inserted_slice : tensor<256x256xf32>819 }820 %2 = tensor.empty() : tensor<256x256x24xf32>821 %3 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel"]} ins(%1, %arg2 : tensor<256x256xf32>, tensor<24xf32>) outs(%2 : tensor<256x256x24xf32>) {822 ^bb0(%in: f32, %in_0: f32, %out: f32):823 %4 = arith.addf %in, %in_0 : f32824 linalg.yield %4 : f32825 } -> tensor<256x256x24xf32>826 return %3 : tensor<256x256x24xf32>827 }828}829 830module attributes {transform.with_named_sequence} {831 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {832 %consumer = transform.structured.match ops{["linalg.generic"]} in %arg1833 : (!transform.any_op) -> !transform.any_op834 %loop = transform.structured.match ops{["scf.for"]} in %arg1835 : (!transform.any_op) -> !transform.any_op836 %a, %b = transform.test.fuse_consumer %consumer into (%loop)837 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)838 transform.yield839 }840}841// CHECK: func.func @fuse_with_tilable_consumer_with_projected_permutations(842// CHECK-SAME: %[[VAL_0:.*]]: tensor<256x256xf32>, %[[VAL_1:.*]]: tensor<256x256xf32>, %[[VAL_2:.*]]: tensor<24xf32>) -> tensor<256x256x24xf32> {843// CHECK: %[[VAL_3:.*]] = arith.constant 0 : index844// CHECK: %[[VAL_4:.*]] = arith.constant 64 : index845// CHECK: %[[VAL_5:.*]] = arith.constant 256 : index846// CHECK: %[[VAL_6:.*]] = tensor.empty() : tensor<256x256xf32>847// CHECK: %[[VAL_7:.*]] = tensor.empty() : tensor<256x256x24xf32>848// 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>) {849// CHECK: %[[VAL_12:.*]] = tensor.extract_slice %[[VAL_10]]{{\[}}%[[VAL_9]], 0] [64, 256] [1, 1]850// CHECK: %[[VAL_13:.*]] = tensor.extract_slice %[[VAL_0]]{{\[}}%[[VAL_9]], 0] [64, 256] [1, 1]851// CHECK: %[[VAL_14:.*]] = tensor.extract_slice %[[VAL_1]]{{\[}}%[[VAL_9]], 0] [64, 256] [1, 1]852// CHECK: %[[VAL_15:.*]] = linalg.add ins(%[[VAL_13]], %[[VAL_14]] : tensor<64x256xf32>, tensor<64x256xf32>) outs(%[[VAL_12]] : tensor<64x256xf32>) -> tensor<64x256xf32>853// CHECK: %[[VAL_16:.*]] = tensor.insert_slice %[[VAL_15]] into %[[VAL_10]]{{\[}}%[[VAL_9]], 0] [64, 256] [1, 1]854// CHECK: %[[VAL_17:.*]] = tensor.extract_slice %[[VAL_2]][0] [24] [1] : tensor<24xf32> to tensor<24xf32>855// CHECK: %[[VAL_18:.*]] = tensor.extract_slice %[[VAL_11]]{{\[}}%[[VAL_9]], 0, 0] [64, 256, 24] [1, 1, 1]856// CHECK: %[[VAL_19:.*]] = linalg.generic857// CHECK-SAME: ins(%[[VAL_15]], %[[VAL_17]] : tensor<64x256xf32>, tensor<24xf32>) outs(%[[VAL_18]] : tensor<64x256x24xf32>) {858// CHECK: ^bb0(%[[VAL_20:.*]]: f32, %[[VAL_21:.*]]: f32, %[[VAL_22:.*]]: f32):859// CHECK: %[[VAL_23:.*]] = arith.addf %[[VAL_20]], %[[VAL_21]] : f32860// CHECK: linalg.yield %[[VAL_23]] : f32861// CHECK: } -> tensor<64x256x24xf32>862// CHECK: %[[VAL_24:.*]] = tensor.insert_slice %[[VAL_25:.*]] into %[[VAL_11]]{{\[}}%[[VAL_9]], 0, 0] [64, 256, 24] [1, 1, 1]863// CHECK: scf.yield %[[VAL_16]], %[[VAL_24]] : tensor<256x256xf32>, tensor<256x256x24xf32>864// CHECK: }865 866// -----867 868func.func @multi_slice_fusion1(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>, %arg3 : index) -> tensor<?xf32> {869 %c0 = arith.constant 0 : index870 %c1 = arith.constant 1 : index871 %dim0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>872 %dim1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>873 %loop:2 = scf.forall (%iv0) = (%c0) to (%dim0) step (%arg3) shared_outs(%init0 = %arg1, %init1 = %arg2) -> (tensor<?xf32>, tensor<?xf32>) {874 %tilesize = affine.min affine_map<(d0)[s0, s1] -> (s1, s0 - d0)>(%iv0)[%dim0, %arg3]875 %arg0_slice = tensor.extract_slice %arg0[%iv0, 0] [%tilesize, %dim1] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>876 %init0_slice = tensor.extract_slice %init0[%iv0] [%tilesize] [1] : tensor<?xf32> to tensor<?xf32>877 %init1_slice = tensor.extract_slice %init1[%iv0] [%tilesize] [1] : tensor<?xf32> to tensor<?xf32>878 %generic:2 = linalg.generic {879 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0)>],880 iterator_types = ["parallel", "reduction"]}881 ins(%arg0_slice : tensor<?x?xf32>) outs(%init0_slice, %init1_slice : tensor<?xf32>, tensor<?xf32>) {882 ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):883 %0 = arith.mulf %b0, %b1 : f32884 %1 = arith.addf %b0, %b2 : f32885 linalg.yield %0, %1 : f32, f32886 } -> (tensor<?xf32>, tensor<?xf32>)887 scf.forall.in_parallel {888 tensor.parallel_insert_slice %generic#0 into %init0[%iv0] [%tilesize] [1] : tensor<?xf32> into tensor<?xf32>889 tensor.parallel_insert_slice %generic#1 into %init1[%iv0] [%tilesize] [1] : tensor<?xf32> into tensor<?xf32>890 }891 }892 %empty = tensor.empty(%dim0) : tensor<?xf32>893 %result = linalg.generic {894 indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],895 iterator_types = ["parallel"]}896 ins(%loop#0, %loop#1 : tensor<?xf32>, tensor<?xf32>) outs(%empty : tensor<?xf32>) {897 ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):898 %0 = arith.addf %b0, %b1 : f32899 linalg.yield %0 : f32900 } -> tensor<?xf32>901 return %result : tensor<?xf32>902}903 904module attributes {transform.with_named_sequence} {905 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {906 %generics = transform.structured.match ops{["linalg.generic"]} in %arg1907 : (!transform.any_op) -> !transform.any_op908 %loop = transform.structured.match ops{["scf.forall"]} in %arg1909 : (!transform.any_op) -> !transform.any_op910 %producer, %consumer = transform.split_handle %generics : (!transform.any_op) -> (!transform.any_op, !transform.any_op)911 %a, %b = transform.test.fuse_consumer %consumer into (%loop)912 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)913 transform.yield914 }915}916// CHECK-LABEL: func @multi_slice_fusion1(917// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32>918// CHECK: %[[C0:.+]] = arith.constant 0919// CHECK: %[[DIM0:.+]] = tensor.dim %[[ARG0]], %[[C0]]920// CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM0]])921// CHECK: %[[RESULT:.+]]:3 = scf.forall (%[[IV:.+]]) =922// CHECK-SAME: , %[[INIT:[a-zA-Z0-9]+]] = %[[EMPTY]])923// CHECK: %[[TILESIZE:.+]] = affine.min924// CHECK-DAG: %[[GENERIC:.+]]:2 = linalg.generic925// CHECK-DAG: %[[INIT_SLICE:.+]] = tensor.extract_slice %[[INIT]][%[[IV]]] [%[[TILESIZE]]]926// CHECK: %[[FUSED:.+]] = linalg.generic927// CHECK-SAME: ins(%[[GENERIC]]#0, %[[GENERIC]]#1 :928// CHECK: tensor.parallel_insert_slice %[[FUSED]] into %[[INIT]][%[[IV]]] [%[[TILESIZE]]]929// CHECK: return %[[RESULT]]#2930 931 932// -----933 934func.func @multi_slice_fusion2(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>, %arg3 : index) -> tensor<?xf32> {935 %c0 = arith.constant 0 : index936 %c1 = arith.constant 1 : index937 %dim0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>938 %dim1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>939 %loop:2 = scf.forall (%iv0) = (%c0) to (%dim0) step (%arg3) shared_outs(%init0 = %arg1, %init1 = %arg2) -> (tensor<?xf32>, tensor<?xf32>) {940 %tilesize = affine.min affine_map<(d0)[s0, s1] -> (s1, s0 - d0)>(%iv0)[%dim0, %arg3]941 %arg0_slice = tensor.extract_slice %arg0[%iv0, 0] [%tilesize, %dim1] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>942 %init0_slice = tensor.extract_slice %init0[%iv0] [%tilesize] [1] : tensor<?xf32> to tensor<?xf32>943 %generic0 = linalg.generic {944 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>],945 iterator_types = ["parallel", "reduction"]}946 ins(%arg0_slice : tensor<?x?xf32>) outs(%init0_slice : tensor<?xf32>) {947 ^bb0(%b0 : f32, %b1 : f32):948 %0 = arith.mulf %b0, %b1 : f32949 linalg.yield %0 : f32950 } -> tensor<?xf32>951 %init1_slice = tensor.extract_slice %init1[%iv0] [%tilesize] [1] : tensor<?xf32> to tensor<?xf32>952 %generic1 = linalg.generic {953 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>],954 iterator_types = ["parallel", "reduction"]}955 ins(%arg0_slice : tensor<?x?xf32>) outs(%init1_slice: tensor<?xf32>) {956 ^bb0(%b0 : f32, %b1 : f32):957 %0 = arith.addf %b0, %b1 : f32958 linalg.yield %0: f32959 } -> tensor<?xf32>960 scf.forall.in_parallel {961 tensor.parallel_insert_slice %generic0 into %init0[%iv0] [%tilesize] [1] : tensor<?xf32> into tensor<?xf32>962 tensor.parallel_insert_slice %generic1 into %init1[%iv0] [%tilesize] [1] : tensor<?xf32> into tensor<?xf32>963 }964 }965 %empty = tensor.empty(%dim0) : tensor<?xf32>966 %result = linalg.generic {967 indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],968 iterator_types = ["parallel"]}969 ins(%loop#0, %loop#1 : tensor<?xf32>, tensor<?xf32>) outs(%empty : tensor<?xf32>) {970 ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):971 %0 = arith.addf %b0, %b1 : f32972 linalg.yield %0 : f32973 } -> tensor<?xf32>974 return %result : tensor<?xf32>975}976module attributes {transform.with_named_sequence} {977 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {978 %loop = transform.structured.match ops{["scf.forall"]} in %arg1979 : (!transform.any_op) -> !transform.any_op980 %generics = transform.structured.match ops{["linalg.generic"]} in %arg1981 : (!transform.any_op) -> !transform.any_op982 %producer1, %producer2, %consumer = transform.split_handle %generics : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)983 %a, %b = transform.test.fuse_consumer %consumer into (%loop)984 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)985 transform.yield986 }987}988 989// CHECK-LABEL: func @multi_slice_fusion2(990// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32>991// CHECK: %[[C0:.+]] = arith.constant 0992// CHECK: %[[DIM0:.+]] = tensor.dim %[[ARG0]], %[[C0]]993// CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM0]])994// CHECK: %[[RESULT:.+]]:3 = scf.forall (%[[IV:.+]]) =995// CHECK-SAME: , %[[INIT:[a-zA-Z0-9]+]] = %[[EMPTY]])996// CHECK: %[[TILESIZE:.+]] = affine.min997// CHECK: %[[GENERIC0:.+]] = linalg.generic998// CHECK: %[[GENERIC1:.+]] = linalg.generic999// CHECK-DAG: %[[INIT_SLICE:.+]] = tensor.extract_slice %[[INIT]][%[[IV]]] [%[[TILESIZE]]]1000// CHECK: %[[FUSED:.+]] = linalg.generic1001// CHECK-SAME: ins(%[[GENERIC0]], %[[GENERIC1]] :1002// CHECK: tensor.parallel_insert_slice %[[FUSED]] into %[[INIT]][%[[IV]]] [%[[TILESIZE]]]1003// CHECK: return %[[RESULT]]#21004 1005// -----1006 1007func.func @multi_slice_fusion_with_broadcast(%arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?xf32>,1008 %arg3 : index, %arg4 : index) -> tensor<?x?xf32> {1009 %c0 = arith.constant 0 : index1010 %c1 = arith.constant 1 : index1011 %c2 = arith.constant 2 : index1012 %dim0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>1013 %dim1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>1014 %dim2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32>1015 %loop:2 = scf.forall (%iv0, %iv1) = (%c0, %c0) to (%dim0, %dim1) step (%arg3, %arg4)1016 shared_outs(%init0 = %arg1, %init1 = %arg2) -> (tensor<?x?xf32>, tensor<?xf32>) {1017 %tilesize0 = affine.min affine_map<(d0)[s0, s1] -> (s1, s0 - d0)>(%iv0)[%dim0, %arg3]1018 %tilesize1 = affine.min affine_map<(d0)[s0, s1] -> (s1, s0 - d0)>(%iv1)[%dim1, %arg4]1019 %arg0_slice = tensor.extract_slice %arg0[%iv0, %iv1, 0] [%tilesize0, %tilesize1, %dim2] [1, 1, 1]1020 : tensor<?x?x?xf32> to tensor<?x?x?xf32>1021 %init0_slice = tensor.extract_slice %init0[%iv0, %iv1] [%tilesize0, %tilesize1] [1, 1]1022 : tensor<?x?xf32> to tensor<?x?xf32>1023 %generic0 = linalg.generic {1024 indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>],1025 iterator_types = ["parallel", "parallel", "reduction"]}1026 ins(%arg0_slice : tensor<?x?x?xf32>) outs(%init0_slice : tensor<?x?xf32>) {1027 ^bb0(%b0 : f32, %b1 : f32):1028 %0 = arith.mulf %b0, %b1 : f321029 linalg.yield %0 : f321030 } -> tensor<?x?xf32>1031 %init1_slice = tensor.extract_slice %init1[%iv0] [%tilesize0] [1] : tensor<?xf32> to tensor<?xf32>1032 %generic1 = linalg.generic {1033 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>],1034 iterator_types = ["parallel", "reduction"]}1035 ins(%generic0 : tensor<?x?xf32>) outs(%init1_slice: tensor<?xf32>) {1036 ^bb0(%b0 : f32, %b1 : f32):1037 %0 = arith.addf %b0, %b1 : f321038 linalg.yield %0: f321039 } -> tensor<?xf32>1040 scf.forall.in_parallel {1041 tensor.parallel_insert_slice %generic0 into %init0[%iv0, %iv1] [%tilesize0, %tilesize1] [1, 1]1042 : tensor<?x?xf32> into tensor<?x?xf32>1043 tensor.parallel_insert_slice %generic1 into %init1[%iv0] [%tilesize0] [1] : tensor<?xf32> into tensor<?xf32>1044 }1045 }1046 %empty = tensor.empty(%dim0, %dim1) : tensor<?x?xf32>1047 %result = linalg.generic {1048 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0, d1)>],1049 iterator_types = ["parallel", "parallel"]}1050 ins(%loop#0, %loop#1 : tensor<?x?xf32>, tensor<?xf32>) outs(%empty : tensor<?x?xf32>) {1051 ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):1052 %0 = arith.addf %b0, %b1 : f321053 linalg.yield %0 : f321054 } -> tensor<?x?xf32>1055 return %result : tensor<?x?xf32>1056}1057module attributes {transform.with_named_sequence} {1058 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {1059 %loop = transform.structured.match ops{["scf.forall"]} in %arg11060 : (!transform.any_op) -> !transform.any_op1061 %generics = transform.structured.match ops{["linalg.generic"]} in %arg11062 : (!transform.any_op) -> !transform.any_op1063 %producer_1, %producer_2, %consumer = transform.split_handle %generics : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)1064 %a, %b = transform.test.fuse_consumer %consumer into (%loop)1065 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)1066 transform.yield1067 }1068}1069// CHECK-LABEL: func @multi_slice_fusion_with_broadcast(1070// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32>1071// CHECK-DAG: %[[C0:.+]] = arith.constant 01072// CHECK-DAG: %[[C1:.+]] = arith.constant 11073// CHECK-DAG: %[[DIM0:.+]] = tensor.dim %[[ARG0]], %[[C0]]1074// CHECK-DAG: %[[DIM1:.+]] = tensor.dim %[[ARG0]], %[[C1]]1075// CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM0]], %[[DIM1]])1076// CHECK: %[[RESULT:.+]]:3 = scf.forall (%[[IV0:[a-zA-Z0-9]+]], %[[IV1:[a-zA-Z0-9]+]]) =1077// CHECK-SAME: , %[[INIT:[a-zA-Z0-9]+]] = %[[EMPTY]])1078// CHECK-DAG: %[[TILESIZE0:.+]] = affine.min {{.+}}(%[[IV0]])1079// CHECK-DAG: %[[TILESIZE1:.+]] = affine.min {{.+}}(%[[IV1]])1080// CHECK: %[[GENERIC0:.+]] = linalg.generic1081// CHECK: %[[GENERIC1:.+]] = linalg.generic1082// CHECK-DAG: %[[INIT_SLICE:.+]] = tensor.extract_slice %[[INIT]][%[[IV0]], %[[IV1]]] [%[[TILESIZE0]], %[[TILESIZE1]]]1083// CHECK: %[[FUSED:.+]] = linalg.generic1084// CHECK-SAME: ins(%[[GENERIC0]], %[[GENERIC1]] :1085// CHECK: tensor.parallel_insert_slice %[[FUSED]] into %[[INIT]][%[[IV0]], %[[IV1]]] [%[[TILESIZE0]], %[[TILESIZE1]]]1086// CHECK: return %[[RESULT]]#21087 1088// -----1089 1090func.func @multi_slice_fusion_invalid(%arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>,1091 %arg3 : index, %arg4 : index) -> tensor<?x?xf32> {1092 %c0 = arith.constant 0 : index1093 %c1 = arith.constant 1 : index1094 %c2 = arith.constant 2 : index1095 %dim0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>1096 %dim1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>1097 %dim2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32>1098 %loop:2 = scf.forall (%iv0, %iv1) = (%c0, %c0) to (%dim0, %dim1) step (%arg3, %arg4)1099 shared_outs(%init0 = %arg1, %init1 = %arg2) -> (tensor<?x?xf32>, tensor<?x?xf32>) {1100 %tilesize0 = affine.min affine_map<(d0)[s0, s1] -> (s1, s0 - d0)>(%iv0)[%dim0, %arg3]1101 %tilesize1 = affine.min affine_map<(d0)[s0, s1] -> (s1, s0 - d0)>(%iv1)[%dim1, %arg4]1102 %arg0_slice = tensor.extract_slice %arg0[%iv0, %iv1, 0] [%tilesize0, %tilesize1, %dim2] [1, 1, 1]1103 : tensor<?x?x?xf32> to tensor<?x?x?xf32>1104 %init0_slice = tensor.extract_slice %init0[%iv0, %iv1] [%tilesize0, %tilesize1] [1, 1]1105 : tensor<?x?xf32> to tensor<?x?xf32>1106 %generic0 = linalg.generic {1107 indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>],1108 iterator_types = ["parallel", "parallel", "reduction"]}1109 ins(%arg0_slice : tensor<?x?x?xf32>) outs(%init0_slice : tensor<?x?xf32>) {1110 ^bb0(%b0 : f32, %b1 : f32):1111 %0 = arith.mulf %b0, %b1 : f321112 linalg.yield %0 : f321113 } -> tensor<?x?xf32>1114 %init1_slice = tensor.extract_slice %init1[%iv0, %iv1] [%tilesize0, %tilesize1] [1, 1]1115 : tensor<?x?xf32> to tensor<?x?xf32>1116 %generic1 = linalg.generic {1117 indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>],1118 iterator_types = ["parallel", "parallel", "reduction"]}1119 ins(%arg0_slice : tensor<?x?x?xf32>) outs(%init1_slice: tensor<?x?xf32>) {1120 ^bb0(%b0 : f32, %b1 : f32):1121 %0 = arith.addf %b0, %b1 : f321122 linalg.yield %0: f321123 } -> tensor<?x?xf32>1124 scf.forall.in_parallel {1125 tensor.parallel_insert_slice %generic0 into %init0[%iv0, %iv1] [%tilesize0, %tilesize1] [1, 1]1126 : tensor<?x?xf32> into tensor<?x?xf32>1127 tensor.parallel_insert_slice %generic1 into %init1[%iv0, %iv1] [%tilesize0, %tilesize1] [1, 1]1128 : tensor<?x?xf32> into tensor<?x?xf32>1129 }1130 }1131 %empty = tensor.empty(%dim0, %dim1) : tensor<?x?xf32>1132 // expected-error @below {{failed to fuse consumer of slice}}1133 %result = linalg.generic {1134 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d0, d1)>],1135 iterator_types = ["parallel", "parallel"]}1136 ins(%loop#0, %loop#1 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%empty : tensor<?x?xf32>) {1137 ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):1138 %0 = arith.addf %b0, %b1 : f321139 linalg.yield %0 : f321140 } -> tensor<?x?xf32>1141 return %result : tensor<?x?xf32>1142}1143module attributes {transform.with_named_sequence} {1144 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {1145 %loop = transform.structured.match ops{["scf.forall"]} in %arg11146 : (!transform.any_op) -> !transform.any_op1147 %generics = transform.structured.match ops{["linalg.generic"]} in %arg11148 : (!transform.any_op) -> !transform.any_op1149 %producer_1, %producer_2, %consumer = transform.split_handle %generics : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)1150 %a, %b = transform.test.fuse_consumer %consumer into (%loop)1151 : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)1152 transform.yield1153 }1154}1155