98 lines · plain
1// RUN: mlir-opt %s -convert-linalg-to-parallel-loops -split-input-file | FileCheck %s2 3#map0 = affine_map<(d0, d1) -> (d0, d1)>4func.func @linalg_generic_sum(%lhs: memref<2x2xf32>,5 %rhs: memref<2x2xf32>,6 %sum: memref<2x2xf32>) {7 linalg.generic {8 indexing_maps = [#map0, #map0, #map0],9 iterator_types = ["parallel", "parallel"]}10 ins(%lhs, %rhs : memref<2x2xf32>, memref<2x2xf32>)11 outs(%sum : memref<2x2xf32>) {12 ^bb0(%lhs_in: f32, %rhs_in: f32, %sum_out: f32):13 %0 = arith.addf %lhs_in, %rhs_in : f3214 linalg.yield %0 : f3215 }16 return17}18// CHECK-LABEL: @linalg_generic_sum19// CHECK: (%[[LHS:.*]]:{{.*}}, %[[RHS:.*]]:{{.*}}, %[[SUM:.*]]:{{.*}})20// CHECK-DAG: %[[C2:.*]] = arith.constant 221// CHECK-DAG: %[[C0:.*]] = arith.constant 022// CHECK-DAG: %[[C1:.*]] = arith.constant 123// CHECK: scf.parallel (%[[I:.*]], %[[J:.*]]) = {{.*}}24// CHECK: %[[LHS_ELEM:.*]] = memref.load %[[LHS]][%[[I]], %[[J]]]25// CHECK: %[[RHS_ELEM:.*]] = memref.load %[[RHS]][%[[I]], %[[J]]]26// CHECK: %[[SUM:.*]] = arith.addf %[[LHS_ELEM]], %[[RHS_ELEM]] : f3227// CHECK: store %[[SUM]], %{{.*}}[%[[I]], %[[J]]]28// CHECK: scf.reduce29 30// -----31 32#accesses = [33 affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>,34 affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>35]36#trait = {37 iterator_types = ["parallel", "parallel", "reduction", "parallel"],38 indexing_maps = #accesses39}40 41func.func @lower_outer_parallel(%A: memref<?x?x?x?xf32>, %B: memref<?x?x?xf32>) {42 linalg.generic #trait43 ins(%A : memref<?x?x?x?xf32>)44 outs(%B : memref<?x?x?xf32>) {45 ^bb0(%a: f32, %b: f32):46 linalg.yield %a: f3247 }48 return49}50// CHECK-LABEL: @lower_outer_parallel51// CHECK-DAG: %[[C0:.*]] = arith.constant 052// CHECK-DAG: %[[C1:.*]] = arith.constant 153// CHECK-DAG: %[[D0:.*]] = memref.dim %{{.*}}, %c054// CHECK-DAG: %[[D1:.*]] = memref.dim %{{.*}}, %c155// CHECK-DAG: %[[D2:.*]] = memref.dim %{{.*}}, %c256// CHECK-DAG: %[[D3:.*]] = memref.dim %{{.*}}, %c357// CHECK: scf.parallel (%[[IV0:.*]], %[[IV1:.*]]) = (%[[C0]], %[[C0]]) to (%[[D0]], %[[D1]]) step (%[[C1]], %[[C1]])58// CHECK: scf.for %[[IV2:.*]] = %[[C0]] to %[[D2]] step %[[C1]]59// CHECK: scf.parallel (%[[IV3:.*]]) = (%[[C0]]) to (%[[D3]]) step (%[[C1]])60// CHECK: memref.load %{{.*}}[%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]]]61// CHECK: store %{{.*}}, %{{.*}}[%[[IV0]], %[[IV1]], %[[IV3]]]62 63// -----64 65#accesses = [66 affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3, d4, d5)>,67 affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d4, d3)>68]69#trait = {70 iterator_types = ["parallel", "parallel", "reduction", "parallel", "parallel", "reduction"],71 indexing_maps = #accesses72}73 74func.func @lower_mixed_parallel(%A: memref<?x?x?x?x?x?xf32>, %B: memref<?x?x?x?xf32>) {75 linalg.generic #trait76 ins(%A : memref<?x?x?x?x?x?xf32>)77 outs(%B : memref<?x?x?x?xf32>) {78 ^bb0(%a: f32, %b: f32):79 linalg.yield %a: f3280 }81 return82}83// CHECK-LABEL: @lower_mixed_parallel84// CHECK-DAG: %[[C0:.*]] = arith.constant 085// CHECK-DAG: %[[C1:.*]] = arith.constant 186// CHECK-DAG: %[[D0:.*]] = memref.dim %{{.*}}, %c087// CHECK-DAG: %[[D1:.*]] = memref.dim %{{.*}}, %c188// CHECK-DAG: %[[D2:.*]] = memref.dim %{{.*}}, %c289// CHECK-DAG: %[[D3:.*]] = memref.dim %{{.*}}, %c390// CHECK-DAG: %[[D4:.*]] = memref.dim %{{.*}}, %c491// CHECK-DAG: %[[D5:.*]] = memref.dim %{{.*}}, %c592// CHECK: scf.parallel (%[[IV0:.*]], %[[IV1:.*]]) = (%[[C0]], %[[C0]]) to (%[[D0]], %[[D1]]) step (%[[C1]], %[[C1]])93// CHECK: scf.for %[[IV2:.*]] = %[[C0]] to %[[D2]] step %[[C1]]94// CHECK: scf.parallel (%[[IV3:.*]], %[[IV4:.*]]) = (%[[C0]], %[[C0]]) to (%[[D3]], %[[D4]]) step (%[[C1]], %[[C1]])95// CHECK: scf.for %[[IV5:.*]] = %[[C0]] to %[[D5]] step %[[C1]]96// CHECK: memref.load %{{.*}}[%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]], %[[IV4]], %[[IV5]]]97// CHECK: store %{{.*}}, %{{.*}}[%[[IV0]], %[[IV1]], %[[IV4]], %[[IV3]]]98