cubecl_ir/
type.rs

1use super::{ConstantValue, Value, ValueKind};
2use crate::{BarrierLevel, ClampMode, Id, MatrixType, TypeHash};
3use core::fmt::Display;
4use cubecl_common::{
5    e2m1, e2m1x2, e2m3, e3m2, e4m3, e5m2, flex32,
6    quant::scheme::{QuantParam, QuantValue},
7    tf32, ue8m0,
8};
9use derive_more::{Display, From};
10use half::{bf16, f16};
11
12pub use internment::Intern;
13
14#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
15#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)]
16#[allow(missing_docs)]
17pub enum FloatKind {
18    /// FP4, 2 bit exponent, 1 bit mantissa
19    E2M1,
20    /// FP6, 2 bit exponent, 3 bit mantissa
21    /// Note: represented by an 8-bit value, with the upper two bits being insignificant
22    E2M3,
23    /// FP6, 3 bit exponent, 2 bit mantissa
24    /// Note: represented by an 8-bit value, with the upper two bits being insignificant
25    E3M2,
26    /// FP8, 4 bit exponent, 3 bit mantissa
27    E4M3,
28    /// FP8, 5 bit exponent, 2 bit mantissa
29    E5M2,
30    /// FP8, unsigned, 8 bit exponent, 0 bit mantissa
31    UE8M0,
32    F16,
33    BF16,
34    Flex32,
35    F32,
36    TF32,
37    F64,
38}
39
40#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
41#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)]
42#[allow(missing_docs)]
43pub enum IntKind {
44    I8,
45    I16,
46    I32,
47    I64,
48}
49
50#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
51#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)]
52#[allow(missing_docs)]
53pub enum UIntKind {
54    U8,
55    U16,
56    U32,
57    U64,
58}
59
60/// Conceptual element type, not necessarily the physical type used in the code
61#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
62#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord, From)]
63#[allow(missing_docs)]
64pub enum ElemType {
65    Float(FloatKind),
66    Int(IntKind),
67    UInt(UIntKind),
68    Bool,
69}
70
71#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
72#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)]
73pub enum OpaqueType {
74    Barrier(BarrierLevel),
75    BarrierToken(BarrierLevel),
76    TensorMap,
77}
78
79#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
80#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)]
81pub enum SemanticType {
82    TensorLayout(usize, ClampMode),
83    TensorView(usize, bool, [u32; 5]),
84}
85
86/// Physical type containing one or more elements
87#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
88#[derive(Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)]
89pub enum StorageType {
90    /// `ElemType` is the same as the physical type
91    Scalar(ElemType),
92    /// Packed values of type `ElemType`
93    Packed(ElemType, usize),
94}
95
96impl core::fmt::Debug for StorageType {
97    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
98        // Ensure debug is not spread into multiple lines because it makes kernel ids very hard
99        // to read.
100        struct Dummy<'a>(&'a StorageType);
101
102        impl<'a> core::fmt::Debug for Dummy<'a> {
103            fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
104                match self.0 {
105                    StorageType::Scalar(f0) => f.debug_tuple("Scalar").field(&f0).finish(),
106                    StorageType::Packed(f0, f1) => {
107                        f.debug_tuple("Packed").field(&f0).field(&f1).finish()
108                    }
109                }
110            }
111        }
112
113        write!(f, "{:?}", Dummy(self))
114    }
115}
116
117impl ElemType {
118    /// Creates an elem type that correspond to the given [`QuantParam`].
119    pub fn from_quant_param(quant_param: QuantParam) -> Self {
120        match quant_param {
121            QuantParam::F32 => Self::Float(FloatKind::F32),
122            QuantParam::F16 => Self::Float(FloatKind::F16),
123            QuantParam::BF16 => Self::Float(FloatKind::BF16),
124            QuantParam::UE8M0 => Self::Float(FloatKind::UE8M0),
125            QuantParam::UE4M3 => Self::Float(FloatKind::E4M3),
126        }
127    }
128
129    /// Creates an elem type that correspond to the given [`QuantValue`].
130    pub fn from_quant_value(quant_value: QuantValue) -> Self {
131        match quant_value {
132            QuantValue::E5M2 => Self::Float(FloatKind::E5M2),
133            QuantValue::E4M3 => Self::Float(FloatKind::E4M3),
134            QuantValue::E2M1 => Self::Float(FloatKind::E2M1),
135            QuantValue::Q8F | QuantValue::Q8S => Self::Int(IntKind::I8),
136            other => panic!("Unsupported quant value {other:?}"),
137        }
138    }
139
140    /// Create a constant from a constant value.
141    ///
142    /// The output will have the same type as the element.
143    pub fn constant(&self, val: ConstantValue) -> Value {
144        Value::constant(val, Type::scalar(*self))
145    }
146
147    /// Get the size in bytes.
148    pub const fn size(&self) -> usize {
149        match self {
150            ElemType::Float(kind) => match kind {
151                FloatKind::E2M1
152                | FloatKind::E2M3
153                | FloatKind::E3M2
154                | FloatKind::E4M3
155                | FloatKind::E5M2
156                | FloatKind::UE8M0 => core::mem::size_of::<u8>(),
157                FloatKind::F16 => core::mem::size_of::<half::f16>(),
158                FloatKind::BF16 => core::mem::size_of::<half::bf16>(),
159                FloatKind::F32 => core::mem::size_of::<f32>(),
160                FloatKind::F64 => core::mem::size_of::<f64>(),
161                FloatKind::Flex32 => core::mem::size_of::<f32>(),
162                FloatKind::TF32 => core::mem::size_of::<f32>(),
163            },
164            ElemType::Int(kind) => match kind {
165                IntKind::I8 => core::mem::size_of::<i8>(),
166                IntKind::I16 => core::mem::size_of::<i16>(),
167                IntKind::I32 => core::mem::size_of::<i32>(),
168                IntKind::I64 => core::mem::size_of::<i64>(),
169            },
170            ElemType::UInt(kind) => match kind {
171                UIntKind::U8 => core::mem::size_of::<u8>(),
172                UIntKind::U16 => core::mem::size_of::<u16>(),
173                UIntKind::U32 => core::mem::size_of::<u32>(),
174                UIntKind::U64 => core::mem::size_of::<u64>(),
175            },
176            ElemType::Bool => core::mem::size_of::<bool>(),
177        }
178    }
179
180    /// Get the size in bits.
181    pub const fn size_bits(&self) -> usize {
182        match self {
183            ElemType::Float(kind) => match kind {
184                FloatKind::E2M3
185                | FloatKind::E3M2
186                | FloatKind::E4M3
187                | FloatKind::E5M2
188                | FloatKind::UE8M0
189                | FloatKind::F16
190                | FloatKind::BF16
191                | FloatKind::F32
192                | FloatKind::F64
193                | FloatKind::Flex32
194                | FloatKind::TF32 => self.size() * 8,
195                FloatKind::E2M1 => 4,
196            },
197            ElemType::Int(_) | ElemType::UInt(_) | ElemType::Bool => self.size() * 8,
198        }
199    }
200
201    pub const fn min_vector_size(&self) -> u8 {
202        match self {
203            ElemType::Float(FloatKind::E2M1) => 2,
204            _ => 1,
205        }
206    }
207
208    pub fn is_int(&self) -> bool {
209        matches!(self, ElemType::Int(_) | ElemType::UInt(_) | ElemType::Bool)
210    }
211
212    pub fn is_signed_int(&self) -> bool {
213        matches!(self, ElemType::Int(_))
214    }
215
216    pub fn is_unsigned_int(&self) -> bool {
217        matches!(self, ElemType::UInt(_) | ElemType::Bool)
218    }
219
220    pub fn is_float(&self) -> bool {
221        matches!(self, ElemType::Float(_))
222    }
223
224    pub fn is_bool(&self) -> bool {
225        matches!(self, ElemType::Bool)
226    }
227
228    pub fn as_float(&self) -> Option<FloatKind> {
229        match self {
230            ElemType::Float(kind) => Some(*kind),
231            _ => None,
232        }
233    }
234
235    pub fn max_variable(&self) -> Value {
236        let value = match self {
237            ElemType::Float(kind) => match kind {
238                FloatKind::E2M1 => e2m1::MAX,
239                FloatKind::E2M3 => e2m3::MAX,
240                FloatKind::E3M2 => e3m2::MAX,
241                FloatKind::E4M3 => e4m3::MAX.to_f64(),
242                FloatKind::E5M2 => e5m2::MAX.to_f64(),
243                FloatKind::UE8M0 => ue8m0::MAX,
244                FloatKind::F16 => half::f16::MAX.to_f64(),
245                FloatKind::BF16 => half::bf16::MAX.to_f64(),
246                FloatKind::Flex32 | FloatKind::TF32 | FloatKind::F32 => f32::MAX as f64,
247                FloatKind::F64 => f64::MAX,
248            }
249            .into(),
250            ElemType::Int(kind) => match kind {
251                IntKind::I8 => i8::MAX as i64,
252                IntKind::I16 => i16::MAX as i64,
253                IntKind::I32 => i32::MAX as i64,
254                IntKind::I64 => i64::MAX,
255            }
256            .into(),
257            ElemType::UInt(kind) => match kind {
258                UIntKind::U8 => u8::MAX as u64,
259                UIntKind::U16 => u16::MAX as u64,
260                UIntKind::U32 => u32::MAX as u64,
261                UIntKind::U64 => u64::MAX,
262            }
263            .into(),
264            ElemType::Bool => true.into(),
265        };
266
267        Value {
268            kind: ValueKind::Constant(value),
269            ty: Type::scalar(*self),
270        }
271    }
272
273    pub fn min_variable(&self) -> Value {
274        let value = match self {
275            ElemType::Float(kind) => match kind {
276                FloatKind::E2M1 => e2m1::MIN,
277                FloatKind::E2M3 => e2m3::MIN,
278                FloatKind::E3M2 => e3m2::MIN,
279                FloatKind::E4M3 => e4m3::MIN.to_f64(),
280                FloatKind::E5M2 => e5m2::MIN.to_f64(),
281                FloatKind::UE8M0 => ue8m0::MIN,
282                FloatKind::F16 => half::f16::MIN.to_f64(),
283                FloatKind::BF16 => half::bf16::MIN.to_f64(),
284                FloatKind::Flex32 | FloatKind::TF32 | FloatKind::F32 => f32::MIN as f64,
285                FloatKind::F64 => f64::MIN,
286            }
287            .into(),
288            ElemType::Int(kind) => match kind {
289                IntKind::I8 => i8::MIN as i64,
290                IntKind::I16 => i16::MIN as i64,
291                IntKind::I32 => i32::MIN as i64,
292                IntKind::I64 => i64::MIN,
293            }
294            .into(),
295            ElemType::UInt(kind) => match kind {
296                UIntKind::U8 => u8::MIN as u64,
297                UIntKind::U16 => u16::MIN as u64,
298                UIntKind::U32 => u32::MIN as u64,
299                UIntKind::U64 => u64::MIN,
300            }
301            .into(),
302            ElemType::Bool => false.into(),
303        };
304
305        Value {
306            kind: ValueKind::Constant(value),
307            ty: Type::scalar(*self),
308        }
309    }
310
311    pub fn epsilon(&self) -> f64 {
312        match self {
313            ElemType::Float(kind) => match kind {
314                FloatKind::E2M1 => 0.5 * (e2m1::MAX - e2m1::MIN),
315                FloatKind::E2M3 => 0.5 * (e2m3::MAX - e2m3::MIN),
316                FloatKind::E3M2 => 0.5 * (e3m2::MAX - e3m2::MIN),
317                FloatKind::E4M3 => 0.5 * (e4m3::MAX.to_f64() - e4m3::MIN.to_f64()),
318                FloatKind::E5M2 => 0.5 * (e5m2::MAX.to_f64() - e5m2::MIN.to_f64()),
319                FloatKind::UE8M0 => 0.5 * (ue8m0::MAX - ue8m0::MIN),
320                FloatKind::F16 => half::f16::EPSILON.to_f64(),
321                FloatKind::BF16 => 0.0078125, // bf16 epsilon ≈ 2^-7
322                FloatKind::Flex32 | FloatKind::F32 | FloatKind::TF32 => f32::EPSILON.into(),
323                FloatKind::F64 => f64::EPSILON,
324            },
325            ElemType::Int(_) | ElemType::UInt(_) => 1.0, // step of 1
326            ElemType::Bool => 1.0,
327        }
328    }
329}
330
331impl OpaqueType {
332    /// Get the size in bytes.
333    pub const fn size(&self) -> usize {
334        match self {
335            OpaqueType::Barrier(_) => 8,
336            OpaqueType::BarrierToken(_) => 8,
337            OpaqueType::TensorMap => 128,
338        }
339    }
340
341    /// Get the size in bits.
342    pub const fn size_bits(&self) -> usize {
343        self.size() * 8
344    }
345}
346
347impl StorageType {
348    pub fn elem_type(&self) -> ElemType {
349        match self {
350            StorageType::Scalar(ty) | StorageType::Packed(ty, _) => *ty,
351        }
352    }
353
354    pub fn packing_factor(&self) -> usize {
355        match self {
356            StorageType::Packed(_, factor) => *factor,
357            _ => 1,
358        }
359    }
360
361    pub fn size(&self) -> usize {
362        self.size_bits().div_ceil(8)
363    }
364
365    pub fn size_bits(&self) -> usize {
366        match self {
367            StorageType::Packed(ty, factor) => ty.size_bits() * *factor,
368            StorageType::Scalar(ty) => ty.size_bits(),
369        }
370    }
371
372    pub fn is_int(&self) -> bool {
373        self.elem_type().is_int()
374    }
375
376    pub fn is_signed_int(&self) -> bool {
377        self.elem_type().is_signed_int()
378    }
379
380    pub fn is_unsigned_int(&self) -> bool {
381        self.elem_type().is_unsigned_int()
382    }
383
384    pub fn is_float(&self) -> bool {
385        self.elem_type().is_float()
386    }
387
388    pub fn is_bool(&self) -> bool {
389        self.elem_type().is_bool()
390    }
391
392    /// Returns an empirical epsilon for this storage type, taking quantization into account.
393    pub fn epsilon(&self) -> f64 {
394        match self {
395            StorageType::Scalar(ty) => ty.epsilon(),
396            StorageType::Packed(ty, factor) => {
397                // For packed types, we can conservatively scale epsilon by the number of packed elements
398                ty.epsilon() * (*factor as f64)
399            }
400        }
401    }
402
403    pub fn constant(&self, value: ConstantValue) -> Value {
404        Value::constant(value, Type::new(*self))
405    }
406}
407
408macro_rules! storage_from_elem {
409    ($($ty: ty),*) => {
410        $(impl From<$ty> for StorageType {
411            fn from(value: $ty) -> Self {
412                StorageType::Scalar(value.into())
413            }
414        })*
415    };
416}
417
418storage_from_elem!(FloatKind, IntKind, UIntKind, ElemType);
419
420impl From<OpaqueType> for Type {
421    fn from(val: OpaqueType) -> Self {
422        Type::Opaque(val)
423    }
424}
425
426impl<T: Into<StorageType>> From<T> for Type {
427    fn from(val: T) -> Self {
428        Type::new(val.into())
429    }
430}
431
432impl From<SemanticType> for Type {
433    fn from(val: SemanticType) -> Self {
434        Type::semantic(val)
435    }
436}
437
438/// Class of a pointer. For `Global`, the ID contains the underlying buffer ID.
439/// The ID can be used to determine more detailed buffer properties, i.e. for Metal where readability
440/// is part of the pointer class.
441/// For ``CubeCL`` semantics, pointers classes to different buffer IDs should be treated as entirely
442/// separate types.
443#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
444#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)]
445pub enum AddressSpace {
446    Global(Id),
447    Shared,
448    Local,
449}
450
451#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
452#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, PartialOrd, Ord)]
453pub enum Type {
454    /// Scalar type containing a single storage element
455    Scalar(StorageType),
456    /// Opaque types that can be stored but not interacted with normally. i.e. barrier,
457    /// arrival tokens and tensor map descriptor.
458    Opaque(OpaqueType),
459    /// Vector wrapping `n` storage elements
460    Vector(Intern<Type>, VectorSize),
461    /// No defined physical representation, purely semantic. i.e. barrier, pipeline
462    Semantic(SemanticType),
463    /// Atomically accessed version of `Type`
464    Atomic(Intern<Type>),
465    /// Pointer of `Type` into a `PointerClass`
466    Pointer(Intern<Type>, AddressSpace),
467    /// Statically sized array of `Type`s
468    Array(Intern<Type>, usize),
469    /// Dynamically sized array of `Type`s
470    DynamicArray(Intern<Type>),
471    /// Cooperative Matrix
472    Matrix(MatrixType),
473    Aggregate(AggregateKind),
474}
475
476/// `Intern` hashes the pointer, not the values, leading to unstable hashes across runs.
477/// Fix this by manually hashing the value.
478impl core::hash::Hash for Type {
479    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
480        core::mem::discriminant(self).hash(state);
481        match self {
482            Type::Scalar(storage_type) => storage_type.hash(state),
483            Type::Opaque(opaque) => opaque.hash(state),
484            Type::Vector(intern, size) => {
485                intern.as_ref().hash(state);
486                size.hash(state);
487            }
488            Type::Semantic(semantic_type) => semantic_type.hash(state),
489            Type::Atomic(intern) => intern.as_ref().hash(state),
490            Type::Pointer(intern, addr_space) => {
491                intern.as_ref().hash(state);
492                addr_space.hash(state);
493            }
494            Type::Array(intern, size) => {
495                intern.as_ref().hash(state);
496                size.hash(state);
497            }
498            Type::DynamicArray(intern) => {
499                intern.as_ref().hash(state);
500            }
501            Type::Matrix(matrix_type) => {
502                matrix_type.hash(state);
503            }
504            Type::Aggregate(aggregate_kind) => {
505                aggregate_kind.hash(state);
506            }
507        }
508    }
509}
510
511pub type VectorSize = usize;
512
513impl Type {
514    pub fn intern(self) -> Intern<Type> {
515        Intern::new(self)
516    }
517
518    /// Fetch the elem of the item.
519    pub fn elem_type(&self) -> ElemType {
520        self.storage_type().elem_type()
521    }
522
523    /// Create a new item
524    pub fn new(storage: StorageType) -> Self {
525        Type::Scalar(storage)
526    }
527
528    pub fn scalar(elem: ElemType) -> Self {
529        Self::new(StorageType::Scalar(elem))
530    }
531
532    pub fn semantic(ty: SemanticType) -> Self {
533        Self::Semantic(ty)
534    }
535
536    pub fn atomic(ty: impl Into<Type>) -> Self {
537        Self::Atomic(ty.into().intern())
538    }
539
540    pub fn with_vector_size(self, vector_size: VectorSize) -> Self {
541        match self {
542            Type::Scalar(inner) if vector_size > 1 => {
543                Type::Vector(Type::new(inner).intern(), vector_size)
544            }
545            Type::Opaque(opaque) => Type::Opaque(opaque),
546            Type::Vector(inner, _) if vector_size <= 1 => *inner,
547            Type::Vector(inner, _) => Type::Vector(inner, vector_size),
548            Type::Atomic(inner) => Type::Atomic(inner.with_vector_size(vector_size).intern()),
549            Type::Pointer(inner, class) => {
550                Type::Pointer(inner.with_vector_size(vector_size).intern(), class)
551            }
552            Type::Array(inner, size) => {
553                Type::Array(inner.with_vector_size(vector_size).intern(), size)
554            }
555            Type::DynamicArray(inner) => {
556                Type::DynamicArray(inner.with_vector_size(vector_size).intern())
557            }
558            Type::Aggregate(AggregateKind::Ptr { inner_ty, meta }) => {
559                Type::Aggregate(AggregateKind::Ptr {
560                    inner_ty: inner_ty.with_vector_size(vector_size).intern(),
561                    meta,
562                })
563            }
564            this @ (Type::Scalar(_) | Type::Semantic(_) | Type::Matrix(_)) => this,
565        }
566    }
567
568    pub fn pointer(ty: impl Into<Type>, class: AddressSpace) -> Self {
569        Self::Pointer(ty.into().intern(), class)
570    }
571
572    pub fn array(ty: impl Into<Type>, size: usize) -> Self {
573        Self::Array(ty.into().intern(), size)
574    }
575
576    pub fn vector_size(&self) -> VectorSize {
577        match self {
578            Type::Scalar(_) => 1,
579            Type::Opaque(_) => 1,
580            Type::Vector(inner, vector_size) => inner.vector_size() * *vector_size,
581            Type::Array(inner, ..)
582            | Type::DynamicArray(inner, ..)
583            | Type::Atomic(inner)
584            | Type::Pointer(inner, _) => inner.vector_size(),
585            Type::Semantic(_) => 0,
586            Type::Matrix(_) => 1,
587            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.vector_size(),
588        }
589    }
590
591    pub fn array_size(&self) -> usize {
592        match self {
593            Type::Array(_, size) => *size,
594            Type::Scalar(_) => 1,
595            Type::Opaque(_) => 1,
596            Type::Vector(inner, _) | Type::Atomic(inner) | Type::Pointer(inner, _) => {
597                inner.array_size()
598            }
599            Type::Semantic(_) | Type::DynamicArray(..) => 0,
600            Type::Matrix(_) => 1,
601            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.array_size(),
602        }
603    }
604
605    pub fn align(&self) -> usize {
606        match self {
607            Type::Scalar(ty) => ty.size(),
608            Type::Opaque(opaque) => opaque.size(),
609            Type::Vector(ty, vector_size) => ty.size() * *vector_size,
610            Type::Atomic(inner) => inner.align(),
611            Type::Array(inner, _) => inner.align(),
612            Type::DynamicArray(inner, ..) => inner.align(),
613            // All platforms use at least conceptually 64-bit pointers
614            Type::Pointer(..) => align_of::<u64>(),
615            Type::Semantic(_) => 0,
616            Type::Matrix(mat) => mat.storage.size(),
617            Type::Aggregate(..) => panic!("Can't get size of opaque type `Aggregate`"),
618        }
619    }
620
621    pub fn size(&self) -> usize {
622        match self {
623            Type::Scalar(ty) => ty.size(),
624            Type::Opaque(opaque) => opaque.size(),
625            Type::Vector(ty, vector_size) => ty.size() * *vector_size,
626            Type::Atomic(inner) => inner.size(),
627            Type::Array(inner, size) => inner.size() * *size,
628            Type::DynamicArray(inner, ..) => inner.size(),
629            // All platforms use at least conceptually 64-bit pointers
630            Type::Pointer(..) => size_of::<u64>(),
631            Type::Semantic(_) => 0,
632            Type::Matrix(..) => panic!("Can't get size of opaque type `Matrix`"),
633            Type::Aggregate(..) => panic!("Can't get size of opaque type `Aggregate`"),
634        }
635    }
636
637    pub fn size_bits(&self) -> usize {
638        match self {
639            Type::Scalar(ty) => ty.size_bits(),
640            Type::Opaque(opaque) => opaque.size_bits(),
641            Type::Vector(ty, vector_size) => ty.size_bits() * *vector_size,
642            Type::Atomic(inner) => inner.size_bits(),
643            Type::Array(inner, ..) => inner.size_bits(),
644            Type::DynamicArray(inner, ..) => inner.size_bits(),
645            // All platforms use at least conceptually 64-bit pointers
646            Type::Pointer(..) => u64::BITS as usize,
647            Type::Semantic(_) => 0,
648            Type::Matrix(..) => panic!("Can't get size of opaque type `Matrix`"),
649            Type::Aggregate(..) => panic!("Can't get size of opaque type `Aggregate`"),
650        }
651    }
652
653    pub fn packing_factor(&self) -> usize {
654        match self {
655            Type::Scalar(ty) => ty.packing_factor(),
656            Type::Opaque(_) => 1,
657            Type::Vector(ty, _)
658            | Type::Atomic(ty)
659            | Type::Pointer(ty, _)
660            | Type::Array(ty, ..)
661            | Type::DynamicArray(ty, ..) => ty.packing_factor(),
662            Type::Semantic(_) => 1,
663            Type::Matrix(mat) => mat.storage.packing_factor(),
664            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.packing_factor(),
665        }
666    }
667
668    pub fn is_atomic(&self) -> bool {
669        match self {
670            Type::Semantic(_) | Type::Scalar(_) | Type::Matrix(_) | Type::Opaque(_) => false,
671            Type::Atomic(_) => true,
672            Type::Pointer(inner, _)
673            | Type::Vector(inner, _)
674            | Type::Array(inner, ..)
675            | Type::DynamicArray(inner, ..) => inner.is_atomic(),
676            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.is_atomic(),
677        }
678    }
679
680    pub fn is_ptr(&self) -> bool {
681        matches!(self, Type::Pointer(..))
682    }
683
684    pub fn is_int(&self) -> bool {
685        match self {
686            Type::Scalar(ty) => ty.is_int(),
687            Type::Semantic(_) | Type::Opaque(_) => false,
688            Type::Atomic(inner)
689            | Type::Pointer(inner, _)
690            | Type::Vector(inner, _)
691            | Type::Array(inner, ..)
692            | Type::DynamicArray(inner, ..) => inner.is_int(),
693            Type::Matrix(matrix_type) => matrix_type.storage.is_int(),
694            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.is_int(),
695        }
696    }
697
698    pub fn is_signed_int(&self) -> bool {
699        match self {
700            Type::Scalar(ty) => ty.is_signed_int(),
701            Type::Semantic(_) | Type::Opaque(_) => false,
702            Type::Atomic(inner)
703            | Type::Pointer(inner, _)
704            | Type::Vector(inner, _)
705            | Type::Array(inner, ..)
706            | Type::DynamicArray(inner, ..) => inner.is_signed_int(),
707            Type::Matrix(matrix_type) => matrix_type.storage.is_signed_int(),
708            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.is_signed_int(),
709        }
710    }
711
712    pub fn is_unsigned_int(&self) -> bool {
713        match self {
714            Type::Scalar(ty) => ty.is_unsigned_int(),
715            Type::Semantic(_) | Type::Opaque(_) => false,
716            Type::Atomic(inner)
717            | Type::Pointer(inner, _)
718            | Type::Vector(inner, _)
719            | Type::Array(inner, ..)
720            | Type::DynamicArray(inner, ..) => inner.is_unsigned_int(),
721            Type::Matrix(matrix_type) => matrix_type.storage.is_unsigned_int(),
722            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.is_unsigned_int(),
723        }
724    }
725
726    pub fn is_float(&self) -> bool {
727        match self {
728            Type::Scalar(ty) => ty.is_float(),
729            Type::Semantic(_) | Type::Opaque(_) => false,
730            Type::Atomic(inner)
731            | Type::Pointer(inner, _)
732            | Type::Vector(inner, _)
733            | Type::Array(inner, ..)
734            | Type::DynamicArray(inner, ..) => inner.is_float(),
735            Type::Matrix(matrix_type) => matrix_type.storage.is_float(),
736            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.is_float(),
737        }
738    }
739
740    pub fn is_bool(&self) -> bool {
741        match self {
742            Type::Scalar(ty) => ty.is_bool(),
743            Type::Semantic(_) | Type::Opaque(_) => false,
744            Type::Atomic(inner)
745            | Type::Pointer(inner, _)
746            | Type::Vector(inner, _)
747            | Type::Array(inner, ..)
748            | Type::DynamicArray(inner, ..) => inner.is_bool(),
749            Type::Matrix(matrix_type) => matrix_type.storage.is_bool(),
750            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.is_bool(),
751        }
752    }
753
754    pub fn storage_type(&self) -> StorageType {
755        match self {
756            Type::Scalar(ty) => *ty,
757            Type::Semantic(_) | Type::Opaque(_) => {
758                unimplemented!("Can't get storage for semantic type")
759            }
760            Type::Atomic(inner)
761            | Type::Pointer(inner, _)
762            | Type::Vector(inner, _)
763            | Type::Array(inner, ..)
764            | Type::DynamicArray(inner, ..) => inner.storage_type(),
765            Type::Matrix(matrix_type) => matrix_type.storage,
766            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.storage_type(),
767        }
768    }
769
770    pub fn as_scalar(&self) -> Self {
771        match self {
772            Type::Scalar(_) => *self,
773            Type::Vector(inner, _) => inner.as_scalar(),
774            Type::Atomic(inner) => Type::Atomic(inner.as_scalar().intern()),
775            Type::Pointer(inner, class) => Type::Pointer(inner.as_scalar().intern(), *class),
776            Type::Array(inner, size) => Type::Array(inner.as_scalar().intern(), *size),
777            Type::Opaque(opaque_type) => Type::Opaque(*opaque_type),
778            Type::Semantic(semantic_type) => Type::Semantic(*semantic_type),
779            Type::DynamicArray(inner) => Type::DynamicArray(inner.as_scalar().intern()),
780            Type::Matrix(matrix_type) => Type::Matrix(*matrix_type),
781            Type::Aggregate(aggregate_kind) => Type::Aggregate(*aggregate_kind),
782        }
783    }
784
785    /// Utility mainly for use in `cubecl-cpu`
786    pub fn scalar_value_type(&self) -> Self {
787        self.value_type().as_scalar()
788    }
789
790    pub fn is_semantic(&self) -> bool {
791        matches!(self, Type::Semantic(_))
792    }
793
794    pub fn constant(&self, value: ConstantValue) -> Value {
795        Value::constant(value, *self)
796    }
797
798    pub fn unwrap_ptr(&self) -> Type {
799        match self {
800            Type::Pointer(inner, _) => **inner,
801            other => *other,
802        }
803    }
804
805    pub fn address_space(&self) -> Option<AddressSpace> {
806        match self {
807            Type::Scalar(..)
808            | Type::Opaque(..)
809            | Type::Vector(..)
810            | Type::Semantic(..)
811            | Type::Atomic(..)
812            | Type::Matrix(..)
813            | Type::Array(..)
814            | Type::DynamicArray(..)
815            | Type::Aggregate(..) => None,
816            Type::Pointer(.., address_space) => Some(*address_space),
817        }
818    }
819
820    pub fn value_type(&self) -> Type {
821        match self {
822            Type::Pointer(inner, _) | Type::Array(inner, ..) | Type::DynamicArray(inner, ..) => {
823                inner.value_type()
824            }
825            this @ (Type::Scalar(..)
826            | Type::Vector(..)
827            | Type::Semantic(..)
828            | Type::Atomic(..)
829            | Type::Matrix(..)
830            | Type::Opaque(_)) => *this,
831            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.value_type(),
832        }
833    }
834
835    pub fn is_array_like(&self) -> bool {
836        matches!(self, Type::Array(..) | Type::DynamicArray(..))
837    }
838
839    /// Whether a type is destructurable. This implies that
840    /// * it does not have dynamic field offsets (i.e. `Array`)
841    /// * it can exist in registers (i.e. no `Barrier` or `Atomic`)
842    pub fn is_destructurable(&self) -> bool {
843        match self {
844            Type::Scalar(..) | Type::Vector(..) => true,
845            // Should be `true`, but semantics are too dodgy right now. They're registers, but CUDA
846            // wmma uses pointers for all matrix ops. So we need to keep them in memory for now.
847            Type::Matrix(..) => false,
848            Type::Pointer(..)
849            | Type::Array(..)
850            | Type::DynamicArray(..)
851            | Type::Semantic(..)
852            | Type::Atomic(..)
853            | Type::Aggregate(..) => false,
854            Type::Opaque(opaque) => match opaque {
855                // Can only exist in memory
856                OpaqueType::Barrier(..) | OpaqueType::TensorMap => false,
857                OpaqueType::BarrierToken(..) => true,
858            },
859        }
860    }
861
862    pub fn is_value(&self) -> bool {
863        self.value_type() == *self
864    }
865}
866
867impl Display for Type {
868    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
869        match self {
870            Type::Semantic(ty) => write!(f, "{ty}"),
871            Type::Opaque(ty) => write!(f, "{ty}"),
872            Type::Scalar(ty) => write!(f, "{ty}"),
873            Type::Vector(ty, vector_size) => write!(f, "vector<{ty}, {vector_size}>"),
874            Type::Atomic(ty) => write!(f, "atomic<{ty}>"),
875            Type::Pointer(ty, addr_space) => write!(f, "ptr<{ty}, {addr_space}>"),
876            Type::Array(ty, size) => write!(f, "array<{ty}, {size}>"),
877            Type::DynamicArray(ty) => write!(f, "array<{ty}>"),
878            Type::Matrix(mat) => write!(
879                f,
880                "matrix<{}, m{}xn{}xk{}x{}, {}, {}>",
881                mat.ident, mat.m, mat.n, mat.k, mat.storage, mat.layout, mat.storage
882            ),
883            Type::Aggregate(aggregate_kind) => write!(f, "{aggregate_kind}"),
884        }
885    }
886}
887
888impl Display for StorageType {
889    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
890        match self {
891            StorageType::Scalar(ty) => write!(f, "{ty}"),
892            StorageType::Packed(ty, factor) => write!(f, "packed<{ty}, {factor}>"),
893        }
894    }
895}
896
897impl Display for ElemType {
898    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
899        match self {
900            Self::Float(kind) => match kind {
901                FloatKind::E2M1 => f.write_str("e2m1"),
902                FloatKind::E2M3 => f.write_str("e2m3"),
903                FloatKind::E3M2 => f.write_str("e3m2"),
904                FloatKind::E4M3 => f.write_str("e4m3"),
905                FloatKind::E5M2 => f.write_str("e5m2"),
906                FloatKind::UE8M0 => f.write_str("ue8m0"),
907                FloatKind::F16 => f.write_str("f16"),
908                FloatKind::BF16 => f.write_str("bf16"),
909                FloatKind::Flex32 => f.write_str("flex32"),
910                FloatKind::TF32 => f.write_str("tf32"),
911                FloatKind::F32 => f.write_str("f32"),
912                FloatKind::F64 => f.write_str("f64"),
913            },
914            Self::Int(kind) => match kind {
915                IntKind::I8 => f.write_str("i8"),
916                IntKind::I16 => f.write_str("i16"),
917                IntKind::I32 => f.write_str("i32"),
918                IntKind::I64 => f.write_str("i64"),
919            },
920            Self::UInt(kind) => match kind {
921                UIntKind::U8 => f.write_str("u8"),
922                UIntKind::U16 => f.write_str("u16"),
923                UIntKind::U32 => f.write_str("u32"),
924                UIntKind::U64 => f.write_str("u64"),
925            },
926            Self::Bool => f.write_str("bool"),
927        }
928    }
929}
930
931impl Display for SemanticType {
932    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
933        match self {
934            SemanticType::TensorLayout(dims, _) => write!(f, "tensor_layout<{dims}>"),
935            SemanticType::TensorView(dims, has_dims, permutation) => {
936                write!(
937                    f,
938                    "tensor_layout<{:?}, has_dims: {has_dims}>",
939                    &permutation[..*dims]
940                )
941            }
942        }
943    }
944}
945
946impl Display for OpaqueType {
947    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
948        match self {
949            OpaqueType::Barrier(level) => write!(f, "barrier<{level}>"),
950            OpaqueType::BarrierToken(level) => write!(f, "barrier_token<{level}>"),
951            OpaqueType::TensorMap => f.write_str("tensor_map"),
952        }
953    }
954}
955
956impl Display for AddressSpace {
957    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
958        match self {
959            AddressSpace::Global(id) => write!(f, "global<{id}>"),
960            AddressSpace::Shared => write!(f, "shared"),
961            AddressSpace::Local => f.write_str("local"),
962        }
963    }
964}
965
966#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
967#[derive(Debug, Clone, Copy, PartialEq, Eq, TypeHash, PartialOrd, Ord, Display)]
968pub enum AggregateKind {
969    #[display("ptr<{meta}, {inner_ty}>")]
970    Ptr {
971        inner_ty: Intern<Type>,
972        meta: MetadataKind,
973    },
974}
975
976/// Hashed by value rather than derived, for the same reason as [`Type`]: an [`Intern`] hashes the
977/// pointer it holds, which moves between runs.
978impl core::hash::Hash for AggregateKind {
979    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
980        core::mem::discriminant(self).hash(state);
981        match self {
982            AggregateKind::Ptr { inner_ty, meta } => {
983                inner_ty.as_ref().hash(state);
984                meta.hash(state);
985            }
986        }
987    }
988}
989
990impl AggregateKind {
991    pub fn ptr(inner_ty: Type, meta: MetadataKind) -> Self {
992        AggregateKind::Ptr {
993            inner_ty: inner_ty.intern(),
994            meta,
995        }
996    }
997}
998
999#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
1000#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, TypeHash, PartialOrd, Ord, Display)]
1001pub enum MetadataKind {
1002    /// Slice metadata (offset and length)
1003    #[display("slice")]
1004    Slice,
1005    /// Bounds check (in bounds)
1006    #[display("bounds_checked")]
1007    BoundsCheck,
1008}
1009
1010pub struct BoundsCheckMetadata;
1011impl BoundsCheckMetadata {
1012    pub const POINTER: usize = 0;
1013    pub const IS_IN_BOUNDS: usize = 1;
1014}
1015
1016pub struct SliceMetadata;
1017impl SliceMetadata {
1018    pub const LIST: usize = 0;
1019    pub const OFFSET: usize = 1;
1020    pub const LENGTH: usize = 2;
1021}
1022
1023impl From<e2m1x2> for Value {
1024    fn from(_value: e2m1x2) -> Self {
1025        unimplemented!("Can't currently construct e2m1x2")
1026    }
1027}
1028
1029impl From<e2m3> for Value {
1030    fn from(_value: e2m3) -> Self {
1031        unimplemented!("Can't currently construct fp6")
1032    }
1033}
1034
1035impl From<e3m2> for Value {
1036    fn from(_value: e3m2) -> Self {
1037        unimplemented!("Can't currently construct fp6")
1038    }
1039}
1040
1041impl From<i8> for ConstantValue {
1042    fn from(value: i8) -> Self {
1043        ConstantValue::Int(value as i64)
1044    }
1045}
1046
1047impl From<i16> for ConstantValue {
1048    fn from(value: i16) -> Self {
1049        ConstantValue::Int(value as i64)
1050    }
1051}
1052
1053impl From<i32> for ConstantValue {
1054    fn from(value: i32) -> Self {
1055        ConstantValue::Int(value as i64)
1056    }
1057}
1058
1059impl From<isize> for ConstantValue {
1060    fn from(value: isize) -> Self {
1061        ConstantValue::Int(value as i64)
1062    }
1063}
1064
1065impl From<u8> for ConstantValue {
1066    fn from(value: u8) -> Self {
1067        ConstantValue::UInt(value as u64)
1068    }
1069}
1070
1071impl From<u16> for ConstantValue {
1072    fn from(value: u16) -> Self {
1073        ConstantValue::UInt(value as u64)
1074    }
1075}
1076
1077impl From<u32> for ConstantValue {
1078    fn from(value: u32) -> Self {
1079        ConstantValue::UInt(value as u64)
1080    }
1081}
1082
1083impl From<usize> for ConstantValue {
1084    fn from(value: usize) -> Self {
1085        ConstantValue::UInt(value as u64)
1086    }
1087}
1088
1089impl From<e2m1> for ConstantValue {
1090    fn from(value: e2m1) -> Self {
1091        ConstantValue::Float(value.to_f64())
1092    }
1093}
1094
1095impl From<e4m3> for ConstantValue {
1096    fn from(value: e4m3) -> Self {
1097        ConstantValue::Float(value.to_f64())
1098    }
1099}
1100
1101impl From<e5m2> for ConstantValue {
1102    fn from(value: e5m2) -> Self {
1103        ConstantValue::Float(value.to_f64())
1104    }
1105}
1106
1107impl From<ue8m0> for ConstantValue {
1108    fn from(value: ue8m0) -> Self {
1109        ConstantValue::Float(value.to_f64())
1110    }
1111}
1112
1113impl From<half::f16> for ConstantValue {
1114    fn from(value: half::f16) -> Self {
1115        ConstantValue::Float(value.to_f64())
1116    }
1117}
1118
1119impl From<half::bf16> for ConstantValue {
1120    fn from(value: half::bf16) -> Self {
1121        ConstantValue::Float(value.to_f64())
1122    }
1123}
1124
1125impl From<flex32> for ConstantValue {
1126    fn from(value: flex32) -> Self {
1127        ConstantValue::Float(value.to_f64())
1128    }
1129}
1130
1131impl From<tf32> for ConstantValue {
1132    fn from(value: tf32) -> Self {
1133        ConstantValue::Float(value.to_f64())
1134    }
1135}
1136
1137impl From<f32> for ConstantValue {
1138    fn from(value: f32) -> Self {
1139        ConstantValue::Float(value as f64)
1140    }
1141}
1142
1143macro_rules! impl_into_value {
1144    ($($ty: ty => $kind: path,)*) => {
1145        $(
1146            impl From<$ty> for Value {
1147                fn from(value: $ty) -> Self {
1148                    Value {kind: ValueKind::Constant(value.into()), ty: $kind.into()}
1149                }
1150            }
1151        )*
1152    };
1153}
1154
1155impl_into_value!(
1156    bool => ElemType::Bool,
1157
1158    i8 => IntKind::I8,
1159    i16 => IntKind::I16,
1160    i32 => IntKind::I32,
1161    i64 => IntKind::I64,
1162
1163    u8 => UIntKind::U8,
1164    u16 => UIntKind::U16,
1165    u32 => UIntKind::U32,
1166    u64 => UIntKind::U64,
1167
1168    e2m1 => FloatKind::E2M1,
1169    e4m3 => FloatKind::E4M3,
1170    e5m2 => FloatKind::E5M2,
1171    ue8m0 => FloatKind::UE8M0,
1172    f16 => FloatKind::F16,
1173    bf16 => FloatKind::BF16,
1174    f32 => FloatKind::F32,
1175    flex32 => FloatKind::Flex32,
1176    tf32 => FloatKind::TF32,
1177    f64 => FloatKind::F64,
1178
1179    usize => UIntKind::U32,
1180    isize => IntKind::I32,
1181);
1182
1183#[cfg(test)]
1184mod tests {
1185    use super::*;
1186    use core::hash::{Hash, Hasher};
1187
1188    fn hash(ty: Type) -> u64 {
1189        let mut hasher = fnv::FnvHasher::default();
1190        ty.hash(&mut hasher);
1191        hasher.finish()
1192    }
1193
1194    /// `Intern` hashes the pointer it holds, which moves between runs. Every arm that holds one
1195    /// has to reach through it, or a persistent cache keyed on the IR never hits.
1196    #[test]
1197    fn interned_types_hash_by_value() {
1198        let f32_ty = Type::scalar(ElemType::Float(FloatKind::F32));
1199
1200        // Distinct `Intern` allocations of an equal type must agree.
1201        assert_eq!(
1202            hash(Type::Pointer(f32_ty.intern(), AddressSpace::Local)),
1203            hash(Type::Pointer(f32_ty.intern(), AddressSpace::Local))
1204        );
1205        assert_eq!(
1206            hash(Type::Aggregate(AggregateKind::ptr(
1207                f32_ty,
1208                MetadataKind::Slice
1209            ))),
1210            hash(Type::Aggregate(AggregateKind::ptr(
1211                f32_ty,
1212                MetadataKind::Slice
1213            )))
1214        );
1215    }
1216
1217    #[test]
1218    fn vector_size_is_part_of_the_hash() {
1219        let f32_ty = Type::scalar(ElemType::Float(FloatKind::F32));
1220
1221        assert_ne!(
1222            hash(Type::Vector(f32_ty.intern(), 2)),
1223            hash(Type::Vector(f32_ty.intern(), 4))
1224        );
1225    }
1226
1227    #[test]
1228    fn aggregate_inner_type_is_part_of_the_hash() {
1229        let f32_ty = Type::scalar(ElemType::Float(FloatKind::F32));
1230        let u32_ty = Type::scalar(ElemType::UInt(UIntKind::U32));
1231
1232        assert_ne!(
1233            hash(Type::Aggregate(AggregateKind::ptr(
1234                f32_ty,
1235                MetadataKind::Slice
1236            ))),
1237            hash(Type::Aggregate(AggregateKind::ptr(
1238                u32_ty,
1239                MetadataKind::Slice
1240            )))
1241        );
1242    }
1243}
cubecl_ir/type.rs

cubecl_ir/
type.rs
