refactor: Scalars (#1064)

Author: wingertge

crates/cubecl-common/src/format.rs

@@ -104,7 +104,7 @@ mod tests {
     #[test]
     fn test_format_debug() {
         let test = Test {
-            map: HashMap::from_iter([("Hey with space".to_string(), 8)].into_iter()),
+            map: HashMap::from_iter([("Hey with space".to_string(), 8)]),
         };
 
         let formatted = format_debug(&test);

crates/cubecl-core/src/compute/launcher.rs

@@ -1,11 +1,11 @@
-use std::marker::PhantomData;
+use std::{collections::BTreeMap, marker::PhantomData};
 
-use crate::MetadataBuilder;
+use crate::KernelSettings;
 use crate::Runtime;
 use crate::compute::KernelTask;
 use crate::prelude::{ArrayArg, TensorArg, TensorMapArg};
-use crate::{KernelSettings, prelude::CubePrimitive};
-use bytemuck::{AnyBitPattern, NoUninit};
+use crate::{CubeScalar, MetadataBuilder};
+use cubecl_ir::StorageType;
 use cubecl_runtime::server::{Binding, CubeCount, ScalarBinding, TensorMapBinding};
 use cubecl_runtime::{client::ComputeClient, server::Bindings};
 
@@ -14,18 +14,7 @@ use super::CubeKernel;
 /// Prepare a kernel for [launch](KernelLauncher::launch).
 pub struct KernelLauncher<R: Runtime> {
     tensors: TensorState<R>,
-    scalar_bf16: ScalarState<half::bf16>,
-    scalar_f16: ScalarState<half::f16>,
-    scalar_f32: ScalarState<f32>,
-    scalar_f64: ScalarState<f64>,
-    scalar_u64: ScalarState<u64>,
-    scalar_u32: ScalarState<u32>,
-    scalar_u16: ScalarState<u16>,
-    scalar_u8: ScalarState<u8>,
-    scalar_i64: ScalarState<i64>,
-    scalar_i32: ScalarState<i32>,
-    scalar_i16: ScalarState<i16>,
-    scalar_i8: ScalarState<i8>,
+    scalars: ScalarState,
     pub settings: KernelSettings,
     runtime: PhantomData<R>,
 }
@@ -46,64 +35,14 @@ impl<R: Runtime> KernelLauncher<R> {
         self.tensors.push_array(array);
     }
 
-    /// Register a u8 scalar to be launched.
-    pub fn register_u8(&mut self, scalar: u8) {
-        self.scalar_u8.push(scalar);
+    /// Register a scalar to be launched.
+    pub fn register_scalar<C: CubeScalar>(&mut self, scalar: C) {
+        self.scalars.push(scalar);
     }
 
-    /// Register a u16 scalar to be launched.
-    pub fn register_u16(&mut self, scalar: u16) {
-        self.scalar_u16.push(scalar);
-    }
-
-    /// Register a u32 scalar to be launched.
-    pub fn register_u32(&mut self, scalar: u32) {
-        self.scalar_u32.push(scalar);
-    }
-
-    /// Register a u64 scalar to be launched.
-    pub fn register_u64(&mut self, scalar: u64) {
-        self.scalar_u64.push(scalar);
-    }
-
-    /// Register a i8 scalar to be launched.
-    pub fn register_i8(&mut self, scalar: i8) {
-        self.scalar_i8.push(scalar);
-    }
-
-    /// Register a i16 scalar to be launched.
-    pub fn register_i16(&mut self, scalar: i16) {
-        self.scalar_i16.push(scalar);
-    }
-
-    /// Register a i32 scalar to be launched.
-    pub fn register_i32(&mut self, scalar: i32) {
-        self.scalar_i32.push(scalar);
-    }
-
-    /// Register a i64 scalar to be launched.
-    pub fn register_i64(&mut self, scalar: i64) {
-        self.scalar_i64.push(scalar);
-    }
-
-    /// Register a bf16 scalar to be launched.
-    pub fn register_bf16(&mut self, scalar: half::bf16) {
-        self.scalar_bf16.push(scalar);
-    }
-
-    /// Register a f16 scalar to be launched.
-    pub fn register_f16(&mut self, scalar: half::f16) {
-        self.scalar_f16.push(scalar);
-    }
-
-    /// Register a f32 scalar to be launched.
-    pub fn register_f32(&mut self, scalar: f32) {
-        self.scalar_f32.push(scalar);
-    }
-
-    /// Register a f64 scalar to be launched.
-    pub fn register_f64(&mut self, scalar: f64) {
-        self.scalar_f64.push(scalar);
+    /// Register a scalar to be launched from raw data.
+    pub fn register_scalar_raw(&mut self, bytes: &[u8], dtype: StorageType) {
+        self.scalars.push_raw(bytes, dtype);
     }
 
     /// Launch the kernel.
@@ -156,19 +95,7 @@ impl<R: Runtime> KernelLauncher<R> {
         let mut bindings = Bindings::new();
 
         self.tensors.register(&mut bindings);
-
-        self.scalar_u8.register(&mut bindings);
-        self.scalar_u16.register(&mut bindings);
-        self.scalar_u32.register(&mut bindings);
-        self.scalar_u64.register(&mut bindings);
-        self.scalar_i8.register(&mut bindings);
-        self.scalar_i16.register(&mut bindings);
-        self.scalar_i32.register(&mut bindings);
-        self.scalar_i64.register(&mut bindings);
-        self.scalar_f16.register(&mut bindings);
-        self.scalar_bf16.register(&mut bindings);
-        self.scalar_f32.register(&mut bindings);
-        self.scalar_f64.register(&mut bindings);
+        self.scalars.register(&mut bindings);
 
         bindings
     }
@@ -190,13 +117,14 @@ pub enum TensorState<R: Runtime> {
 /// Handles the scalar state of an element type
 ///
 /// The scalars are grouped to reduce the number of buffers needed to send data to the compute device.
-pub enum ScalarState<T> {
-    /// No scalar of that type is registered yet.
-    Empty,
-    /// The registered scalars.
-    Some(Vec<T>),
+#[derive(Default, Clone)]
+pub struct ScalarState {
+    data: BTreeMap<StorageType, ScalarValues>,
 }
 
+/// Stores the data and type for a scalar arg
+pub type ScalarValues = Vec<u8>;
+
 impl<R: Runtime> TensorState<R> {
     fn maybe_init(&mut self) {
         if matches!(self, TensorState::Empty) {
@@ -316,26 +244,36 @@ impl<R: Runtime> TensorState<R> {
     }
 }
 
-impl<T: NoUninit + AnyBitPattern + CubePrimitive> ScalarState<T> {
+impl ScalarState {
     /// Add a new scalar value to the state.
-    pub fn push(&mut self, val: T) {
-        match self {
-            ScalarState::Empty => *self = Self::Some(vec![val]),
-            ScalarState::Some(values) => values.push(val),
-        }
+    pub fn push<T: CubeScalar>(&mut self, val: T) {
+        let val = [val];
+        let bytes = T::as_bytes(&val);
+        self.data
+            .entry(T::cube_type())
+            .or_default()
+            .extend(bytes.iter().copied());
+    }
+
+    /// Add a new raw value to the state.
+    pub fn push_raw(&mut self, bytes: &[u8], dtype: StorageType) {
+        self.data
+            .entry(dtype)
+            .or_default()
+            .extend(bytes.iter().copied());
     }
 
     fn register(&self, bindings: &mut Bindings) {
-        if let ScalarState::Some(values) = self {
-            let len = values.len();
-            let len_u64 = len.div_ceil(size_of::<u64>() / size_of::<T>());
+        for (ty, values) in self.data.iter() {
+            let len = values.len() / ty.size();
+            let len_u64 = len.div_ceil(size_of::<u64>() / ty.size());
+
             let mut data = vec![0; len_u64];
-            let slice = bytemuck::cast_slice_mut::<u64, T>(&mut data);
+            let slice = bytemuck::cast_slice_mut::<u64, u8>(&mut data);
             slice[0..values.len()].copy_from_slice(values);
-            let elem = T::as_type_native_unchecked();
             bindings
                 .scalars
-                .insert(elem, ScalarBinding::new(elem, len, data));
+                .insert(*ty, ScalarBinding::new(*ty, len, data));
         }
     }
 }
@@ -344,18 +282,7 @@ impl<R: Runtime> Default for KernelLauncher<R> {
     fn default() -> Self {
         Self {
             tensors: TensorState::Empty,
-            scalar_bf16: ScalarState::Empty,
-            scalar_f16: ScalarState::Empty,
-            scalar_f32: ScalarState::Empty,
-            scalar_f64: ScalarState::Empty,
-            scalar_u64: ScalarState::Empty,
-            scalar_u32: ScalarState::Empty,
-            scalar_u16: ScalarState::Empty,
-            scalar_u8: ScalarState::Empty,
-            scalar_i64: ScalarState::Empty,
-            scalar_i32: ScalarState::Empty,
-            scalar_i16: ScalarState::Empty,
-            scalar_i8: ScalarState::Empty,
+            scalars: Default::default(),
             settings: Default::default(),
             runtime: PhantomData,
         }

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

@@ -79,7 +79,7 @@ mod fill {
                 let length = self.expand.ty.line_size();
                 let output = scope.create_local(Type::new(P::as_type(scope)).line(length));
 
-                cast::expand::<P>(scope, value, output.clone().into());
+                cast::expand::<P, Line<P>>(scope, value, output.clone().into());
 
                 output.into()
             })

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

@@ -14,15 +14,15 @@ pub trait Cast: CubePrimitive {
         scope: &mut Scope,
         value: ExpandElementTyped<From>,
     ) -> <Self as CubeType>::ExpandType {
-        if core::any::TypeId::of::<Self>() == core::any::TypeId::of::<From>() {
+        if Self::as_type(scope) == From::as_type(scope) {
             return value.expand.into();
         }
         let line_size_in = value.expand.ty.line_size();
         let line_size_out = line_size_in * value.expand.ty.storage_type().packing_factor()
             / Self::as_type(scope).packing_factor();
         let new_var = scope
             .create_local(Type::new(<Self as CubePrimitive>::as_type(scope)).line(line_size_out));
-        cast::expand(scope, value, new_var.clone().into());
+        cast::expand::<From, Self>(scope, value, new_var.clone().into());
         new_var.into()
     }
 }

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

@@ -164,27 +164,3 @@ impl_float!(half f16, F16);
 impl_float!(half bf16, BF16);
 impl_float!(f32, F32);
 impl_float!(f64, F64);
-
-impl ScalarArgSettings for f16 {
-    fn register<R: Runtime>(&self, settings: &mut KernelLauncher<R>) {
-        settings.register_f16(*self);
-    }
-}
-
-impl ScalarArgSettings for bf16 {
-    fn register<R: Runtime>(&self, settings: &mut KernelLauncher<R>) {
-        settings.register_bf16(*self);
-    }
-}
-
-impl ScalarArgSettings for f32 {
-    fn register<R: Runtime>(&self, settings: &mut KernelLauncher<R>) {
-        settings.register_f32(*self);
-    }
-}
-
-impl ScalarArgSettings for f64 {
-    fn register<R: Runtime>(&self, settings: &mut KernelLauncher<R>) {
-        settings.register_f64(*self);
-    }
-}

crates/cubecl-core/src/frontend/element/float/fp4.rs

@@ -1,13 +1,9 @@
 use cubecl_common::{e2m1, e2m1x2};
 use cubecl_ir::{ConstantScalarValue, ElemType, ExpandElement, FloatKind, Scope, StorageType};
 
-use crate::{
-    Runtime,
-    compute::KernelLauncher,
-    prelude::{
-        CubePrimitive, CubeType, ExpandElementIntoMut, ExpandElementTyped, IntoRuntime,
-        ScalarArgSettings, into_mut_expand_element, into_runtime_expand_element,
-    },
+use crate::prelude::{
+    CubePrimitive, CubeType, ExpandElementIntoMut, ExpandElementTyped, IntoRuntime,
+    into_mut_expand_element, into_runtime_expand_element,
 };
 
 impl CubeType for e2m1 {
@@ -73,9 +69,3 @@ impl ExpandElementIntoMut for e2m1x2 {
         into_mut_expand_element(scope, elem)
     }
 }
-
-impl ScalarArgSettings for e2m1x2 {
-    fn register<R: Runtime>(&self, _settings: &mut KernelLauncher<R>) {
-        todo!("Not yet supported for scalars")
-    }
-}

crates/cubecl-core/src/frontend/element/float/fp8.rs

@@ -1,13 +1,9 @@
 use cubecl_common::{e4m3, e5m2, ue8m0};
 use cubecl_ir::{ConstantScalarValue, ElemType, ExpandElement, FloatKind, Scope, StorageType};
 
-use crate::{
-    Runtime,
-    compute::KernelLauncher,
-    prelude::{
-        CubePrimitive, CubeType, ExpandElementIntoMut, ExpandElementTyped, IntoRuntime, Numeric,
-        ScalarArgSettings, into_mut_expand_element, into_runtime_expand_element,
-    },
+use crate::prelude::{
+    CubePrimitive, CubeType, ExpandElementIntoMut, ExpandElementTyped, IntoRuntime, Numeric,
+    into_mut_expand_element, into_runtime_expand_element,
 };
 
 impl CubeType for e4m3 {
@@ -50,12 +46,6 @@ impl ExpandElementIntoMut for e4m3 {
     }
 }
 
-impl ScalarArgSettings for e4m3 {
-    fn register<R: Runtime>(&self, _settings: &mut KernelLauncher<R>) {
-        todo!("Not yet supported for scalars")
-    }
-}
-
 impl CubeType for e5m2 {
     type ExpandType = ExpandElementTyped<e5m2>;
 }
@@ -96,12 +86,6 @@ impl ExpandElementIntoMut for e5m2 {
     }
 }
 
-impl ScalarArgSettings for e5m2 {
-    fn register<R: Runtime>(&self, _settings: &mut KernelLauncher<R>) {
-        todo!("Not yet supported for scalars")
-    }
-}
-
 impl CubeType for ue8m0 {
     type ExpandType = ExpandElementTyped<ue8m0>;
 }
@@ -141,9 +125,3 @@ impl ExpandElementIntoMut for ue8m0 {
         into_mut_expand_element(scope, elem)
     }
 }
-
-impl ScalarArgSettings for ue8m0 {
-    fn register<R: Runtime>(&self, _settings: &mut KernelLauncher<R>) {
-        todo!("Not yet supported for scalars")
-    }
-}

crates/cubecl-core/src/frontend/element/float/relaxed.rs

@@ -5,7 +5,7 @@ use crate::prelude::{Numeric, into_runtime_expand_element};
 
 use super::{
     CubePrimitive, CubeType, ExpandElementIntoMut, ExpandElementTyped, Float, IntoRuntime,
-    KernelLauncher, Runtime, ScalarArgSettings, into_mut_expand_element,
+    into_mut_expand_element,
 };
 
 impl CubeType for flex32 {
@@ -79,9 +79,3 @@ impl Float for flex32 {
         flex32::from_f32(val)
     }
 }
-
-impl ScalarArgSettings for flex32 {
-    fn register<R: Runtime>(&self, settings: &mut KernelLauncher<R>) {
-        settings.register_f32(self.to_f32());
-    }
-}

crates/cubecl-core/src/frontend/element/float/tensor_float.rs

@@ -6,7 +6,7 @@ use crate::prelude::{Numeric, into_runtime_expand_element};
 
 use super::{
     CubePrimitive, CubeType, ExpandElementIntoMut, ExpandElementTyped, Float, IntoRuntime,
-    KernelLauncher, Runtime, ScalarArgSettings, into_mut_expand_element,
+    into_mut_expand_element,
 };
 
 impl CubeType for tf32 {
@@ -81,9 +81,3 @@ impl Float for tf32 {
         tf32::from_f32(val)
     }
 }
-
-impl ScalarArgSettings for tf32 {
-    fn register<R: Runtime>(&self, settings: &mut KernelLauncher<R>) {
-        settings.register_f32((*self).to_f32());
-    }
-}