cubecl_core/frontend/
scalar.rs

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

cubecl_core/frontend/
scalar.rs
