cubecl_core/frontend/
barrier.rs

1//! This module exposes barrier for asynchronous data transfer
2
3use alloc::vec;
4
5use crate as cubecl;
6use cubecl_ir::{Instruction, OpaqueType, Value};
7use cubecl_macros::intrinsic;
8use paste::paste;
9
10use crate::{
11    ir::{BarrierOps, Scope},
12    prelude::*,
13    unexpanded,
14};
15
16use super::{CubePrimitive, CubeType, NativeExpand, SliceExpand, TensorMap};
17
18/// A mechanism for awaiting on asynchronous data transfers
19/// Behavior is defined by its ``BarrierLevel``.
20#[derive(Clone, Copy, PartialEq, Eq)]
21pub struct Barrier;
22pub type BarrierExpand = NativeExpand<Barrier>;
23
24#[derive(Clone, Copy, PartialEq)]
25pub struct BarrierToken;
26
27impl CubeType for Barrier {
28    type ExpandType = NativeExpand<Barrier>;
29}
30
31impl CubeDebug for Barrier {}
32
33impl CubePrimitive for Barrier {
34    type Scalar = u32; // Dummy, maybe we need another trait for non-standard primitives
35    type Size = Const<1>;
36    type WithScalar<S: Scalar> = S;
37    fn from_const_value(_value: cubecl_ir::ConstantValue) -> Self {
38        unreachable!("Can't create from const value")
39    }
40}
41
42impl NativeAssign for Barrier {
43    fn elem_init_mut(_scope: &Scope, elem: Value) -> Value {
44        elem
45    }
46}
47
48impl CubeType for BarrierToken {
49    type ExpandType = NativeExpand<BarrierToken>;
50}
51
52impl NativeAssign for BarrierToken {
53    fn elem_init_mut(_scope: &Scope, elem: Value) -> Value {
54        elem
55    }
56}
57
58impl AsMutExpand for NativeExpand<BarrierToken> {
59    fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
60        self
61    }
62}
63
64macro_rules! tensor_map_load {
65    ($dim: literal, $($arg: expr),*) => {
66        paste! {
67            impl Barrier {
68                /// Copy a tile from a global memory `source` to a shared memory `destination`, with
69                /// the provided offsets.
70                #[allow(unused, clippy::too_many_arguments)]
71                pub fn [<tma_load_ $dim d>]<C1: CubePrimitive, C2: CubePrimitive<Scalar = C1::Scalar>>(
72                    &self,
73                    source: &TensorMap<C1, Tiled>,
74                    destination: &mut [C2],
75                    $($arg: i32),*
76                ) {
77                    unexpanded!()
78                }
79
80                #[allow(clippy::too_many_arguments)]
81                pub fn [<__expand_tma_load_ $dim d>]<C1: CubePrimitive, C2: CubePrimitive<Scalar = C1::Scalar>>(
82                    scope: &Scope,
83                    expand: &NativeExpand<Barrier>,
84                    source: &NativeExpand<TensorMap<C1, Tiled>>,
85                    destination: &mut SliceExpand<C2>,
86                    $($arg: NativeExpand<i32>),*
87                ) {
88                    expand.[<__expand_tma_load_ $dim d_method>](scope, source, destination, $($arg),*);
89                }
90            }
91
92            impl NativeExpand<Barrier> {
93                #[allow(clippy::too_many_arguments)]
94                pub fn [<__expand_tma_load_ $dim d_method>]<C1: CubePrimitive, C2: CubePrimitive<Scalar = C1::Scalar>>(
95                    &self,
96                    scope: &Scope,
97                    source: &NativeExpand<TensorMap<C1, Tiled>>,
98                    destination: &mut SliceExpand<C2>,
99                    $($arg: NativeExpand<i32>),*
100                ) {
101                    let barrier = self.expand;
102                    let source = source.expand;
103                    let destination = unsafe { *destination.__expand_as_ptr_method(scope) }.expand;
104
105                    let mem_copy = BarrierOps::TmaLoad {
106                        barrier,
107                        tensor_map: source,
108                        destination,
109                        indices: vec![$($arg.expand),*],
110                    };
111
112                    scope.register(Instruction::no_out(mem_copy));
113                }
114            }
115        }
116    };
117}
118
119macro_rules! tensor_map_load_im2col {
120    ($dim: literal, $($arg: expr),*; $($offset: expr),*) => {
121        paste! {
122            impl Barrier {
123                /// Copy a tile from a global memory `source` to a shared memory `destination`, with
124                /// the provided offsets.
125                #[allow(unused, clippy::too_many_arguments)]
126                pub fn [<tma_load_im2col_ $dim d>]<C1: CubePrimitive, C2: CubePrimitive<Scalar = C1::Scalar>>(
127                    &self,
128                    source: &TensorMap<C1, Im2col>,
129                    destination: &mut [C2],
130                    $($arg: i32,)*
131                    $($offset: u16),*
132                ) {
133                    unexpanded!()
134                }
135
136                #[allow(clippy::too_many_arguments)]
137                pub fn [<__expand_tma_load_im2col_ $dim d>]<C1: CubePrimitive, C2: CubePrimitive<Scalar = C1::Scalar>>(
138                    scope: &Scope,
139                    expand: &NativeExpand<Barrier>,
140                    source: &NativeExpand<TensorMap<C1, Im2col>>,
141                    destination: &mut SliceExpand<C2>,
142                    $($arg: NativeExpand<i32>,)*
143                    $($offset: NativeExpand<u16>),*
144                ) {
145                    expand.[<__expand_tma_load_im2col_ $dim d_method>](scope, source, destination, $($arg),*, $($offset),*);
146                }
147            }
148
149            impl NativeExpand<Barrier> {
150                #[allow(clippy::too_many_arguments)]
151                pub fn [<__expand_tma_load_im2col_ $dim d_method>]<C1: CubePrimitive, C2: CubePrimitive<Scalar = C1::Scalar>>(
152                    &self,
153                    scope: &Scope,
154                    source: &NativeExpand<TensorMap<C1, Im2col>>,
155                    destination: &mut SliceExpand<C2>,
156                    $($arg: NativeExpand<i32>,)*
157                    $($offset: NativeExpand<u16>),*
158                ) {
159                    let barrier = self.expand;
160                    let source = source.expand;
161                    let destination = unsafe { *destination.__expand_as_ptr_method(scope) }.expand;
162
163                    let mem_copy = BarrierOps::TmaLoadIm2col {
164                        barrier,
165                        tensor_map: source,
166                        destination,
167                        indices: vec![$($arg.expand),*],
168                        offsets: vec![$($offset.expand),*],
169                    };
170
171                    scope.register(Instruction::no_out(mem_copy));
172                }
173            }
174        }
175    };
176}
177
178tensor_map_load!(1, x);
179tensor_map_load!(2, y, x);
180tensor_map_load!(3, z, y, x);
181tensor_map_load!(4, w, z, y, x);
182tensor_map_load!(5, v, w, z, y, x);
183
184tensor_map_load_im2col!(3, n, w, c; w_offset);
185tensor_map_load_im2col!(4, n, h, w, c; h_offset, w_offset);
186tensor_map_load_im2col!(5, n, d, h, w, c; d_offset, h_offset, w_offset);
187
188#[cube]
189impl Barrier {
190    /// Create a local barrier object for the current unit. Automatically initialized with an
191    /// arrival count of `1`.
192    pub fn local() -> Self {
193        intrinsic!(|scope| {
194            let value = scope.create_local_mut(OpaqueType::Barrier(cubecl_ir::BarrierLevel::Unit));
195            scope.register(BarrierOps::Init {
196                barrier: value,
197                is_elected: true.into(),
198                arrival_count: 1.into(),
199            });
200            value.into()
201        })
202    }
203
204    /// Create a shared memory barrier that can be accesses by all units in the cube. Initialized
205    /// by the `is_elected` unit with an arrival count of `arrival_count`. This is the number of
206    /// times `arrive` or one of its variants needs to be called before the barrier advances.
207    ///
208    /// If all units in the cube arrive on the barrier, use `CUBE_DIM` as the arrival count. For
209    /// other purposes, only a subset may need to arrive.
210    pub fn shared(arrival_count: u32, is_elected: bool) -> Shared<Barrier> {
211        intrinsic!(|scope| {
212            let value =
213                scope.create_shared(OpaqueType::Barrier(cubecl_ir::BarrierLevel::Cube), None);
214            scope.register(BarrierOps::Init {
215                barrier: value,
216                is_elected: is_elected.expand,
217                arrival_count: arrival_count.expand,
218            });
219            value.into()
220        })
221    }
222
223    /// Create a shared memory barrier that can be accesses by all units in the cube. Only declared,
224    /// but not initialized.
225    pub fn shared_uninit() -> Shared<Barrier> {
226        intrinsic!(|scope| {
227            let value =
228                scope.create_shared(OpaqueType::Barrier(cubecl_ir::BarrierLevel::Cube), None);
229            value.into()
230        })
231    }
232
233    /// Initializes a barrier with a given `arrival_count`. This is the number of
234    /// times `arrive` or one of its variants needs to be called before the barrier advances.
235    ///
236    /// If all units in the cube arrive on the barrier, use `CUBE_DIM` as the arrival count. For
237    /// other purposes, only a subset may need to arrive.
238    ///
239    /// # Note
240    ///
241    /// No synchronization or election is performed, this is raw initialization. For shared barriers
242    /// ensure only one unit performs the initialization, and synchronize the cube afterwards. There
243    /// may also be additional synchronization requirements for bulk copy operations, like
244    /// [`sync_async_proxy_shared()`].
245    pub fn init_manual(&self, arrival_count: u32) {
246        intrinsic!(|scope| {
247            let barrier = self.expand;
248
249            scope.register(BarrierOps::InitManual {
250                barrier,
251                arrival_count: arrival_count.expand,
252            });
253        })
254    }
255}
256
257// MemcpyAsync
258
259#[cube]
260impl Barrier {
261    /// Copy the source slice to destination
262    ///
263    /// # Safety
264    ///
265    /// This will try to copy the whole source slice, so
266    /// make sure source length <= destination length
267    pub fn memcpy_async<C: CubePrimitive>(&self, source: &[C], destination: &mut [C]) {
268        intrinsic!(|scope| {
269            let barrier = self.expand;
270            let source_length = source.__extract_length(scope).expand;
271            let source = unsafe { *source.__expand_as_ptr_method(scope) }.expand;
272            let destination = unsafe { *destination.__expand_as_ptr_method(scope) }.expand;
273
274            let mem_copy = BarrierOps::MemCopyAsync {
275                barrier,
276                destination,
277                source,
278                source_length,
279            };
280
281            scope.register(Instruction::no_out(mem_copy));
282        })
283    }
284
285    /// Copy the source slice to destination
286    ///
287    /// # Safety
288    ///
289    /// This will try to copy the whole source slice, so
290    /// make sure source length <= destination length
291    pub fn memcpy_async_cooperative<C: CubePrimitive>(&self, source: &[C], destination: &mut [C]) {
292        intrinsic!(|scope| {
293            let barrier = self.expand;
294            let source_length = source.__extract_length(scope).expand;
295            let source = unsafe { *source.__expand_as_ptr_method(scope) }.expand;
296            let destination = unsafe { *destination.__expand_as_ptr_method(scope) }.expand;
297
298            let mem_copy = BarrierOps::MemCopyAsyncCooperative {
299                barrier,
300                source,
301                destination,
302                source_length,
303            };
304
305            scope.register(Instruction::no_out(mem_copy));
306        })
307    }
308
309    /// Copy the source slice to destination. Uses transaction count like TMA, so use with
310    /// `expect_tx` or `arrive_and_expect_tx`.
311    ///
312    /// # Safety
313    ///
314    /// This will try to copy the whole source slice, so
315    /// make sure source length <= destination length
316    pub fn memcpy_async_tx<C: CubePrimitive>(&self, source: &[C], destination: &mut [C]) {
317        intrinsic!(|scope| {
318            let barrier = self.expand;
319            let source_length = source.__extract_length(scope).expand;
320            let source = unsafe { *source.__expand_as_ptr_method(scope) }.expand;
321            let destination = unsafe { *destination.__expand_as_ptr_method(scope) }.expand;
322
323            let mem_copy = BarrierOps::MemCopyAsyncTx {
324                barrier,
325                source,
326                destination,
327                source_length,
328            };
329
330            scope.register(Instruction::no_out(mem_copy));
331        })
332    }
333}
334
335// Arrival and Wait
336
337#[cube]
338impl Barrier {
339    /// Arrive at the barrier, decrementing arrival count
340    pub fn arrive(&self) -> BarrierToken {
341        intrinsic!(|scope| {
342            let barrier = self.expand;
343            let Type::Opaque(OpaqueType::Barrier(level)) = barrier.ty.unwrap_ptr() else {
344                unreachable!()
345            };
346            let token = scope.create_value(Type::Opaque(OpaqueType::BarrierToken(level)));
347            scope.register(Instruction::new(BarrierOps::Arrive { barrier }, token));
348            token.into()
349        })
350    }
351
352    /// Arrive at the barrier, decrementing arrival count. Additionally increments expected count.
353    pub fn arrive_and_expect_tx(&self, arrival_count: u32, transaction_count: u32) -> BarrierToken {
354        intrinsic!(|scope| {
355            let barrier = self.expand;
356            let Type::Opaque(OpaqueType::Barrier(level)) = barrier.ty.unwrap_ptr() else {
357                unreachable!()
358            };
359            let token = scope.create_value(Type::Opaque(OpaqueType::BarrierToken(level)));
360            let arrival_count: Value = arrival_count.into();
361            let transaction_count: Value = transaction_count.into();
362            scope.register(Instruction::new(
363                BarrierOps::ArriveTx {
364                    barrier,
365                    arrive_count_update: arrival_count,
366                    transaction_count_update: transaction_count,
367                },
368                token,
369            ));
370            token.into()
371        })
372    }
373
374    /// Increments the expected count of the barrier.
375    pub fn expect_tx(&self, expected_count: u32) {
376        intrinsic!(|scope| {
377            let barrier = self.expand;
378            let transaction_count: Value = expected_count.into();
379            scope.register(BarrierOps::ExpectTx {
380                barrier,
381                transaction_count_update: transaction_count,
382            });
383        })
384    }
385
386    /// Wait until all data is loaded
387    pub fn arrive_and_wait(&self) {
388        intrinsic!(|scope| {
389            let barrier = self.expand;
390            scope.register(BarrierOps::ArriveAndWait { barrier });
391        })
392    }
393
394    /// Wait at the barrier until all arrivals are done
395    pub fn wait(&self, token: BarrierToken) {
396        intrinsic!(|scope| {
397            let barrier = self.expand;
398            let token = token.expand;
399            scope.register(BarrierOps::Wait { barrier, token });
400        })
401    }
402
403    /// Wait at the barrier until the `phase` is completed. Doesn't require a token, but needs phase
404    /// to be managed manually.
405    pub fn wait_parity(&self, phase: u32) {
406        intrinsic!(|scope| {
407            let barrier = self.expand;
408            let phase = phase.expand;
409            scope.register(BarrierOps::WaitParity { barrier, phase });
410        })
411    }
412}
413
414// Copy async
415
416/// Copy the source slice in global memory to destination in shared memory with a low level async
417/// copy. This only copies up to 128 bits/16 bytes, and does not synchronize. Use
418/// `barrier.copy_async_arrive` to make the reads visible.
419/// `copy_size` is in terms of elements to simplify copying between different vector sizes.
420///
421/// # Safety
422///
423/// This will try to copy the entire `copy_size`, so make sure the full width is in bounds.
424/// Starting address must be aligned to the full copy size.
425pub fn copy_async<C: CubePrimitive>(_source: &[C], _destination: &mut [C], _copy_size: u32) {
426    unexpanded!()
427}
428
429pub mod copy_async {
430    use super::*;
431
432    pub fn expand<C: CubePrimitive>(
433        scope: &Scope,
434        source: &SliceExpand<C>,
435        destination: &mut SliceExpand<C>,
436        copy_length: u32,
437    ) {
438        let source_length = copy_length.into();
439        let source = unsafe { *source.__expand_as_ptr_method(scope) }.expand;
440        let destination = unsafe { *destination.__expand_as_ptr_method(scope) }.expand;
441        let scalar_size = C::__expand_as_type(scope).storage_type().size();
442
443        let mem_copy = BarrierOps::CopyAsync {
444            source,
445            destination,
446            source_length,
447            copy_length: copy_length * scalar_size as u32,
448            checked: false,
449        };
450
451        scope.register(Instruction::no_out(mem_copy));
452    }
453}
454
455/// Copy the source slice in global memory to destination in shared memory with a low level async
456/// copy. This only copies up to 128 bits/16 bytes, and does not synchronize. Use
457/// `barrier.copy_async_arrive` to make the reads visible.
458/// `copy_size` is in terms of elements to simplify copying between different vector sizes.
459///
460/// Will only copy the length of the source slice, and zero fill the rest. Source length must be
461/// <= copy size.
462///
463/// # Safety
464/// Starting address must be aligned to the full copy size.
465/// **This will silently fail if the address is only aligned to the source length and not the copy size!**
466pub fn copy_async_checked<C: CubePrimitive>(
467    _source: &[C],
468    _destination: &mut [C],
469    _copy_size: u32,
470) {
471    unexpanded!();
472}
473
474pub mod copy_async_checked {
475    use super::*;
476
477    pub fn expand<C: CubePrimitive>(
478        scope: &Scope,
479        source: &SliceExpand<C>,
480        destination: &mut SliceExpand<C>,
481        copy_length: u32,
482    ) {
483        let source_length = source.__extract_length(scope).expand;
484
485        // OOB pointer is allowed as long as length is 0
486        let source = unsafe { (*source.__expand_as_ptr_unchecked_method(scope)).expand };
487        let destination = unsafe { (*destination.__expand_as_ptr_unchecked_method(scope)).expand };
488        let scalar_size = C::__expand_as_type(scope).storage_type().size();
489
490        let mem_copy = BarrierOps::CopyAsync {
491            source,
492            destination,
493            source_length,
494            copy_length: copy_length * scalar_size as u32,
495            checked: true,
496        };
497
498        scope.register(Instruction::no_out(mem_copy));
499    }
500}
501
502#[cube]
503impl Barrier {
504    /// Makes all previous `copy_async` operations visible on the barrier.
505    /// Should be called once after all copies have been dispatched, before reading from the shared
506    /// memory.
507    ///
508    /// Does *not* count as an arrive in terms of the barrier arrival count. So `arrive` or
509    /// `arrive_and_wait` should still be called afterwards.
510    pub fn commit_copy_async(&self) {
511        intrinsic!(|scope| {
512            let barrier = self.expand;
513            let Type::Opaque(OpaqueType::Barrier(level)) = barrier.ty.unwrap_ptr() else {
514                unreachable!()
515            };
516            let token = scope.create_value(Type::Opaque(OpaqueType::BarrierToken(level)));
517            scope.register(Instruction::new(
518                BarrierOps::CommitCopyAsync { barrier },
519                token,
520            ));
521        })
522    }
523}
524
525impl From<SharedExpand<Barrier>> for BarrierExpand {
526    fn from(value: SharedExpand<Barrier>) -> Self {
527        value.expand.into()
528    }
529}
cubecl_core/frontend/barrier.rs

cubecl_core/frontend/
barrier.rs
