cubecl_core/frontend/
barrier.rs

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

cubecl_core/frontend/
barrier.rs
