[Draft] OneSweep Radix Sort

crates/cubecl-core/src/frontend/barrier.rs

@@ -50,16 +50,43 @@ pub struct BarrierExpand<C: CubePrimitive> {
     _c: PhantomData<C>,
 }
 
+#[derive(Clone, PartialEq, Eq)]
 pub struct BarrierLevel(InnerBarrierLevel);
 
-#[derive(Clone)]
-/// Defines how many units must reach the barrier to allow continuation
+impl CubeType for BarrierLevel {
+    type ExpandType = Self;
+}
+
+impl Init for BarrierLevel {
+    fn init(self, _scope: &mut Scope) -> Self {
+        self
+    }
+}
+
+impl CubeDebug for BarrierLevel {
+    fn set_debug_name(&self, _scope: &mut Scope, _name: &'static str) {}
+}
+
+#[derive(Clone, Eq, PartialEq)]
+/// Defines how many units must reach the barrier before execution can continue.
+/// This also determines how `memcpy_async` operations should be handled.
 enum InnerBarrierLevel {
-    /// Only waits for the unit who declared this barrier.
-    /// This may be useful for waiting upon async data loading
+    /// Waits only for the unit that declared this barrier.
+    /// Useful for synchronizing after async data loading.
     Unit,
-    /// All units in the Cube must reach the barrier before continuing
-    Cube(u32),
+
+    /// All units in the Cube must reach the barrier before continuing.
+    /// The argument is the ID of the unit elected for initialization.
+    ///
+    /// `memcpy_async` is **cooperative**, so all units in the Cube must call `memcpy_async` with the same arguments.
+    /// The called is not elected by default, so it must be done manually if wanted
+    CubeCoop(u32),
+
+    /// All units in the Cube must reach the barrier before continuing.
+    /// The argument is the ID of the unit elected for initialization.
+    ///
+    /// `memcpy_async` is **not cooperative**, so each unit must manually handle its own data slice.
+    CubeManual(u32),
 }
 
 impl BarrierLevel {
@@ -68,28 +95,43 @@ impl BarrierLevel {
         BarrierLevel(InnerBarrierLevel::Unit)
     }
 
-    /// Creates a Cube barrier level
+    /// Creates a CubeCoop barrier level
     ///
-    /// The field elected_unit is the UNIT_POS of the unit that will
-    /// perform the underlying initialization. Typically, 0 should work
-    pub fn cube(elected_unit: u32) -> Self {
-        BarrierLevel(InnerBarrierLevel::Cube(elected_unit))
+    /// Will sync all units
+    pub fn cube_coop(elected_unit: u32) -> Self {
+        BarrierLevel(InnerBarrierLevel::CubeCoop(elected_unit))
+    }
+
+    /// Creates a CubeManual barrier level
+    ///
+    /// Will sync all units
+    pub fn cube_manual(elected_unit: u32) -> Self {
+        BarrierLevel(InnerBarrierLevel::CubeManual(elected_unit))
     }
 
     pub fn __expand_unit(_scope: &mut Scope) -> BarrierLevel {
         BarrierLevel(InnerBarrierLevel::Unit)
     }
 
-    pub fn __expand_cube(_scope: &mut Scope, elected_unit: u32) -> Self {
-        BarrierLevel(InnerBarrierLevel::Cube(elected_unit))
+    pub fn __expand_cube_coop(_scope: &mut Scope, elected_unit: u32) -> Self {
+        BarrierLevel(InnerBarrierLevel::CubeCoop(elected_unit))
+    }
+
+    pub fn __expand_cube_manual(_scope: &mut Scope, elected_unit: u32) -> Self {
+        BarrierLevel(InnerBarrierLevel::CubeManual(elected_unit))
     }
 }
 
 impl From<InnerBarrierLevel> for cubecl_ir::BarrierLevel {
     fn from(val: InnerBarrierLevel) -> Self {
         match val {
             InnerBarrierLevel::Unit => cubecl_ir::BarrierLevel::Unit,
-            InnerBarrierLevel::Cube(elected_unit) => cubecl_ir::BarrierLevel::Cube(elected_unit),
+            InnerBarrierLevel::CubeCoop(elected_unit) => {
+                cubecl_ir::BarrierLevel::CubeCoop(elected_unit)
+            }
+            InnerBarrierLevel::CubeManual(elected_unit) => {
+                cubecl_ir::BarrierLevel::CubeManual(elected_unit)
+            }
         }
     }
 }
@@ -106,7 +148,7 @@ impl<C: CubePrimitive> Barrier<C> {
     ///
     /// This will try to copy the whole source slice, so
     /// make sure source length <= destination length
-    pub fn memcpy_async(&self, _source: Slice<Line<C>>, _destination: SliceMut<Line<C>>) {
+    pub fn memcpy_async(&self, _source: &Slice<Line<C>>, _destination: &mut SliceMut<Line<C>>) {
         unexpanded!()
     }
 

crates/cubecl-core/src/frontend/branch.rs

@@ -363,7 +363,7 @@ impl<C: CubePrimitive> IfElseExprExpand<C> {
             } => {
                 let mut else_child = scope.child();
                 let ret = else_block(&mut else_child);
-                assign::expand(&mut else_child, ret, out.clone());
+                assign::expand::<C>(&mut else_child, ret, out.clone());
 
                 scope.register(Branch::IfElse(Box::new(IfElse {
                     cond: *runtime_cond,
@@ -394,7 +394,7 @@ pub fn if_else_expr_expand<C: CubePrimitive>(
             let mut then_child = scope.child();
             let ret = then_block(&mut then_child);
             let out: ExpandElementTyped<C> = scope.create_local_mut(ret.expand.item).into();
-            assign::expand(&mut then_child, ret, out.clone());
+            assign::expand::<C>(&mut then_child, ret, out.clone());
 
             IfElseExprExpand::Runtime {
                 runtime_cond,
@@ -471,7 +471,7 @@ impl<I: Int, C: CubePrimitive> SwitchExpandExpr<I, C> {
         let value = I::from(value).unwrap();
         let mut case_child = scope.child();
         let ret = block(&mut case_child);
-        assign::expand(&mut case_child, ret, self.out.clone());
+        assign::expand::<C>(&mut case_child, ret, self.out.clone());
         self.cases.push((value.into(), case_child));
         self
     }
@@ -499,7 +499,7 @@ pub fn switch_expand_expr<I: Int, C: CubePrimitive>(
     let mut default_child = scope.child();
     let default = default_block(&mut default_child);
     let out: ExpandElementTyped<C> = scope.create_local_mut(default.expand.item).into();
-    assign::expand(&mut default_child, default, out.clone());
+    assign::expand::<C>(&mut default_child, default, out.clone());
 
     SwitchExpandExpr {
         value,

crates/cubecl-core/src/frontend/container/array/base.rs

@@ -183,7 +183,7 @@ mod vectorization {
                     self.clone(),
                     ExpandElementTyped::from_lit(scope, 0u32),
                 );
-                assign::expand(scope, element, new_var.clone().into());
+                assign::expand::<C>(scope, element, new_var.clone().into());
                 new_var
             } else {
                 let new_var = scope.create_local_mut(item);

crates/cubecl-core/src/frontend/container/option/launch.rs

@@ -3,14 +3,8 @@ use crate::{
     prelude::{ArgSettings, CompilationArg, LaunchArg, LaunchArgExpand},
     Runtime,
 };
-use serde::{Deserialize, Serialize};
 use std::marker::PhantomData;
 
-#[derive(Clone, PartialEq, Eq, Hash, Debug, Serialize, Deserialize)]
-pub struct CubeOptionCompilationArg<C> {
-    pub inner: Option<C>,
-}
-
 impl<C: CompilationArg> CompilationArg for Option<C> {}
 
 pub struct CubeOptionHandleRef<'a, H, R: Runtime> {
@@ -43,7 +37,6 @@ impl<R: Runtime, A: ArgSettings<R>> ArgSettings<R> for Option<A> {
 
 impl<T: LaunchArg> LaunchArg for Option<T> {
     type RuntimeArg<'a, R: Runtime> = Option<T::RuntimeArg<'a, R>>;
-
     fn compilation_arg<R: Runtime>(runtime_arg: &Self::RuntimeArg<'_, R>) -> Self::CompilationArg {
         runtime_arg.as_ref().map(T::compilation_arg)
     }

crates/cubecl-core/src/frontend/element/numeric.rs

@@ -124,7 +124,7 @@ pub trait ScalarArgSettings: Send + Sync {
 
 #[derive(new)]
 pub struct ScalarArg<T: Numeric> {
-    elem: T,
+    pub elem: T,
 }
 
 impl<T: Numeric, R: Runtime> ArgSettings<R> for ScalarArg<T> {

crates/cubecl-core/src/frontend/pipeline.rs

@@ -135,7 +135,7 @@ impl<C: CubePrimitive> Pipeline<C> {
     ///
     /// This will try to copy the whole source slice, so
     /// make sure source length <= destination length
-    pub fn memcpy_async(&self, _source: Slice<Line<C>>, _destination: SliceMut<Line<C>>) {
+    pub fn memcpy_async(&self, _source: &Slice<Line<C>>, _destination: &mut SliceMut<Line<C>>) {
         unexpanded!()
     }
 

crates/cubecl-core/src/frontend/plane.rs

@@ -393,3 +393,33 @@ pub mod plane_ballot {
         output.into()
     }
 }
+
+#[allow(unused_variables)]
+pub fn plane_shuffle<E: CubePrimitive>(value: E, index: u32) -> E {
+    unexpanded!()
+}
+
+/// Module containing the expand function for [plane_shuffle()].
+pub mod plane_shuffle {
+
+    use super::*;
+
+    /// Expand method of [plane_shuffle()].
+    pub fn expand<E: CubePrimitive>(
+        scope: &mut Scope,
+        value: ExpandElementTyped<E>,
+        id: ExpandElementTyped<u32>,
+    ) -> ExpandElementTyped<E> {
+        let output = scope.create_local(value.expand.item);
+        let out = *output;
+        let lhs = *value.expand;
+        let rhs = *id.expand;
+
+        scope.register(Instruction::new(
+            Plane::Shuffle(crate::ir::BinaryOperator { lhs, rhs }),
+            out,
+        ));
+
+        output.into()
+    }
+}

crates/cubecl-core/src/runtime_tests/barrier.rs

@@ -8,9 +8,9 @@ pub fn async_copy_test<F: Float>(input: &Array<Line<F>>, output: &mut Array<Line
     let mut smem = SharedMemory::<F>::new_lined(1u32, 1u32);
 
     let source = input.slice(2, 3);
-    let destination = smem.slice_mut(0, 1);
+    let mut destination = smem.slice_mut(0, 1);
 
-    barrier.memcpy_async(source, destination);
+    barrier.memcpy_async(&source, &mut destination);
 
     barrier.wait();
     output[0] = smem[0];
@@ -47,11 +47,11 @@ pub fn test_async_copy<R: Runtime, F: Float + CubeElement>(
 fn one_load<F: Float>(lhs: &Tensor<Line<F>>, output: &mut Tensor<Line<F>>) {
     let mut lhs_smem = SharedMemory::<F>::new_lined(4u32, 1u32);
 
-    let barrier = Barrier::<F>::new(BarrierLevel::cube(0u32));
+    let barrier = Barrier::<F>::new(BarrierLevel::cube_manual(0u32));
     sync_units();
 
     // Can't use lhs.to_slice() because then generated input_length will not exist
-    barrier.memcpy_async(lhs.slice(0u32, 4u32), lhs_smem.to_slice_mut());
+    barrier.memcpy_async(&lhs.slice(0u32, 4u32), &mut lhs_smem.to_slice_mut());
 
     barrier.wait();
 
@@ -72,14 +72,14 @@ fn two_loads<F: Float>(
     let mut lhs_smem = SharedMemory::<F>::new_lined(num_data, 1u32);
     let mut rhs_smem = SharedMemory::<F>::new_lined(num_data, 1u32);
 
-    let barrier = Barrier::<F>::new(BarrierLevel::cube(0u32));
+    let barrier = Barrier::<F>::new(BarrierLevel::cube_manual(0u32));
     sync_units();
 
     let start = UNIT_POS_X * num_data / 2;
     let end = start + num_data / 2;
 
-    barrier.memcpy_async(lhs.slice(start, end), lhs_smem.slice_mut(start, end));
-    barrier.memcpy_async(rhs.slice(start, end), rhs_smem.slice_mut(start, end));
+    barrier.memcpy_async(&lhs.slice(start, end), &mut lhs_smem.slice_mut(start, end));
+    barrier.memcpy_async(&rhs.slice(start, end), &mut rhs_smem.slice_mut(start, end));
 
     barrier.wait();
     let mut dot = Line::cast_from(0u32);
@@ -100,8 +100,8 @@ fn two_independant_loads<F: Float>(
     let mut lhs_smem = SharedMemory::<F>::new_lined(num_data, 1u32);
     let mut rhs_smem = SharedMemory::<F>::new_lined(num_data, 1u32);
 
-    let barrier_0 = barrier::Barrier::new(BarrierLevel::cube(0u32));
-    let barrier_1 = barrier::Barrier::new(BarrierLevel::cube(0u32));
+    let barrier_0 = barrier::Barrier::new(BarrierLevel::cube_manual(0u32));
+    let barrier_1 = barrier::Barrier::new(BarrierLevel::cube_manual(0u32));
     // At the Cube level, we must sync after barrier creation to make sure they
     // exist for all units
     sync_units();
@@ -115,8 +115,8 @@ fn two_independant_loads<F: Float>(
         output[i] = Line::cast_from(0u32);
     }
 
-    barrier_0.memcpy_async(lhs.slice(start, end), lhs_smem.slice_mut(start, end));
-    barrier_1.memcpy_async(rhs.slice(start, end), rhs_smem.slice_mut(start, end));
+    barrier_0.memcpy_async(&lhs.slice(start, end), &mut lhs_smem.slice_mut(start, end));
+    barrier_1.memcpy_async(&rhs.slice(start, end), &mut rhs_smem.slice_mut(start, end));
 
     let mut dot = Line::cast_from(0u32);