refactor: Launch (#944)

Author: wingertge

Cargo.toml

@@ -35,6 +35,7 @@ serde_json = { version = "1.0.119", default-features = false }
 toml = "0.9.1"
 variadics_please = "1"
 
+# no_std compatiblity
 dashmap = "6.1.0"
 foldhash = { version = "0.1.2", default-features = false }
 hashbrown = "0.15.5"
@@ -86,11 +87,10 @@ tracy-client = { version = "0.18.0" }
 strum = { version = "0.27.1", features = ["derive"] }
 tracel-xtask = { version = "=2.1.8" }
 
-
-portable-atomic-util = { version = "0.2.4", features = ["alloc"] }
 portable-atomic = { version = "1.11", default-features = false, features = [
     "serde",
 ] }
+portable-atomic-util = { version = "0.2.4", features = ["alloc"] }
 pretty_assertions = "1.4"
 
 # Async

crates/cubecl-convolution/src/kernels/layered/selector/base.rs

@@ -1,5 +1,5 @@
 use cubecl_core as cubecl;
-use cubecl_core::{CubeLaunch, CubeType, prelude::*};
+use cubecl_core::{CubeType, prelude::*};
 use cubecl_std::FastDivmod;
 
 #[derive(CubeType, CubeLaunch, Clone)]

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

@@ -7,8 +7,7 @@ use crate::{
     compute::{KernelBuilder, KernelLauncher},
     ir::{Id, LineSize, Type},
     prelude::{
-        ArgSettings, CompilationArg, CubePrimitive, ExpandElementTyped, LaunchArg, LaunchArgExpand,
-        TensorHandleRef,
+        ArgSettings, CompilationArg, CubePrimitive, ExpandElementTyped, LaunchArg, TensorHandleRef,
     },
 };
 
@@ -30,30 +29,6 @@ pub struct ArrayHandleRef<'a, R: Runtime> {
     runtime: PhantomData<R>,
 }
 
