cubecl_core/frontend/element/
base.rs

1use super::{CubePrimitive, Numeric};
2use crate::{
3    frontend::ConstRuntimeAssign,
4    ir::{ConstantValue, ExpandValue, Scope},
5    prelude::{DynamicSize, KernelBuilder, KernelLauncher, Scalar, assign},
6    unexpanded,
7};
8use alloc::{boxed::Box, vec::Vec};
9use core::{fmt::Debug, marker::PhantomData};
10use cubecl_common::{e2m1, e2m1x2, e2m3, e3m2, e4m3, e5m2, flex32, tf32, ue8m0};
11use cubecl_ir::{
12    VectorSize,
13    attributes::ZeroAttr,
14    ident,
15    interfaces::TypedExt,
16    pliron::{printable::Printable, r#type::Typed, value::Value},
17    types::PointerType,
18};
19use half::{bf16, f16};
20use pliron::r#type::TypeHandle;
21use variadics_please::{all_tuples, all_tuples_enumerated};
22
23/// Types used in a cube function must implement this trait
24///
25/// Values whose values will be known at runtime must
26/// have `Value` as associated type
27/// Values whose values will be known at compile time
28/// must have the primitive type as associated type
29///
30/// Note: Cube functions should be written using `CubeTypes`,
31/// so that the code generated uses the associated `ExpandType`.
32/// This allows Cube code to not necessitate cloning, which is cumbersome
33/// in algorithmic code. The necessary cloning will automatically appear in
34/// the generated code.
35#[diagnostic::on_unimplemented(note = "Consider using `#[derive(CubeType)]` on `{Self}`")]
36pub trait CubeType {
37    type ExpandType: IntoExpand<Expand = Self::ExpandType>
38        + ExpandTypeClone
39        + IntoMut
40        + CubeDebug
41        + AsRefExpand
42        + AsMutExpand;
43}
44
45pub trait NativeCubeType: CubeType<ExpandType = NativeExpand<Self>> {}
46
47impl<'a, T: CubeType + ?Sized> CubeType for &'a T {
48    type ExpandType = &'a T::ExpandType;
49}
50
51impl<'a, T: CubeType + ?Sized> CubeType for &'a mut T {
52    type ExpandType = &'a mut T::ExpandType;
53}
54
55impl<T: CubeType + ?Sized> CubeType for *const T {
56    type ExpandType = *const T::ExpandType;
57}
58
59impl<T: CubeType + ?Sized> CubeType for *mut T {
60    type ExpandType = *mut T::ExpandType;
61}
62
63impl<T: CubeType<ExpandType = NativeExpand<T>> + ?Sized> NativeCubeType for T {}
64
65pub trait IntoExpand {
66    type Expand;
67    fn into_expand(self, scope: &Scope) -> Self::Expand;
68}
69
70impl<'a, T: IntoExpand<Expand = T> + ?Sized> IntoExpand for &'a T {
71    type Expand = &'a T;
72
73    fn into_expand(self, _: &Scope) -> Self::Expand {
74        self
75    }
76}
77
78impl<'a, T: IntoExpand<Expand = T> + ?Sized> IntoExpand for &'a mut T {
79    type Expand = &'a mut T;
80
81    fn into_expand(self, _: &Scope) -> Self::Expand {
82        self
83    }
84}
85
86impl<T: IntoExpand<Expand = T> + ?Sized> IntoExpand for *const T {
87    type Expand = *const T;
88
89    fn into_expand(self, _: &Scope) -> Self::Expand {
90        self
91    }
92}
93
94impl<T: IntoExpand<Expand = T> + ?Sized> IntoExpand for *mut T {
95    type Expand = *mut T;
96
97    fn into_expand(self, _: &Scope) -> Self::Expand {
98        self
99    }
100}
101
102pub trait ExpandTypeClone {
103    /// Unchecked clone that only clones the conceptual runtime value. Should only be used in cases
104    /// where each copy is used in a mutually exclusive branch (i.e. match, runtime enums). This is
105    /// intentionally separated from Rust's `Clone` semantics and should only be used for the
106    /// conceptual expand values, never real data. Using two values in the same branch is undefined
107    /// behaviour.
108    fn clone_unchecked(&self) -> Self;
109}
110
111impl<T: ExpandTypeClone + ?Sized> ExpandTypeClone for &T {
112    fn clone_unchecked(&self) -> Self {
113        self
114    }
115}
116
117impl<T: ExpandTypeClone + ?Sized> ExpandTypeClone for &mut T {
118    #[allow(mutable_transmutes)]
119    fn clone_unchecked(&self) -> Self {
120        unsafe { core::mem::transmute(&**self) }
121    }
122}
123
124impl<T: ExpandTypeClone + ?Sized> ExpandTypeClone for *const T {
125    fn clone_unchecked(&self) -> Self {
126        *self
127    }
128}
129
130impl<T: ExpandTypeClone + ?Sized> ExpandTypeClone for *mut T {
131    fn clone_unchecked(&self) -> Self {
132        *self
133    }
134}
135
136/// Expand version of [`AsRef`](core::convert::AsRef). Like [`AsRef<Self>`](core::convert::AsRef)
137/// it's implemented for all [`ExpandType`](CubeType::ExpandType)s. This is called when the Rust
138/// code uses `&x`.
139pub trait AsRefExpand<T: ?Sized = Self> {
140    fn __expand_as_ref_method(&self, scope: &Scope) -> &T {
141        self.__expand_ref_method(scope)
142    }
143    fn __expand_ref_method(&self, scope: &Scope) -> &T;
144}
145
146impl<T: AsRefExpand + ?Sized> AsRefExpand for &T {
147    fn __expand_ref_method(&self, _: &Scope) -> &Self {
148        self
149    }
150}
151
152impl<T: AsRefExpand + ?Sized> AsRefExpand for &mut T {
153    fn __expand_ref_method(&self, _: &Scope) -> &Self {
154        self
155    }
156}
157
158impl<T: AsRefExpand + ?Sized> AsRefExpand for *const T {
159    fn __expand_ref_method(&self, _: &Scope) -> &Self {
160        self
161    }
162}
163
164impl<T: AsRefExpand + ?Sized> AsRefExpand for *mut T {
165    fn __expand_ref_method(&self, _: &Scope) -> &Self {
166        self
167    }
168}
169
170/// Expand version of [`AsMut`](core::convert::AsMut). The `Self` version must be implemented by
171/// all [`ExpandType`](CubeType::ExpandType)s, since `CubeCL` also uses it to implement `&mut x`.
172pub trait AsMutExpand<T: ?Sized = Self> {
173    fn __expand_as_mut_method(&mut self, scope: &Scope) -> &mut T {
174        self.__expand_ref_mut_method(scope)
175    }
176    fn __expand_ref_mut_method(&mut self, scope: &Scope) -> &mut T;
177}
178
179impl<T: AsMutExpand + ?Sized> AsMutExpand for &T {
180    fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
181        self
182    }
183}
184
185impl<T: AsMutExpand + ?Sized> AsMutExpand for &mut T {
186    fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
187        self
188    }
189}
190
191impl<T: AsMutExpand + ?Sized> AsMutExpand for *const T {
192    fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
193        self
194    }
195}
196
197impl<T: AsMutExpand + ?Sized> AsMutExpand for *mut T {
198    fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
199        self
200    }
201}
202
203/// `CubeCL` version of [`Deref`](core::ops::Deref). Unlike those traits, this trait produces owned
204/// values directly. Maps to `*x`.
205pub trait DerefExpand {
206    type Target;
207
208    fn __expand_deref_method(&self, scope: &Scope) -> Self::Target;
209}
210
211pub fn __expand_deref<T: DerefExpand<Target = T>>(scope: &Scope, value: &T) -> T {
212    value.__expand_deref_method(scope)
213}
214
215pub trait AsDerefExpand {
216    type Target;
217    fn __expand_as_deref_method(&self, scope: &Scope) -> &Self::Target;
218}
219
220pub trait AsDerefMutExpand: AsDerefExpand {
221    fn __expand_as_deref_mut_method(&mut self, scope: &Scope) -> &mut Self::Target;
222}
223
224impl<T> AsDerefExpand for &mut T {
225    type Target = T;
226    fn __expand_as_deref_method(&self, _: &Scope) -> &T {
227        self
228    }
229}
230
231pub trait CubeEnum: Sized {
232    type RuntimeValue: ExpandTypeClone + CubeDebug;
233
234    fn discriminant(&self) -> NativeExpand<i32>;
235
236    /// Return the runtime value of this enum, if only one variant has a value.
237    /// Should return () for all other cases.
238    fn runtime_value(self) -> Self::RuntimeValue;
239
240    fn discriminant_of_value(&self, variant_name: &'static str) -> i32 {
241        Self::discriminant_of(variant_name)
242    }
243
244    fn discriminant_of(variant_name: &'static str) -> i32;
245}
246
247pub trait Assign<T = Self> {
248    /// Assign `value` to `self` in `scope`.
249    fn __expand_assign_method(&mut self, scope: &Scope, value: T);
250}
251
252pub trait RuntimeAssign<T = <Self as IntoExpand>::Expand>: IntoExpand<Expand: Assign<T>> {
253    /// Create a new mutable variable of this type in `scope`.
254    fn init_mut(&self, scope: &Scope) -> Self::Expand;
255}
256
257pub fn __expand_assign<T: Assign<T>>(scope: &Scope, target: &mut T, value: T) {
258    target.__expand_assign_method(scope, value);
259}
260
261impl<T: CubePrimitive> Assign for T {
262    fn __expand_assign_method(&mut self, _scope: &Scope, value: Self) {
263        *self = value;
264    }
265}
266
267impl<T: CubePrimitive + IntoExpand<Expand = NativeExpand<T>>> RuntimeAssign for T {
268    fn init_mut(&self, scope: &Scope) -> NativeExpand<T> {
269        init_mut_of_type(scope, T::__expand_as_type(scope)).into()
270    }
271}
272#[diagnostic::do_not_recommend]
273impl<T: CubePrimitive + IntoExpand<Expand = NativeExpand<T>>> ConstRuntimeAssign for T {
274    fn const_init_mut(scope: &Scope) -> NativeExpand<T> {
275        init_mut_of_type(scope, T::__expand_as_type(scope)).into()
276    }
277}
278
279impl<T: NativeAssign + NativeCubeType + CanReadValue> Assign for NativeExpand<T> {
280    fn __expand_assign_method(&mut self, scope: &Scope, value: Self) {
281        let value = value.read_value(scope);
282        assign::expand(scope, value.into(), self);
283    }
284}
285
286impl<T: NativeAssign + NativeCubeType + CanReadValue> RuntimeAssign for NativeExpand<T> {
287    fn init_mut(&self, scope: &Scope) -> Self::Expand {
288        T::elem_init_mut(scope, self.expand).into()
289    }
290}
291impl<T: CubePrimitive> ConstRuntimeAssign for NativeExpand<T> {
292    fn const_init_mut(scope: &Scope) -> Self::Expand {
293        init_mut_of_type(scope, T::__expand_as_type(scope)).into()
294    }
295}
296
297impl<T: Assign> Assign for Option<T> {
298    fn __expand_assign_method(&mut self, scope: &Scope, value: Self) {
299        match (self, value) {
300            (Some(this), Some(other)) => this.__expand_assign_method(scope, other),
301            (None, None) => {}
302            _ => panic!("Can't assign mismatched enum variants"),
303        }
304    }
305}
306
307impl<T: Assign> Assign for Vec<T> {
308    fn __expand_assign_method(&mut self, scope: &Scope, value: Self) {
309        assert!(
310            self.len() == value.len(),
311            "Can't assign mismatched vector lengths"
312        );
313        for (this, other) in self.iter_mut().zip(value) {
314            this.__expand_assign_method(scope, other);
315        }
316    }
317}
318
319pub trait CloneExpand {
320    fn __expand_clone_method(&self, scope: &Scope) -> Self;
321}
322impl<T: Clone> CloneExpand for T {
323    fn __expand_clone_method(&self, _: &Scope) -> Self {
324        self.clone()
325    }
326}
327
328/// Trait useful to convert a comptime value into runtime value.
329pub trait IntoRuntime:
330    IntoExpand<Expand = <Self as CubeType>::ExpandType> + CubeType + Sized
331{
332    fn runtime(self) -> Self {
333        self
334    }
335
336    fn __expand_runtime_method(self, scope: &Scope) -> Self::ExpandType;
337}
338
339/// Trait for marking a function return value as comptime when the compiler can't infer it.
340pub trait IntoComptime: Sized {
341    #[allow(clippy::wrong_self_convention)]
342    fn comptime(self) -> Self {
343        self
344    }
345}
346
347impl<T: Sized> IntoComptime for T {}
348
349/// Convert an expand type to a version with mutable registers when necessary.
350pub trait IntoMut: Sized {
351    /// Convert the variable into a potentially new mutable variable in `scope`, copying if needed.
352    fn into_mut(self, scope: &Scope) -> Self;
353}
354
355impl<T: IntoMut> IntoMut for &T {
356    fn into_mut(self, _: &Scope) -> Self {
357        self
358    }
359}
360
361impl<T: IntoMut> IntoMut for &mut T {
362    fn into_mut(self, _: &Scope) -> Self {
363        self
364    }
365}
366
367impl<T: IntoMut> IntoMut for *const T {
368    fn into_mut(self, _: &Scope) -> Self {
369        self
370    }
371}
372
373impl<T: IntoMut> IntoMut for *mut T {
374    fn into_mut(self, _: &Scope) -> Self {
375        self
376    }
377}
378
379pub fn into_mut_assign<T: RuntimeAssign>(value: T, scope: &Scope) -> T::Expand {
380    let mut out = value.init_mut(scope);
381    out.__expand_assign_method(scope, value.into_expand(scope));
382    out
383}
384
385pub trait CubeDebug {
386    /// Set the debug name of this type's expansion. Should do nothing for types that don't appear
387    /// at runtime
388    #[allow(unused)]
389    fn set_debug_name(&self, scope: &Scope, name: &'static str) {}
390}
391
392impl<T: CubeDebug + ?Sized> CubeDebug for &T {
393    fn set_debug_name(&self, scope: &Scope, name: &'static str) {
394        T::set_debug_name(self, scope, name);
395    }
396}
397
398impl<T: CubeDebug + ?Sized> CubeDebug for &mut T {
399    fn set_debug_name(&self, scope: &Scope, name: &'static str) {
400        T::set_debug_name(self, scope, name);
401    }
402}
403
404impl<T: CubeDebug + ?Sized> CubeDebug for *const T {
405    fn set_debug_name(&self, scope: &Scope, name: &'static str) {
406        T::set_debug_name(unsafe { &**self }, scope, name);
407    }
408}
409
410impl<T: CubeDebug + ?Sized> CubeDebug for *mut T {
411    fn set_debug_name(&self, scope: &Scope, name: &'static str) {
412        T::set_debug_name(unsafe { &**self }, scope, name);
413    }
414}
415
416impl CubeDebug for i128 {}
417
418/// A type that can be used as a kernel comptime argument.
419/// Note that a type doesn't need to implement `CubeComptime` to be used as
420/// a comptime argument. However, this facilitate the declaration of generic cube types.
421///
422/// # Example
423///
424/// ```ignore
425/// #[derive(CubeType)]
426/// pub struct Example<A: CubeType, B: CubeComptime> {
427///     a: A,
428///     #[cube(comptime)]
429///     b: B
430/// }
431/// ```
432pub trait CubeComptime: core::fmt::Debug + core::hash::Hash + Eq + Clone + Copy {}
433impl<T> CubeComptime for T where T: core::fmt::Debug + core::hash::Hash + Eq + Clone + Copy {}
434
435/// Argument used during the compilation of kernels.
436pub trait CompilationArg:
437    Clone + PartialEq + Eq + core::hash::Hash + core::fmt::Debug + Send + Sync + 'static
438{
439    /// Compilation args should be the same even with different element types. However, it isn't
440    /// possible to enforce it with the type system. So, we make the compilation args serializable
441    /// and dynamically cast them.
442    ///
443    /// Without this, the compilation time is unreasonable. The performance drop isn't a concern
444    /// since this is only done once when compiling a kernel for the first time.
445    fn dynamic_cast<Arg: CompilationArg>(&self) -> Arg {
446        // Dynamic cast, unlike transmute it does not require statically proving the types are the
447        // same size. We assert at runtime to avoid undefined behaviour and help the compiler optimize.
448        assert!(size_of::<Arg>() == size_of::<Self>());
449        let this = Box::new(self.clone());
450        unsafe { *Box::from_raw(Box::into_raw(this) as *mut Arg) }
451    }
452}
453
454impl<T: Clone + PartialEq + Eq + core::hash::Hash + core::fmt::Debug + Send + Sync + 'static>
455    CompilationArg for T
456{
457}
458
459/// Defines how a [launch argument](LaunchArg) can be expanded.
460///
461/// TODO Verify the accuracy of the next comment.
462///
463/// Normally this type should be implemented two times for an argument.
464/// Once for the reference and the other for the mutable reference. Often time, the reference
465/// should expand the argument as an input while the mutable reference should expand the argument
466/// as an output.
467#[diagnostic::on_unimplemented(note = "Consider using `#[derive(CubeLaunch)]` on `{Self}`")]
468pub trait LaunchArg: CubeType + 'static {
469    /// The runtime argument for the kernel.
470    type RuntimeArg: Send + Sync;
471    /// Compilation argument.
472    type CompilationArg: CompilationArg;
473
474    fn register(arg: Self::RuntimeArg, launcher: &mut KernelLauncher) -> Self::CompilationArg;
475
476    /// Register a variable during compilation that fill the [`KernelBuilder`].
477    fn expand(
478        arg: &Self::CompilationArg,
479        builder: &mut KernelBuilder,
480    ) -> <Self as CubeType>::ExpandType;
481}
482
483macro_rules! impl_launch_arg_ref {
484    ($ty: ty) => {
485        impl<T: LaunchArg + ?Sized + 'static> LaunchArg for $ty {
486            type RuntimeArg = T::RuntimeArg;
487            type CompilationArg = T::CompilationArg;
488
489            fn register(
490                arg: Self::RuntimeArg,
491                launcher: &mut KernelLauncher,
492            ) -> Self::CompilationArg {
493                T::register(arg, launcher)
494            }
495
496            fn expand(
497                arg: &Self::CompilationArg,
498                builder: &mut KernelBuilder,
499            ) -> <Self as CubeType>::ExpandType {
500                let value = T::expand(arg, builder);
501                builder.scope.create_kernel_ref(value)
502            }
503        }
504    };
505}
506
507impl_launch_arg_ref!(&'static T);
508impl_launch_arg_ref!(&'static mut T);
509impl_launch_arg_ref!(*const T);
510impl_launch_arg_ref!(*mut T);
511
512macro_rules! launch_tuple {
513    ($(($T:ident, $t:ident)),*) => {
514        impl<$($T: LaunchArg),*> LaunchArg for ($($T,)*) {
515            type RuntimeArg = ($($T::RuntimeArg,)*);
516            type CompilationArg = ($($T::CompilationArg,)*);
517
518            fn register(runtime_arg: Self::RuntimeArg, launcher: &mut KernelLauncher) -> Self::CompilationArg {
519                let ($($t,)*) = runtime_arg;
520                ($($T::register($t, launcher),)*)
521            }
522
523            fn expand(arg: &Self::CompilationArg, builder: &mut KernelBuilder) -> ($(<$T as CubeType>::ExpandType,)*) {
524                let ($($t,)*) = arg;
525                ($($T::expand($t, builder),)*)
526            }
527        }
528    };
529}
530
531all_tuples!(launch_tuple, 1, 12, T, t);
532
533macro_rules! as_ref_tuple {
534    ($(($T:ident, $t:ident)),*) => {
535        impl<$($T: AsRefExpand),*> AsRefExpand for ($($T,)*) {
536            fn __expand_ref_method(&self, _: &Scope) -> &($($T,)*) {
537                self
538            }
539        }
540    };
541}
542
543all_tuples!(as_ref_tuple, 1, 12, T, t);
544
545macro_rules! as_mut_tuple {
546    ($(($T:ident, $t:ident)),*) => {
547        impl<$($T: AsMutExpand),*> AsMutExpand for ($($T,)*) {
548            fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut ($($T,)*) {
549                self
550            }
551        }
552    };
553}
554
555all_tuples!(as_mut_tuple, 1, 12, T, t);
556
557macro_rules! deref_tuple {
558    ($(($T:ident, $t:ident)),*) => {
559        impl<$($T: DerefExpand),*> DerefExpand for ($($T,)*) {
560            type Target = ($($T::Target,)*);
561
562            fn __expand_deref_method(&self, scope: &Scope) -> Self::Target {
563                let ($($t,)*) = self;
564                ($($t.__expand_deref_method(scope),)*)
565            }
566        }
567    };
568}
569
570all_tuples!(deref_tuple, 1, 12, T, t);
571
572/// Expand type of a native GPU type, i.e. scalar primitives, arrays, shared memory.
573#[derive(new, Clone, Copy, Debug)]
574pub struct NativeExpand<T: ?Sized> {
575    pub expand: ExpandValue,
576    pub(crate) _type: PhantomData<T>,
577}
578
579impl<T: ?Sized> IntoExpand for NativeExpand<T> {
580    type Expand = Self;
581
582    fn into_expand(self, _: &Scope) -> Self::Expand {
583        self
584    }
585}
586
587impl<T: ?Sized> ExpandTypeClone for NativeExpand<T> {
588    fn clone_unchecked(&self) -> Self {
589        NativeExpand {
590            expand: self.expand,
591            _type: PhantomData,
592        }
593    }
594}
595
596impl<T: ?Sized> NativeExpand<T> {
597    /// Casts a reference of this expand element to a different type.
598    /// # Safety
599    /// There's no guarantee the new type is valid for the `Value`
600    pub unsafe fn as_type_ref_unchecked<E: ?Sized>(&self) -> &NativeExpand<E> {
601        unsafe { core::mem::transmute::<&NativeExpand<T>, &NativeExpand<E>>(self) }
602    }
603
604    /// Casts a mutable reference of this expand element to a different type.
605    /// # Safety
606    /// There's no guarantee the new type is valid for the `Value`
607    pub unsafe fn as_type_mut_unchecked<E: ?Sized>(&mut self) -> &mut NativeExpand<E> {
608        unsafe { core::mem::transmute::<&mut NativeExpand<T>, &mut NativeExpand<E>>(self) }
609    }
610}
611
612/// Read a value into registers. Should only be implemented for types that can exist outside of
613/// memory like primitives or `Array`, but not for things like `Barrier` that must exist behind a
614/// pointer.
615pub trait ReadValue {
616    fn read_value(&self, scope: &Scope) -> Value;
617}
618
619impl<T: CubePrimitive> ReadValue for NativeExpand<T> {
620    fn read_value(&self, scope: &Scope) -> Value {
621        self.expand.read_value(scope)
622    }
623}
624
625pub trait CanReadValue: CubeType<ExpandType: ReadValue> {}
626impl<T: CubeType<ExpandType: ReadValue>> CanReadValue for T {}
627
628pub trait HasValue {
629    fn value(&self, scope: &Scope) -> Value;
630}
631
632impl<T: ?Sized> HasValue for NativeExpand<T> {
633    fn value(&self, scope: &Scope) -> Value {
634        self.expand.value(scope)
635    }
636}
637
638impl<T: ?Sized> HasValue for &NativeExpand<T> {
639    fn value(&self, scope: &Scope) -> Value {
640        self.expand.value(scope)
641    }
642}
643
644impl<T: ?Sized> HasValue for &mut NativeExpand<T> {
645    fn value(&self, scope: &Scope) -> Value {
646        self.expand.value(scope)
647    }
648}
649
650impl<T: ?Sized> HasValue for *const NativeExpand<T> {
651    fn value(&self, scope: &Scope) -> Value {
652        unsafe { &**self }.value(scope)
653    }
654}
655
656impl<T: ?Sized> HasValue for *mut NativeExpand<T> {
657    fn value(&self, scope: &Scope) -> Value {
658        unsafe { &**self }.value(scope)
659    }
660}
661
662impl<T: ?Sized> AsRefExpand for NativeExpand<T> {
663    fn __expand_ref_method(&self, _: &Scope) -> &Self {
664        self
665    }
666}
667
668#[diagnostic::do_not_recommend]
669impl<T: CubePrimitive> AsMutExpand for NativeExpand<T> {
670    fn __expand_ref_mut_method(&mut self, _scope: &Scope) -> &mut Self {
671        self
672    }
673}
674
675impl<T: CubePrimitive> DerefExpand for NativeExpand<T> {
676    type Target = Self;
677
678    fn __expand_deref_method(&self, scope: &Scope) -> NativeExpand<T> {
679        self.read_value(scope).into()
680    }
681}
682
683impl<T: ?Sized> From<NativeExpand<T>> for ExpandValue {
684    fn from(value: NativeExpand<T>) -> Self {
685        value.expand
686    }
687}
688
689macro_rules! from_const {
690    ($lit:ty) => {
691        impl From<$lit> for NativeExpand<$lit> {
692            fn from(value: $lit) -> Self {
693                let variable: ExpandValue = value.into();
694                variable.into()
695            }
696        }
697    };
698}
699
700from_const!(u8);
701from_const!(u16);
702from_const!(u32);
703from_const!(u64);
704from_const!(usize);
705from_const!(isize);
706from_const!(i64);
707from_const!(i8);
708from_const!(i16);
709from_const!(i32);
710from_const!(f64);
711from_const!(f16);
712from_const!(bf16);
713from_const!(flex32);
714from_const!(tf32);
715from_const!(f32);
716from_const!(e2m1);
717from_const!(e2m1x2);
718from_const!(e2m3);
719from_const!(e3m2);
720from_const!(e4m3);
721from_const!(e5m2);
722from_const!(ue8m0);
723from_const!(bool);
724from_const!(num_complex::Complex<f32>);
725from_const!(num_complex::Complex<f64>);
726
727macro_rules! tuple_cube_type {
728    ($($P:ident),*) => {
729        impl<$($P: CubeType),*> CubeType for ($($P,)*) {
730            type ExpandType = ($($P::ExpandType,)*);
731        }
732
733        impl<$($P: IntoExpand),*> IntoExpand for ($($P,)*) {
734            type Expand = ($($P::Expand,)*);
735
736            #[allow(non_snake_case, unused, clippy::unused_unit)]
737            fn into_expand(self, scope: &Scope) -> Self::Expand {
738                let ($($P,)*) = self;
739                ($(
740                    $P.into_expand(scope),
741                )*)
742            }
743        }
744
745        impl<$($P: ExpandTypeClone),*> ExpandTypeClone for ($($P,)*) {
746            #[allow(non_snake_case, unused, clippy::unused_unit)]
747            fn clone_unchecked(&self) -> Self {
748                let ($($P,)*) = self;
749                ($(
750                    $P.clone_unchecked(),
751                )*)
752            }
753        }
754    }
755}
756macro_rules! tuple_init {
757    ($($P:ident),*) => {
758        impl<$($P: IntoMut),*> IntoMut for ($($P,)*) {
759            #[allow(non_snake_case, unused, clippy::unused_unit)]
760            fn into_mut(self, scope: &Scope) -> Self {
761                let ($($P,)*) = self;
762                ($(
763                    $P.into_mut(scope),
764                )*)
765            }
766        }
767    }
768}
769macro_rules! tuple_debug {
770    ($($P:ident),*) => {
771        impl<$($P: CubeDebug),*> CubeDebug for ($($P,)*) {}
772    }
773}
774macro_rules! tuple_runtime {
775    ($($P:ident),*) => {
776        impl<$($P: IntoRuntime),*> IntoRuntime for ($($P,)*) {
777            #[allow(non_snake_case, unused, clippy::unused_unit)]
778            fn __expand_runtime_method(self, scope: &Scope) -> Self::ExpandType {
779                let ($($P,)*) = self;
780                ($(
781                    $P.__expand_runtime_method(scope),
782                )*)
783            }
784        }
785    }
786}
787macro_rules! tuple_assign {
788    ($(($n: tt, $P:ident)),*) => {
789        impl<$($P: Assign),*> Assign for ($($P,)*) {
790            #[allow(non_snake_case, unused, clippy::unused_unit)]
791            fn __expand_assign_method(&mut self, scope: &Scope, value: Self) {
792                let ($($P,)*) = self;
793                $(
794                    $P.__expand_assign_method(scope, value.$n);
795                )*
796            }
797        }
798
799        impl<$($P: RuntimeAssign),*> RuntimeAssign for ($($P,)*) {
800            #[allow(non_snake_case, unused, clippy::unused_unit)]
801            fn init_mut(&self, scope: &Scope) -> Self::Expand {
802                let ($($P,)*) = self;
803                ($(
804                    $P.init_mut(scope),
805                )*)
806            }
807        }
808
809        impl<$($P: ConstRuntimeAssign),*> ConstRuntimeAssign for ($($P,)*) {
810            #[allow(non_snake_case, unused, clippy::unused_unit)]
811            fn const_init_mut(scope: &Scope) -> Self::Expand {
812                ($(
813                    $P::const_init_mut(scope),
814                )*)
815            }
816        }
817    }
818}
819
820all_tuples!(tuple_cube_type, 1, 12, P);
821all_tuples!(tuple_debug, 1, 12, P);
822all_tuples!(tuple_init, 1, 12, P);
823all_tuples!(tuple_runtime, 1, 12, P);
824all_tuples_enumerated!(tuple_assign, 1, 12, P);
825
826/// Trait for native types that can be assigned. For non-native composites, use the normal [`Assign`].
827pub trait NativeAssign: CubeType {
828    fn elem_init_mut(scope: &Scope, elem: ExpandValue) -> ExpandValue {
829        init_mut_of_type(scope, elem.value(scope).get_type(scope.ctx()))
830    }
831}
832
833impl<T: NativeAssign + NativeCubeType + CanReadValue> IntoMut for NativeExpand<T> {
834    fn into_mut(self, scope: &Scope) -> Self {
835        into_mut_assign(self, scope)
836    }
837}
838
839impl<T: ?Sized> CubeDebug for NativeExpand<T> {
840    fn set_debug_name(&self, scope: &Scope, name: &'static str) {
841        self.value(scope).set_name(scope.ctx(), Some(ident(name)));
842    }
843}
844
845impl<T: CubePrimitive> NativeExpand<T> {
846    // Expanded version of vectorization factor.
847    pub fn __expand_vector_size_method(&self, scope: &Scope) -> VectorSize {
848        self.value(scope).vector_size(scope.ctx())
849    }
850}
851
852impl<T: ?Sized> From<ExpandValue> for NativeExpand<T> {
853    fn from(expand: ExpandValue) -> Self {
854        Self {
855            expand,
856            _type: PhantomData,
857        }
858    }
859}
860
861impl<T: ?Sized> From<Value> for NativeExpand<T> {
862    fn from(expand: Value) -> Self {
863        Self {
864            expand: expand.into(),
865            _type: PhantomData,
866        }
867    }
868}
869
870impl<T: Scalar + Into<ConstantValue>> NativeExpand<T> {
871    /// Create an [`NativeExpand`] from a value that is normally a literal.
872    pub fn from_lit(scope: &Scope, lit: T) -> Self {
873        T::elem_type(scope).constant(lit.into()).into()
874    }
875
876    /// Get the [`ConstantValue`] from the variable.
877    pub fn constant(&self) -> Option<ConstantValue> {
878        match self.expand {
879            ExpandValue::Constant { value, .. } => Some(value),
880            _ => None,
881        }
882    }
883
884    pub fn __expand_into_lit_unchecked_method(self, _scope: &Scope) -> T {
885        let value = self.constant().unwrap();
886        T::from_const_value(value)
887    }
888}
889
890pub(crate) fn zeroed_mut_of_type(scope: &Scope, mut ty: TypeHandle) -> ExpandValue {
891    let ctx = scope.ctx();
892    if let Some(PointerType { inner, .. }) = ty.deref(ctx).downcast_ref() {
893        ty = *inner;
894    }
895    if ty.is_ptr(ctx) {
896        panic!("tried initializing mut for ptr {}", ty.disp(ctx));
897    }
898    let zero = ZeroAttr::new(ty);
899    scope.create_local_mut(ty, Some(zero.into())).into()
900}
901
902pub(crate) fn init_mut_of_type(scope: &Scope, mut ty: TypeHandle) -> ExpandValue {
903    let ctx = scope.ctx();
904    if let Some(PointerType { inner, .. }) = ty.deref(ctx).downcast_ref() {
905        ty = *inner;
906    }
907    if ty.is_ptr(ctx) {
908        panic!("tried initializing mut for ptr {}", ty.disp(ctx));
909    }
910    scope.create_local_mut(ty, None).into()
911}
912
913impl<T: IntoMut> IntoMut for Option<T> {
914    fn into_mut(self, scope: &Scope) -> Self {
915        self.map(|o| IntoMut::into_mut(o, scope))
916    }
917}
918
919impl<T: CubeType> CubeType for Vec<T> {
920    type ExpandType = Vec<T::ExpandType>;
921}
922
923impl<T: IntoExpand> IntoExpand for Vec<T> {
924    type Expand = Self;
925
926    fn into_expand(self, _: &Scope) -> Self::Expand {
927        self
928    }
929}
930
931impl<T: ExpandTypeClone> ExpandTypeClone for Vec<T> {
932    fn clone_unchecked(&self) -> Self {
933        self.iter().map(|it| it.clone_unchecked()).collect()
934    }
935}
936
937impl<T: IntoMut> IntoMut for Vec<T> {
938    fn into_mut(self, scope: &Scope) -> Self {
939        self.into_iter().map(|e| e.into_mut(scope)).collect()
940    }
941}
942impl<T: CubeDebug> CubeDebug for Vec<T> {}
943
944impl<T: AsRefExpand> AsRefExpand for Vec<T> {
945    fn __expand_ref_method(&self, _: &Scope) -> &Self {
946        self
947    }
948}
949impl<T: AsMutExpand> AsMutExpand for Vec<T> {
950    fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
951        self
952    }
953}
954
955/// Create a constant element of the correct type during expansion.
956pub(crate) fn __expand_new<C: Numeric, Out: Numeric>(scope: &Scope, val: C) -> NativeExpand<Out> {
957    let input: ConstantValue = val.into();
958    Out::elem_type(scope).constant(input).into()
959}
960
961impl CubeType for () {
962    type ExpandType = ();
963}
964
965impl LaunchArg for () {
966    type RuntimeArg = ();
967    type CompilationArg = ();
968
969    fn register(_runtime_arg: Self::RuntimeArg, _launcher: &mut KernelLauncher) {
970        // nothing to do
971    }
972
973    fn expand(
974        _: &Self::CompilationArg,
975        _builder: &mut KernelBuilder,
976    ) -> <Self as CubeType>::ExpandType {
977    }
978}
979
980impl Assign for () {
981    fn __expand_assign_method(&mut self, _: &Scope, _: Self) {}
982}
983
984impl RuntimeAssign for () {
985    fn init_mut(&self, _: &Scope) {}
986}
987impl ConstRuntimeAssign for () {
988    fn const_init_mut(_: &Scope) -> Self::Expand {}
989}
990
991impl IntoRuntime for () {
992    fn __expand_runtime_method(self, _: &Scope) -> Self::ExpandType {
993        self
994    }
995}
996
997impl IntoExpand for () {
998    type Expand = ();
999
1000    fn into_expand(self, _: &Scope) -> Self::Expand {
1001        self
1002    }
1003}
1004
1005impl CubeDebug for () {}
1006
1007impl ExpandTypeClone for () {
1008    fn clone_unchecked(&self) -> Self {
1009        *self
1010    }
1011}
1012
1013impl IntoMut for () {
1014    fn into_mut(self, _: &Scope) -> Self {
1015        self
1016    }
1017}
1018
1019impl AsRefExpand for () {
1020    fn __expand_ref_method(&self, _: &Scope) -> &Self {
1021        self
1022    }
1023}
1024impl AsMutExpand for () {
1025    fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
1026        self
1027    }
1028}
1029
1030impl<T: CubePrimitive + IntoExpand<Expand = NativeExpand<T>> + Debug, const N: usize> IntoExpand
1031    for [T; N]
1032{
1033    type Expand = NativeExpand<[T; N]>;
1034
1035    fn into_expand(self, scope: &Scope) -> Self::Expand {
1036        let values = self
1037            .into_iter()
1038            .map(|it| it.into_expand(scope))
1039            .collect::<Vec<_>>()
1040            .try_into()
1041            .unwrap();
1042        NativeExpand::__expand_construct(scope, values)
1043    }
1044}
1045
1046impl<T: CubePrimitive + IntoRuntime + Debug, const N: usize> IntoRuntime for [T; N] {
1047    fn __expand_runtime_method(self, scope: &Scope) -> Self::ExpandType {
1048        self.into_expand(scope)
1049    }
1050}
1051
1052pub trait DefaultExpand: CubeType {
1053    fn __expand_default(scope: &Scope) -> Self::ExpandType;
1054}
1055
1056impl<T: CubeType + Default + IntoRuntime> DefaultExpand for T {
1057    fn __expand_default(scope: &Scope) -> T::ExpandType {
1058        T::default().__expand_runtime_method(scope)
1059    }
1060}
1061
1062#[derive(Clone, Copy, Debug)]
1063pub struct Const<const N: usize>;
1064
1065pub trait Size: core::fmt::Debug + Clone + Copy + Send + Sync + 'static {
1066    fn __expand_value(scope: &Scope) -> usize;
1067    fn value() -> usize {
1068        unexpanded!()
1069    }
1070    fn try_value_const() -> Option<usize> {
1071        None
1072    }
1073}
1074
1075impl<const VALUE: usize> Size for Const<VALUE> {
1076    fn __expand_value(_scope: &Scope) -> usize {
1077        VALUE
1078    }
1079    fn value() -> usize {
1080        VALUE
1081    }
1082    fn try_value_const() -> Option<usize> {
1083        Some(VALUE)
1084    }
1085}
1086
1087impl<Marker: 'static> Size for DynamicSize<Marker> {
1088    fn __expand_value(scope: &Scope) -> usize {
1089        scope.resolve_size::<Self>().expect("Size to be registered")
1090    }
1091    fn value() -> usize {
1092        unexpanded!()
1093    }
1094}
1095
1096/// Define a custom type to be used for a comptime scalar type.
1097/// Useful for cases where generics can't work.
1098#[macro_export]
1099macro_rules! define_scalar {
1100    ($vis: vis $name: ident) => {
1101        $crate::__private::paste! {
1102            $vis struct [<__ $name>];
1103            $vis type $name = $crate::prelude::DynamicScalar<[<__ $name>]>;
1104        }
1105    };
1106}
1107
1108/// Define a custom type to be used for a comptime size. Useful for cases where generics can't work.
1109#[macro_export]
1110macro_rules! define_size {
1111    ($vis: vis $name: ident) => {
1112        $crate::__private::paste! {
1113            $vis struct [<__ $name>];
1114            $vis type $name = $crate::prelude::DynamicSize<[<__ $name>]>;
1115        }
1116    };
1117}
cubecl_core/frontend/element/base.rs

cubecl_core/frontend/element/
base.rs
