cubecl_core/frontend/
scalar.rs

1use alloc::vec::Vec;
2use cubecl::prelude::*;
3use cubecl_common::{e4m3, e5m2, ue8m0};
4use cubecl_ir::{Instruction, Operator, Value};
5use serde::{Deserialize, Serialize};
6
7use crate::{
8    self as cubecl, ScalarArgType, intrinsic,
9    ir::{ElemType, FloatKind, IntKind, UIntKind},
10};
11
12#[derive(Clone, Copy, Debug)]
13/// A way to define an input scalar without a generic attached to it.
14///
15/// It uses comptime enum with zero-cost runtime abstraction for kernel generation.
16pub struct InputScalar {
17    data: [u8; 8],
18    dtype: StorageType,
19}
20
21#[derive(Clone)]
22pub struct InputScalarExpand {
23    pub expand: Value,
24}
25
26impl CubeType for InputScalar {
27    type ExpandType = InputScalarExpand;
28}
29
30impl ExpandTypeClone for InputScalarExpand {
31    fn clone_unchecked(&self) -> Self {
32        self.clone()
33    }
34}
35
36impl IntoExpand for InputScalarExpand {
37    type Expand = Self;
38
39    fn into_expand(self, _scope: &Scope) -> Self::Expand {
40        self
41    }
42}
43
44impl IntoMut for InputScalarExpand {
45    fn into_mut(self, _scope: &Scope) -> Self {
46        self
47    }
48}
49
50impl CubeDebug for InputScalarExpand {}
51
52impl AsRefExpand for InputScalarExpand {
53    fn __expand_ref_method(&self, _: &Scope) -> &Self {
54        self
55    }
56}
57impl AsMutExpand for InputScalarExpand {
58    fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
59        self
60    }
61}
62
63impl InputScalar {
64    /// Creates an [`InputScalar`] from the given element and dtype.
65    ///
66    /// # Panics
67    ///
68    /// If the given numeric element can't be transformed into the passed [`ElemType`].
69    pub fn new<E: num_traits::ToPrimitive>(val: E, dtype: impl Into<StorageType>) -> Self {
70        let dtype: StorageType = dtype.into();
71        let mut out = InputScalar {
72            data: Default::default(),
73            dtype,
74        };
75        fn write<E: ScalarArgType>(val: impl num_traits::ToPrimitive, out: &mut [u8]) {
76            let val = [E::from(val).unwrap()];
77            let bytes = E::as_bytes(&val);
78            out[..bytes.len()].copy_from_slice(bytes);
79        }
80        match dtype {
81            StorageType::Scalar(elem) => match elem {
82                ElemType::Float(float_kind) => match float_kind {
83                    FloatKind::F16 => write::<half::f16>(val, &mut out.data),
84                    FloatKind::BF16 => write::<half::bf16>(val, &mut out.data),
85                    FloatKind::Flex32 | FloatKind::F32 | FloatKind::TF32 => {
86                        write::<f32>(val, &mut out.data)
87                    }
88                    FloatKind::F64 => write::<f64>(val, &mut out.data),
89                    FloatKind::E2M1 | FloatKind::E2M3 | FloatKind::E3M2 => {
90                        unimplemented!("fp6 CPU conversion not yet implemented")
91                    }
92                    FloatKind::E4M3 => write::<e4m3>(val, &mut out.data),
93                    FloatKind::E5M2 => write::<e5m2>(val, &mut out.data),
94                    FloatKind::UE8M0 => write::<ue8m0>(val, &mut out.data),
95                },
96                ElemType::Int(int_kind) => match int_kind {
97                    IntKind::I8 => write::<i8>(val, &mut out.data),
98                    IntKind::I16 => write::<i16>(val, &mut out.data),
99                    IntKind::I32 => write::<i32>(val, &mut out.data),
100                    IntKind::I64 => write::<i64>(val, &mut out.data),
101                },
102                ElemType::UInt(uint_kind) => match uint_kind {
103                    UIntKind::U8 => write::<u8>(val, &mut out.data),
104                    UIntKind::U16 => write::<u16>(val, &mut out.data),
105                    UIntKind::U32 => write::<u32>(val, &mut out.data),
106                    UIntKind::U64 => write::<u64>(val, &mut out.data),
107                },
108                ElemType::Bool => panic!("Bool isn't a scalar"),
109            },
110            other => unimplemented!("{other} not supported for scalars"),
111        };
112        out
113    }
114}
115
116#[cube]
117impl InputScalar {
118    /// Reads the scalar with the given element type.
119    ///
120    /// Performs casting if necessary.
121    pub fn get<C: Scalar>(&self) -> C {
122        intrinsic!(|scope| {
123            let dtype = C::__expand_as_type(scope);
124            cast_expand_elem(scope, self.expand, dtype).into()
125        })
126    }
127}
128
129impl InputScalar {
130    pub fn as_bytes(&self) -> Vec<u8> {
131        self.data[..self.dtype.size()].to_vec()
132    }
133}
134
135impl LaunchArg for InputScalar {
136    type RuntimeArg<R: Runtime> = InputScalar;
137    type CompilationArg = InputScalarCompilationArg;
138
139    fn register<R: Runtime>(
140        arg: Self::RuntimeArg<R>,
141        launcher: &mut KernelLauncher<R>,
142    ) -> Self::CompilationArg {
143        let dtype = arg.dtype;
144        launcher.register_scalar_raw(&arg.data[..dtype.size()], dtype);
145        InputScalarCompilationArg::new(arg.dtype)
146    }
147
148    fn expand(
149        arg: &Self::CompilationArg,
150        builder: &mut KernelBuilder,
151    ) -> <Self as CubeType>::ExpandType {
152        let expand = builder.create_value(arg.ty.into());
153        let id = builder.scalar(arg.ty);
154        builder.register(Instruction::new(Operator::ReadScalar(id), expand));
155        InputScalarExpand { expand }
156    }
157}
158
159#[derive(Serialize, Deserialize, Clone, PartialEq, Eq, Hash, Debug)]
160pub struct InputScalarCompilationArg {
161    ty: StorageType,
162}
163
164impl InputScalarCompilationArg {
165    pub fn new(ty: StorageType) -> Self {
166        Self { ty }
167    }
168}
cubecl_core/frontend/scalar.rs

cubecl_core/frontend/
scalar.rs
