You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
102 lines
3.0 KiB
102 lines
3.0 KiB
4 months ago
|
// RUN: mlir-opt -gpu-async-region %s | FileCheck %s
|
||
|
|
|
||
|
|
// CHECK: module attributes {gpu.container_module}
|
||
|
|
module attributes {gpu.container_module} {
|
||
|
|
|
||
|
|
gpu.module @kernels {
|
||
|
|
gpu.func @kernel() kernel { gpu.return }
|
||
|
|
}
|
||
|
|
|
||
|
|
func private @foo() -> ()
|
||
|
|
|
||
|
|
// CHECK-LABEL:func @async(%{{.*}}: index)
|
||
|
|
func @async(%sz : index) {
|
||
|
|
// CHECK: %[[t0:.*]] = gpu.wait async
|
||
|
|
// CHECK: %[[t1:.*]] = gpu.launch_func async [%[[t0]]]
|
||
|
|
gpu.launch_func @kernels::@kernel
|
||
|
|
blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
|
||
|
|
// CHECK: %[[t2:.*]] = gpu.launch_func async [%[[t1]]]
|
||
|
|
gpu.launch_func @kernels::@kernel
|
||
|
|
blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
|
||
|
|
// CHECK: gpu.wait [%[[t2]]]
|
||
|
|
// CHECK: call @foo
|
||
|
|
call @foo() : () -> ()
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
// CHECK-LABEL:func @defer_wait(%{{.*}}: index)
|
||
|
|
func @defer_wait(%sz : index) {
|
||
|
|
// CHECK: %[[a0:.*]], %[[f0:.*]] = async.execute
|
||
|
|
%a0 = async.execute {
|
||
|
|
// CHECK: %[[t:.*]] = gpu.launch_func async
|
||
|
|
gpu.launch_func @kernels::@kernel
|
||
|
|
blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
|
||
|
|
// CHECK-NOT: gpu.wait
|
||
|
|
// CHECK: async.yield %[[t]]
|
||
|
|
async.yield
|
||
|
|
}
|
||
|
|
|
||
|
|
// CHECK: %[[a1:.*]], %[[f1:.*]] = async.execute
|
||
|
|
// CHECK-SAME: %[[f0]]
|
||
|
|
%a1 = async.execute [%a0] {
|
||
|
|
// CHECK: %[[t:.*]] = gpu.launch_func async
|
||
|
|
gpu.launch_func @kernels::@kernel
|
||
|
|
blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
|
||
|
|
// CHECK-NOT: gpu.wait
|
||
|
|
// CHECK: async.yield %[[t]]
|
||
|
|
async.yield
|
||
|
|
}
|
||
|
|
|
||
|
|
// CHECK: async.await %[[a1]]
|
||
|
|
// CHECK: %[[t:.*]] = async.await %[[f1]]
|
||
|
|
// CHECK: gpu.wait [%[[t]]]
|
||
|
|
async.await %a1 : !async.token
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
// CHECK-LABEL:func @defer_wait_blocked_by_side_effect(%{{.*}}: index)
|
||
|
|
func @defer_wait_blocked_by_side_effect(%sz : index) {
|
||
|
|
// CHECK: %[[a:.*]] = async.execute
|
||
|
|
%a = async.execute {
|
||
|
|
// CHECK: %[[t:.*]] = gpu.launch_func async
|
||
|
|
gpu.launch_func @kernels::@kernel
|
||
|
|
blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
|
||
|
|
// CHECK: gpu.wait [%[[t]]]
|
||
|
|
call @foo() : () -> ()
|
||
|
|
async.yield
|
||
|
|
}
|
||
|
|
|
||
|
|
// CHECK: async.await %[[a]]
|
||
|
|
// CHECK-NOT: gpu.wait
|
||
|
|
async.await %a : !async.token
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
// CHECK-LABEL:func @defer_wait_pass_through(%{{.*}}: index)
|
||
|
|
func @defer_wait_pass_through(%sz : index) {
|
||
|
|
// CHECK: %[[a0:.*]], %[[f0:.*]] = async.execute
|
||
|
|
%a0 = async.execute {
|
||
|
|
// CHECK: %[[t:.*]] = gpu.launch_func async
|
||
|
|
gpu.launch_func @kernels::@kernel
|
||
|
|
blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
|
||
|
|
// CHECK-NOT: gpu.wait
|
||
|
|
// CHECK: async.yield %[[t]]
|
||
|
|
async.yield
|
||
|
|
}
|
||
|
|
|
||
|
|
// CHECK: %[[a1:.*]], %[[f1:.*]] = async.execute
|
||
|
|
// CHECK-SAME: %[[f0]]
|
||
|
|
%a1 = async.execute [%a0] {
|
||
|
|
// CHECK-NOT: gpu.wait
|
||
|
|
// CHECK: async.yield %{{.*}}
|
||
|
|
async.yield
|
||
|
|
}
|
||
|
|
|
||
|
|
// CHECK: async.await %[[a1]]
|
||
|
|
// CHECK: %[[t:.*]] = async.await %[[f1]]
|
||
|
|
// CHECK: gpu.wait [%[[t]]]
|
||
|
|
async.await %a1 : !async.token
|
||
|
|
return
|
||
|
|
}
|
||
|
|
}
|