cubecl_cpp/cuda/
barrier.rs

1use cubecl_core::ir::{
2    dialect::{barrier::*, base::ptr_value_ty},
3    interfaces::TypedExt,
4};
5
6use crate::{
7    cuda::{cuda_op, cuda_op_with_out},
8    shared::CppValue,
9};
10
11cuda_op!(InitOp, |op, ctx| format!(
12    "init({}, {});",
13    op.barrier(ctx).name(ctx),
14    op.arrival_count(ctx).name(ctx)
15));
16
17cuda_op!(MemCopyAsyncOp, |op, ctx| {
18    let source = op.source(ctx).name(ctx);
19    let barrier = op.barrier(ctx).name(ctx);
20    let destination = op.destination(ctx).name(ctx);
21    let source_length = op.source_length(ctx).name(ctx);
22    let size = ptr_value_ty(ctx, &op.source(ctx)).size(ctx);
23    match op.cooperative(ctx).0 {
24        false => format!(
25            "cuda::memcpy_async({destination}, {source}, {source_length} * {size}, *{barrier});"
26        ),
27        true => format!(
28            "cuda::memcpy_async(cooperative_groups::this_thread_block(), {destination}, {source}, {source_length} * {size}, *{barrier});"
29        ),
30    }
31});
32
33cuda_op!(MemCopyAsyncTxOp, |op, ctx| {
34    let barrier = op.barrier(ctx).name(ctx);
35    let source = op.source(ctx).name(ctx);
36    let destination = op.destination(ctx).name(ctx);
37    let source_length = op.source_length(ctx).name(ctx);
38    let size = ptr_value_ty(ctx, &op.source(ctx)).size(ctx);
39    format!(
40        "cuda::device::memcpy_async_tx({destination}, {source}, {source_length} * {size}, *{barrier});"
41    )
42});
43
44cuda_op_with_out!(ArriveOp, |op, ctx| {
45    let barrier = op.barrier(ctx).name(ctx);
46    format!("{barrier}->arrive()")
47});
48
49cuda_op_with_out!(ArriveAndExpectTxOp, |op, ctx| {
50    let barrier = op.barrier(ctx).name(ctx);
51    let arrive_count_update = op.arrive_count_update(ctx).name(ctx);
52    let transaction_count_update = op.transaction_count_update(ctx).name(ctx);
53    format!(
54        "cuda::device::barrier_arrive_tx(*{barrier}, {arrive_count_update}, {transaction_count_update})"
55    )
56});
57
58cuda_op!(ExpectTxOp, |op, ctx| {
59    let barrier = op.barrier(ctx).name(ctx);
60    let transaction_count_update = op.transaction_count_update(ctx).name(ctx);
61    format!("cuda::device::barrier_expect_tx(*{barrier}, {transaction_count_update});")
62});
63
64cuda_op!(WaitOp, |op, ctx| {
65    let barrier = op.barrier(ctx).name(ctx);
66    let token = op.token(ctx).name(ctx);
67    format!("{barrier}->wait(std::move({token}));")
68});
69
70cuda_op!(WaitParityOp, |op, ctx| {
71    let barrier = op.barrier(ctx).name(ctx);
72    let phase = op.phase(ctx).name(ctx);
73    format!("{barrier}->wait_parity({phase});")
74});
75
76cuda_op!(ArriveAndWaitOp, |op, ctx| {
77    let barrier = op.barrier(ctx).name(ctx);
78    format!("{barrier}->arrive_and_wait();")
79});
cubecl_cpp/cuda/barrier.rs

cubecl_cpp/cuda/
barrier.rs
