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 accesses 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    /// If all units in the cube arrive on the barrier, use `CUBE_DIM` as the arrival count. For
214    /// other purposes, only a subset may need to arrive.
215    pub fn shared(arrival_count: u32, is_elected: bool) -> Shared<Barrier> {
216        intrinsic!(|scope| {
217            let value =
218                scope.create_shared(BarrierType::get(scope.ctx(), BarrierLevel::Cube), None);
219            if_expand(scope, is_elected, |scope| {
220                let arrival_count = arrival_count.read_value(scope);
221                let op = InitOp::new(scope.ctx_mut(), value, arrival_count);
222                scope.register(&op);
223            });
224            sync_cube::expand(scope);
225            value.into()
226        })
227    }
228
229    /// Create a shared memory barrier that can be accesses by all units in the cube. Only declared,
230    /// but not initialized.
231    pub fn shared_uninit() -> Shared<Barrier> {
232        intrinsic!(|scope| {
233            let value =
234                scope.create_shared(BarrierType::get(scope.ctx(), BarrierLevel::Cube), None);
235            value.into()
236        })
237    }
238
239    /// Initializes a barrier with a given `arrival_count`. This is the number of
240    /// times `arrive` or one of its variants needs to be called before the barrier advances.
241    ///
242    /// If all units in the cube arrive on the barrier, use `CUBE_DIM` as the arrival count. For
243    /// other purposes, only a subset may need to arrive.
244    ///
245    /// # Note
246    ///
247    /// No synchronization or election is performed, this is raw initialization. For shared barriers
248    /// ensure only one unit performs the initialization, and synchronize the cube afterwards. There
249    /// may also be additional synchronization requirements for bulk copy operations, like
250    /// [`sync_async_proxy_shared()`].
251    pub fn init_manual(&self, arrival_count: u32) {
252        intrinsic!(|scope| {
253            let barrier = self.value(scope);
254            let arrival_count = arrival_count.read_value(scope);
255            let op = InitOp::new(scope.ctx_mut(), barrier, arrival_count);
256            scope.register(&op);
257        })
258    }
259}
260
261// MemcpyAsync
262
263#[cube]
264impl Barrier {
265    /// Copy the source slice to destination
266    ///
267    /// # Safety
268    ///
269    /// This will try to copy the whole source slice, so
270    /// make sure source length <= destination length
271    pub fn memcpy_async<C: CubePrimitive>(&self, source: &[C], destination: &mut [C]) {
272        intrinsic!(|scope| {
273            let barrier = self.value(scope);
274            let source_length = source.__extract_length(scope).value(scope);
275            let source = unsafe { *source.__expand_as_ptr_method(scope) }.value(scope);
276            let destination = unsafe { *destination.__expand_as_ptr_method(scope) }.value(scope);
277
278            let mem_copy = MemCopyAsyncOp::new(
279                scope.ctx_mut(),
280                barrier,
281                source,
282                destination,
283                source_length,
284                false,
285            );
286
287            scope.register(&mem_copy);
288        })
289    }
290
291    /// Copy the source slice to destination
292    ///
293    /// # Safety
294    ///
295    /// This will try to copy the whole source slice, so
296    /// make sure source length <= destination length
297    pub fn memcpy_async_cooperative<C: CubePrimitive>(&self, source: &[C], destination: &mut [C]) {
298        intrinsic!(|scope| {
299            let barrier = self.value(scope);
300            let source_length = source.__extract_length(scope).value(scope);
301            let source = unsafe { *source.__expand_as_ptr_method(scope) }.value(scope);
302            let destination = unsafe { *destination.__expand_as_ptr_method(scope) }.value(scope);
303
304            let mem_copy = MemCopyAsyncOp::new(
305                scope.ctx_mut(),
306                barrier,
307                source,
308                destination,
309                source_length,
310                true,
311            );
312
313            scope.register(&mem_copy);
314        })
315    }
316
317    /// Copy the source slice to destination. Uses transaction count like TMA, so use with
318    /// `expect_tx` or `arrive_and_expect_tx`.
319    ///
320    /// # Safety
321    ///
322    /// This will try to copy the whole source slice, so
323    /// make sure source length <= destination length
324    pub fn memcpy_async_tx<C: CubePrimitive>(&self, source: &[C], destination: &mut [C]) {
325        intrinsic!(|scope| {
326            let barrier = self.value(scope);
327            let source_length = source.__extract_length(scope).value(scope);
328            let source = unsafe { *source.__expand_as_ptr_method(scope) }.value(scope);
329            let destination = unsafe { *destination.__expand_as_ptr_method(scope) }.value(scope);
330
331            let mem_copy =
332                MemCopyAsyncTxOp::new(scope.ctx_mut(), barrier, source, destination, source_length);
333
334            scope.register(&mem_copy);
335        })
336    }
337}
338
339// Arrival and Wait
340
341#[cube]
342impl Barrier {
343    /// Arrive at the barrier, decrementing arrival count
344    pub fn arrive(&self) -> BarrierToken {
345        intrinsic!(|scope| {
346            let barrier = self.value(scope);
347            let arrive = ArriveOp::new(scope.ctx_mut(), barrier);
348            scope.register_with_result(&arrive).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.value(scope);
356            let arrival_count = arrival_count.read_value(scope);
357            let transaction_count = transaction_count.read_value(scope);
358            let op = ArriveAndExpectTxOp::new(
359                scope.ctx_mut(),
360                barrier,
361                arrival_count,
362                transaction_count,
363            );
364            scope.register_with_result(&op).into()
365        })
366    }
367
368    /// Increments the expected count of the barrier.
369    pub fn expect_tx(&self, transaction_count_update: u32) {
370        intrinsic!(|scope| {
371            let barrier = self.value(scope);
372            let transaction_count_update = transaction_count_update.value(scope);
373            scope.register(&ExpectTxOp::new(
374                scope.ctx_mut(),
375                barrier,
376                transaction_count_update,
377            ));
378        })
379    }
380
381    /// Wait until all data is loaded
382    pub fn arrive_and_wait(&self) {
383        intrinsic!(|scope| {
384            let barrier = self.value(scope);
385            scope.register(&ArriveAndWaitOp::new(scope.ctx_mut(), barrier));
386        })
387    }
388
389    /// Wait at the barrier until all arrivals are done
390    pub fn wait(&self, token: BarrierToken) {
391        intrinsic!(|scope| {
392            let barrier = self.value(scope);
393            let token = token.value(scope);
394            scope.register(&WaitOp::new(scope.ctx_mut(), barrier, token));
395        })
396    }
397
398    /// Wait at the barrier until the `phase` is completed. Doesn't require a token, but needs phase
399    /// to be managed manually.
400    pub fn wait_parity(&self, phase: u32) {
401        intrinsic!(|scope| {
402            let barrier = self.value(scope);
403            let phase = phase.read_value(scope);
404            scope.register(&WaitParityOp::new(scope.ctx_mut(), barrier, phase));
405        })
406    }
407}
408
409// Copy async
410
411/// Copy the source slice in global memory to destination in shared memory with a low level async
412/// copy. This only copies up to 128 bits/16 bytes, and does not synchronize. Use
413/// `barrier.copy_async_arrive` to make the reads visible.
414/// `copy_size` is in terms of elements to simplify copying between different vector sizes.
415///
416/// # Safety
417///
418/// This will try to copy the entire `copy_size`, so make sure the full width is in bounds.
419/// Starting address must be aligned to the full copy size.
420pub fn copy_async<C: CubePrimitive>(_source: &[C], _destination: &mut [C], _copy_size: u32) {
421    unexpanded!()
422}
423
424pub mod copy_async {
425    use super::*;
426
427    pub fn expand<C: CubePrimitive>(
428        scope: &Scope,
429        source: &SliceExpand<C>,
430        destination: &mut SliceExpand<C>,
431        copy_length: u32,
432    ) {
433        let source = unsafe { *source.__expand_as_ptr_method(scope) }.value(scope);
434        let destination = unsafe { *destination.__expand_as_ptr_method(scope) }.value(scope);
435        let scalar_size = C::Scalar::__expand_size(scope);
436        let copy_length_bytes = copy_length as usize * scalar_size;
437        let source_length = ExpandValue::from(copy_length_bytes).read_value(scope);
438
439        let mem_copy = CopyAsyncOp::new(
440            scope.ctx_mut(),
441            source,
442            destination,
443            source_length,
444            copy_length_bytes,
445            false,
446        );
447
448        scope.register(&mem_copy);
449    }
450}
451
452/// Copy the source slice in global memory to destination in shared memory with a low level async
453/// copy. This only copies up to 128 bits/16 bytes, and does not synchronize. Use
454/// `barrier.copy_async_arrive` to make the reads visible.
455/// `copy_size` is in terms of elements to simplify copying between different vector sizes.
456///
457/// Will only copy the length of the source slice, and zero fill the rest. Source length must be
458/// <= copy size.
459///
460/// # Safety
461/// Starting address must be aligned to the full copy size.
462/// **This will silently fail if the address is only aligned to the source length and not the copy size!**
463pub fn copy_async_checked<C: CubePrimitive>(
464    _source: &[C],
465    _destination: &mut [C],
466    _copy_size: u32,
467) {
468    unexpanded!();
469}
470
471pub mod copy_async_checked {
472    use super::*;
473
474    pub fn expand<C: CubePrimitive>(
475        scope: &Scope,
476        source: &SliceExpand<C>,
477        destination: &mut SliceExpand<C>,
478        copy_length: u32,
479    ) {
480        let source_length = source.__extract_length(scope);
481
482        let source = unsafe { *source.__expand_as_ptr_unchecked_method(scope) }.value(scope);
483        let destination =
484            unsafe { *destination.__expand_as_ptr_unchecked_method(scope) }.value(scope);
485        let scalar_size = C::Scalar::__expand_size(scope);
486        let vector_size = C::__expand_size(scope).__expand_runtime_method(scope);
487        let source_length_bytes = source_length.__expand_mul_method(scope, vector_size);
488
489        let mem_copy = CopyAsyncOp::new(
490            scope.ctx_mut(),
491            source,
492            destination,
493            source_length_bytes.read_value(scope),
494            copy_length as usize * scalar_size,
495            true,
496        );
497
498        scope.register(&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.value(scope);
513            scope.register(&CommitCopyAsyncOp::new(scope.ctx_mut(), barrier));
514        })
515    }
516}
517
518impl From<SharedExpand<Barrier>> for BarrierExpand {
519    fn from(value: SharedExpand<Barrier>) -> Self {
520        value.expand.into()
521    }
522}
cubecl_core/frontend/barrier.rs

cubecl_core/frontend/
barrier.rs
