695 lines · plain
1// RUN: mlir-opt %s -transform-interpreter -split-input-file -canonicalize -cse -verify-diagnostics | FileCheck %s2 3func.func @reduction_tile(%arg0: tensor<?x?xf32>, %out: tensor<?xf32>) -> tensor<?xf32> {4 %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,5 affine_map<(d0, d1) -> (d0)>],6 iterator_types = ["parallel", "reduction"]}7 ins(%arg0 : tensor<?x?xf32>)8 outs(%out : tensor<?xf32>) {9 ^bb0(%arg7: f32, %arg9: f32):10 %1 = arith.mulf %arg7, %arg7 : f3211 %2 = arith.addf %1, %arg9 : f3212 linalg.yield %2 : f3213 } -> tensor<?xf32>14 return %red : tensor<?xf32>15}16 17module attributes {transform.with_named_sequence} {18 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {19 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op20 %1, %2, %3, %loop = transform.structured.tile_reduction_using_for %021 by tile_sizes = [0, 5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)22 transform.yield23 }24}25 26// CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0)[s0] -> (-d0 + s0, 5)>27// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1) -> (d0, d1)>28// CHECK: func @reduction_tile(%[[ARG0:.+]]: tensor<?x?xf32>, %[[ARG1:.+]]: tensor<?xf32>29// CHECK-DAG: %[[I:.*]] = arith.constant 0.000000e+00 : f3230// CHECK-DAG: %[[C5:.*]] = arith.constant 5 : index31// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index32// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index33// CHECK-DAG: %[[D0:.*]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x?xf32>34// CHECK-DAG: %[[D1:.*]] = tensor.dim %[[ARG0]], %[[C1]] : tensor<?x?xf32>35// CHECK: %[[E:.*]] = tensor.empty(%[[D0]]) : tensor<?x5xf32>36// CHECK: %[[F:.*]] = linalg.fill ins(%[[I]] : f32) outs(%[[E]] : tensor<?x5xf32>) -> tensor<?x5xf32>37// CHECK: %[[L:.*]] = scf.for %[[K:.*]] = %[[C0]] to %[[D1]] step %[[C5]] iter_args(%[[ARG3:.*]] = %[[F]]) -> (tensor<?x5xf32>) {38// CHECK: %[[PS:.*]] = affine.min #[[MAP0]](%[[K]])[%[[D1]]]39// CHECK: %[[EXT2:.*]] = tensor.extract_slice %[[ARG0]][0, %[[K:.*]]] [%[[D0]], %[[PS]]] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>40// CHECK: %[[EXT:.*]] = tensor.extract_slice %[[ARG3]][0, 0] [%[[D0]], %[[PS]]] [1, 1] : tensor<?x5xf32> to tensor<?x?xf32>41// CHECK: %[[PR:.*]] = linalg.generic {indexing_maps = [#[[MAP1]], #[[MAP1]]], iterator_types = ["parallel", "parallel"]} ins(%[[EXT2]] : tensor<?x?xf32>) outs(%[[EXT]] : tensor<?x?xf32>) {42// CHECK: arith.mulf43// CHECK: arith.addf44// CHECK: linalg.yield45// CHECK: } -> tensor<?x?xf32>46// CHECK: %[[INS:.*]] = tensor.insert_slice %[[PR]] into %[[ARG3]][0, 0] [%[[D0]], %[[PS]]] [1, 1] : tensor<?x?xf32> into tensor<?x5xf32>47// CHECK: scf.yield %[[INS]] : tensor<?x5xf32>48// CHECK: }49// CHECK: %[[R:.*]] = linalg.reduce ins(%[[L]] : tensor<?x5xf32>) outs(%[[ARG1]] : tensor<?xf32>) dimensions = [1]50// CHECK: arith.addf51// CHECK: linalg.yield52// CHECK: }53// CHECK: return %[[R]] : tensor<?xf32>54 55// -----56 57func.func @reduction_tile_transpose(%arg0: tensor<?x?xf32>, %out: tensor<?xf32>) -> tensor<?xf32> {58 %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,59 affine_map<(d0, d1) -> (d1)>],60 iterator_types = ["reduction", "parallel"]}61 ins(%arg0 : tensor<?x?xf32>)62 outs(%out : tensor<?xf32>) {63 ^bb0(%arg7: f32, %arg9: f32):64 %42 = arith.addf %arg7, %arg9 : f3265 linalg.yield %42 : f3266 } -> tensor<?xf32>67 return %red : tensor<?xf32>68}69 70module attributes {transform.with_named_sequence} {71 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {72 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op73 %1, %2, %3, %loop = transform.structured.tile_reduction_using_for %074 by tile_sizes = [5, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)75 transform.yield76 }77}78 79// CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0)[s0] -> (-d0 + s0, 5)>80// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1) -> (d0, d1)>81// CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0, d1) -> (d1, d0)>82// CHECK: func @reduction_tile_transpose83// CHECK: tensor.empty(%{{.*}}) : tensor<?x5xf32>84// CHECK: linalg.fill {{.*}} : tensor<?x5xf32>) -> tensor<?x5xf32>85// CHECK: scf.for86// CHECK: %[[EXT:.*]] = tensor.extract_slice %[[ARG3:.*]][0, 0] [%[[D0:.*]], %[[D1:.*]]] [1, 1] : tensor<?x5xf32> to tensor<?x?xf32>87// CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[MAP1]], #[[MAP2]]], iterator_types = ["parallel", "parallel"]} ins(%[[L:.*]] : tensor<?x?xf32>) outs(%[[EXT]] : tensor<?x?xf32>)88// CHECK: %[[INS:.*]] = tensor.insert_slice %[[R]] into %[[ARG3]][0, 0] [%[[D0]], %[[D1]]] [1, 1] : tensor<?x?xf32> into tensor<?x5xf32>89// CHECK: scf.yield {{.*}} : tensor<?x5xf32>90// CHECK: }91// CHECK: linalg.reduce92// CHECK: return93 94// -----95 96func.func @reduction_tile_parallel(97 %arg0: tensor<?x?xf32>, %out: tensor<?xf32>) -> tensor<?xf32> {98 %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,99 affine_map<(d0, d1) -> (d0)>],100 iterator_types = ["parallel", "reduction"]}101 ins(%arg0 : tensor<?x?xf32>)102 outs(%out : tensor<?xf32>) {103 ^bb0(%arg7: f32, %arg9: f32):104 %1 = arith.mulf %arg7, %arg7 : f32105 %2 = arith.addf %1, %arg9 : f32106 linalg.yield %2 : f32107 } -> tensor<?xf32>108 return %red : tensor<?xf32>109}110 111module attributes {transform.with_named_sequence} {112 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {113 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op114 %1, %2, %3, %loop = transform.structured.tile_reduction_using_forall %0115 by num_threads = [0, 5] tile_sizes = [] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)116 transform.yield117 }118}119 120// CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0)[s0] -> (-(d0 * (s0 ceildiv 5)) + s0, s0 ceildiv 5)>121// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0) -> (0, d0)>122// CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0)[s0] -> (d0 * (s0 ceildiv 5))>123// CHECK-DAG: #[[MAP3:.*]] = affine_map<(d0, d1) -> (d0, d1)>124// CHECK-DAG: #[[MAP4:.*]] = affine_map<(d0, d1) -> (d0)>125// CHECK: func @reduction_tile_parallel(%[[ARG0:.+]]: tensor<?x?xf32>, %[[ARG1:.+]]: tensor<?xf32>126// CHECK-DAG: %[[I:.*]] = arith.constant 0.000000e+00 : f32127// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index128// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index129// CHECK-DAG: %[[D0:.*]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x?xf32>130// CHECK-DAG: %[[D1:.*]] = tensor.dim %[[ARG0]], %[[C1]] : tensor<?x?xf32>131// CHECK: %[[E:.*]] = tensor.empty(%[[D0]]) : tensor<?x5xf32>132// CHECK: %[[F:.*]] = linalg.fill ins(%[[I]] : f32) outs(%[[E]] : tensor<?x5xf32>) -> tensor<?x5xf32>133// CHECK: %[[L:.*]] = scf.forall (%[[IV:.+]]) in (5) shared_outs(%[[ARG3:.+]] = %[[F]]) -> (tensor<?x5xf32>) {134// CHECK-DAG: %[[TS0:.+]] = affine.min #[[MAP0]](%[[IV]])[%[[D1]]]135// CHECK-DAG: %[[TS1:.+]] = affine.max #[[MAP1]](%[[TS0]])136// CHECK-DAG: %[[ET:.+]] = tensor.extract_slice %[[ARG3:.+]][0, %[[IV]]] [%[[D0]], 1] [1, 1] : tensor<?x5xf32> to tensor<?xf32>137// CHECK-DAG: %[[TINDEX:.+]] = affine.apply #[[MAP2]](%[[IV]])[%[[D1]]]138// CHECK-DAG: %[[INCHUNK:.+]] = tensor.extract_slice %[[ARG0]][0, %[[TINDEX]]] [%[[D0]], %[[TS1]]] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>139// CHECK: %[[PARTIAL:.+]] = linalg.generic {indexing_maps = [#[[MAP3]], #[[MAP4]]], iterator_types = ["parallel", "reduction"]} ins(%[[INCHUNK]] : tensor<?x?xf32>) outs(%[[ET]] : tensor<?xf32>) {140// CHECK: arith.mulf141// CHECK: arith.addf142// CHECK: linalg.yield143// CHECK: } -> tensor<?xf32>144// CHECK: scf.forall.in_parallel {145// CHECK: tensor.parallel_insert_slice %[[PARTIAL]] into %[[ARG3]][0, %[[IV]]] [%[[D0]], 1] [1, 1] : tensor<?xf32> into tensor<?x5xf32>146// CHECK: }147// CHECK: }148// CHECK: %[[R:.*]] = linalg.reduce ins(%[[L]] : tensor<?x5xf32>) outs(%[[ARG1]] : tensor<?xf32>) dimensions = [1]149// CHECK: {150// CHECK: arith.addf151// CHECK: linalg.yield152// CHECK: }153// CHECK: return %[[R]] : tensor<?xf32>154 155// -----156 157func.func @matmul_tile_parallel(158 %A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %out: tensor<?x?xf32>) -> tensor<?x?xf32> {159 %matmul = linalg.matmul ins(%A, %B: tensor<?x?xf32>, tensor<?x?xf32>)160 outs(%out: tensor<?x?xf32>) -> tensor<?x?xf32>161 return %matmul : tensor<?x?xf32>162}163 164module attributes {transform.with_named_sequence} {165 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {166 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op167 %1, %2, %3, %loop = transform.structured.tile_reduction_using_forall %0168 by num_threads = [0, 0, 5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)169 transform.yield170 }171}172 173// CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0)[s0] -> (-(d0 * (s0 ceildiv 5)) + s0, s0 ceildiv 5)>174// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0) -> (0, d0)>175// CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0)[s0] -> (d0 * (s0 ceildiv 5))>176// CHECK: func @matmul_tile_parallel(%[[ARG0:.+]]: tensor<?x?xf32>, %[[ARG1:.+]]: tensor<?x?xf32>, %[[ARG2:.+]]: tensor<?x?xf32>177// CHECK-DAG: %[[I:.*]] = arith.constant 0.000000e+00 : f32178// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index179// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index180// CHECK-DAG: %[[D0:.*]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x?xf32>181// CHECK-DAG: %[[D1:.*]] = tensor.dim %[[ARG0]], %[[C1]] : tensor<?x?xf32>182// CHECK-DAG: %[[D2:.*]] = tensor.dim %[[ARG1]], %[[C1]] : tensor<?x?xf32>183// CHECK: %[[E:.*]] = tensor.empty(%[[D0]], %[[D2]]) : tensor<?x?x5xf32>184// CHECK: %[[F:.*]] = linalg.fill ins(%[[I]] : f32) outs(%[[E]] : tensor<?x?x5xf32>) -> tensor<?x?x5xf32>185// CHECK: %[[L:.*]] = scf.forall (%[[IV:.+]]) in (5) shared_outs(%[[ARG3:.+]] = %[[F]]) -> (tensor<?x?x5xf32>) {186// CHECK-DAG: %[[TS0:.+]] = affine.min #[[MAP0]](%[[IV]])[%[[D1]]]187// CHECK-DAG: %[[TS1:.+]] = affine.max #[[MAP1]](%[[TS0]])188// CHECK-DAG: %[[ET:.+]] = tensor.extract_slice %[[ARG3:.+]][0, 0, %[[IV]]] [%[[D0]], %[[D2]], 1] [1, 1, 1] : tensor<?x?x5xf32> to tensor<?x?xf32>189// CHECK-DAG: %[[TINDEX:.+]] = affine.apply #[[MAP2]](%[[IV]])[%[[D1]]]190// CHECK-DAG: %[[INCHUNKA:.+]] = tensor.extract_slice %[[ARG0]][0, %[[TINDEX]]] [%[[D0]], %[[TS1]]] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>191// CHECK-DAG: %[[INCHUNKB:.+]] = tensor.extract_slice %[[ARG1]][%[[TINDEX]], 0] [%[[TS1]], %[[D2]]] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>192// CHECK: %[[PARTIAL:.+]] = linalg.matmul ins(%[[INCHUNKA]], %[[INCHUNKB]] : tensor<?x?xf32>, tensor<?x?xf32>) outs(%[[ET]] : tensor<?x?xf32>) -> tensor<?x?xf32>193// CHECK: scf.forall.in_parallel {194// CHECK: tensor.parallel_insert_slice %[[PARTIAL]] into %[[ARG3]][0, 0, %[[IV]]] [%[[D0]], %[[D2]], 1] [1, 1, 1] : tensor<?x?xf32> into tensor<?x?x5xf32>195// CHECK: }196// CHECK: }197// CHECK: %[[R:.*]] = linalg.reduce ins(%[[L]] : tensor<?x?x5xf32>) outs(%[[ARG2]] : tensor<?x?xf32>) dimensions = [2]198// CHECK: arith.addf199// CHECK: linalg.yield200// CHECK: }201// CHECK: return %[[R]] : tensor<?x?xf32>202 203// -----204 205func.func @reduction_untiled_forall(206 %arg0: tensor<?x?xf32>, %out: tensor<?xf32>) -> tensor<?xf32> {207 // expected-error @below {{tiling parallel dimensions is not supported with partial reduction tiling strategies}}208 %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,209 affine_map<(d0, d1) -> (d0)>],210 iterator_types = ["parallel", "reduction"]}211 ins(%arg0 : tensor<?x?xf32>)212 outs(%out : tensor<?xf32>) {213 ^bb0(%arg7: f32, %arg9: f32):214 %1 = arith.mulf %arg7, %arg7 : f32215 %2 = arith.addf %1, %arg9 : f32216 linalg.yield %2 : f32217 } -> tensor<?xf32>218 return %red : tensor<?xf32>219}220 221module attributes {transform.with_named_sequence} {222 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {223 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op224 // expected-error @below {{could not tile reduction}}225 %1, %2, %3, %loop = transform.structured.tile_reduction_using_forall %0226 by num_threads = [5] tile_sizes = [3] mapping = [#gpu.thread<x>] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)227 transform.yield228 }229}230 231// -----232 233#map = affine_map<(d0, d1) -> (d0, d1)>234#map1 = affine_map<(d0, d1) -> (d0)>235 236module {237 func.func @fail_for_float_neutral(%arg0: tensor<?x?xf32>, %arg1: tensor<?xf32>) -> tensor<?xf32> {238 // expected-error @below {{'linalg.generic' op Failed to get an identity value for the reduction operation.}}239 %0 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<?x?xf32>) outs(%arg1 : tensor<?xf32>) {240 ^bb0(%in: f32, %out: f32):241 %1 = llvm.fmul %in, %in : f32242 %2 = llvm.fadd %1, %out : f32243 linalg.yield %2 : f32244 } -> tensor<?xf32>245 return %0 : tensor<?xf32>246 }247 module attributes {transform.with_named_sequence} {248 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {249 %0 = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op250 // expected-error @below {{failed to tile using partial reduction}}251 %fill_op, %split_linalg_op, %combining_linalg_op, %for_op = transform.structured.tile_reduction_using_for %0 by tile_sizes = [0, 5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)252 transform.yield253 }254 }255}256 257// -----258 259#map = affine_map<(d0, d1, d2) -> (d1, d2)>260#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>261#map2 = affine_map<(d0, d1, d2) -> (d0)>262module {263 func.func @reduction_tile_multiple_reduction(%arg0: tensor<86x128xf32>, %arg1: tensor<4096x86x128xf32>, %arg2: tensor<4096xf32>) -> tensor<4096xf32> {264 %0 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "reduction", "reduction"]} ins(%arg0, %arg1 : tensor<86x128xf32>, tensor<4096x86x128xf32>) outs(%arg2 : tensor<4096xf32>) {265 ^bb0(%in: f32, %in_0: f32, %out: f32):266 %1 = arith.mulf %in, %in_0 : f32267 %2 = arith.addf %1, %out : f32268 linalg.yield %2 : f32269 } -> tensor<4096xf32>270 return %0 : tensor<4096xf32>271 }272 module attributes {transform.with_named_sequence} {273 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {274 %0 = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op275 %fill_op, %split_linalg_op, %combining_linalg_op, %for_op = transform.structured.tile_reduction_using_for %0 by tile_sizes = [0, 2, 64] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)276 transform.yield277 }278 }279}280 281// CHECK: func @reduction_tile_multiple_reduction(%[[ARG0:.+]]: tensor<86x128xf32>, %[[ARG1:.+]]: tensor<4096x86x128xf32>, %[[ARG2:.+]]: tensor<4096xf32>282// CHECK: %[[F:.*]] = linalg.fill ins(%{{.*}} : f32) outs(%{{.*}} : tensor<4096x2x64xf32>) -> tensor<4096x2x64xf32>283// CHECK: %[[L0:.*]] = scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[ARG3:.*]] = %[[F]]) -> (tensor<4096x2x64xf32>)284// CHECK: %[[L1:.*]] = scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[ARG4:.*]] = %[[ARG3]]) -> (tensor<4096x2x64xf32>)285// CHECK: %[[OUT:.*]] = linalg.generic {indexing_maps = [{{.*}}, {{.*}}, {{.*}}], iterator_types = ["parallel", "parallel", "parallel"]} ins(%{{.*}}, %{{.*}}: tensor<2x64xf32>, tensor<4096x2x64xf32>) outs(%{{.*}}: tensor<4096x2x64xf32>)286// CHECK: scf.yield %[[OUT]] : tensor<4096x2x64xf32>287// CHECK: scf.yield %[[L1]] : tensor<4096x2x64xf32>288// CHECK: %[[OUT2:.*]] = linalg.reduce ins(%{{.*}} : tensor<4096x2x64xf32>) outs(%{{.*}} : tensor<4096xf32>)289// CHECK: return %[[OUT2]] : tensor<4096xf32>290 291// -----292 293func.func @reduction_tile_multiple_results(%arg0: tensor<?x?xf32>, %out: tensor<?xf32>, %out2: tensor<?xf32>) -> (tensor<?xf32>, tensor<?xf32>) {294 %red:2 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,295 affine_map<(d0, d1) -> (d0)>,296 affine_map<(d0, d1) -> (d0)>],297 iterator_types = ["parallel", "reduction"]}298 ins(%arg0 : tensor<?x?xf32>)299 outs(%out, %out2 : tensor<?xf32>, tensor<?xf32>) {300 ^bb0(%arg7: f32, %arg9: f32, %arg9_1: f32):301 %1 = arith.mulf %arg7, %arg7 : f32302 %2 = arith.addf %1, %arg9 : f32303 %3 = arith.maximumf %1, %arg9_1 : f32304 linalg.yield %2, %3 : f32, f32305 } -> (tensor<?xf32>, tensor<?xf32>)306 return %red#0, %red#1 : tensor<?xf32>, tensor<?xf32>307}308 309module attributes {transform.with_named_sequence} {310 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {311 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op312 %1, %12, %2, %3, %4, %loop = transform.structured.tile_reduction_using_for %0313 by tile_sizes = [0, 5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)314 transform.yield315 }316}317 318// CHECK: func @reduction_tile_multiple_results319// CHECK-DAG: %[[SUM_ID:.+]] = arith.constant 0.000000e+00 : f32320// CHECK-DAG: %[[MAX_ID:.+]] = arith.constant 0xFF800000 : f32321// CHECK-DAG: %[[SUM_INIT:.+]] = linalg.fill ins(%[[SUM_ID]] : f32) outs(%{{.*}} : tensor<?x5xf32>) -> tensor<?x5xf32>322// CHECK-DAG: %[[MAX_INIT:.+]] = linalg.fill ins(%[[MAX_ID]] : f32) outs(%{{.*}} : tensor<?x5xf32>) -> tensor<?x5xf32>323// CHECK: %[[OUT:.+]]:2 = scf.for324// CHECK-SAME: iter_args(%[[SUM:.+]] = %[[SUM_INIT]], %[[MAX:.+]] = %[[MAX_INIT]])325// CHECK: %[[UPDATED:.*]]:2 = linalg.generic326// CHECK: arith.mulf327// CHECK: arith.addf328// CHECK: arith.maximumf329// CHECK: %[[INSERT1:.+]] = tensor.insert_slice %[[UPDATED]]#0 into %[[SUM]]330// CHECK: %[[INSERT2:.+]] = tensor.insert_slice %[[UPDATED]]#1 into %[[MAX]]331// CHECK: scf.yield %[[INSERT1]], %[[INSERT1]]332// CHECK: linalg.reduce333// CHECK: arith.addf334// CHECK: linalg.reduce335// CHECK: arith.maximumf336 337// -----338 339func.func @reduction_tile_multi_dim_transpose(%arg0: tensor<?x?x?xf32>, %out: tensor<?x?xf32>) -> tensor<?x?xf32> {340 %red = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,341 affine_map<(d0, d1, d2) -> (d2, d0)>],342 iterator_types = ["parallel", "reduction", "parallel"]}343 ins(%arg0 : tensor<?x?x?xf32>)344 outs(%out : tensor<?x?xf32>) {345 ^bb0(%arg7: f32, %arg9: f32):346 %42 = arith.addf %arg7, %arg9 : f32347 linalg.yield %42 : f32348 } -> tensor<?x?xf32>349 return %red : tensor<?x?xf32>350}351 352module attributes {transform.with_named_sequence} {353 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {354 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op355 %1, %2, %3, %loop = transform.structured.tile_reduction_using_for %0356 by tile_sizes = [0, 5, 0] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)357 transform.yield358 }359}360 361// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>362// CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0, d1, d2) -> (d2, d0, d1)>363// CHECK: func @reduction_tile_multi_dim_transpose364// CHECK: tensor.empty(%{{.*}}) : tensor<?x?x5xf32>365// CHECK: linalg.fill {{.*}} : tensor<?x?x5xf32>) -> tensor<?x?x5xf32>366// CHECK: scf.for367// CHECK: %[[K:.*]] = affine.min368// CHECK: %[[EXT:.*]] = tensor.extract_slice %[[ARG3:.*]][0, 0, 0] [%[[D2:.*]], %[[D0:.*]], %[[K]]] [1, 1, 1] : tensor<?x?x5xf32> to tensor<?x?x?xf32>369// CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[MAP1]], #[[MAP2]]], iterator_types = ["parallel", "parallel", "parallel"]} ins(%[[L:.*]] : tensor<?x?x?xf32>) outs(%[[EXT]] : tensor<?x?x?xf32>)370// CHECK: %[[INS:.*]] = tensor.insert_slice %[[R]] into %[[ARG3]][0, 0, 0] [%[[D2]], %[[D0]], %[[K]]] [1, 1, 1] : tensor<?x?x?xf32> into tensor<?x?x5xf32>371// CHECK: scf.yield {{.*}} : tensor<?x?x5xf32>372// CHECK: }373// CHECK: linalg.reduce374// CHECK: return375 376// -----377 378// Check that only one of the reduction dimension can be tiled (in this case outer).379 380#map = affine_map<(d0, d1, d2) -> (d1, d2)>381#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>382#map2 = affine_map<(d0, d1, d2) -> (d0)>383module {384 func.func @reduction_tile_single_of_multiple_reduction_outer(385 %arg0: tensor<86x128xf32>, %arg1: tensor<4096x86x128xf32>, %arg2: tensor<4096xf32>) -> tensor<4096xf32> {386 %0 = linalg.generic {387 indexing_maps = [#map, #map1, #map2],388 iterator_types = ["parallel", "reduction", "reduction"]}389 ins(%arg0, %arg1 : tensor<86x128xf32>, tensor<4096x86x128xf32>) outs(%arg2 : tensor<4096xf32>) {390 ^bb0(%in: f32, %in_0: f32, %out: f32):391 %1 = arith.mulf %in, %in_0 : f32392 %2 = arith.addf %1, %out : f32393 linalg.yield %2 : f32394 } -> tensor<4096xf32>395 return %0 : tensor<4096xf32>396 }397 module attributes {transform.with_named_sequence} {398 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {399 %0 = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op400 %fill_op, %split_linalg_op, %combining_linalg_op, %for_op =401 transform.structured.tile_reduction_using_for %0 reduction_dims = [1] by tile_sizes = [0, 2]402 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)403 transform.yield404 }405 }406}407// CHECK: #[[INIT_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>408// CHECK: @reduction_tile_single_of_multiple_reduction_outer(409// CHECK-SAME: %[[INIT:[a-zA-Z0-9]+]]: tensor<4096xf32>410// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index411// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index412// CHECK-DAG: %[[C86:.+]] = arith.constant 86 : index413// CHECK-DAG: %[[EMPTY:.+]] = tensor.empty() : tensor<4096x2xf32>414// CHECK: %[[FILL:.+]] = linalg.fill415// CHECK-SAME: outs(%[[EMPTY]] :416// CHECK: %[[RESULT:.+]] = scf.for %[[IV:[a-zA-Z0-9]+]] = %[[C0]] to %[[C86]] step %[[C2]]417// CHECK-SAME: iter_args(%[[ITER_ARG:.+]] = %[[FILL]])418// CHECK: %[[PARTIAL_RESULT:.+]] = linalg.generic419// CHECK-SAME: indexing_maps = [#{{.+}}, #{{.+}}, #[[INIT_MAP]]]420// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]421// CHECK-SAME: outs(%[[ITER_ARG]] :422// CHECK: scf.yield %[[PARTIAL_RESULT]]423// CHECK: %[[REDUCE:.+]] = linalg.reduce424// CHECK-SAME: ins(%[[RESULT]] :425// CHECK-SAME: outs(%[[INIT]] :426// CHECK-SAME: dimensions = [1]427// CHECK: return %[[REDUCE]]428 429// -----430 431// Check that only one of the reduction dimension can be tiled (in this case inner).432 433#map = affine_map<(d0, d1, d2) -> (d1, d2)>434#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>435#map2 = affine_map<(d0, d1, d2) -> (d0)>436module {437 func.func @reduction_tile_single_of_multiple_reduction_inner(438 %arg0: tensor<86x128xf32>, %arg1: tensor<4096x86x128xf32>, %arg2: tensor<4096xf32>) -> tensor<4096xf32> {439 %0 = linalg.generic {440 indexing_maps = [#map, #map1, #map2],441 iterator_types = ["parallel", "reduction", "reduction"]}442 ins(%arg0, %arg1 : tensor<86x128xf32>, tensor<4096x86x128xf32>) outs(%arg2 : tensor<4096xf32>) {443 ^bb0(%in: f32, %in_0: f32, %out: f32):444 %1 = arith.mulf %in, %in_0 : f32445 %2 = arith.addf %1, %out : f32446 linalg.yield %2 : f32447 } -> tensor<4096xf32>448 return %0 : tensor<4096xf32>449 }450 module attributes {transform.with_named_sequence} {451 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {452 %0 = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op453 %fill_op, %split_linalg_op, %combining_linalg_op, %for_op =454 transform.structured.tile_reduction_using_for %0 reduction_dims = [2] by tile_sizes = [0, 0, 64]455 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)456 transform.yield457 }458 }459}460// CHECK: #[[INIT_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>461// CHECK: @reduction_tile_single_of_multiple_reduction_inner(462// CHECK-SAME: %[[INIT:[a-zA-Z0-9]+]]: tensor<4096xf32>463// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index464// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index465// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index466// CHECK-DAG: %[[EMPTY:.+]] = tensor.empty() : tensor<4096x64xf32>467// CHECK: %[[FILL:.+]] = linalg.fill468// CHECK-SAME: outs(%[[EMPTY]] :469// CHECK: %[[RESULT:.+]] = scf.for %[[IV:[a-zA-Z0-9]+]] = %[[C0]] to %[[C128]] step %[[C64]]470// CHECK-SAME: iter_args(%[[ITER_ARG:.+]] = %[[FILL]])471// CHECK: %[[PARTIAL_RESULT:.+]] = linalg.generic472// CHECK-SAME: indexing_maps = [#{{.+}}, #{{.+}}, #[[INIT_MAP]]]473// CHECK-SAME: iterator_types = ["parallel", "reduction", "parallel"]474// CHECK-SAME: outs(%[[ITER_ARG]] :475// CHECK: scf.yield %[[PARTIAL_RESULT]]476// CHECK: %[[REDUCE:.+]] = linalg.reduce477// CHECK-SAME: ins(%[[RESULT]] :478// CHECK-SAME: outs(%[[INIT]] :479// CHECK-SAME: dimensions = [1]480// CHECK: return %[[REDUCE]]481 482// -----483 484// Check that both the reduction dimensions are tiled but the dimensions in the output are swapped.485 486#map = affine_map<(d0, d1, d2) -> (d1, d2)>487#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>488#map2 = affine_map<(d0, d1, d2) -> (d0)>489module {490 func.func @reduction_tile_single_of_multiple_reduction_reversed(491 %arg0: tensor<86x128xf32>, %arg1: tensor<4096x86x128xf32>, %arg2: tensor<4096xf32>) -> tensor<4096xf32> {492 %0 = linalg.generic {493 indexing_maps = [#map, #map1, #map2],494 iterator_types = ["parallel", "reduction", "reduction"]}495 ins(%arg0, %arg1 : tensor<86x128xf32>, tensor<4096x86x128xf32>) outs(%arg2 : tensor<4096xf32>) {496 ^bb0(%in: f32, %in_0: f32, %out: f32):497 %1 = arith.mulf %in, %in_0 : f32498 %2 = arith.addf %1, %out : f32499 linalg.yield %2 : f32500 } -> tensor<4096xf32>501 return %0 : tensor<4096xf32>502 }503 module attributes {transform.with_named_sequence} {504 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {505 %0 = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op506 %fill_op, %split_linalg_op, %combining_linalg_op, %for_op =507 transform.structured.tile_reduction_using_for %0 reduction_dims = [2, 1] by tile_sizes = [0, 2, 64]508 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)509 transform.yield510 }511 }512}513// CHECK: #[[INIT_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d2, d1)>514// CHECK: @reduction_tile_single_of_multiple_reduction_reversed(515// CHECK-SAME: %[[INIT:[a-zA-Z0-9]+]]: tensor<4096xf32>516// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index517// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index518// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index519// CHECK-DAG: %[[C86:.+]] = arith.constant 86 : index520// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index521// CHECK-DAG: %[[EMPTY:.+]] = tensor.empty() : tensor<4096x64x2xf32>522// CHECK: %[[FILL:.+]] = linalg.fill523// CHECK-SAME: outs(%[[EMPTY]] :524// CHECK: %[[RESULT:.+]] = scf.for %[[IV0:[a-zA-Z0-9]+]] = %[[C0]] to %[[C86]] step %[[C2]]525// CHECK-SAME: iter_args(%[[ITER_ARG:.+]] = %[[FILL]])526// CHECK: %[[RESULT0:.+]] = scf.for %[[IV1:[a-zA-Z0-9]+]] = %[[C0]] to %[[C128]] step %[[C64]]527// CHECK-SAME: iter_args(%[[ITER_ARG0:.+]] = %[[ITER_ARG]])528// CHECK: %[[PARTIAL_RESULT:.+]] = linalg.generic529// CHECK-SAME: indexing_maps = [#{{.+}}, #{{.+}}, #[[INIT_MAP]]]530// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]531// CHECK-SAME: outs(%[[ITER_ARG0]] :532// CHECK: scf.yield %[[PARTIAL_RESULT]]533// CHECK scf.yield %[[RESULT0]]534// CHECK: %[[REDUCE:.+]] = linalg.reduce535// CHECK-SAME: ins(%[[RESULT]] :536// CHECK-SAME: outs(%[[INIT]] :537// CHECK-SAME: dimensions = [1, 2]538// CHECK: return %[[REDUCE]]539 540// -----541 542func.func @reduction_tile_parallel_using_tile_sizes(543 %arg0: tensor<?x?xf32>, %out: tensor<?xf32>) -> tensor<?xf32> {544 %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,545 affine_map<(d0, d1) -> (d0)>],546 iterator_types = ["parallel", "reduction"]}547 ins(%arg0 : tensor<?x?xf32>)548 outs(%out : tensor<?xf32>) {549 ^bb0(%arg7: f32, %arg9: f32):550 %1 = arith.mulf %arg7, %arg7 : f32551 %2 = arith.addf %1, %arg9 : f32552 linalg.yield %2 : f32553 } -> tensor<?xf32>554 return %red : tensor<?xf32>555}556// CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> (s0 ceildiv 5)>557// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0)[s0] -> (-d0 + s0, 5)>558// CHECK-DAG: #[[MAP2:.*]] = affine_map<()[s0] -> (s0 floordiv 5)>559// CHECK: func @reduction_tile_parallel_using_tile_sizes(%[[ARG0:.+]]: tensor<?x?xf32>, %[[ARG1:.+]]: tensor<?xf32>560// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index561// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index562// CHECK-DAG: %[[D0:.*]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x?xf32>563// CHECK-DAG: %[[D1:.*]] = tensor.dim %[[ARG0]], %[[C1]] : tensor<?x?xf32>564// CHECK-DAG: %[[PARALLEL_DIM:.+]] = affine.apply #[[MAP0]]()[%[[D1]]]565// CHECK: %[[E:.*]] = tensor.empty(%[[D0]], %[[PARALLEL_DIM]]) : tensor<?x?xf32>566// CHECK: %[[F:.*]] = linalg.fill567// CHECK-SAME: outs(%[[E]] :568// CHECK: %[[L:.*]] = scf.forall (%[[IV:.+]]) = (0) to (%[[D1]]) step (5) shared_outs(%[[ARG3:.+]] = %[[F]])569// CHECK-DAG: %[[TS0:.+]] = affine.min #[[MAP1]](%[[IV]])[%[[D1]]]570// CHECK-DAG: %[[INIT_OFFSET:.+]] = affine.apply #[[MAP2]]()[%[[IV]]]571// CHECK-DAG: %[[INCHUNK:.+]] = tensor.extract_slice %[[ARG0]][0, %[[IV]]] [%[[D0]], %[[TS0]]] [1, 1]572// CHECK-DAG: %[[ET:.+]] = tensor.extract_slice %[[ARG3]][0, %[[INIT_OFFSET]]] [%[[D0]], 1] [1, 1]573// CHECK: %[[PARTIAL:.+]] = linalg.generic574// CHECK-SAME: ins(%[[INCHUNK]] :575// CHECK-SAME: outs(%[[ET]] :576// CHECK: scf.forall.in_parallel {577// CHECK: tensor.parallel_insert_slice %[[PARTIAL]] into %[[ARG3]][0, %[[INIT_OFFSET]]] [%[[D0]], 1] [1, 1]578// CHECK: }579// CHECK: }580// CHECK: %[[R:.*]] = linalg.reduce ins(%[[L]]581// CHECK-SAME: outs(%[[ARG1]] :582// CHECK: return %[[R]] : tensor<?xf32>583module attributes {transform.with_named_sequence} {584 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {585 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op586 %1, %2, %3, %loop = transform.structured.tile_reduction_using_forall %0587 by tile_sizes = [0, 5] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)588 transform.yield589 }590}591 592// -----593 594// Check that only one of the reduction dimension can be tiled (in this case inner).595 596#map = affine_map<(d0, d1, d2) -> (d1, d2)>597#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>598#map2 = affine_map<(d0, d1, d2) -> (d0)>599module {600 func.func @reduction_using_forall_tile_single_of_multiple_reduction_inner(601 %arg0: tensor<86x128xf32>, %arg1: tensor<4096x86x128xf32>, %arg2: tensor<4096xf32>) -> tensor<4096xf32> {602 %0 = linalg.generic {603 indexing_maps = [#map, #map1, #map2],604 iterator_types = ["parallel", "reduction", "reduction"]}605 ins(%arg0, %arg1 : tensor<86x128xf32>, tensor<4096x86x128xf32>) outs(%arg2 : tensor<4096xf32>) {606 ^bb0(%in: f32, %in_0: f32, %out: f32):607 %1 = arith.mulf %in, %in_0 : f32608 %2 = arith.addf %1, %out : f32609 linalg.yield %2 : f32610 } -> tensor<4096xf32>611 return %0 : tensor<4096xf32>612 }613 module attributes {transform.with_named_sequence} {614 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {615 %0 = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op616 %fill_op, %split_linalg_op, %combining_linalg_op, %for_op =617 transform.structured.tile_reduction_using_forall %0 reduction_dims = [2] by tile_sizes = [0, 0, 64]618 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)619 transform.yield620 }621 }622}623// CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> (s0 floordiv 64)>624// CHECK: func @reduction_using_forall_tile_single_of_multiple_reduction_inner(%[[ARG0:.+]]: tensor<86x128xf32>, %[[ARG1:.+]]: tensor<4096x86x128xf32>, %[[ARG2:.+]]: tensor<4096xf32>)625// CHECK: %[[E:.*]] = tensor.empty() : tensor<4096x2xf32>626// CHECK: %[[F:.*]] = linalg.fill627// CHECK-SAME: outs(%[[E]] :628// CHECK: %[[L:.*]] = scf.forall (%[[IV:.+]]) = (0) to (128) step (64) shared_outs(%[[ARG3:.+]] = %[[F]])629// CHECK-DAG: %[[INIT_OFFSET:.+]] = affine.apply #[[MAP0]]()[%[[IV]]]630// CHECK-DAG: %[[ARG0_SLICE:.+]] = tensor.extract_slice %[[ARG0]][0, %[[IV]]] [86, 64] [1, 1]631// CHECK-DAG: %[[ARG1_SLICE:.+]] = tensor.extract_slice %[[ARG1]][0, 0, %[[IV]]] [4096, 86, 64] [1, 1, 1]632// CHECK-DAG: %[[ET:.+]] = tensor.extract_slice %[[ARG3]][0, %[[INIT_OFFSET]]] [4096, 1] [1, 1]633// CHECK: %[[PARTIAL:.+]] = linalg.generic634// CHECK-SAME: ins(%[[ARG0_SLICE]], %[[ARG1_SLICE]] :635// CHECK-SAME: outs(%[[ET]] :636// CHECK: scf.forall.in_parallel {637// CHECK: tensor.parallel_insert_slice %[[PARTIAL]] into %[[ARG3]][0, %[[INIT_OFFSET]]] [4096, 1] [1, 1]638// CHECK: }639// CHECK: }640// CHECK: %[[R:.*]] = linalg.reduce ins(%[[L]]641// CHECK-SAME: outs(%[[ARG2]] :642// CHECK: return %[[R]]643 644// -----645 646// Check that specifying both reduction dimensions, but setting tile size to 0 for one of them behaves consistent with specifying single reduction dimension.647 648#map = affine_map<(d0, d1, d2) -> (d1, d2)>649#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>650#map2 = affine_map<(d0, d1, d2) -> (d0)>651module {652 func.func @reduction_using_forall_tilesize_0_of_multiple_reduction_inner(653 %arg0: tensor<86x128xf32>, %arg1: tensor<4096x86x128xf32>, %arg2: tensor<4096xf32>) -> tensor<4096xf32> {654 %0 = linalg.generic {655 indexing_maps = [#map, #map1, #map2],656 iterator_types = ["parallel", "reduction", "reduction"]}657 ins(%arg0, %arg1 : tensor<86x128xf32>, tensor<4096x86x128xf32>) outs(%arg2 : tensor<4096xf32>) {658 ^bb0(%in: f32, %in_0: f32, %out: f32):659 %1 = arith.mulf %in, %in_0 : f32660 %2 = arith.addf %1, %out : f32661 linalg.yield %2 : f32662 } -> tensor<4096xf32>663 return %0 : tensor<4096xf32>664 }665 module attributes {transform.with_named_sequence} {666 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {667 %0 = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op668 %fill_op, %split_linalg_op, %combining_linalg_op, %for_op =669 transform.structured.tile_reduction_using_forall %0 reduction_dims = [1, 2] by tile_sizes = [0, 0, 64]670 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)671 transform.yield672 }673 }674}675// CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0] -> (s0 floordiv 64)>676// CHECK: func @reduction_using_forall_tilesize_0_of_multiple_reduction_inner(%[[ARG0:.+]]: tensor<86x128xf32>, %[[ARG1:.+]]: tensor<4096x86x128xf32>, %[[ARG2:.+]]: tensor<4096xf32>)677// CHECK: %[[E:.*]] = tensor.empty() : tensor<4096x2xf32>678// CHECK: %[[F:.*]] = linalg.fill679// CHECK-SAME: outs(%[[E]] :680// CHECK: %[[L:.*]] = scf.forall (%[[IV:.+]]) = (0) to (128) step (64) shared_outs(%[[ARG3:.+]] = %[[F]])681// CHECK-DAG: %[[INIT_OFFSET:.+]] = affine.apply #[[MAP0]]()[%[[IV]]]682// CHECK-DAG: %[[ARG0_SLICE:.+]] = tensor.extract_slice %[[ARG0]][0, %[[IV]]] [86, 64] [1, 1]683// CHECK-DAG: %[[ARG1_SLICE:.+]] = tensor.extract_slice %[[ARG1]][0, 0, %[[IV]]] [4096, 86, 64] [1, 1, 1]684// CHECK-DAG: %[[ET:.+]] = tensor.extract_slice %[[ARG3]][0, %[[INIT_OFFSET]]] [4096, 1] [1, 1]685// CHECK: %[[PARTIAL:.+]] = linalg.generic686// CHECK-SAME: ins(%[[ARG0_SLICE]], %[[ARG1_SLICE]] :687// CHECK-SAME: outs(%[[ET]] :688// CHECK: scf.forall.in_parallel {689// CHECK: tensor.parallel_insert_slice %[[PARTIAL]] into %[[ARG3]][0, %[[INIT_OFFSET]]] [4096, 1] [1, 1]690// CHECK: }691// CHECK: }692// CHECK: %[[R:.*]] = linalg.reduce ins(%[[L]]693// CHECK-SAME: outs(%[[ARG2]] :694// CHECK: return %[[R]]695