cubecl_cpp/cuda/
tma.rs

1use cubecl_core::ir::dialect::tma::*;
2use itertools::Itertools;
3
4use crate::{
5    cuda::cuda_op,
6    shared::{CppValue, signature::op_includes},
7    target::Cuda,
8};
9
10op_includes!(Cuda, [TmaStoreOp, CommitGroupOp, WaitGroupOp, WaitGroupReadOp] => "cuda/barrier");
11
12cuda_op!(TmaLoadOp, |op, ctx| {
13    let barrier = op.barrier(ctx).name(ctx);
14    let tensor_map = op.tensor_map(ctx).name(ctx);
15    let smem_ptr = op.destination(ctx).name(ctx);
16    let indices = op.indices(ctx);
17    let indices = indices.iter().map(|it| it.name(ctx)).rev().join(", ");
18    let rank = op.rank(ctx);
19    format!(
20        "cuda::device::experimental::cp_async_bulk_tensor_{rank}d_global_to_shared({smem_ptr}, &{tensor_map}, {indices}, *{barrier});"
21    )
22});
23
24cuda_op!(TmaStoreOp, |op, ctx| {
25    let tensor_map = op.tensor_map(ctx).name(ctx);
26    let smem_ptr = op.source(ctx).name(ctx);
27    let indices = op.indices(ctx);
28    let indices = indices.iter().map(|it| it.name(ctx)).rev().join(", ");
29    let rank = op.rank(ctx);
30    format!(
31        "cuda::device::experimental::cp_async_bulk_tensor_{rank}d_shared_to_global(&{tensor_map}, {indices}, {smem_ptr});"
32    )
33});
34
35cuda_op!(CommitGroupOp, |_, _| {
36    "cuda::device::experimental::cp_async_bulk_commit_group();".into()
37});
38cuda_op!(WaitGroupOp, |op, ctx| {
39    let max_pending = op.max_pending(ctx).0;
40    format!("cuda::device::experimental::cp_async_bulk_wait_group<{max_pending}>();")
41});
42cuda_op!(WaitGroupReadOp, |op, ctx| {
43    let max_pending = op.max_pending(ctx).0;
44    format!("cuda::device::experimental::cp_async_bulk_wait_group_read<{max_pending}>();")
45});
cubecl_cpp/cuda/tma.rs

cubecl_cpp/cuda/
tma.rs
