cubecl_core/frontend/element/
base.rs

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

cubecl_core/frontend/element/
base.rs