-impl<C: CubePrimitive> LaunchArgExpand for Array<C> {
-    type CompilationArg = ArrayCompilationArg;
-
-    fn expand(
-        arg: &Self::CompilationArg,
-        builder: &mut KernelBuilder,
-    ) -> ExpandElementTyped<Array<C>> {
-        builder
-            .input_array(Type::new(C::as_type(&builder.scope)).line(arg.line_size))
-            .into()
-    }
-    fn expand_output(
-        arg: &Self::CompilationArg,
-        builder: &mut KernelBuilder,
-    ) -> ExpandElementTyped<Array<C>> {
-        match arg.inplace {
-            Some(id) => builder.inplace_output(id).into(),
-            None => builder
-                .output_array(Type::new(C::as_type(&builder.scope)).line(arg.line_size))
-                .into(),
-        }
-    }
-}
-
 pub enum ArrayArg<'a, R: Runtime> {
     /// The array is passed with an array handle.
     Handle {
@@ -153,6 +128,7 @@ impl<'a, R: Runtime> ArrayHandleRef<'a, R> {
 
 impl<C: CubePrimitive> LaunchArg for Array<C> {
     type RuntimeArg<'a, R: Runtime> = ArrayArg<'a, R>;
+    type CompilationArg = ArrayCompilationArg;
 
     fn compilation_arg<R: Runtime>(runtime_arg: &Self::RuntimeArg<'_, R>) -> Self::CompilationArg {
         match runtime_arg {
@@ -166,4 +142,24 @@ impl<C: CubePrimitive> LaunchArg for Array<C> {
             },
         }
     }
+
+    fn expand(
+        arg: &Self::CompilationArg,
+        builder: &mut KernelBuilder,
+    ) -> ExpandElementTyped<Array<C>> {
+        builder
+            .input_array(Type::new(C::as_type(&builder.scope)).line(arg.line_size))
+            .into()
+    }
+    fn expand_output(
+        arg: &Self::CompilationArg,
+        builder: &mut KernelBuilder,
+    ) -> ExpandElementTyped<Array<C>> {
+        match arg.inplace {
+            Some(id) => builder.inplace_output(id).into(),
+            None => builder
+                .output_array(Type::new(C::as_type(&builder.scope)).line(arg.line_size))
+                .into(),
+        }
+    }
 }

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

@@ -1,11 +1,9 @@
 use std::{cell::RefCell, rc::Rc};
 
-use serde::{Deserialize, Serialize};
-
 use crate::{
     Runtime,
     compute::KernelBuilder,
-    prelude::{ArgSettings, CompilationArg, LaunchArg, LaunchArgExpand},
+    prelude::{ArgSettings, CompilationArg, LaunchArg},
 };
 
 use super::{Sequence, SequenceExpand};
@@ -29,7 +27,6 @@ impl<'a, R: Runtime, T: LaunchArg> SequenceArg<'a, R, T> {
     }
 }
 
-#[derive(Serialize, Deserialize)]
 pub struct SequenceCompilationArg<C: LaunchArg> {
     pub values: Vec<C::CompilationArg>,
 }
@@ -65,6 +62,7 @@ impl<C: LaunchArg> core::cmp::Eq for SequenceCompilationArg<C> {}
 
 impl<C: LaunchArg> LaunchArg for Sequence<C> {
     type RuntimeArg<'a, R: Runtime> = SequenceArg<'a, R, C>;
+    type CompilationArg = SequenceCompilationArg<C>;
 
     fn compilation_arg<R: Runtime>(runtime_arg: &Self::RuntimeArg<'_, R>) -> Self::CompilationArg {
         SequenceCompilationArg {
@@ -75,16 +73,6 @@ impl<C: LaunchArg> LaunchArg for Sequence<C> {
                 .collect(),
         }
     }
-}
-
-impl<R: Runtime, T: LaunchArg> ArgSettings<R> for SequenceArg<'_, R, T> {
-    fn register(&self, launcher: &mut crate::prelude::KernelLauncher<R>) {
-        self.values.iter().for_each(|arg| arg.register(launcher));
-    }
-}
-
-impl<C: LaunchArg> LaunchArgExpand for Sequence<C> {
-    type CompilationArg = SequenceCompilationArg<C>;
 
     fn expand(arg: &Self::CompilationArg, builder: &mut KernelBuilder) -> SequenceExpand<C> {
         let values = arg
@@ -110,3 +98,9 @@ impl<C: LaunchArg> LaunchArgExpand for Sequence<C> {
         }
     }
 }
+
+impl<R: Runtime, T: LaunchArg> ArgSettings<R> for SequenceArg<'_, R, T> {
+    fn register(&self, launcher: &mut crate::prelude::KernelLauncher<R>) {
+        self.values.iter().for_each(|arg| arg.register(launcher));
+    }
+}

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

@@ -48,7 +48,7 @@ impl<E: CubePrimitive, IO: SliceVisibility> Iterator for Slice<E, IO> {
     }
 }
 
-pub trait SliceVisibility: Clone + Copy {}
+pub trait SliceVisibility: Clone + Copy + Send + Sync + 'static {}
 
 impl SliceVisibility for ReadOnly {}
 

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

@@ -6,9 +6,7 @@ use crate::{
     Runtime,
     compute::{KernelBuilder, KernelLauncher},
     ir::{Id, LineSize, Type},
-    prelude::{
-        ArgSettings, CompilationArg, CubePrimitive, ExpandElementTyped, LaunchArg, LaunchArgExpand,
-    },
+    prelude::{ArgSettings, CompilationArg, CubePrimitive, ExpandElementTyped, LaunchArg},
 };
 
 use super::Tensor;
