cubecl_core/frontend/element/
base.rs

1use super::{CubePrimitive, Numeric};
2use crate::{
3    ir::{ConstantValue, ExpandValue, Scope},
4    prelude::{DynamicSize, KernelBuilder, KernelLauncher, Scalar, assign},
5    unexpanded,
6};
7use alloc::{boxed::Box, vec::Vec};
8use core::{fmt::Debug, marker::PhantomData};
9use cubecl_common::{e2m1, e2m1x2, e2m3, e3m2, e4m3, e5m2, flex32, tf32, ue8m0};
10use cubecl_ir::{
11    VectorSize, ident,
12    interfaces::TypedExt,
13    pliron::{
14        debug_info::set_operation_result_name, printable::Printable, r#type::Typed, value::Value,
15    },
16    types::PointerType,
17};
18use cubecl_runtime::runtime::Runtime;
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
273impl<T: NativeAssign + NativeCubeType + CanReadValue> Assign for NativeExpand<T> {
274    fn __expand_assign_method(&mut self, scope: &Scope, value: Self) {
275        let value = value.read_value(scope);
276        assign::expand(scope, value.into(), self);
277    }
278}
279
280impl<T: NativeAssign + NativeCubeType + CanReadValue> RuntimeAssign for NativeExpand<T> {
281    fn init_mut(&self, scope: &Scope) -> Self::Expand {
282        T::elem_init_mut(scope, self.expand).into()
283    }
284}
285
286impl<T: Assign> Assign for Option<T> {
287    fn __expand_assign_method(&mut self, scope: &Scope, value: Self) {
288        match (self, value) {
289            (Some(this), Some(other)) => this.__expand_assign_method(scope, other),
290            (None, None) => {}
291            _ => panic!("Can't assign mismatched enum variants"),
292        }
293    }
294}
295
296impl<T: Assign> Assign for Vec<T> {
297    fn __expand_assign_method(&mut self, scope: &Scope, value: Self) {
298        assert!(
299            self.len() == value.len(),
300            "Can't assign mismatched vector lengths"
301        );
302        for (this, other) in self.iter_mut().zip(value) {
303            this.__expand_assign_method(scope, other);
304        }
305    }
306}
307
308pub trait CloneExpand {
309    fn __expand_clone_method(&self, scope: &Scope) -> Self;
310}
311impl<T: Clone> CloneExpand for T {
312    fn __expand_clone_method(&self, _: &Scope) -> Self {
313        self.clone()
314    }
315}
316
317/// Trait useful to convert a comptime value into runtime value.
318pub trait IntoRuntime:
319    IntoExpand<Expand = <Self as CubeType>::ExpandType> + CubeType + Sized
320{
321    fn runtime(self) -> Self {
322        self
323    }
324
325    fn __expand_runtime_method(self, scope: &Scope) -> Self::ExpandType;
326}
327
328/// Trait for marking a function return value as comptime when the compiler can't infer it.
329pub trait IntoComptime: Sized {
330    #[allow(clippy::wrong_self_convention)]
331    fn comptime(self) -> Self {
332        self
333    }
334}
335
336impl<T: Sized> IntoComptime for T {}
337
338/// Convert an expand type to a version with mutable registers when necessary.
339pub trait IntoMut: Sized {
340    /// Convert the variable into a potentially new mutable variable in `scope`, copying if needed.
341    fn into_mut(self, scope: &Scope) -> Self;
342}
343
344impl<T: IntoMut> IntoMut for &T {
345    fn into_mut(self, _: &Scope) -> Self {
346        self
347    }
348}
349
350impl<T: IntoMut> IntoMut for &mut T {
351    fn into_mut(self, _: &Scope) -> Self {
352        self
353    }
354}
355
356impl<T: IntoMut> IntoMut for *const T {
357    fn into_mut(self, _: &Scope) -> Self {
358        self
359    }
360}
361
362impl<T: IntoMut> IntoMut for *mut T {
363    fn into_mut(self, _: &Scope) -> Self {
364        self
365    }
366}
367
368pub fn into_mut_assign<T: RuntimeAssign>(value: T, scope: &Scope) -> T::Expand {
369    let mut out = value.init_mut(scope);
370    out.__expand_assign_method(scope, value.into_expand(scope));
371    out
372}
373
374pub trait CubeDebug {
375    /// Set the debug name of this type's expansion. Should do nothing for types that don't appear
376    /// at runtime
377    #[allow(unused)]
378    fn set_debug_name(&self, scope: &Scope, name: &'static str) {}
379}
380
381impl<T: CubeDebug + ?Sized> CubeDebug for &T {
382    fn set_debug_name(&self, scope: &Scope, name: &'static str) {
383        T::set_debug_name(self, scope, name);
384    }
385}
386
387impl<T: CubeDebug + ?Sized> CubeDebug for &mut T {
388    fn set_debug_name(&self, scope: &Scope, name: &'static str) {
389        T::set_debug_name(self, scope, name);
390    }
391}
392
393impl<T: CubeDebug + ?Sized> CubeDebug for *const T {
394    fn set_debug_name(&self, scope: &Scope, name: &'static str) {
395        T::set_debug_name(unsafe { &**self }, scope, name);
396    }
397}
398
399impl<T: CubeDebug + ?Sized> CubeDebug for *mut T {
400    fn set_debug_name(&self, scope: &Scope, name: &'static str) {
401        T::set_debug_name(unsafe { &**self }, scope, name);
402    }
403}
404
405impl CubeDebug for i128 {}
406
407/// A type that can be used as a kernel comptime argument.
408/// Note that a type doesn't need to implement `CubeComptime` to be used as
409/// a comptime argument. However, this facilitate the declaration of generic cube types.
410///
411/// # Example
412///
413/// ```ignore
414/// #[derive(CubeType)]
415/// pub struct Example<A: CubeType, B: CubeComptime> {
416///     a: A,
417///     #[cube(comptime)]
418///     b: B
419/// }
420/// ```
421pub trait CubeComptime: core::fmt::Debug + core::hash::Hash + Eq + Clone + Copy {}
422impl<T> CubeComptime for T where T: core::fmt::Debug + core::hash::Hash + Eq + Clone + Copy {}
423
424/// Argument used during the compilation of kernels.
425pub trait CompilationArg:
426    Clone + PartialEq + Eq + core::hash::Hash + core::fmt::Debug + Send + Sync + 'static
427{
428    /// Compilation args should be the same even with different element types. However, it isn't
429    /// possible to enforce it with the type system. So, we make the compilation args serializable
430    /// and dynamically cast them.
431    ///
432    /// Without this, the compilation time is unreasonable. The performance drop isn't a concern
433    /// since this is only done once when compiling a kernel for the first time.
434    fn dynamic_cast<Arg: CompilationArg>(&self) -> Arg {
435        // Dynamic cast, unlike transmute it does not require statically proving the types are the
436        // same size. We assert at runtime to avoid undefined behaviour and help the compiler optimize.
437        assert!(size_of::<Arg>() == size_of::<Self>());
438        let this = Box::new(self.clone());
439        unsafe { *Box::from_raw(Box::into_raw(this) as *mut Arg) }
440    }
441}
442
443impl<T: Clone + PartialEq + Eq + core::hash::Hash + core::fmt::Debug + Send + Sync + 'static>
444    CompilationArg for T
445{
446}
447
448/// Defines how a [launch argument](LaunchArg) can be expanded.
449///
450/// TODO Verify the accuracy of the next comment.
451///
452/// Normally this type should be implemented two times for an argument.
453/// Once for the reference and the other for the mutable reference. Often time, the reference
454/// should expand the argument as an input while the mutable reference should expand the argument
455/// as an output.
456#[diagnostic::on_unimplemented(note = "Consider using `#[derive(CubeLaunch)]` on `{Self}`")]
457pub trait LaunchArg: CubeType + 'static {
458    /// The runtime argument for the kernel.
459    type RuntimeArg<R: Runtime>: Send + Sync;
460    /// Compilation argument.
461    type CompilationArg: CompilationArg;
462
463    fn register<R: Runtime>(
464        arg: Self::RuntimeArg<R>,
465        launcher: &mut KernelLauncher<R>,
466    ) -> Self::CompilationArg;
467
468    /// Register a variable during compilation that fill the [`KernelBuilder`].
469    fn expand(
470        arg: &Self::CompilationArg,
471        builder: &mut KernelBuilder,
472    ) -> <Self as CubeType>::ExpandType;
473}
474
475macro_rules! impl_launch_arg_ref {
476    ($ty: ty) => {
477        impl<T: LaunchArg + ?Sized + 'static> LaunchArg for $ty {
478            type RuntimeArg<R: Runtime> = T::RuntimeArg<R>;
479            type CompilationArg = T::CompilationArg;
480
481            fn register<R: Runtime>(
482                arg: Self::RuntimeArg<R>,
483                launcher: &mut KernelLauncher<R>,
484            ) -> Self::CompilationArg {
485                T::register(arg, launcher)
486            }
487
488            fn expand(
489                arg: &Self::CompilationArg,
490                builder: &mut KernelBuilder,
491            ) -> <Self as CubeType>::ExpandType {
492                let value = T::expand(arg, builder);
493                builder.scope.create_kernel_ref(value)
494            }
495        }
496    };
497}
498
499impl_launch_arg_ref!(&'static T);
500impl_launch_arg_ref!(&'static mut T);
501impl_launch_arg_ref!(*const T);
502impl_launch_arg_ref!(*mut T);
503
504macro_rules! launch_tuple {
505    ($(($T:ident, $t:ident)),*) => {
506        impl<$($T: LaunchArg),*> LaunchArg for ($($T,)*) {
507            type RuntimeArg<R: Runtime> = ($($T::RuntimeArg<R>,)*);
508            type CompilationArg = ($($T::CompilationArg,)*);
509
510            fn register<R: Runtime>(runtime_arg: Self::RuntimeArg<R>, launcher: &mut KernelLauncher<R>) -> Self::CompilationArg {
511                let ($($t,)*) = runtime_arg;
512                ($($T::register($t, launcher),)*)
513            }
514
515            fn expand(arg: &Self::CompilationArg, builder: &mut KernelBuilder) -> ($(<$T as CubeType>::ExpandType,)*) {
516                let ($($t,)*) = arg;
517                ($($T::expand($t, builder),)*)
518            }
519        }
520    };
521}
522
523all_tuples!(launch_tuple, 1, 12, T, t);
524
525macro_rules! as_ref_tuple {
526    ($(($T:ident, $t:ident)),*) => {
527        impl<$($T: AsRefExpand),*> AsRefExpand for ($($T,)*) {
528            fn __expand_ref_method(&self, _: &Scope) -> &($($T,)*) {
529                self
530            }
531        }
532    };
533}
534
535all_tuples!(as_ref_tuple, 1, 12, T, t);
536
537macro_rules! as_mut_tuple {
538    ($(($T:ident, $t:ident)),*) => {
539        impl<$($T: AsMutExpand),*> AsMutExpand for ($($T,)*) {
540            fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut ($($T,)*) {
541                self
542            }
543        }
544    };
545}
546
547all_tuples!(as_mut_tuple, 1, 12, T, t);
548
549macro_rules! deref_tuple {
550    ($(($T:ident, $t:ident)),*) => {
551        impl<$($T: DerefExpand),*> DerefExpand for ($($T,)*) {
552            type Target = ($($T::Target,)*);
553
554            fn __expand_deref_method(&self, scope: &Scope) -> Self::Target {
555                let ($($t,)*) = self;
556                ($($t.__expand_deref_method(scope),)*)
557            }
558        }
559    };
560}
561
562all_tuples!(deref_tuple, 1, 12, T, t);
563
564/// Expand type of a native GPU type, i.e. scalar primitives, arrays, shared memory.
565#[derive(new, Clone, Copy, Debug)]
566pub struct NativeExpand<T: ?Sized> {
567    pub expand: ExpandValue,
568    pub(crate) _type: PhantomData<T>,
569}
570
571impl<T: ?Sized> IntoExpand for NativeExpand<T> {
572    type Expand = Self;
573
574    fn into_expand(self, _: &Scope) -> Self::Expand {
575        self
576    }
577}
578
579impl<T: ?Sized> ExpandTypeClone for NativeExpand<T> {
580    fn clone_unchecked(&self) -> Self {
581        NativeExpand {
582            expand: self.expand,
583            _type: PhantomData,
584        }
585    }
586}
587
588impl<T: ?Sized> NativeExpand<T> {
589    /// Casts a reference of this expand element to a different type.
590    /// # Safety
591    /// There's no guarantee the new type is valid for the `Value`
592    pub unsafe fn as_type_ref_unchecked<E: ?Sized>(&self) -> &NativeExpand<E> {
593        unsafe { core::mem::transmute::<&NativeExpand<T>, &NativeExpand<E>>(self) }
594    }
595
596    /// Casts a mutable reference of this expand element to a different type.
597    /// # Safety
598    /// There's no guarantee the new type is valid for the `Value`
599    pub unsafe fn as_type_mut_unchecked<E: ?Sized>(&mut self) -> &mut NativeExpand<E> {
600        unsafe { core::mem::transmute::<&mut NativeExpand<T>, &mut NativeExpand<E>>(self) }
601    }
602}
603
604/// Read a value into registers. Should only be implemented for types that can exist outside of
605/// memory like primitives or `Array`, but not for things like `Barrier` that must exist behind a
606/// pointer.
607pub trait ReadValue {
608    fn read_value(&self, scope: &Scope) -> Value;
609}
610
611impl<T: CubePrimitive> ReadValue for NativeExpand<T> {
612    fn read_value(&self, scope: &Scope) -> Value {
613        self.expand.read_value(scope)
614    }
615}
616
617pub trait CanReadValue: CubeType<ExpandType: ReadValue> {}
618impl<T: CubeType<ExpandType: ReadValue>> CanReadValue for T {}
619
620pub trait HasValue {
621    fn value(&self, scope: &Scope) -> Value;
622}
623
624impl<T: ?Sized> HasValue for NativeExpand<T> {
625    fn value(&self, scope: &Scope) -> Value {
626        self.expand.value(scope)
627    }
628}
629
630impl<T: ?Sized> HasValue for &NativeExpand<T> {
631    fn value(&self, scope: &Scope) -> Value {
632        self.expand.value(scope)
633    }
634}
635
636impl<T: ?Sized> HasValue for &mut NativeExpand<T> {
637    fn value(&self, scope: &Scope) -> Value {
638        self.expand.value(scope)
639    }
640}
641
642impl<T: ?Sized> HasValue for *const NativeExpand<T> {
643    fn value(&self, scope: &Scope) -> Value {
644        unsafe { &**self }.value(scope)
645    }
646}
647
648impl<T: ?Sized> HasValue for *mut NativeExpand<T> {
649    fn value(&self, scope: &Scope) -> Value {
650        unsafe { &**self }.value(scope)
651    }
652}
653
654impl<T: ?Sized> AsRefExpand for NativeExpand<T> {
655    fn __expand_ref_method(&self, _: &Scope) -> &Self {
656        self
657    }
658}
659
660#[diagnostic::do_not_recommend]
661impl<T: CubePrimitive> AsMutExpand for NativeExpand<T> {
662    fn __expand_ref_mut_method(&mut self, _scope: &Scope) -> &mut Self {
663        self
664    }
665}
666
667impl<T: CubePrimitive> DerefExpand for NativeExpand<T> {
668    type Target = Self;
669
670    fn __expand_deref_method(&self, scope: &Scope) -> NativeExpand<T> {
671        self.read_value(scope).into()
672    }
673}
674
675impl<T: ?Sized> From<NativeExpand<T>> for ExpandValue {
676    fn from(value: NativeExpand<T>) -> Self {
677        value.expand
678    }
679}
680
681macro_rules! from_const {
682    ($lit:ty) => {
683        impl From<$lit> for NativeExpand<$lit> {
684            fn from(value: $lit) -> Self {
685                let variable: ExpandValue = value.into();
686                variable.into()
687            }
688        }
689    };
690}
691
692from_const!(u8);
693from_const!(u16);
694from_const!(u32);
695from_const!(u64);
696from_const!(usize);
697from_const!(isize);
698from_const!(i64);
699from_const!(i8);
700from_const!(i16);
701from_const!(i32);
702from_const!(f64);
703from_const!(f16);
704from_const!(bf16);
705from_const!(flex32);
706from_const!(tf32);
707from_const!(f32);
708from_const!(e2m1);
709from_const!(e2m1x2);
710from_const!(e2m3);
711from_const!(e3m2);
712from_const!(e4m3);
713from_const!(e5m2);
714from_const!(ue8m0);
715from_const!(bool);
716
717macro_rules! tuple_cube_type {
718    ($($P:ident),*) => {
719        impl<$($P: CubeType),*> CubeType for ($($P,)*) {
720            type ExpandType = ($($P::ExpandType,)*);
721        }
722
723        impl<$($P: IntoExpand),*> IntoExpand for ($($P,)*) {
724            type Expand = ($($P::Expand,)*);
725
726            #[allow(non_snake_case, unused, clippy::unused_unit)]
727            fn into_expand(self, scope: &Scope) -> Self::Expand {
728                let ($($P,)*) = self;
729                ($(
730                    $P.into_expand(scope),
731                )*)
732            }
733        }
734
735        impl<$($P: ExpandTypeClone),*> ExpandTypeClone for ($($P,)*) {
736            #[allow(non_snake_case, unused, clippy::unused_unit)]
737            fn clone_unchecked(&self) -> Self {
738                let ($($P,)*) = self;
739                ($(
740                    $P.clone_unchecked(),
741                )*)
742            }
743        }
744    }
745}
746macro_rules! tuple_init {
747    ($($P:ident),*) => {
748        impl<$($P: IntoMut),*> IntoMut for ($($P,)*) {
749            #[allow(non_snake_case, unused, clippy::unused_unit)]
750            fn into_mut(self, scope: &Scope) -> Self {
751                let ($($P,)*) = self;
752                ($(
753                    $P.into_mut(scope),
754                )*)
755            }
756        }
757    }
758}
759macro_rules! tuple_debug {
760    ($($P:ident),*) => {
761        impl<$($P: CubeDebug),*> CubeDebug for ($($P,)*) {}
762    }
763}
764macro_rules! tuple_runtime {
765    ($($P:ident),*) => {
766        impl<$($P: IntoRuntime),*> IntoRuntime for ($($P,)*) {
767            #[allow(non_snake_case, unused, clippy::unused_unit)]
768            fn __expand_runtime_method(self, scope: &Scope) -> Self::ExpandType {
769                let ($($P,)*) = self;
770                ($(
771                    $P.__expand_runtime_method(scope),
772                )*)
773            }
774        }
775    }
776}
777macro_rules! tuple_assign {
778    ($(($n: tt, $P:ident)),*) => {
779        impl<$($P: Assign),*> Assign for ($($P,)*) {
780            #[allow(non_snake_case, unused, clippy::unused_unit)]
781            fn __expand_assign_method(&mut self, scope: &Scope, value: Self) {
782                let ($($P,)*) = self;
783                $(
784                    $P.__expand_assign_method(scope, value.$n);
785                )*
786            }
787        }
788
789        impl<$($P: RuntimeAssign),*> RuntimeAssign for ($($P,)*) {
790            #[allow(non_snake_case, unused, clippy::unused_unit)]
791            fn init_mut(&self, scope: &Scope) -> Self::Expand {
792                let ($($P,)*) = self;
793                ($(
794                    $P.init_mut(scope),
795                )*)
796            }
797        }
798    }
799}
800
801all_tuples!(tuple_cube_type, 1, 12, P);
802all_tuples!(tuple_debug, 1, 12, P);
803all_tuples!(tuple_init, 1, 12, P);
804all_tuples!(tuple_runtime, 1, 12, P);
805all_tuples_enumerated!(tuple_assign, 1, 12, P);
806
807/// Trait for native types that can be assigned. For non-native composites, use the normal [`Assign`].
808pub trait NativeAssign: CubeType {
809    fn elem_init_mut(scope: &Scope, elem: ExpandValue) -> ExpandValue {
810        init_mut_of_type(scope, elem.value(scope).get_type(scope.ctx()))
811    }
812}
813
814impl<T: NativeAssign + NativeCubeType + CanReadValue> IntoMut for NativeExpand<T> {
815    fn into_mut(self, scope: &Scope) -> Self {
816        into_mut_assign(self, scope)
817    }
818}
819
820impl<T: ?Sized> CubeDebug for NativeExpand<T> {
821    fn set_debug_name(&self, scope: &Scope, name: &'static str) {
822        let op = self.value(scope).defining_op().unwrap();
823        set_operation_result_name(scope.ctx(), op, 0, Some(ident(name)));
824    }
825}
826
827impl<T: CubePrimitive> NativeExpand<T> {
828    // Expanded version of vectorization factor.
829    pub fn __expand_vector_size_method(&self, scope: &Scope) -> VectorSize {
830        self.value(scope).vector_size(scope.ctx())
831    }
832}
833
834impl<T: ?Sized> From<ExpandValue> for NativeExpand<T> {
835    fn from(expand: ExpandValue) -> Self {
836        Self {
837            expand,
838            _type: PhantomData,
839        }
840    }
841}
842
843impl<T: ?Sized> From<Value> for NativeExpand<T> {
844    fn from(expand: Value) -> Self {
845        Self {
846            expand: expand.into(),
847            _type: PhantomData,
848        }
849    }
850}
851
852impl<T: Scalar + Into<ConstantValue>> NativeExpand<T> {
853    /// Create an [`NativeExpand`] from a value that is normally a literal.
854    pub fn from_lit(scope: &Scope, lit: T) -> Self {
855        T::elem_type(scope).constant(lit.into()).into()
856    }
857
858    /// Get the [`ConstantValue`] from the variable.
859    pub fn constant(&self) -> Option<ConstantValue> {
860        match self.expand {
861            ExpandValue::Constant { value, .. } => Some(value),
862            _ => None,
863        }
864    }
865
866    pub fn __expand_into_lit_unchecked_method(self, _scope: &Scope) -> T {
867        let value = self.constant().unwrap();
868        T::from_const_value(value)
869    }
870}
871
872pub(crate) fn init_mut_of_type(scope: &Scope, mut ty: TypeHandle) -> ExpandValue {
873    let ctx = scope.ctx();
874    if let Some(PointerType { inner, .. }) = ty.deref(ctx).downcast_ref() {
875        ty = *inner;
876    }
877    if ty.is_ptr(ctx) {
878        panic!("tried initializing mut for ptr {}", ty.disp(ctx));
879    }
880    scope.create_local_mut(ty, None).into()
881}
882
883impl<T: IntoMut> IntoMut for Option<T> {
884    fn into_mut(self, scope: &Scope) -> Self {
885        self.map(|o| IntoMut::into_mut(o, scope))
886    }
887}
888
889impl<T: CubeType> CubeType for Vec<T> {
890    type ExpandType = Vec<T::ExpandType>;
891}
892
893impl<T: IntoExpand> IntoExpand for Vec<T> {
894    type Expand = Self;
895
896    fn into_expand(self, _: &Scope) -> Self::Expand {
897        self
898    }
899}
900
901impl<T: ExpandTypeClone> ExpandTypeClone for Vec<T> {
902    fn clone_unchecked(&self) -> Self {
903        self.iter().map(|it| it.clone_unchecked()).collect()
904    }
905}
906
907impl<T: IntoMut> IntoMut for Vec<T> {
908    fn into_mut(self, scope: &Scope) -> Self {
909        self.into_iter().map(|e| e.into_mut(scope)).collect()
910    }
911}
912impl<T: CubeDebug> CubeDebug for Vec<T> {}
913
914impl<T: AsRefExpand> AsRefExpand for Vec<T> {
915    fn __expand_ref_method(&self, _: &Scope) -> &Self {
916        self
917    }
918}
919impl<T: AsMutExpand> AsMutExpand for Vec<T> {
920    fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
921        self
922    }
923}
924
925/// Create a constant element of the correct type during expansion.
926pub(crate) fn __expand_new<C: Numeric, Out: Numeric>(scope: &Scope, val: C) -> NativeExpand<Out> {
927    let input: ConstantValue = val.into();
928    Out::elem_type(scope).constant(input).into()
929}
930
931impl CubeType for () {
932    type ExpandType = ();
933}
934
935impl LaunchArg for () {
936    type RuntimeArg<R: Runtime> = ();
937    type CompilationArg = ();
938
939    fn register<R: Runtime>(_runtime_arg: Self::RuntimeArg<R>, _launcher: &mut KernelLauncher<R>) {
940        // nothing to do
941    }
942
943    fn expand(
944        _: &Self::CompilationArg,
945        _builder: &mut KernelBuilder,
946    ) -> <Self as CubeType>::ExpandType {
947    }
948}
949
950impl Assign for () {
951    fn __expand_assign_method(&mut self, _: &Scope, _: Self) {}
952}
953
954impl RuntimeAssign for () {
955    fn init_mut(&self, _: &Scope) {}
956}
957
958impl IntoRuntime for () {
959    fn __expand_runtime_method(self, _: &Scope) -> Self::ExpandType {
960        self
961    }
962}
963
964impl IntoExpand for () {
965    type Expand = ();
966
967    fn into_expand(self, _: &Scope) -> Self::Expand {
968        self
969    }
970}
971
972impl CubeDebug for () {}
973
974impl ExpandTypeClone for () {
975    fn clone_unchecked(&self) -> Self {
976        *self
977    }
978}
979
980impl IntoMut for () {
981    fn into_mut(self, _: &Scope) -> Self {
982        self
983    }
984}
985
986impl AsRefExpand for () {
987    fn __expand_ref_method(&self, _: &Scope) -> &Self {
988        self
989    }
990}
991impl AsMutExpand for () {
992    fn __expand_ref_mut_method(&mut self, _: &Scope) -> &mut Self {
993        self
994    }
995}
996
997pub trait DefaultExpand: CubeType {
998    fn __expand_default(scope: &Scope) -> Self::ExpandType;
999}
1000
1001impl<T: CubeType + Default + IntoRuntime> DefaultExpand for T {
1002    fn __expand_default(scope: &Scope) -> T::ExpandType {
1003        T::default().__expand_runtime_method(scope)
1004    }
1005}
1006
1007#[derive(Clone, Copy, Debug)]
1008pub struct Const<const N: usize>;
1009
1010pub trait Size: core::fmt::Debug + Clone + Copy + Send + Sync + 'static {
1011    fn __expand_value(scope: &Scope) -> usize;
1012    fn value() -> usize {
1013        unexpanded!()
1014    }
1015    fn try_value_const() -> Option<usize> {
1016        None
1017    }
1018}
1019
1020impl<const VALUE: usize> Size for Const<VALUE> {
1021    fn __expand_value(_scope: &Scope) -> usize {
1022        VALUE
1023    }
1024    fn value() -> usize {
1025        VALUE
1026    }
1027    fn try_value_const() -> Option<usize> {
1028        Some(VALUE)
1029    }
1030}
1031
1032impl<Marker: 'static> Size for DynamicSize<Marker> {
1033    fn __expand_value(scope: &Scope) -> usize {
1034        scope.resolve_size::<Self>().expect("Size to be registered")
1035    }
1036    fn value() -> usize {
1037        unexpanded!()
1038    }
1039}
1040
1041/// Define a custom type to be used for a comptime scalar type.
1042/// Useful for cases where generics can't work.
1043#[macro_export]
1044macro_rules! define_scalar {
1045    ($vis: vis $name: ident) => {
1046        $crate::__private::paste! {
1047            $vis struct [<__ $name>];
1048            $vis type $name = $crate::prelude::DynamicScalar<[<__ $name>]>;
1049        }
1050    };
1051}
1052
1053/// Define a custom type to be used for a comptime size. Useful for cases where generics can't work.
1054#[macro_export]
1055macro_rules! define_size {
1056    ($vis: vis $name: ident) => {
1057        $crate::__private::paste! {
1058            $vis struct [<__ $name>];
1059            $vis type $name = $crate::prelude::DynamicSize<[<__ $name>]>;
1060        }
1061    };
1062}
cubecl_core/frontend/element/base.rs

cubecl_core/frontend/element/
base.rs
