59 lines · plain
1// RUN: mlir-opt -transform-interpreter -cse -mlir-print-local-scope -split-input-file -verify-diagnostics %s | FileCheck %s2 3// Check tile+ fuse works with partial reduction outer parallel strategy.4 5module{6 func.func @tile_and_fuse_with_partial_reduction_outer_parallel(7 %arg0 : tensor<?x?xf32>) -> tensor<?xf32> {8 %c0 = arith.constant 0 : index9 %cst = arith.constant 0.0 : f3210 %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>11 %empty = tensor.empty(%d0) : tensor<?xf32>12 %fill = linalg.fill ins(%cst : f32) outs(%empty : tensor<?xf32>) -> tensor<?xf32>13 %generic = linalg.generic {14 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>],15 iterator_types = ["parallel", "reduction"]}16 ins(%arg0 : tensor<?x?xf32>) outs(%fill : tensor<?xf32>) {17 ^bb0(%b0 : f32, %b1 : f32):18 %0 = arith.addf %b0, %b1 : f3219 linalg.yield %0 : f3220 } -> tensor<?xf32>21 return %generic : tensor<?xf32>22 }23}24module attributes {transform.with_named_sequence} {25 transform.named_sequence @__transform_main(%arg1 : !transform.any_op {transform.readonly}) {26 %generic = transform.structured.match ops{["linalg.generic"]} in %arg127 : (!transform.any_op) -> !transform.any_op28 %a, %loop = transform.test.tile_and_fuse_outer_parallel_partial_reduction29 %generic tile_sizes = [128] 30 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)31 transform.yield32 }33}34// CHECK-LABEL: func @tile_and_fuse_with_partial_reduction_outer_parallel(35// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32>)36// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index37// CHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]]38// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index39// CHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG0]], %[[C1]]40// CHECK: %[[REDUCTION_NUM:.+]] = affine.apply affine_map<()[s0] -> (s0 ceildiv 128)>()[%[[D1]]]41// CHECK: %[[EMPTY:.+]] = tensor.empty(%[[D0]], %[[REDUCTION_NUM]])42// CHECK: %[[FORALL:.+]] = scf.forall (%[[IV0:[a-zA-Z0-9]+]]) =43// CHECK-SAME: shared_outs(%[[ITER_ARG:.+]] = %[[EMPTY]])44// CHECK-DAG: %[[TILESIZE:.+]] = affine.min affine_map<(d0)[s0] -> (-d0 + s0, 128)>(%[[IV0]])[%[[D1]]]45// CHECK-DAG: %[[REDUCTION_IV:.+]] = affine.apply affine_map<()[s0] -> (s0 floordiv 128)>()[%[[IV0]]]46// CHECK-DAG: %[[ARG0_SLICE:.+]] = tensor.extract_slice %[[ARG0]][0, %[[IV0]]] [%[[D0]], %[[TILESIZE]]] [1, 1]47// CHECK: %[[ITER_ARG_SLICE:.+]] = tensor.extract_slice %[[ITER_ARG]][0, %[[REDUCTION_IV]]] [%[[D0]], 1] [1, 1]48// CHECK: %[[FILL:.+]] = linalg.fill49// CHECK-SAME: outs(%[[ITER_ARG_SLICE]] : tensor<?x1xf32>)50// CHECK: %[[REDUCING_SLICE:.+]] = tensor.extract_slice %[[FILL]][0, 0] [%[[D0]], 1] [1, 1] : tensor<?x1xf32> to tensor<?xf32>51// CHECK: %[[GENERIC:.+]] = linalg.generic52// CHECK-SAME: ins(%[[ARG0_SLICE]] :53// CHECK-SAME: outs(%[[REDUCING_SLICE]] :54// CHECK: tensor.parallel_insert_slice %[[GENERIC]] into %[[ITER_ARG]]55// CHECK-SAME: [0, %[[REDUCTION_IV]]] [%[[D0]], 1] [1, 1]56// CHECK: %[[REDUCE:.+]] = linalg.reduce57// CHECK-SAME: ins(%[[FORALL]] :58// CHECK: return %[[REDUCE]]59