1// RUN: mlir-opt %s -async-parallel-for | FileCheck %s
2
3// CHECK-LABEL: @loop_1d
4func @loop_1d(%arg0: index, %arg1: index, %arg2: index, %arg3: memref<?xf32>) {
5  // CHECK: %[[GROUP:.*]] = async.create_group
6  // CHECK: scf.for
7  // CHECK:   %[[TOKEN:.*]] = async.execute {
8  // CHECK:     scf.for
9  // CHECK:       store
10  // CHECK:     async.yield
11  // CHECK:   }
12  // CHECK:   async.add_to_group %[[TOKEN]], %[[GROUP]]
13  // CHECK: async.await_all %[[GROUP]]
14  scf.parallel (%i) = (%arg0) to (%arg1) step (%arg2) {
15    %one = constant 1.0 : f32
16    store %one, %arg3[%i] : memref<?xf32>
17  }
18
19  return
20}
21
22// CHECK-LABEL: @loop_2d
23func @loop_2d(%arg0: index, %arg1: index, %arg2: index, // lb, ub, step
24              %arg3: index, %arg4: index, %arg5: index, // lb, ub, step
25              %arg6: memref<?x?xf32>) {
26  // CHECK: %[[GROUP:.*]] = async.create_group
27  // CHECK: scf.for
28  // CHECK:   scf.for
29  // CHECK:     %[[TOKEN:.*]] = async.execute {
30  // CHECK:       scf.for
31  // CHECK:         scf.for
32  // CHECK:           store
33  // CHECK:       async.yield
34  // CHECK:     }
35  // CHECK:     async.add_to_group %[[TOKEN]], %[[GROUP]]
36  // CHECK: async.await_all %[[GROUP]]
37  scf.parallel (%i0, %i1) = (%arg0, %arg3) to (%arg1, %arg4)
38                            step (%arg2, %arg5) {
39    %one = constant 1.0 : f32
40    store %one, %arg6[%i0, %i1] : memref<?x?xf32>
41  }
42
43  return
44}
45