@@ -73,9 +71,23 @@ pub struct TensorCompilationArg {
 
 impl CompilationArg for TensorCompilationArg {}
 
-impl<C: CubePrimitive> LaunchArgExpand for Tensor<C> {
+impl<C: CubePrimitive> LaunchArg for Tensor<C> {
+    type RuntimeArg<'a, R: Runtime> = TensorArg<'a, R>;
     type CompilationArg = TensorCompilationArg;
 
+    fn compilation_arg<R: Runtime>(runtime_arg: &Self::RuntimeArg<'_, R>) -> Self::CompilationArg {
+        match runtime_arg {
+            TensorArg::Handle { line_size, .. } => TensorCompilationArg {
+                inplace: None,
+                line_size: *line_size as u32,
+            },
+            TensorArg::Alias { input_pos } => TensorCompilationArg {
+                inplace: Some(*input_pos as Id),
+                line_size: 0,
+            },
+        }
+    }
+
     fn expand(
         arg: &Self::CompilationArg,
         builder: &mut KernelBuilder,
@@ -97,23 +109,6 @@ impl<C: CubePrimitive> LaunchArgExpand for Tensor<C> {
     }
 }
 
-impl<C: CubePrimitive> LaunchArg for Tensor<C> {
-    type RuntimeArg<'a, R: Runtime> = TensorArg<'a, R>;
-
-    fn compilation_arg<R: Runtime>(runtime_arg: &Self::RuntimeArg<'_, R>) -> Self::CompilationArg {
-        match runtime_arg {
-            TensorArg::Handle { line_size, .. } => TensorCompilationArg {
-                inplace: None,
-                line_size: *line_size as u32,
-            },
-            TensorArg::Alias { input_pos } => TensorCompilationArg {
-                inplace: Some(*input_pos as Id),
-                line_size: 0,
-            },
-        }
-    }
-}
-
 impl<'a, R: Runtime> TensorArg<'a, R> {
     /// Create a new tensor argument specified with its vectorization factor.
     ///

crates/cubecl-core/src/frontend/container/tensor/tensormap.rs

@@ -119,9 +119,14 @@ pub struct TensorMapCompilationArg;
 
 impl CompilationArg for TensorMapCompilationArg {}
 
-impl<E: CubePrimitive> LaunchArgExpand for TensorMap<E> {
+impl<E: CubePrimitive> LaunchArg for TensorMap<E> {
+    type RuntimeArg<'a, R: Runtime> = TensorMapArg<'a, R>;
     type CompilationArg = TensorMapCompilationArg;
 
+    fn compilation_arg<R: Runtime>(_runtime_arg: &Self::RuntimeArg<'_, R>) -> Self::CompilationArg {
+        TensorMapCompilationArg
+    }
+
     fn expand(
         _arg: &Self::CompilationArg,
         builder: &mut KernelBuilder,
@@ -138,14 +143,6 @@ impl<E: CubePrimitive> LaunchArgExpand for TensorMap<E> {
     }
 }
 
-impl<E: CubePrimitive> LaunchArg for TensorMap<E> {
-    type RuntimeArg<'a, R: Runtime> = TensorMapArg<'a, R>;
-
-    fn compilation_arg<R: Runtime>(_runtime_arg: &Self::RuntimeArg<'_, R>) -> Self::CompilationArg {
-        TensorMapCompilationArg
-    }
-}
-
 /// Commit an async tensor operation. Not sure how this works, poor docs. But you need to call it
 /// after a write, but not after reads.
 pub fn tma_group_commit() {

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

@@ -1,13 +1,9 @@
 use cubecl_ir::{AtomicOp, ExpandElement, StorageType};
 
-use super::{
-    ExpandElementIntoMut, ExpandElementTyped, Int, LaunchArgExpand, Numeric,
-    into_mut_expand_element,
-};
+use super::{ExpandElementIntoMut, ExpandElementTyped, Int, Numeric, into_mut_expand_element};
 use crate::{
     frontend::{CubePrimitive, CubeType},
     ir::{BinaryOperator, CompareAndSwapOperator, Instruction, Scope, Type, UnaryOperator},
-    prelude::KernelBuilder,
     unexpanded,
 };
 
@@ -317,11 +313,3 @@ impl<Inner: CubePrimitive> ExpandElementIntoMut for Atomic<Inner> {
         into_mut_expand_element(scope, elem)
     }
 }
-
-impl<Inner: CubePrimitive> LaunchArgExpand for Atomic<Inner> {
-    type CompilationArg = ();
-
-    fn expand(_: &Self::CompilationArg, builder: &mut KernelBuilder) -> ExpandElementTyped<Self> {
-        builder.scalar(Self::as_type_native_unchecked()).into()
-    }
-}