cubecl_core/frontend/element/
numeric.rs

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

cubecl_core/frontend/element/
numeric.rs
