146 lines · plain
1// RUN: mlir-opt %s -split-input-file \2// RUN: -async-parallel-for=async-dispatch=true \3// RUN: | FileCheck %s4 5// RUN: mlir-opt %s -split-input-file \6// RUN: -async-parallel-for=async-dispatch=false \7// RUN: -canonicalize -inline -symbol-dce \8// RUN: | FileCheck %s9 10// Check that constants defined outside of the `scf.parallel` body will be11// sunk into the parallel compute function to avoid blowing up the number12// of parallel compute function arguments.13 14// CHECK-LABEL: func @clone_constant(15func.func @clone_constant(%arg0: memref<?xf32>, %lb: index, %ub: index, %st: index) {16 %one = arith.constant 1.0 : f3217 18 scf.parallel (%i) = (%lb) to (%ub) step (%st) {19 memref.store %one, %arg0[%i] : memref<?xf32>20 }21 22 return23}24 25// CHECK-LABEL: func private @parallel_compute_fn(26// CHECK-SAME: %[[BLOCK_INDEX:arg[0-9]+]]: index,27// CHECK-SAME: %[[BLOCK_SIZE:arg[0-9]+]]: index,28// CHECK-SAME: %[[TRIP_COUNT:arg[0-9]+]]: index,29// CHECK-SAME: %[[LB:arg[0-9]+]]: index,30// CHECK-SAME: %[[UB:arg[0-9]+]]: index,31// CHECK-SAME: %[[STEP:arg[0-9]+]]: index,32// CHECK-SAME: %[[MEMREF:arg[0-9]+]]: memref<?xf32>33// CHECK-SAME: ) {34// CHECK: %[[CST:.*]] = arith.constant 1.0{{.*}} : f3235// CHECK: scf.for36// CHECK: memref.store %[[CST]], %[[MEMREF]]37 38// -----39 40// Check that constant loop bound sunk into the parallel compute function.41 42// CHECK-LABEL: func @sink_constant_step(43func.func @sink_constant_step(%arg0: memref<?xf32>, %lb: index, %ub: index) {44 %one = arith.constant 1.0 : f3245 %st = arith.constant 123 : index46 47 scf.parallel (%i) = (%lb) to (%ub) step (%st) {48 memref.store %one, %arg0[%i] : memref<?xf32>49 }50 51 return52}53 54// CHECK-LABEL: func private @parallel_compute_fn(55// CHECK-SAME: %[[BLOCK_INDEX:arg[0-9]+]]: index,56// CHECK-SAME: %[[BLOCK_SIZE:arg[0-9]+]]: index,57// CHECK-SAME: %[[TRIP_COUNT:arg[0-9]+]]: index,58// CHECK-SAME: %[[LB:arg[0-9]+]]: index,59// CHECK-SAME: %[[UB:arg[0-9]+]]: index,60// CHECK-SAME: %[[STEP:arg[0-9]+]]: index,61// CHECK-SAME: %[[MEMREF:arg[0-9]+]]: memref<?xf32>62// CHECK-SAME: ) {63// CHECK: %[[CSTEP:.*]] = arith.constant 123 : index64// CHECK-NOT: %[[STEP]]65// CHECK: scf.for %[[I:arg[0-9]+]]66// CHECK: %[[TMP:.*]] = arith.muli %[[I]], %[[CSTEP]]67// CHECK: %[[IDX:.*]] = arith.addi %[[LB]], %[[TMP]]68// CHECK: memref.store69 70// -----71 72// Smoke test that parallel for doesn't crash when func dialect is not used.73 74// CHECK-LABEL: llvm.func @without_func_dialect()75llvm.func @without_func_dialect() {76 %cst = arith.constant 0.0 : f3277 78 %c0 = arith.constant 0 : index79 %c22 = arith.constant 22 : index80 %c1 = arith.constant 1 : index81 %54 = memref.alloc() : memref<22xf32>82 %alloc_4 = memref.alloc() : memref<22xf32>83 scf.parallel (%arg0) = (%c0) to (%c22) step (%c1) {84 memref.store %cst, %alloc_4[%arg0] : memref<22xf32>85 }86 llvm.return87}88 89// -----90 91// Check that for statically known inner loop bound block size is aligned and92// inner loop uses statically known loop trip counts.93 94// CHECK-LABEL: func @sink_constant_step(95func.func @sink_constant_step(%arg0: memref<?x10xf32>, %lb: index, %ub: index) {96 %one = arith.constant 1.0 : f3297 98 %c0 = arith.constant 0 : index99 %c1 = arith.constant 1 : index100 %c10 = arith.constant 10 : index101 102 scf.parallel (%i, %j) = (%lb, %c0) to (%ub, %c10) step (%c1, %c1) {103 memref.store %one, %arg0[%i, %j] : memref<?x10xf32>104 }105 106 return107}108 109// CHECK-LABEL: func private @parallel_compute_fn_with_aligned_loops(110// CHECK-SAME: %[[BLOCK_INDEX:arg[0-9]+]]: index,111// CHECK-SAME: %[[BLOCK_SIZE:arg[0-9]+]]: index,112// CHECK-SAME: %[[TRIP_COUNT0:arg[0-9]+]]: index,113// CHECK-SAME: %[[TRIP_COUNT1:arg[0-9]+]]: index,114// CHECK-SAME: %[[LB0:arg[0-9]+]]: index,115// CHECK-SAME: %[[LB1:arg[0-9]+]]: index,116// CHECK-SAME: %[[UB0:arg[0-9]+]]: index,117// CHECK-SAME: %[[UB1:arg[0-9]+]]: index,118// CHECK-SAME: %[[STEP0:arg[0-9]+]]: index,119// CHECK-SAME: %[[STEP1:arg[0-9]+]]: index,120// CHECK-SAME: %[[MEMREF:arg[0-9]+]]: memref<?x10xf32>121// CHECK-SAME: ) {122// CHECK: %[[C0:.*]] = arith.constant 0 : index123// CHECK: %[[C1:.*]] = arith.constant 1 : index124// CHECK: %[[C10:.*]] = arith.constant 10 : index125// CHECK: scf.for %[[I:arg[0-9]+]]126// CHECK-NOT: arith.select127// CHECK: scf.for %[[J:arg[0-9]+]] = %c0 to %c10 step %c1128 129// CHECK-LABEL: func private @parallel_compute_fn(130// CHECK-SAME: %[[BLOCK_INDEX:arg[0-9]+]]: index,131// CHECK-SAME: %[[BLOCK_SIZE:arg[0-9]+]]: index,132// CHECK-SAME: %[[TRIP_COUNT0:arg[0-9]+]]: index,133// CHECK-SAME: %[[TRIP_COUNT1:arg[0-9]+]]: index,134// CHECK-SAME: %[[LB0:arg[0-9]+]]: index,135// CHECK-SAME: %[[LB1:arg[0-9]+]]: index,136// CHECK-SAME: %[[UB0:arg[0-9]+]]: index,137// CHECK-SAME: %[[UB1:arg[0-9]+]]: index,138// CHECK-SAME: %[[STEP0:arg[0-9]+]]: index,139// CHECK-SAME: %[[STEP1:arg[0-9]+]]: index,140// CHECK-SAME: %[[MEMREF:arg[0-9]+]]: memref<?x10xf32>141// CHECK-SAME: ) {142// CHECK: scf.for %[[I:arg[0-9]+]]143// CHECK: arith.select144// CHECK: scf.for %[[J:arg[0-9]+]]145// CHECK: memref.store146