cubecl_core/frontend/element/
numeric.rs

1use cubecl_ir::{ConstantValue, ElemType, ExpandValue, dialect::general::ReadScalarOp};
2use cubecl_runtime::runtime::Runtime;
3use num_traits::{NumCast, One, Zero};
4
5use crate::unexpanded;
6use crate::{IntoRuntime, ScalarArgType, compute::KernelBuilder, frontend::*};
7use crate::{compute::KernelLauncher, frontend::AtomicNumeric};
8use crate::{frontend::CubeType, prelude::InputScalar};
9use crate::{ir::Scope, prelude::Scalar};
10
11use super::{LaunchArg, NativeAssign, NativeExpand};
12
13/// Type that encompasses both (unsigned or signed) integers and floats
14/// Used in kernels that should work for both.
15pub trait Numeric:
16    Copy
17    + ScalarAbs
18    + ScalarVectorSum
19    + ScalarAdd
20    + ScalarSub
21    + ScalarMul
22    + ScalarDiv
23    + ScalarRem
24    + ScalarModFloor
25    + Scalar
26    + AtomicNumeric
27    + PlaneNumeric
28    + NativeAssign
29    + Into<NativeExpand<Self>>
30    + Into<ConstantValue>
31    + ScalarPartialOrd
32    + num_traits::NumCast
33    + num_traits::NumAssign
34    + core::fmt::Debug
35    + bytemuck::Zeroable
36{
37    fn min_value() -> Self;
38    fn max_value() -> Self;
39
40    fn __expand_min_value(scope: &Scope) -> <Self as CubeType>::ExpandType {
41        let elem = Self::elem_type(scope);
42        let val = elem.min_variable();
43        val.into()
44    }
45
46    fn __expand_max_value(scope: &Scope) -> <Self as CubeType>::ExpandType {
47        let elem = Self::elem_type(scope);
48        let val = elem.max_variable(scope);
49        val.into()
50    }
51
52    /// Create a new constant numeric.
53    ///
54    /// Note: since this must work for both integer and float
55    /// only the less expressive of both can be created (int)
56    /// If a number with decimals is needed, use `Float::new`.
57    ///
58    /// This method panics when unexpanded. For creating an element
59    /// with a val, use the new method of the sub type.
60    fn from_int(val: i64) -> Self {
61        <Self as NumCast>::from(val).unwrap()
62    }
63
64    /// Create a new constant numeric. Uses `i128` to be able to represent both signed integers, and
65    /// `u64::MAX`.
66    ///
67    /// Note: since this must work for both integer and float
68    /// only the less expressive of both can be created (int)
69    /// If a number with decimals is needed, use `Float::new`.
70    ///
71    /// This method panics when unexpanded. For creating an element
72    /// with a val, use the new method of the sub type.
73    fn from_int_128(val: i128) -> Self {
74        <Self as NumCast>::from(val).unwrap()
75    }
76
77    fn from_vec<const D: usize>(_vec: [u32; D]) -> Self {
78        unexpanded!()
79    }
80
81    fn __expand_from_int(scope: &Scope, val: NativeExpand<i64>) -> <Self as CubeType>::ExpandType {
82        let elem = Self::elem_type(scope);
83        let val: ExpandValue = elem.constant(val.constant().unwrap());
84
85        val.into()
86    }
87}
88
89/// Similar to [`ArgSettings`], however only for scalar types that don't depend on the [Runtime]
90/// trait.
91pub trait ScalarArgSettings: Send + Sync + Scalar {
92    /// Register the information to the [`KernelLauncher`].
93    fn register<R: Runtime>(&self, launcher: &mut KernelLauncher<R>);
94    fn expand_scalar(builder: &mut KernelBuilder) -> NativeExpand<Self> {
95        let storage_ty = Self::elem_type(&builder.scope);
96        let id = builder.scalar(storage_ty);
97        let ty = storage_ty.to_type(builder.ctx_mut());
98        let op = ReadScalarOp::new(builder.ctx_mut(), ty, id);
99        builder.register_with_result(&op).into()
100    }
101}
102
103impl<E: ScalarArgType> ScalarArgSettings for E {
104    fn register<R: Runtime>(&self, launcher: &mut KernelLauncher<R>) {
105        launcher.register_scalar(*self);
106    }
107}
108
109impl ScalarArgSettings for usize {
110    fn register<R: Runtime>(&self, launcher: &mut KernelLauncher<R>) {
111        let value = InputScalar::new(*self, launcher.settings.address_type.unsigned_type());
112        launcher.register_scalar_raw(value.as_bytes(), ElemType::Index);
113    }
114}
115
116impl ScalarArgSettings for isize {
117    fn register<R: Runtime>(&self, launcher: &mut KernelLauncher<R>) {
118        let value = InputScalar::new(*self, launcher.settings.address_type.signed_type());
119        InputScalar::register(value, launcher);
120    }
121}
122
123macro_rules! impl_scalar_launch {
124    ($ty: ty) => {
125        impl LaunchArg for $ty {
126            type RuntimeArg<R: Runtime> = $ty;
127            type CompilationArg = ();
128
129            fn register<R: Runtime>(arg: Self::RuntimeArg<R>, launcher: &mut KernelLauncher<R>) {
130                arg.register(launcher);
131            }
132
133            fn expand(_: &(), builder: &mut KernelBuilder) -> NativeExpand<Self> {
134                <$ty>::expand_scalar(builder)
135            }
136        }
137    };
138}
139pub(crate) use impl_scalar_launch;
140
141pub trait ZeroExpand: CubeType + Zero {
142    fn __expand_zero(scope: &Scope) -> Self::ExpandType;
143}
144
145pub trait OneExpand: CubeType + One {
146    fn __expand_one(scope: &Scope) -> Self::ExpandType;
147}
148
149impl<T: CubeType + Zero + IntoRuntime> ZeroExpand for T {
150    fn __expand_zero(scope: &Scope) -> Self::ExpandType {
151        T::zero().__expand_runtime_method(scope)
152    }
153}
154
155impl<T: CubeType + One + IntoRuntime> OneExpand for T {
156    fn __expand_one(scope: &Scope) -> Self::ExpandType {
157        T::one().__expand_runtime_method(scope)
158    }
159}
cubecl_core/frontend/element/numeric.rs

cubecl_core/frontend/element/
numeric.rs
