1// RUN: mlir-opt %s -split-input-file -async-parallel-for=async-dispatch=true  \
2// RUN: | FileCheck %s
3
4// CHECK-LABEL: @loop_1d
5func @loop_1d(%arg0: index, %arg1: index, %arg2: index, %arg3: memref<?xf32>) {
6  // CHECK: %[[C0:.*]] = constant 0 : index
7  // CHECK: %[[GROUP:.*]] = async.create_group
8  // CHECK: scf.if {{.*}} {
9  // CHECK:   call @parallel_compute_fn(%[[C0]]
10  // CHECK: } else {
11  // CHECK:   call @async_dispatch_fn
12  // CHECK: }
13  // CHECK: async.await_all %[[GROUP]]
14  scf.parallel (%i) = (%arg0) to (%arg1) step (%arg2) {
15    %one = constant 1.0 : f32
16    memref.store %one, %arg3[%i] : memref<?xf32>
17  }
18  return
19}
20
21// CHECK-LABEL: func private @parallel_compute_fn
22// CHECK:       scf.for
23// CHECK:         memref.store
24
25// CHECK-LABEL: func private @async_dispatch_fn
26// CHECK-SAME:  (
27// CHECK-SAME:    %[[GROUP:arg0]]: !async.group,
28// CHECK-SAME:    %[[BLOCK_START:arg1]]: index
29// CHECK-SAME:    %[[BLOCK_END:arg2]]: index
30// CHECK-SAME:  )
31// CHECK:         %[[C1:.*]] = constant 1 : index
32// CHECK:         %[[C2:.*]] = constant 2 : index
33// CHECK:         scf.while (%[[S0:.*]] = %[[BLOCK_START]],
34// CHECK-SAME:               %[[E0:.*]] = %[[BLOCK_END]])
35// While loop `before` block decides if we need to dispatch more tasks.
36// CHECK:         {
37// CHECK:           %[[DIFF0:.*]] = subi %[[E0]], %[[S0]]
38// CHECK:           %[[COND:.*]] = cmpi sgt, %[[DIFF0]], %[[C1]]
39// CHECK:           scf.condition(%[[COND]])
40// While loop `after` block splits the range in half and submits async task
41// to process the second half using the call to the same dispatch function.
42// CHECK:         } do {
43// CHECK:         ^bb0(%[[S1:.*]]: index, %[[E1:.*]]: index):
44// CHECK:           %[[DIFF1:.*]] = subi %[[E1]], %[[S1]]
45// CHECK:           %[[HALF:.*]] = divi_signed %[[DIFF1]], %[[C2]]
46// CHECK:           %[[MID:.*]] = addi %[[S1]], %[[HALF]]
47// CHECK:           %[[TOKEN:.*]] = async.execute
48// CHECK:             call @async_dispatch_fn
49// CHECK:           async.add_to_group
50// CHECK:           scf.yield %[[S1]], %[[MID]]
51// CHECK:         }
52// After async dispatch the first block processed in the caller thread.
53// CHECK:         call @parallel_compute_fn(%[[BLOCK_START]]
54
55// -----
56
57// CHECK-LABEL: @loop_2d
58func @loop_2d(%arg0: index, %arg1: index, %arg2: index, // lb, ub, step
59              %arg3: index, %arg4: index, %arg5: index, // lb, ub, step
60              %arg6: memref<?x?xf32>) {
61  // CHECK: %[[GROUP:.*]] = async.create_group
62  // CHECK: call @async_dispatch_fn
63  // CHECK: async.await_all %[[GROUP]]
64  scf.parallel (%i0, %i1) = (%arg0, %arg3) to (%arg1, %arg4)
65                            step (%arg2, %arg5) {
66    %one = constant 1.0 : f32
67    memref.store %one, %arg6[%i0, %i1] : memref<?x?xf32>
68  }
69  return
70}
71
72// CHECK-LABEL: func private @parallel_compute_fn
73// CHECK:       scf.for
74// CHECK:         scf.for
75// CHECK:           memref.store
76
77// CHECK-LABEL: func private @async_dispatch_fn
78