cubecl_cpp/cuda/ptx/
copy_async.rs

1use cubecl_core::{
2    self as cubecl,
3    frontend::barrier::Barrier,
4    ir::dialect::barrier::{CommitCopyAsyncOp, CopyAsyncOp},
5    prelude::*,
6};
7use pliron::{derive::op_interface_impl, value::Value};
8
9use crate::{
10    cuda::ptx::{barrier_native_handle, generic_to_shared},
11    shared::lowering::LowerOp,
12    target::Cuda,
13};
14
15// Ptr type doesn't matter
16
17#[cube]
18pub fn cp_async_global_to_shared(
19    src: *const u32,
20    smem: *const u32,
21    #[comptime] cache: &str,
22    #[comptime] copy_size: usize,
23) {
24    let smem = generic_to_shared::<u32>(smem);
25    gpu_asm!(
26        "cp.async.{cache}.shared::cta.global [{}], [{}], {size}, {size};",
27        mem_out(_) smem, mem_in(_) src, size = const copy_size, options(explicit_mem)
28    );
29}
30
31#[cube]
32pub fn cp_async_global_to_shared_checked(
33    src: *const u32,
34    smem: *const u32,
35    src_size: usize,
36    #[comptime] cache: &str,
37    #[comptime] copy_size: usize,
38) {
39    let smem = generic_to_shared::<u32>(smem);
40    gpu_asm!(
41        "cp.async.{cache}.shared::cta.global [{}], [{}], {copy_size}, {len};",
42        mem_out(_) smem, mem_in(_) src, len = in(_) src_size, options(explicit_mem)
43    );
44}
45
46#[cube]
47pub fn commit_copy_async(bar: &Barrier) {
48    let bar_handle = barrier_native_handle(bar);
49    gpu_asm!("cp.async.mbarrier.arrive.shared::cta.b64 [{}];", mem_inout(_) bar_handle, options(explicit_mem));
50}
51
52#[op_interface_impl]
53impl LowerOp<Cuda> for CopyAsyncOp {
54    fn lower(&self, scope: &Scope) -> Vec<Value> {
55        let copy_size = self.copy_length(scope.ctx()).0;
56        let src = self.source(scope.ctx()).into();
57        let smem = self.destination(scope.ctx()).into();
58        let cache = if copy_size == 16 { "cg" } else { "ca" };
59        if self.checked(scope.ctx()).0 {
60            let src_size = self.source_length(scope.ctx()).into();
61            cp_async_global_to_shared_checked::expand(
62                scope, &src, &smem, src_size, cache, copy_size,
63            );
64        } else {
65            cp_async_global_to_shared::expand(scope, &src, &smem, cache, copy_size);
66        }
67        vec![]
68    }
69}
70
71#[op_interface_impl]
72impl LowerOp<Cuda> for CommitCopyAsyncOp {
73    fn lower(&self, scope: &Scope) -> Vec<Value> {
74        let bar = self.barrier(scope.ctx()).into();
75        commit_copy_async::expand(scope, &bar);
76        vec![]
77    }
78}
cubecl_cpp/cuda/ptx/copy_async.rs

cubecl_cpp/cuda/ptx/
copy_async.rs
