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

cubecl_core/frontend/element/
base.rs
