cubecl_cpp/cuda/ptx/
tma_load_im2col.rs

1use cubecl_core::{
2    self as cubecl, frontend::barrier::Barrier, ir::dialect::tma::TmaLoadIm2colOp, prelude::*,
3};
4use pliron::{derive::op_interface_impl, value::Value};
5
6use crate::{
7    cuda::ptx::{barrier_native_handle, generic_to_shared, tensor_map_address},
8    shared::lowering::LowerOp,
9    target::Cuda,
10};
11
12#[cube]
13pub fn tma_load_im2col_3d(
14    tensor_map: &TensorMap<u32, Im2col>,
15    bar: &Barrier,
16    smem: *const u32,
17    pos: (i32, i32, i32),
18    offset: u16,
19) {
20    let bar_handle = barrier_native_handle(bar);
21    let smem = generic_to_shared::<u32>(smem);
22    let descriptor_address = tensor_map_address(tensor_map);
23    let (n, w, c) = pos;
24    // Note: tensor maps are opaque descriptors not pointers, so memory effects don't apply
25    gpu_asm!(
26        "cp.async.bulk.tensor.3d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes ",
27        "[{smem}], [{tensor_map}, {{{c}, {w}, {n}}}], [{bar}], {{{offs_w}}};",
28        smem = mem_out(_) smem, tensor_map = in(_) descriptor_address,
29        c = in(_) c, w = in(_) w, n = in(_) n, offs_w = in(_) offset,
30        bar = mem_inout(_) bar_handle, options(explicit_mem)
31    );
32}
33
34#[cube]
35pub fn tma_load_im2col_4d(
36    tensor_map: &TensorMap<u32, Im2col>,
37    bar: &Barrier,
38    smem: *const u32,
39    pos: (i32, i32, i32, i32),
40    offset: (u16, u16),
41) {
42    let bar_handle = barrier_native_handle(bar);
43    let smem = generic_to_shared::<u32>(smem);
44    let descriptor_address = tensor_map_address(tensor_map);
45    let (n, h, w, c) = pos;
46    let (offs_h, offs_w) = offset;
47    // Note: tensor maps are opaque descriptors not pointers, so memory effects don't apply
48    gpu_asm!(
49        "cp.async.bulk.tensor.4d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes ",
50        "[{smem}], [{tensor_map}, {{{c}, {w}, {h}, {n}}}], [{bar}], {{{offs_w}, {offs_h}}};",
51        smem = mem_out(_) smem, tensor_map = in(_) descriptor_address,
52        c = in(_) c, w = in(_) w, h = in(_) h, n = in(_) n,
53        offs_w = in(_) offs_w, offs_h = in(_) offs_h,
54        bar = mem_inout(_) bar_handle, options(explicit_mem),
55    );
56}
57
58#[cube]
59pub fn tma_load_im2col_5d(
60    tensor_map: &TensorMap<u32, Im2col>,
61    bar: &Barrier,
62    smem: *const u32,
63    pos: (i32, i32, i32, i32, i32),
64    offset: (u16, u16, u16),
65) {
66    let bar_handle = barrier_native_handle(bar);
67    let smem = generic_to_shared::<u32>(smem);
68    let descriptor_address = tensor_map_address(tensor_map);
69    let (n, d, h, w, c) = pos;
70    let (offs_d, offs_h, offs_w) = offset;
71    // Note: tensor maps are opaque descriptors not pointers, so memory effects don't apply
72    gpu_asm!(
73        "cp.async.bulk.tensor.5d.shared::cluster.global.im2col.mbarrier::complete_tx::bytes ",
74        "[{smem}], [{tensor_map}, {{{c}, {w}, {h}, {d}, {n}}}], [{bar}], {{{offs_w}, {offs_h}, {offs_d}}};",
75        smem = in(mem_out) smem, tensor_map = in(_) descriptor_address,
76        c = in(_) c, w = in(_) w, h = in(_) h, d = in(_) d, n = in(_) n,
77        offs_w = in(_) offs_w, offs_h = in(_) offs_h, offs_d = in(_) offs_d,
78        bar = in(mem_inout) bar_handle, options(explicit_mem)
79    );
80}
81
82#[op_interface_impl]
83impl LowerOp<Cuda> for TmaLoadIm2colOp {
84    fn lower(&self, scope: &Scope) -> Vec<Value> {
85        let ctx = scope.ctx_mut();
86        let tensor_map = self.tensor_map(ctx).into();
87        let bar = self.barrier(ctx).into();
88        let smem = self.destination(ctx).into();
89        let pos = self.indices(ctx);
90        let offsets = self.offsets(ctx);
91
92        let pos_3 = (pos[0].into(), pos[1].into(), pos[2].into());
93
94        match self.rank(ctx) {
95            3 => {
96                let offset = offsets[0].into();
97                tma_load_im2col_3d::expand(scope, &tensor_map, &bar, &smem, pos_3, offset);
98            }
99            4 => {
100                let pos = (pos_3.0, pos_3.1, pos_3.2, pos[3].into());
101                let offset = (offsets[0].into(), offsets[1].into());
102                tma_load_im2col_4d::expand(scope, &tensor_map, &bar, &smem, pos, offset);
103            }
104            5 => {
105                let pos = (pos_3.0, pos_3.1, pos_3.2, pos[3].into(), pos[4].into());
106                let offset = (offsets[0].into(), offsets[1].into(), offsets[2].into());
107                tma_load_im2col_5d::expand(scope, &tensor_map, &bar, &smem, pos, offset);
108            }
109            _ => unreachable!("Should be 3D-5D"),
110        }
111        vec![]
112    }
113}
cubecl_cpp/cuda/ptx/tma_load_im2col.rs

cubecl_cpp/cuda/ptx/
tma_load_im2col.rs
