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::UE8M0),
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, _) => intern.as_ref().hash(state),
485            Type::Semantic(semantic_type) => semantic_type.hash(state),
486            Type::Atomic(intern) => intern.as_ref().hash(state),
487            Type::Pointer(intern, addr_space) => {
488                intern.as_ref().hash(state);
489                addr_space.hash(state);
490            }
491            Type::Array(intern, size) => {
492                intern.as_ref().hash(state);
493                size.hash(state);
494            }
495            Type::DynamicArray(intern) => {
496                intern.as_ref().hash(state);
497            }
498            Type::Matrix(matrix_type) => {
499                matrix_type.hash(state);
500            }
501            Type::Aggregate(aggregate_kind) => {
502                aggregate_kind.hash(state);
503            }
504        }
505    }
506}
507
508pub type VectorSize = usize;
509
510impl Type {
511    pub fn intern(self) -> Intern<Type> {
512        Intern::new(self)
513    }
514
515    /// Fetch the elem of the item.
516    pub fn elem_type(&self) -> ElemType {
517        self.storage_type().elem_type()
518    }
519
520    /// Create a new item
521    pub fn new(storage: StorageType) -> Self {
522        Type::Scalar(storage)
523    }
524
525    pub fn scalar(elem: ElemType) -> Self {
526        Self::new(StorageType::Scalar(elem))
527    }
528
529    pub fn semantic(ty: SemanticType) -> Self {
530        Self::Semantic(ty)
531    }
532
533    pub fn atomic(ty: impl Into<Type>) -> Self {
534        Self::Atomic(ty.into().intern())
535    }
536
537    pub fn with_vector_size(self, vector_size: VectorSize) -> Self {
538        match self {
539            Type::Scalar(inner) if vector_size > 1 => {
540                Type::Vector(Type::new(inner).intern(), vector_size)
541            }
542            Type::Opaque(opaque) => Type::Opaque(opaque),
543            Type::Vector(inner, _) if vector_size <= 1 => *inner,
544            Type::Vector(inner, _) => Type::Vector(inner, vector_size),
545            Type::Atomic(inner) => Type::Atomic(inner.with_vector_size(vector_size).intern()),
546            Type::Pointer(inner, class) => {
547                Type::Pointer(inner.with_vector_size(vector_size).intern(), class)
548            }
549            Type::Array(inner, size) => {
550                Type::Array(inner.with_vector_size(vector_size).intern(), size)
551            }
552            Type::DynamicArray(inner) => {
553                Type::DynamicArray(inner.with_vector_size(vector_size).intern())
554            }
555            Type::Aggregate(AggregateKind::Ptr { inner_ty, meta }) => {
556                Type::Aggregate(AggregateKind::Ptr {
557                    inner_ty: inner_ty.with_vector_size(vector_size).intern(),
558                    meta,
559                })
560            }
561            this @ (Type::Scalar(_) | Type::Semantic(_) | Type::Matrix(_)) => this,
562        }
563    }
564
565    pub fn pointer(ty: impl Into<Type>, class: AddressSpace) -> Self {
566        Self::Pointer(ty.into().intern(), class)
567    }
568
569    pub fn array(ty: impl Into<Type>, size: usize) -> Self {
570        Self::Array(ty.into().intern(), size)
571    }
572
573    pub fn vector_size(&self) -> VectorSize {
574        match self {
575            Type::Scalar(_) => 1,
576            Type::Opaque(_) => 1,
577            Type::Vector(inner, vector_size) => inner.vector_size() * *vector_size,
578            Type::Array(inner, ..)
579            | Type::DynamicArray(inner, ..)
580            | Type::Atomic(inner)
581            | Type::Pointer(inner, _) => inner.vector_size(),
582            Type::Semantic(_) => 0,
583            Type::Matrix(_) => 1,
584            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.vector_size(),
585        }
586    }
587
588    pub fn array_size(&self) -> usize {
589        match self {
590            Type::Array(_, size) => *size,
591            Type::Scalar(_) => 1,
592            Type::Opaque(_) => 1,
593            Type::Vector(inner, _) | Type::Atomic(inner) | Type::Pointer(inner, _) => {
594                inner.array_size()
595            }
596            Type::Semantic(_) | Type::DynamicArray(..) => 0,
597            Type::Matrix(_) => 1,
598            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.array_size(),
599        }
600    }
601
602    pub fn align(&self) -> usize {
603        match self {
604            Type::Scalar(ty) => ty.size(),
605            Type::Opaque(opaque) => opaque.size(),
606            Type::Vector(ty, vector_size) => ty.size() * *vector_size,
607            Type::Atomic(inner) => inner.align(),
608            Type::Array(inner, _) => inner.align(),
609            Type::DynamicArray(inner, ..) => inner.align(),
610            // All platforms use at least conceptually 64-bit pointers
611            Type::Pointer(..) => align_of::<u64>(),
612            Type::Semantic(_) => 0,
613            Type::Matrix(mat) => mat.storage.size(),
614            Type::Aggregate(..) => panic!("Can't get size of opaque type `Aggregate`"),
615        }
616    }
617
618    pub fn size(&self) -> usize {
619        match self {
620            Type::Scalar(ty) => ty.size(),
621            Type::Opaque(opaque) => opaque.size(),
622            Type::Vector(ty, vector_size) => ty.size() * *vector_size,
623            Type::Atomic(inner) => inner.size(),
624            Type::Array(inner, size) => inner.size() * *size,
625            Type::DynamicArray(inner, ..) => inner.size(),
626            // All platforms use at least conceptually 64-bit pointers
627            Type::Pointer(..) => size_of::<u64>(),
628            Type::Semantic(_) => 0,
629            Type::Matrix(..) => panic!("Can't get size of opaque type `Matrix`"),
630            Type::Aggregate(..) => panic!("Can't get size of opaque type `Aggregate`"),
631        }
632    }
633
634    pub fn size_bits(&self) -> usize {
635        match self {
636            Type::Scalar(ty) => ty.size_bits(),
637            Type::Opaque(opaque) => opaque.size_bits(),
638            Type::Vector(ty, vector_size) => ty.size_bits() * *vector_size,
639            Type::Atomic(inner) => inner.size_bits(),
640            Type::Array(inner, ..) => inner.size_bits(),
641            Type::DynamicArray(inner, ..) => inner.size_bits(),
642            // All platforms use at least conceptually 64-bit pointers
643            Type::Pointer(..) => u64::BITS as usize,
644            Type::Semantic(_) => 0,
645            Type::Matrix(..) => panic!("Can't get size of opaque type `Matrix`"),
646            Type::Aggregate(..) => panic!("Can't get size of opaque type `Aggregate`"),
647        }
648    }
649
650    pub fn packing_factor(&self) -> usize {
651        match self {
652            Type::Scalar(ty) => ty.packing_factor(),
653            Type::Opaque(_) => 1,
654            Type::Vector(ty, _)
655            | Type::Atomic(ty)
656            | Type::Pointer(ty, _)
657            | Type::Array(ty, ..)
658            | Type::DynamicArray(ty, ..) => ty.packing_factor(),
659            Type::Semantic(_) => 1,
660            Type::Matrix(mat) => mat.storage.packing_factor(),
661            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.packing_factor(),
662        }
663    }
664
665    pub fn is_atomic(&self) -> bool {
666        match self {
667            Type::Semantic(_) | Type::Scalar(_) | Type::Matrix(_) | Type::Opaque(_) => false,
668            Type::Atomic(_) => true,
669            Type::Pointer(inner, _)
670            | Type::Vector(inner, _)
671            | Type::Array(inner, ..)
672            | Type::DynamicArray(inner, ..) => inner.is_atomic(),
673            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.is_atomic(),
674        }
675    }
676
677    pub fn is_ptr(&self) -> bool {
678        matches!(self, Type::Pointer(..))
679    }
680
681    pub fn is_int(&self) -> bool {
682        match self {
683            Type::Scalar(ty) => ty.is_int(),
684            Type::Semantic(_) | Type::Opaque(_) => false,
685            Type::Atomic(inner)
686            | Type::Pointer(inner, _)
687            | Type::Vector(inner, _)
688            | Type::Array(inner, ..)
689            | Type::DynamicArray(inner, ..) => inner.is_int(),
690            Type::Matrix(matrix_type) => matrix_type.storage.is_int(),
691            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.is_int(),
692        }
693    }
694
695    pub fn is_signed_int(&self) -> bool {
696        match self {
697            Type::Scalar(ty) => ty.is_signed_int(),
698            Type::Semantic(_) | Type::Opaque(_) => false,
699            Type::Atomic(inner)
700            | Type::Pointer(inner, _)
701            | Type::Vector(inner, _)
702            | Type::Array(inner, ..)
703            | Type::DynamicArray(inner, ..) => inner.is_signed_int(),
704            Type::Matrix(matrix_type) => matrix_type.storage.is_signed_int(),
705            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.is_signed_int(),
706        }
707    }
708
709    pub fn is_unsigned_int(&self) -> bool {
710        match self {
711            Type::Scalar(ty) => ty.is_unsigned_int(),
712            Type::Semantic(_) | Type::Opaque(_) => false,
713            Type::Atomic(inner)
714            | Type::Pointer(inner, _)
715            | Type::Vector(inner, _)
716            | Type::Array(inner, ..)
717            | Type::DynamicArray(inner, ..) => inner.is_unsigned_int(),
718            Type::Matrix(matrix_type) => matrix_type.storage.is_unsigned_int(),
719            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.is_unsigned_int(),
720        }
721    }
722
723    pub fn is_float(&self) -> bool {
724        match self {
725            Type::Scalar(ty) => ty.is_float(),
726            Type::Semantic(_) | Type::Opaque(_) => false,
727            Type::Atomic(inner)
728            | Type::Pointer(inner, _)
729            | Type::Vector(inner, _)
730            | Type::Array(inner, ..)
731            | Type::DynamicArray(inner, ..) => inner.is_float(),
732            Type::Matrix(matrix_type) => matrix_type.storage.is_float(),
733            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.is_float(),
734        }
735    }
736
737    pub fn is_bool(&self) -> bool {
738        match self {
739            Type::Scalar(ty) => ty.is_bool(),
740            Type::Semantic(_) | Type::Opaque(_) => false,
741            Type::Atomic(inner)
742            | Type::Pointer(inner, _)
743            | Type::Vector(inner, _)
744            | Type::Array(inner, ..)
745            | Type::DynamicArray(inner, ..) => inner.is_bool(),
746            Type::Matrix(matrix_type) => matrix_type.storage.is_bool(),
747            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.is_bool(),
748        }
749    }
750
751    pub fn storage_type(&self) -> StorageType {
752        match self {
753            Type::Scalar(ty) => *ty,
754            Type::Semantic(_) | Type::Opaque(_) => {
755                unimplemented!("Can't get storage for semantic type")
756            }
757            Type::Atomic(inner)
758            | Type::Pointer(inner, _)
759            | Type::Vector(inner, _)
760            | Type::Array(inner, ..)
761            | Type::DynamicArray(inner, ..) => inner.storage_type(),
762            Type::Matrix(matrix_type) => matrix_type.storage,
763            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.storage_type(),
764        }
765    }
766
767    pub fn as_scalar(&self) -> Self {
768        match self {
769            Type::Scalar(_) => *self,
770            Type::Vector(inner, _) => inner.as_scalar(),
771            Type::Atomic(inner) => Type::Atomic(inner.as_scalar().intern()),
772            Type::Pointer(inner, class) => Type::Pointer(inner.as_scalar().intern(), *class),
773            Type::Array(inner, size) => Type::Array(inner.as_scalar().intern(), *size),
774            Type::Opaque(opaque_type) => Type::Opaque(*opaque_type),
775            Type::Semantic(semantic_type) => Type::Semantic(*semantic_type),
776            Type::DynamicArray(inner) => Type::DynamicArray(inner.as_scalar().intern()),
777            Type::Matrix(matrix_type) => Type::Matrix(*matrix_type),
778            Type::Aggregate(aggregate_kind) => Type::Aggregate(*aggregate_kind),
779        }
780    }
781
782    /// Utility mainly for use in `cubecl-cpu`
783    pub fn scalar_value_type(&self) -> Self {
784        self.value_type().as_scalar()
785    }
786
787    pub fn is_semantic(&self) -> bool {
788        matches!(self, Type::Semantic(_))
789    }
790
791    pub fn constant(&self, value: ConstantValue) -> Value {
792        Value::constant(value, *self)
793    }
794
795    pub fn unwrap_ptr(&self) -> Type {
796        match self {
797            Type::Pointer(inner, _) => **inner,
798            other => *other,
799        }
800    }
801
802    pub fn address_space(&self) -> Option<AddressSpace> {
803        match self {
804            Type::Scalar(..)
805            | Type::Opaque(..)
806            | Type::Vector(..)
807            | Type::Semantic(..)
808            | Type::Atomic(..)
809            | Type::Matrix(..)
810            | Type::Array(..)
811            | Type::DynamicArray(..)
812            | Type::Aggregate(..) => None,
813            Type::Pointer(.., address_space) => Some(*address_space),
814        }
815    }
816
817    pub fn value_type(&self) -> Type {
818        match self {
819            Type::Pointer(inner, _) | Type::Array(inner, ..) | Type::DynamicArray(inner, ..) => {
820                inner.value_type()
821            }
822            this @ (Type::Scalar(..)
823            | Type::Vector(..)
824            | Type::Semantic(..)
825            | Type::Atomic(..)
826            | Type::Matrix(..)
827            | Type::Opaque(_)) => *this,
828            Type::Aggregate(AggregateKind::Ptr { inner_ty, .. }) => inner_ty.value_type(),
829        }
830    }
831
832    pub fn is_array_like(&self) -> bool {
833        matches!(self, Type::Array(..) | Type::DynamicArray(..))
834    }
835
836    /// Whether a type is destructurable. This implies that
837    /// * it does not have dynamic field offsets (i.e. `Array`)
838    /// * it can exist in registers (i.e. no `Barrier` or `Atomic`)
839    pub fn is_destructurable(&self) -> bool {
840        match self {
841            Type::Scalar(..) | Type::Vector(..) => true,
842            // Should be `true`, but semantics are too dodgy right now. They're registers, but CUDA
843            // wmma uses pointers for all matrix ops. So we need to keep them in memory for now.
844            Type::Matrix(..) => false,
845            Type::Pointer(..)
846            | Type::Array(..)
847            | Type::DynamicArray(..)
848            | Type::Semantic(..)
849            | Type::Atomic(..)
850            | Type::Aggregate(..) => false,
851            Type::Opaque(opaque) => match opaque {
852                // Can only exist in memory
853                OpaqueType::Barrier(..) | OpaqueType::TensorMap => false,
854                OpaqueType::BarrierToken(..) => true,
855            },
856        }
857    }
858
859    pub fn is_value(&self) -> bool {
860        self.value_type() == *self
861    }
862}
863
864impl Display for Type {
865    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
866        match self {
867            Type::Semantic(ty) => write!(f, "{ty}"),
868            Type::Opaque(ty) => write!(f, "{ty}"),
869            Type::Scalar(ty) => write!(f, "{ty}"),
870            Type::Vector(ty, vector_size) => write!(f, "vector<{ty}, {vector_size}>"),
871            Type::Atomic(ty) => write!(f, "atomic<{ty}>"),
872            Type::Pointer(ty, addr_space) => write!(f, "ptr<{ty}, {addr_space}>"),
873            Type::Array(ty, size) => write!(f, "array<{ty}, {size}>"),
874            Type::DynamicArray(ty) => write!(f, "array<{ty}>"),
875            Type::Matrix(mat) => write!(
876                f,
877                "matrix<{}, m{}xn{}xk{}x{}, {}, {}>",
878                mat.ident, mat.m, mat.n, mat.k, mat.storage, mat.layout, mat.storage
879            ),
880            Type::Aggregate(aggregate_kind) => write!(f, "{aggregate_kind}"),
881        }
882    }
883}
884
885impl Display for StorageType {
886    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
887        match self {
888            StorageType::Scalar(ty) => write!(f, "{ty}"),
889            StorageType::Packed(ty, factor) => write!(f, "packed<{ty}, {factor}>"),
890        }
891    }
892}
893
894impl Display for ElemType {
895    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
896        match self {
897            Self::Float(kind) => match kind {
898                FloatKind::E2M1 => f.write_str("e2m1"),
899                FloatKind::E2M3 => f.write_str("e2m3"),
900                FloatKind::E3M2 => f.write_str("e3m2"),
901                FloatKind::E4M3 => f.write_str("e4m3"),
902                FloatKind::E5M2 => f.write_str("e5m2"),
903                FloatKind::UE8M0 => f.write_str("ue8m0"),
904                FloatKind::F16 => f.write_str("f16"),
905                FloatKind::BF16 => f.write_str("bf16"),
906                FloatKind::Flex32 => f.write_str("flex32"),
907                FloatKind::TF32 => f.write_str("tf32"),
908                FloatKind::F32 => f.write_str("f32"),
909                FloatKind::F64 => f.write_str("f64"),
910            },
911            Self::Int(kind) => match kind {
912                IntKind::I8 => f.write_str("i8"),
913                IntKind::I16 => f.write_str("i16"),
914                IntKind::I32 => f.write_str("i32"),
915                IntKind::I64 => f.write_str("i64"),
916            },
917            Self::UInt(kind) => match kind {
918                UIntKind::U8 => f.write_str("u8"),
919                UIntKind::U16 => f.write_str("u16"),
920                UIntKind::U32 => f.write_str("u32"),
921                UIntKind::U64 => f.write_str("u64"),
922            },
923            Self::Bool => f.write_str("bool"),
924        }
925    }
926}
927
928impl Display for SemanticType {
929    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
930        match self {
931            SemanticType::TensorLayout(dims, _) => write!(f, "tensor_layout<{dims}>"),
932            SemanticType::TensorView(dims, has_dims, permutation) => {
933                write!(
934                    f,
935                    "tensor_layout<{:?}, has_dims: {has_dims}>",
936                    &permutation[..*dims]
937                )
938            }
939        }
940    }
941}
942
943impl Display for OpaqueType {
944    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
945        match self {
946            OpaqueType::Barrier(level) => write!(f, "barrier<{level}>"),
947            OpaqueType::BarrierToken(level) => write!(f, "barrier_token<{level}>"),
948            OpaqueType::TensorMap => f.write_str("tensor_map"),
949        }
950    }
951}
952
953impl Display for AddressSpace {
954    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
955        match self {
956            AddressSpace::Global(id) => write!(f, "global<{id}>"),
957            AddressSpace::Shared => write!(f, "shared"),
958            AddressSpace::Local => f.write_str("local"),
959        }
960    }
961}
962
963#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
964#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, TypeHash, PartialOrd, Ord, Display)]
965pub enum AggregateKind {
966    #[display("ptr<{meta}, {inner_ty}>")]
967    Ptr {
968        inner_ty: Intern<Type>,
969        meta: MetadataKind,
970    },
971}
972
973impl AggregateKind {
974    pub fn ptr(inner_ty: Type, meta: MetadataKind) -> Self {
975        AggregateKind::Ptr {
976            inner_ty: inner_ty.intern(),
977            meta,
978        }
979    }
980}
981
982#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
983#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, TypeHash, PartialOrd, Ord, Display)]
984pub enum MetadataKind {
985    /// Slice metadata (offset and length)
986    #[display("slice")]
987    Slice,
988    /// Bounds check (in bounds)
989    #[display("bounds_checked")]
990    BoundsCheck,
991}
992
993pub struct BoundsCheckMetadata;
994impl BoundsCheckMetadata {
995    pub const POINTER: usize = 0;
996    pub const IS_IN_BOUNDS: usize = 1;
997}
998
999pub struct SliceMetadata;
1000impl SliceMetadata {
1001    pub const LIST: usize = 0;
1002    pub const OFFSET: usize = 1;
1003    pub const LENGTH: usize = 2;
1004}
1005
1006impl From<e2m1x2> for Value {
1007    fn from(_value: e2m1x2) -> Self {
1008        unimplemented!("Can't currently construct e2m1x2")
1009    }
1010}
1011
1012impl From<e2m3> for Value {
1013    fn from(_value: e2m3) -> Self {
1014        unimplemented!("Can't currently construct fp6")
1015    }
1016}
1017
1018impl From<e3m2> for Value {
1019    fn from(_value: e3m2) -> Self {
1020        unimplemented!("Can't currently construct fp6")
1021    }
1022}
1023
1024impl From<i8> for ConstantValue {
1025    fn from(value: i8) -> Self {
1026        ConstantValue::Int(value as i64)
1027    }
1028}
1029
1030impl From<i16> for ConstantValue {
1031    fn from(value: i16) -> Self {
1032        ConstantValue::Int(value as i64)
1033    }
1034}
1035
1036impl From<i32> for ConstantValue {
1037    fn from(value: i32) -> Self {
1038        ConstantValue::Int(value as i64)
1039    }
1040}
1041
1042impl From<isize> for ConstantValue {
1043    fn from(value: isize) -> Self {
1044        ConstantValue::Int(value as i64)
1045    }
1046}
1047
1048impl From<u8> for ConstantValue {
1049    fn from(value: u8) -> Self {
1050        ConstantValue::UInt(value as u64)
1051    }
1052}
1053
1054impl From<u16> for ConstantValue {
1055    fn from(value: u16) -> Self {
1056        ConstantValue::UInt(value as u64)
1057    }
1058}
1059
1060impl From<u32> for ConstantValue {
1061    fn from(value: u32) -> Self {
1062        ConstantValue::UInt(value as u64)
1063    }
1064}
1065
1066impl From<usize> for ConstantValue {
1067    fn from(value: usize) -> Self {
1068        ConstantValue::UInt(value as u64)
1069    }
1070}
1071
1072impl From<e2m1> for ConstantValue {
1073    fn from(value: e2m1) -> Self {
1074        ConstantValue::Float(value.to_f64())
1075    }
1076}
1077
1078impl From<e4m3> for ConstantValue {
1079    fn from(value: e4m3) -> Self {
1080        ConstantValue::Float(value.to_f64())
1081    }
1082}
1083
1084impl From<e5m2> for ConstantValue {
1085    fn from(value: e5m2) -> Self {
1086        ConstantValue::Float(value.to_f64())
1087    }
1088}
1089
1090impl From<ue8m0> for ConstantValue {
1091    fn from(value: ue8m0) -> Self {
1092        ConstantValue::Float(value.to_f64())
1093    }
1094}
1095
1096impl From<half::f16> for ConstantValue {
1097    fn from(value: half::f16) -> Self {
1098        ConstantValue::Float(value.to_f64())
1099    }
1100}
1101
1102impl From<half::bf16> for ConstantValue {
1103    fn from(value: half::bf16) -> Self {
1104        ConstantValue::Float(value.to_f64())
1105    }
1106}
1107
1108impl From<flex32> for ConstantValue {
1109    fn from(value: flex32) -> Self {
1110        ConstantValue::Float(value.to_f64())
1111    }
1112}
1113
1114impl From<tf32> for ConstantValue {
1115    fn from(value: tf32) -> Self {
1116        ConstantValue::Float(value.to_f64())
1117    }
1118}
1119
1120impl From<f32> for ConstantValue {
1121    fn from(value: f32) -> Self {
1122        ConstantValue::Float(value as f64)
1123    }
1124}
1125
1126macro_rules! impl_into_value {
1127    ($($ty: ty => $kind: path,)*) => {
1128        $(
1129            impl From<$ty> for Value {
1130                fn from(value: $ty) -> Self {
1131                    Value {kind: ValueKind::Constant(value.into()), ty: $kind.into()}
1132                }
1133            }
1134        )*
1135    };
1136}
1137
1138impl_into_value!(
1139    bool => ElemType::Bool,
1140
1141    i8 => IntKind::I8,
1142    i16 => IntKind::I16,
1143    i32 => IntKind::I32,
1144    i64 => IntKind::I64,
1145
1146    u8 => UIntKind::U8,
1147    u16 => UIntKind::U16,
1148    u32 => UIntKind::U32,
1149    u64 => UIntKind::U64,
1150
1151    e2m1 => FloatKind::E2M1,
1152    e4m3 => FloatKind::E4M3,
1153    e5m2 => FloatKind::E5M2,
1154    ue8m0 => FloatKind::UE8M0,
1155    f16 => FloatKind::F16,
1156    bf16 => FloatKind::BF16,
1157    f32 => FloatKind::F32,
1158    flex32 => FloatKind::Flex32,
1159    tf32 => FloatKind::TF32,
1160    f64 => FloatKind::F64,
1161
1162    usize => UIntKind::U32,
1163    isize => IntKind::I32,
1164);
cubecl_ir/type.rs

cubecl_ir/
type.rs
