1// RUN: mlir-opt -gpu-async-region %s | FileCheck %s
2
3// CHECK: module attributes {gpu.container_module}
4module attributes {gpu.container_module} {
5
6  gpu.module @kernels {
7    gpu.func @kernel() kernel { gpu.return }
8  }
9
10  func private @foo() -> ()
11
12  // CHECK-LABEL:func @async(%{{.*}}: index)
13  func @async(%sz : index) {
14    // CHECK: %[[t0:.*]] = gpu.wait async
15    // CHECK: %[[t1:.*]] = gpu.launch_func async [%[[t0]]]
16    gpu.launch_func @kernels::@kernel
17        blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
18    // CHECK: %[[t2:.*]] = gpu.launch_func async [%[[t1]]]
19    gpu.launch_func @kernels::@kernel
20        blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
21    // CHECK: gpu.wait [%[[t2]]]
22    // CHECK: call @foo
23    call @foo() : () -> ()
24    return
25  }
26
27  // CHECK-LABEL:func @defer_wait(%{{.*}}: index)
28  func @defer_wait(%sz : index) {
29    // CHECK: %[[a0:.*]], %[[f0:.*]] = async.execute
30    %a0 = async.execute {
31      // CHECK: %[[t:.*]] = gpu.launch_func async
32      gpu.launch_func @kernels::@kernel
33          blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
34      // CHECK-NOT: gpu.wait
35      // CHECK: async.yield %[[t]]
36      async.yield
37    }
38
39    // CHECK: %[[a1:.*]], %[[f1:.*]] = async.execute
40    // CHECK-SAME: %[[f0]]
41    %a1 = async.execute [%a0] {
42      // CHECK: %[[t:.*]] = gpu.launch_func async
43      gpu.launch_func @kernels::@kernel
44          blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
45      // CHECK-NOT: gpu.wait
46      // CHECK: async.yield %[[t]]
47      async.yield
48    }
49
50    // CHECK: async.await %[[a1]]
51    // CHECK: %[[t:.*]] = async.await %[[f1]]
52    // CHECK: gpu.wait [%[[t]]]
53    async.await %a1 : !async.token
54    return
55  }
56
57  // CHECK-LABEL:func @defer_wait_blocked_by_side_effect(%{{.*}}: index)
58  func @defer_wait_blocked_by_side_effect(%sz : index) {
59    // CHECK: %[[a:.*]] = async.execute
60    %a = async.execute {
61      // CHECK: %[[t:.*]] = gpu.launch_func async
62      gpu.launch_func @kernels::@kernel
63          blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
64      // CHECK: gpu.wait [%[[t]]]
65      call @foo() : () -> ()
66      async.yield
67    }
68
69    // CHECK: async.await %[[a]]
70    // CHECK-NOT: gpu.wait
71    async.await %a : !async.token
72    return
73  }
74
75  // CHECK-LABEL:func @defer_wait_pass_through(%{{.*}}: index)
76  func @defer_wait_pass_through(%sz : index) {
77    // CHECK: %[[a0:.*]], %[[f0:.*]] = async.execute
78    %a0 = async.execute {
79      // CHECK: %[[t:.*]] = gpu.launch_func async
80      gpu.launch_func @kernels::@kernel
81          blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
82      // CHECK-NOT: gpu.wait
83      // CHECK: async.yield %[[t]]
84      async.yield
85    }
86
87    // CHECK: %[[a1:.*]], %[[f1:.*]] = async.execute
88    // CHECK-SAME: %[[f0]]
89    %a1 = async.execute [%a0] {
90      // CHECK-NOT: gpu.wait
91      // CHECK: async.yield %{{.*}}
92      async.yield
93    }
94
95    // CHECK: async.await %[[a1]]
96    // CHECK: %[[t:.*]] = async.await %[[f1]]
97    // CHECK: gpu.wait [%[[t]]]
98    async.await %a1 : !async.token
99    return
100  }
101}
102