cubecl_ir/
barrier.rs

1use crate::{Instruction, TypeHash};
2use alloc::{format, string::String, vec::Vec};
3use core::fmt::{Display, Write};
4
5use crate::OperationReflect;
6
7use super::Value;
8
9#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
10#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, Copy, PartialOrd, Ord)]
11pub enum BarrierLevel {
12    Unit,
13    Cube,
14}
15
16#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
17#[derive(Debug, Clone, TypeHash, PartialEq, Eq, Hash, OperationReflect)]
18#[operation(opcode_name = BarrierOpCode)]
19/// Operations available on a barrier
20pub enum BarrierOps {
21    /// Initialize the barrier, optionally with a cta proxy fence
22    Init {
23        barrier: Value,
24        is_elected: Value,
25        arrival_count: Value,
26    },
27    /// Manually initialize the barrier with an arrival count, without any sync or election handling
28    InitManual {
29        barrier: Value,
30        arrival_count: Value,
31    },
32    /// Copy source to destination
33    MemCopyAsync {
34        barrier: Value,
35        #[args(allow_ptr, ptr_read)]
36        source: Value,
37        #[args(allow_ptr, ptr_write)]
38        destination: Value,
39        source_length: Value,
40    },
41    /// Copy source to destination, with cooperative behaviour
42    MemCopyAsyncCooperative {
43        barrier: Value,
44        #[args(allow_ptr, ptr_read)]
45        source: Value,
46        #[args(allow_ptr, ptr_write)]
47        destination: Value,
48        source_length: Value,
49    },
50    /// Copy source to destination, with transaction count
51    MemCopyAsyncTx {
52        barrier: Value,
53        #[args(allow_ptr, ptr_read)]
54        source: Value,
55        #[args(allow_ptr, ptr_write)]
56        destination: Value,
57        source_length: Value,
58    },
59    /// Copy source to destination
60    CopyAsync {
61        #[args(allow_ptr, ptr_read)]
62        source: Value,
63        #[args(allow_ptr, ptr_write)]
64        destination: Value,
65        source_length: Value,
66        copy_length: u32,
67        checked: bool,
68    },
69    TmaLoad {
70        barrier: Value,
71        tensor_map: Value,
72        #[args(allow_ptr, ptr_write)]
73        destination: Value,
74        indices: Vec<Value>,
75    },
76    TmaLoadIm2col {
77        barrier: Value,
78        tensor_map: Value,
79        #[args(allow_ptr, ptr_write)]
80        destination: Value,
81        indices: Vec<Value>,
82        offsets: Vec<Value>,
83    },
84    /// Arrives at the barrier (decrements barrier count)
85    Arrive {
86        barrier: Value,
87    },
88    ArriveTx {
89        barrier: Value,
90        arrive_count_update: Value,
91        transaction_count_update: Value,
92    },
93    CommitCopyAsync {
94        barrier: Value,
95    },
96    ExpectTx {
97        barrier: Value,
98        transaction_count_update: Value,
99    },
100    Wait {
101        barrier: Value,
102        token: Value,
103    },
104    WaitParity {
105        barrier: Value,
106        phase: Value,
107    },
108    /// Waits until data is loaded
109    ArriveAndWait {
110        barrier: Value,
111    },
112}
113
114impl Display for BarrierOps {
115    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
116        match self {
117            BarrierOps::Init {
118                barrier,
119                arrival_count,
120                ..
121            } => write!(f, "{barrier}.init_barrier({arrival_count})"),
122            BarrierOps::InitManual {
123                barrier,
124                arrival_count,
125            } => write!(f, "{barrier}.init_barrier({arrival_count})"),
126            BarrierOps::MemCopyAsync {
127                barrier, source, ..
128            } => {
129                write!(f, "mem_copy_async({barrier}, source: {source})",)
130            }
131            BarrierOps::MemCopyAsyncCooperative {
132                barrier, source, ..
133            } => {
134                write!(f, "mem_copy_async_cooperative({barrier}, source: {source})",)
135            }
136            BarrierOps::MemCopyAsyncTx {
137                barrier, source, ..
138            } => {
139                write!(f, "mem_copy_async_tx({barrier}, source: {source})",)
140            }
141            BarrierOps::CopyAsync {
142                source,
143                destination,
144                source_length,
145                copy_length,
146                checked,
147            } => {
148                let source_slice = if *checked {
149                    format!("[..{source_length}]")
150                } else {
151                    String::new()
152                };
153                write!(
154                    f,
155                    "copy_async(source: {source}{source_slice}, destination: {destination}, bytes: {copy_length})",
156                )
157            }
158            BarrierOps::ArriveAndWait { barrier } => write!(f, "arrive_and_wait({barrier})"),
159            BarrierOps::TmaLoad {
160                barrier,
161                tensor_map,
162                destination,
163                indices,
164            } => {
165                let rank = indices.len();
166                let indices = indices.iter().fold(String::new(), |mut s, it| {
167                    let _ = write!(s, "{it}, ");
168                    s
169                });
170                write!(
171                    f,
172                    "tma_load::<{rank}>(bar: {barrier}, from: {tensor_map}, to: {destination}, indices: {indices})"
173                )
174            }
175            BarrierOps::TmaLoadIm2col {
176                barrier,
177                tensor_map,
178                destination,
179                indices,
180                offsets,
181            } => {
182                let rank = indices.len();
183                let indices = indices.iter().fold(String::new(), |mut s, it| {
184                    let _ = write!(s, "{it}, ");
185                    s
186                });
187                let offsets = offsets.iter().fold(String::new(), |mut s, it| {
188                    let _ = write!(s, "{it}, ");
189                    s
190                });
191                write!(
192                    f,
193                    "tma_load_im2col::<{rank}>(bar: {barrier}, from: {tensor_map}, to: {destination}, indices: ({indices}), offsets: ({offsets}))"
194                )
195            }
196            BarrierOps::Arrive { barrier } => write!(f, "arrive({barrier})"),
197            BarrierOps::CommitCopyAsync { barrier } => write!(f, "commit_copy_async({barrier})"),
198            BarrierOps::ArriveTx {
199                barrier,
200                arrive_count_update,
201                transaction_count_update,
202            } => write!(
203                f,
204                "arrive_tx({barrier}, {arrive_count_update}, {transaction_count_update})"
205            ),
206            BarrierOps::ExpectTx {
207                barrier,
208                transaction_count_update,
209            } => write!(f, "expect_tx({barrier}, {transaction_count_update})"),
210            BarrierOps::Wait { barrier, token } => write!(f, "wait({barrier}, {token})"),
211            BarrierOps::WaitParity { barrier, phase } => {
212                write!(f, "wait_parity({barrier}, {phase})")
213            }
214        }
215    }
216}
217
218impl Display for BarrierLevel {
219    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
220        match self {
221            BarrierLevel::Unit => f.write_str("unit"),
222            BarrierLevel::Cube => f.write_str("cube"),
223        }
224    }
225}
226
227impl From<BarrierOps> for Instruction {
228    fn from(value: BarrierOps) -> Self {
229        Instruction::no_out(value)
230    }
231}
cubecl_ir/barrier.rs

cubecl_ir/
barrier.rs
