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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

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

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

#[cube]
pub fn tma_load_im2col_3d(
    tensor_map: &TensorMap<u32, Im2col>,
    bar: &Barrier,
    smem: *const u32,
    pos: (i32, i32, i32),
    offset: u16,
) {
    let bar_handle = barrier_native_handle(bar);
    let smem = generic_to_shared::<u32>(smem);
    let descriptor_address = tensor_map_address(tensor_map);
    let (n, w, c) = pos;
    // Note: tensor maps are opaque descriptors not pointers, so memory effects don't apply
    gpu_asm!(
        "cp.async.bulk.tensor.3d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes ",
        "[{smem}], [{tensor_map}, {{{c}, {w}, {n}}}], [{bar}], {{{offs_w}}};",
        smem = mem_out(_) smem, tensor_map = in(_) descriptor_address,
        c = in(_) c, w = in(_) w, n = in(_) n, offs_w = in(_) offset,
        bar = mem_inout(_) bar_handle, options(explicit_mem)
    );
}

#[cube]
pub fn tma_load_im2col_4d(
    tensor_map: &TensorMap<u32, Im2col>,
    bar: &Barrier,
    smem: *const u32,
    pos: (i32, i32, i32, i32),
    offset: (u16, u16),
) {
    let bar_handle = barrier_native_handle(bar);
    let smem = generic_to_shared::<u32>(smem);
    let descriptor_address = tensor_map_address(tensor_map);
    let (n, h, w, c) = pos;
    let (offs_h, offs_w) = offset;
    // Note: tensor maps are opaque descriptors not pointers, so memory effects don't apply
    gpu_asm!(
        "cp.async.bulk.tensor.4d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes ",
        "[{smem}], [{tensor_map}, {{{c}, {w}, {h}, {n}}}], [{bar}], {{{offs_w}, {offs_h}}};",
        smem = mem_out(_) smem, tensor_map = in(_) descriptor_address,
        c = in(_) c, w = in(_) w, h = in(_) h, n = in(_) n,
        offs_w = in(_) offs_w, offs_h = in(_) offs_h,
        bar = mem_inout(_) bar_handle, options(explicit_mem),
    );
}

#[cube]
pub fn tma_load_im2col_5d(
    tensor_map: &TensorMap<u32, Im2col>,
    bar: &Barrier,
    smem: *const u32,
    pos: (i32, i32, i32, i32, i32),
    offset: (u16, u16, u16),
) {
    let bar_handle = barrier_native_handle(bar);
    let smem = generic_to_shared::<u32>(smem);
    let descriptor_address = tensor_map_address(tensor_map);
    let (n, d, h, w, c) = pos;
    let (offs_d, offs_h, offs_w) = offset;
    // Note: tensor maps are opaque descriptors not pointers, so memory effects don't apply
    gpu_asm!(
        "cp.async.bulk.tensor.5d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes ",
        "[{smem}], [{tensor_map}, {{{c}, {w}, {h}, {d}, {n}}}], [{bar}], {{{offs_w}, {offs_h}, {offs_d}}};",
        smem = in(mem_out) smem, tensor_map = in(_) descriptor_address,
        c = in(_) c, w = in(_) w, h = in(_) h, d = in(_) d, n = in(_) n,
        offs_w = in(_) offs_w, offs_h = in(_) offs_h, offs_d = in(_) offs_d,
        bar = in(mem_inout) bar_handle, options(explicit_mem)
    );
}

#[op_interface_impl]
impl LowerOp<Cuda> for TmaLoadIm2colOp {
    fn lower(&self, scope: &Scope) -> Vec<Value> {
        let ctx = scope.ctx_mut();
        let tensor_map = self.tensor_map(ctx).into();
        let bar = self.barrier(ctx).into();
        let smem = self.destination(ctx).into();
        let pos = self.indices(ctx);
        let offsets = self.offsets(ctx);

        let pos_3 = (pos[0].into(), pos[1].into(), pos[2].into());

        match self.rank(ctx) {
            3 => {
                let offset = offsets[0].into();
                tma_load_im2col_3d::expand(scope, &tensor_map, &bar, &smem, pos_3, offset);
            }
            4 => {
                let pos = (pos_3.0, pos_3.1, pos_3.2, pos[3].into());
                let offset = (offsets[0].into(), offsets[1].into());
                tma_load_im2col_4d::expand(scope, &tensor_map, &bar, &smem, pos, offset);
            }
            5 => {
                let pos = (pos_3.0, pos_3.1, pos_3.2, pos[3].into(), pos[4].into());
                let offset = (offsets[0].into(), offsets[1].into(), offsets[2].into());
                tma_load_im2col_5d::expand(scope, &tensor_map, &bar, &smem, pos, offset);
            }
            _ => unreachable!("Should be 3D-5D"),
        }
        vec![]
    }
}