cubecl_ir/dialect/
barrier.rs

1use cubecl_macros_internal::{cube_op, op_traits};
2use pliron::r#type::{TypeHandle, TypedHandle};
3
4use crate::{
5    CanMaterialize, HasSideEffects,
6    attributes::{BoolAttr, IndexAttr},
7    dialect::synchronization::SyncScope,
8    interfaces::Synchronizes,
9    prelude::*,
10    types::{PointerType, barrier::BarrierTokenType},
11};
12
13#[cube_op(name = "barrier.init")]
14#[result_ty(none)]
15#[op_traits(CanMaterialize)]
16pub struct InitOp {
17    // Opaque so we can't know exact memory effects. Treat it as atomic read-update.
18    #[operand(ptr_read, ptr_write)]
19    pub barrier: Value,
20    pub arrival_count: Value,
21}
22
23#[cube_op(name = "barrier.memcpy_async")]
24#[result_ty(none)]
25#[op_traits(CanMaterialize)]
26pub struct MemCopyAsyncOp {
27    // Opaque so we can't know exact memory effects. Treat it as atomic read-update.
28    #[operand(ptr_read, ptr_write)]
29    pub barrier: Value,
30    #[operand(ptr_read)]
31    pub source: Value,
32    #[operand(ptr_write)]
33    pub destination: Value,
34    pub source_length: Value,
35    pub cooperative: BoolAttr,
36}
37
38#[cube_op(name = "barrier.memcpy_async_tx")]
39#[result_ty(none)]
40#[op_traits(CanMaterialize)]
41pub struct MemCopyAsyncTxOp {
42    // Opaque so we can't know exact memory effects. Treat it as atomic read-update.
43    #[operand(ptr_read, ptr_write)]
44    pub barrier: Value,
45    #[operand(ptr_read)]
46    pub source: Value,
47    #[operand(ptr_write)]
48    pub destination: Value,
49    pub source_length: Value,
50}
51
52#[cube_op(name = "barrier.copy_async")]
53#[result_ty(none)]
54#[op_traits(CanMaterialize)]
55pub struct CopyAsyncOp {
56    #[operand(ptr_read)]
57    pub source: Value,
58    #[operand(ptr_write)]
59    pub destination: Value,
60    pub source_length: Value,
61    pub copy_length: IndexAttr,
62    pub checked: BoolAttr,
63}
64
65#[cube_op(name = "barrier.arrive")]
66#[result_ty(from_inputs = token_ty)]
67#[op_traits(CanMaterialize)]
68pub struct ArriveOp {
69    // Opaque so we can't know exact memory effects. Treat it as atomic read-update.
70    #[operand(ptr_read, ptr_write)]
71    pub barrier: Value,
72}
73
74#[cube_op(name = "barrier.arrive_and_expect_tx")]
75#[result_ty(from_inputs = |ctx, bar, _, _| token_ty(ctx, bar))]
76#[op_traits(CanMaterialize)]
77pub struct ArriveAndExpectTxOp {
78    // Opaque so we can't know exact memory effects. Treat it as atomic read-update.
79    #[operand(ptr_read, ptr_write)]
80    pub barrier: Value,
81    pub arrive_count_update: Value,
82    pub transaction_count_update: Value,
83}
84
85fn token_ty(ctx: &Context, barrier: &Value) -> TypeHandle {
86    let bar_ptr = barrier.get_type(ctx).deref(ctx);
87    let bar_ptr = bar_ptr.downcast_ref::<PointerType>().unwrap();
88    let bar = TypedHandle::from_handle(bar_ptr.inner, ctx).expect("Should be barrier");
89    BarrierTokenType::get(ctx, bar).into()
90}
91
92#[cube_op(name = "barrier.commit_copy_async")]
93#[result_ty(none)]
94#[op_traits(CanMaterialize)]
95pub struct CommitCopyAsyncOp {
96    // Opaque so we can't know exact memory effects. Treat it as atomic read-update.
97    #[operand(ptr_read, ptr_write)]
98    pub barrier: Value,
99}
100
101#[cube_op(name = "barrier.expect_tx")]
102#[result_ty(none)]
103#[op_traits(CanMaterialize)]
104pub struct ExpectTxOp {
105    // Opaque so we can't know exact memory effects. Treat it as atomic read-update.
106    #[operand(ptr_read, ptr_write)]
107    pub barrier: Value,
108    pub transaction_count_update: Value,
109}
110
111#[cube_op(name = "barrier.wait")]
112#[result_ty(none)]
113#[op_traits(CanMaterialize, HasSideEffects)]
114pub struct WaitOp {
115    // Opaque so we can't know exact memory effects. Treat it as atomic read-update.
116    #[operand(ptr_read, ptr_write)]
117    pub barrier: Value,
118    pub token: Value,
119}
120
121/// Sync scope depends on init params and use, so can't be trivially analyzed. So just be
122/// conservative.
123#[op_interface_impl]
124impl Synchronizes for WaitOp {
125    fn minimum_scope(&self, _ctx: &Context) -> SyncScope {
126        SyncScope::Unit
127    }
128
129    fn maximum_scope(&self, _ctx: &Context) -> SyncScope {
130        SyncScope::Device
131    }
132}
133
134#[cube_op(name = "barrier.wait_parity")]
135#[result_ty(none)]
136#[op_traits(CanMaterialize, HasSideEffects)]
137pub struct WaitParityOp {
138    // Opaque so we can't know exact memory effects. Treat it as atomic read-update.
139    #[operand(ptr_read, ptr_write)]
140    pub barrier: Value,
141    pub phase: Value,
142}
143
144/// Sync scope depends on init params and use, so can't be trivially analyzed. So just be
145/// conservative.
146#[op_interface_impl]
147impl Synchronizes for WaitParityOp {
148    fn minimum_scope(&self, _ctx: &Context) -> SyncScope {
149        SyncScope::Unit
150    }
151
152    fn maximum_scope(&self, _ctx: &Context) -> SyncScope {
153        SyncScope::Device
154    }
155}
156
157#[cube_op(name = "barrier.arrive_and_wait")]
158#[result_ty(none)]
159#[op_traits(CanMaterialize, HasSideEffects)]
160pub struct ArriveAndWaitOp {
161    // Opaque so we can't know exact memory effects. Treat it as atomic read-update.
162    #[operand(ptr_read, ptr_write)]
163    pub barrier: Value,
164}
165
166/// Sync scope depends on init params and use, so can't be trivially analyzed. So just be
167/// conservative.
168#[op_interface_impl]
169impl Synchronizes for ArriveAndWaitOp {
170    fn minimum_scope(&self, _ctx: &Context) -> SyncScope {
171        SyncScope::Unit
172    }
173
174    fn maximum_scope(&self, _ctx: &Context) -> SyncScope {
175        SyncScope::Device
176    }
177}
cubecl_ir/dialect/barrier.rs

cubecl_ir/dialect/
barrier.rs
