cubecl-cpp 0.11.0-pre.4

CPP transpiler for CubeCL
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78

use cubecl_core::{
    self as cubecl,
    frontend::barrier::Barrier,
    ir::dialect::barrier::{CommitCopyAsyncOp, CopyAsyncOp},
    prelude::*,
};
use pliron::{derive::op_interface_impl, value::Value};

use crate::{
    cuda::ptx::{barrier_native_handle, generic_to_shared},
    shared::lowering::LowerOp,
    target::Cuda,
};

// Ptr type doesn't matter

#[cube]
pub fn cp_async_global_to_shared(
    src: *const u32,
    smem: *const u32,
    #[comptime] cache: &str,
    #[comptime] copy_size: usize,
) {
    let smem = generic_to_shared::<u32>(smem);
    gpu_asm!(
        "cp.async.{cache}.shared::cta.global [{}], [{}], {size}, {size};",
        mem_out(_) smem, mem_in(_) src, size = const copy_size, options(explicit_mem)
    );
}

#[cube]
pub fn cp_async_global_to_shared_checked(
    src: *const u32,
    smem: *const u32,
    src_size: usize,
    #[comptime] cache: &str,
    #[comptime] copy_size: usize,
) {
    let smem = generic_to_shared::<u32>(smem);
    gpu_asm!(
        "cp.async.{cache}.shared::cta.global [{}], [{}], {copy_size}, {len};",
        mem_out(_) smem, mem_in(_) src, len = in(_) src_size, options(explicit_mem)
    );
}

#[cube]
pub fn commit_copy_async(bar: &Barrier) {
    let bar_handle = barrier_native_handle(bar);
    gpu_asm!("cp.async.mbarrier.arrive.shared::cta.b64 [{}];", mem_inout(_) bar_handle, options(explicit_mem));
}

#[op_interface_impl]
impl LowerOp<Cuda> for CopyAsyncOp {
    fn lower(&self, scope: &Scope) -> Vec<Value> {
        let copy_size = self.copy_length(scope.ctx()).0;
        let src = self.source(scope.ctx()).into();
        let smem = self.destination(scope.ctx()).into();
        let cache = if copy_size == 16 { "cg" } else { "ca" };
        if self.checked(scope.ctx()).0 {
            let src_size = self.source_length(scope.ctx()).into();
            cp_async_global_to_shared_checked::expand(
                scope, &src, &smem, src_size, cache, copy_size,
            );
        } else {
            cp_async_global_to_shared::expand(scope, &src, &smem, cache, copy_size);
        }
        vec![]
    }
}

#[op_interface_impl]
impl LowerOp<Cuda> for CommitCopyAsyncOp {
    fn lower(&self, scope: &Scope) -> Vec<Value> {
        let bar = self.barrier(scope.ctx()).into();
        commit_copy_async::expand(scope, &bar);
        vec![]
    }
}