cubecl_ir/
type.rs

1use super::{ConstantValue, ExpandValue};
2use crate::{
3    AddressType, ContextExt, Scope, typed_vec_attr,
4    types::{scalar::*, spirv::ClampMode},
5};
6use alloc::vec::Vec;
7use core::fmt::Display;
8use cubecl_common::{
9    e2m1, e2m1x2, e2m3, e3m2, e4m3, e5m2, flex32,
10    quant::scheme::{QuantValue, ScaleDtype},
11    tf32, ue8m0,
12};
13use cubecl_macros_internal::TypeHash;
14use derive_more::{Display, From};
15use half::{bf16, f16};
16
17pub use internment::Intern;
18use pliron::{
19    builtin::types::{IntegerType, Signedness},
20    context::Context,
21    derive::format,
22    r#type::TypeHandle,
23};
24
25#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
26#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)]
27#[allow(missing_docs)]
28pub enum FloatKind {
29    /// FP4, 2 bit exponent, 1 bit mantissa
30    E2M1,
31    /// `FP4x2`, 2 bit exponent, 1 bit mantissa
32    E2M1x2,
33    /// FP6, 2 bit exponent, 3 bit mantissa
34    /// Note: represented by an 8-bit value, with the upper two bits being insignificant
35    E2M3,
36    /// FP6, 3 bit exponent, 2 bit mantissa
37    /// Note: represented by an 8-bit value, with the upper two bits being insignificant
38    E3M2,
39    /// FP8, 4 bit exponent, 3 bit mantissa
40    E4M3,
41    /// FP8, 5 bit exponent, 2 bit mantissa
42    E5M2,
43    /// FP8, unsigned, 8 bit exponent, 0 bit mantissa
44    UE8M0,
45    F16,
46    BF16,
47    Flex32,
48    F32,
49    TF32,
50    F64,
51}
52
53impl FloatKind {
54    pub fn to_type(&self, ctx: &Context) -> TypeHandle {
55        match self {
56            FloatKind::E2M1 => Float4E2M1Type::get(ctx).into(),
57            FloatKind::E2M1x2 => Float4E2M1x2Type::get(ctx).into(),
58            FloatKind::E2M3 => Float6E2M3Type::get(ctx).into(),
59            FloatKind::E3M2 => Float6E3M2Type::get(ctx).into(),
60            FloatKind::E4M3 => Float8E4M3Type::get(ctx).into(),
61            FloatKind::E5M2 => Float8E5M2Type::get(ctx).into(),
62            FloatKind::UE8M0 => Float8E8M0Type::get(ctx).into(),
63            FloatKind::F16 => Float16Type::get(ctx).into(),
64            FloatKind::BF16 => BFloat16Type::get(ctx).into(),
65            FloatKind::Flex32 => FloatFlex32Type::get(ctx).into(),
66            FloatKind::F32 => Float32Type::get(ctx).into(),
67            FloatKind::TF32 => TFloat32Type::get(ctx).into(),
68            FloatKind::F64 => Float64Type::get(ctx).into(),
69        }
70    }
71}
72
73#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
74#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)]
75#[allow(missing_docs)]
76pub enum IntKind {
77    I8,
78    I16,
79    I32,
80    I64,
81}
82
83impl IntKind {
84    pub fn to_type(&self, ctx: &Context) -> TypeHandle {
85        IntegerType::get(ctx, self.size_bits() as u32, Signedness::Signed).into()
86    }
87
88    pub fn size_bits(&self) -> usize {
89        match self {
90            IntKind::I8 => 8,
91            IntKind::I16 => 16,
92            IntKind::I32 => 32,
93            IntKind::I64 => 64,
94        }
95    }
96}
97
98#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
99#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)]
100#[allow(missing_docs)]
101pub enum UIntKind {
102    U8,
103    U16,
104    U32,
105    U64,
106}
107
108#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
109#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)]
110#[allow(missing_docs)]
111pub enum ComplexKind {
112    C32,
113    C64,
114}
115
116impl UIntKind {
117    pub fn to_type(&self, ctx: &Context) -> TypeHandle {
118        IntegerType::get(ctx, self.size_bits() as u32, Signedness::Unsigned).into()
119    }
120
121    pub fn size_bits(&self) -> usize {
122        match self {
123            UIntKind::U8 => 8,
124            UIntKind::U16 => 16,
125            UIntKind::U32 => 32,
126            UIntKind::U64 => 64,
127        }
128    }
129}
130
131/// Conceptual element type, not necessarily the physical type used in the code
132#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
133#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord, From)]
134#[allow(missing_docs)]
135pub enum ElemType {
136    Index,
137    Float(FloatKind),
138    Int(IntKind),
139    UInt(UIntKind),
140    Complex(ComplexKind),
141    Bool,
142}
143
144#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
145#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)]
146pub enum OpaqueType {
147    Barrier,
148    TensorMap,
149}
150
151#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
152#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)]
153pub enum SemanticType {
154    TensorLayout(usize, ClampMode),
155    TensorView(usize, bool, [u32; 5]),
156}
157
158impl ElemType {
159    /// Creates an elem type that correspond to the given [`ScaleDtype`].
160    pub fn from_scale_dtype(dtype: ScaleDtype) -> Self {
161        match dtype {
162            ScaleDtype::F32 => Self::Float(FloatKind::F32),
163            ScaleDtype::F16 => Self::Float(FloatKind::F16),
164            ScaleDtype::BF16 => Self::Float(FloatKind::BF16),
165            ScaleDtype::UE8M0 => Self::Float(FloatKind::UE8M0),
166            ScaleDtype::UE4M3 => Self::Float(FloatKind::E4M3),
167        }
168    }
169
170    /// Creates an elem type that correspond to the given [`QuantValue`].
171    pub fn from_quant_value(quant_value: QuantValue) -> Self {
172        match quant_value {
173            QuantValue::E5M2 => Self::Float(FloatKind::E5M2),
174            QuantValue::E4M3 => Self::Float(FloatKind::E4M3),
175            QuantValue::E2M1 => Self::Float(FloatKind::E2M1),
176            QuantValue::Q8F | QuantValue::Q8S => Self::Int(IntKind::I8),
177            other => panic!("Unsupported quant value {other:?}"),
178        }
179    }
180
181    pub fn to_type(&self, ctx: &Context) -> TypeHandle {
182        match self {
183            ElemType::Index => IndexType::get(ctx).into(),
184            ElemType::Float(float_kind) => float_kind.to_type(ctx),
185            ElemType::Int(int_kind) => int_kind.to_type(ctx),
186            ElemType::UInt(uint_kind) => uint_kind.to_type(ctx),
187            ElemType::Complex(ComplexKind::C32) => Complex32Type::get(ctx).into(),
188            ElemType::Complex(ComplexKind::C64) => Complex64Type::get(ctx).into(),
189            ElemType::Bool => BoolType::get(ctx).into(),
190        }
191    }
192
193    /// Create a constant from a constant value.
194    ///
195    /// The output will have the same type as the element.
196    pub fn constant(&self, val: ConstantValue) -> ExpandValue {
197        ExpandValue::constant(val, *self)
198    }
199
200    pub fn with_vector_size(self, vector_size: VectorSize) -> Type {
201        let ty = Type::Scalar(self);
202        if vector_size > 1 {
203            Type::Vector(ty.intern(), vector_size)
204        } else {
205            ty
206        }
207    }
208
209    pub fn expand_size(&self, address_type: AddressType) -> usize {
210        match self {
211            ElemType::Index => address_type.size(),
212            other => other.size(),
213        }
214    }
215
216    /// Get the size in bytes.
217    pub fn size(&self) -> usize {
218        match self {
219            ElemType::Index => panic!("Can't get index size outside kernel"),
220            ElemType::Float(kind) => match kind {
221                FloatKind::E2M1
222                | FloatKind::E2M1x2
223                | FloatKind::E2M3
224                | FloatKind::E3M2
225                | FloatKind::E4M3
226                | FloatKind::E5M2
227                | FloatKind::UE8M0 => core::mem::size_of::<u8>(),
228                FloatKind::F16 => core::mem::size_of::<half::f16>(),
229                FloatKind::BF16 => core::mem::size_of::<half::bf16>(),
230                FloatKind::F32 => core::mem::size_of::<f32>(),
231                FloatKind::F64 => core::mem::size_of::<f64>(),
232                FloatKind::Flex32 => core::mem::size_of::<f32>(),
233                FloatKind::TF32 => core::mem::size_of::<f32>(),
234            },
235            ElemType::Int(kind) => match kind {
236                IntKind::I8 => core::mem::size_of::<i8>(),
237                IntKind::I16 => core::mem::size_of::<i16>(),
238                IntKind::I32 => core::mem::size_of::<i32>(),
239                IntKind::I64 => core::mem::size_of::<i64>(),
240            },
241            ElemType::UInt(kind) => match kind {
242                UIntKind::U8 => core::mem::size_of::<u8>(),
243                UIntKind::U16 => core::mem::size_of::<u16>(),
244                UIntKind::U32 => core::mem::size_of::<u32>(),
245                UIntKind::U64 => core::mem::size_of::<u64>(),
246            },
247            ElemType::Complex(ComplexKind::C32) => {
248                core::mem::size_of::<num_complex::Complex<f32>>()
249            }
250            ElemType::Complex(ComplexKind::C64) => {
251                core::mem::size_of::<num_complex::Complex<f64>>()
252            }
253            ElemType::Bool => core::mem::size_of::<bool>(),
254        }
255    }
256
257    /// Get the size in bits.
258    pub fn size_bits(&self) -> usize {
259        match self {
260            ElemType::Index => panic!("Can't get index size outside kernel"),
261            ElemType::Float(kind) => match kind {
262                FloatKind::E2M1x2
263                | FloatKind::E2M3
264                | FloatKind::E3M2
265                | FloatKind::E4M3
266                | FloatKind::E5M2
267                | FloatKind::UE8M0
268                | FloatKind::F16
269                | FloatKind::BF16
270                | FloatKind::F32
271                | FloatKind::F64
272                | FloatKind::Flex32
273                | FloatKind::TF32 => self.size() * 8,
274                FloatKind::E2M1 => 4,
275            },
276            ElemType::Int(_) | ElemType::UInt(_) | ElemType::Complex(_) | ElemType::Bool => {
277                self.size() * 8
278            }
279        }
280    }
281
282    pub const fn min_vector_size(&self) -> u8 {
283        match self {
284            ElemType::Float(FloatKind::E2M1) => 2,
285            _ => 1,
286        }
287    }
288
289    pub fn is_int(&self) -> bool {
290        matches!(self, ElemType::Int(_) | ElemType::UInt(_) | ElemType::Bool)
291    }
292
293    pub fn is_signed_int(&self) -> bool {
294        matches!(self, ElemType::Int(_))
295    }
296
297    pub fn is_unsigned_int(&self) -> bool {
298        matches!(self, ElemType::UInt(_) | ElemType::Bool)
299    }
300
301    pub fn is_float(&self) -> bool {
302        matches!(self, ElemType::Float(_))
303    }
304
305    pub fn is_bool(&self) -> bool {
306        matches!(self, ElemType::Bool)
307    }
308
309    pub fn is_complex(&self) -> bool {
310        matches!(self, ElemType::Complex(_))
311    }
312
313    pub fn as_complex(&self) -> Option<ComplexKind> {
314        match self {
315            ElemType::Complex(kind) => Some(*kind),
316            _ => None,
317        }
318    }
319
320    pub fn as_float(&self) -> Option<FloatKind> {
321        match self {
322            ElemType::Float(kind) => Some(*kind),
323            _ => None,
324        }
325    }
326
327    pub fn max_variable(&self, scope: &Scope) -> ExpandValue {
328        let value = match self {
329            ElemType::Index => {
330                let addr = scope.ctx().address_type().unsigned_type();
331                return addr.max_variable(scope);
332            }
333            ElemType::Float(kind) => match kind {
334                FloatKind::E2M1 => e2m1::MAX.to_f64(),
335                FloatKind::E2M1x2 => e2m1::MAX.to_f64(),
336                FloatKind::E2M3 => e2m3::MAX,
337                FloatKind::E3M2 => e3m2::MAX,
338                FloatKind::E4M3 => e4m3::MAX.to_f64(),
339                FloatKind::E5M2 => e5m2::MAX.to_f64(),
340                FloatKind::UE8M0 => ue8m0::MAX.to_f64(),
341                FloatKind::F16 => half::f16::MAX.to_f64(),
342                FloatKind::BF16 => half::bf16::MAX.to_f64(),
343                FloatKind::Flex32 | FloatKind::TF32 | FloatKind::F32 => f32::MAX as f64,
344                FloatKind::F64 => f64::MAX,
345            }
346            .into(),
347            ElemType::Int(kind) => match kind {
348                IntKind::I8 => i8::MAX as i64,
349                IntKind::I16 => i16::MAX as i64,
350                IntKind::I32 => i32::MAX as i64,
351                IntKind::I64 => i64::MAX,
352            }
353            .into(),
354            ElemType::UInt(kind) => match kind {
355                UIntKind::U8 => u8::MAX as u64,
356                UIntKind::U16 => u16::MAX as u64,
357                UIntKind::U32 => u32::MAX as u64,
358                UIntKind::U64 => u64::MAX,
359            }
360            .into(),
361            ElemType::Complex(_) => panic!("Complex numbers have no maximum"),
362            ElemType::Bool => true.into(),
363        };
364
365        ExpandValue::Constant { value, ty: *self }
366    }
367
368    pub fn min_variable(&self) -> ExpandValue {
369        let value = match self {
370            ElemType::Index => 0u64.into(),
371            ElemType::Float(kind) => match kind {
372                FloatKind::E2M1 => e2m1::MIN.to_f64(),
373                FloatKind::E2M1x2 => e2m1::MIN.to_f64(),
374                FloatKind::E2M3 => e2m3::MIN,
375                FloatKind::E3M2 => e3m2::MIN,
376                FloatKind::E4M3 => e4m3::MIN.to_f64(),
377                FloatKind::E5M2 => e5m2::MIN.to_f64(),
378                FloatKind::UE8M0 => ue8m0::MIN.to_f64(),
379                FloatKind::F16 => half::f16::MIN.to_f64(),
380                FloatKind::BF16 => half::bf16::MIN.to_f64(),
381                FloatKind::Flex32 | FloatKind::TF32 | FloatKind::F32 => f32::MIN as f64,
382                FloatKind::F64 => f64::MIN,
383            }
384            .into(),
385            ElemType::Int(kind) => match kind {
386                IntKind::I8 => i8::MIN as i64,
387                IntKind::I16 => i16::MIN as i64,
388                IntKind::I32 => i32::MIN as i64,
389                IntKind::I64 => i64::MIN,
390            }
391            .into(),
392            ElemType::UInt(kind) => match kind {
393                UIntKind::U8 => u8::MIN as u64,
394                UIntKind::U16 => u16::MIN as u64,
395                UIntKind::U32 => u32::MIN as u64,
396                UIntKind::U64 => u64::MIN,
397            }
398            .into(),
399            ElemType::Complex(_) => panic!("Complex numbers have no minimum"),
400            ElemType::Bool => false.into(),
401        };
402
403        ExpandValue::Constant { value, ty: *self }
404    }
405
406    pub fn epsilon(&self) -> f64 {
407        match self {
408            ElemType::Float(kind) => match kind {
409                FloatKind::E2M1 => 0.5 * (e2m1::MAX.to_f64() - e2m1::MIN.to_f64()),
410                FloatKind::E2M1x2 => 0.5 * (e2m1::MAX.to_f64() - e2m1::MIN.to_f64()),
411                FloatKind::E2M3 => 0.5 * (e2m3::MAX - e2m3::MIN),
412                FloatKind::E3M2 => 0.5 * (e3m2::MAX - e3m2::MIN),
413                FloatKind::E4M3 => 0.5 * (e4m3::MAX.to_f64() - e4m3::MIN.to_f64()),
414                FloatKind::E5M2 => 0.5 * (e5m2::MAX.to_f64() - e5m2::MIN.to_f64()),
415                FloatKind::UE8M0 => 0.5 * (ue8m0::MAX.to_f64() - ue8m0::MIN.to_f64()),
416                FloatKind::F16 => half::f16::EPSILON.to_f64(),
417                FloatKind::BF16 => 0.0078125, // bf16 epsilon ≈ 2^-7
418                FloatKind::Flex32 | FloatKind::F32 | FloatKind::TF32 => f32::EPSILON.into(),
419                FloatKind::F64 => f64::EPSILON,
420            },
421            ElemType::Index | ElemType::Int(_) | ElemType::UInt(_) => 1.0, // step of 1
422            ElemType::Complex(ComplexKind::C32) => f32::EPSILON.into(),
423            ElemType::Complex(ComplexKind::C64) => f64::EPSILON,
424            ElemType::Bool => 1.0,
425        }
426    }
427}
428
429impl From<OpaqueType> for Type {
430    fn from(val: OpaqueType) -> Self {
431        Type::Opaque(val)
432    }
433}
434
435impl<T: Into<ElemType>> From<T> for Type {
436    fn from(val: T) -> Self {
437        Type::new(val.into())
438    }
439}
440
441impl From<SemanticType> for Type {
442    fn from(val: SemanticType) -> Self {
443        Type::semantic(val)
444    }
445}
446
447/// Class of a pointer. For `Global`, the ID contains the underlying buffer ID.
448/// The ID can be used to determine more detailed buffer properties, i.e. for Metal where readability
449/// is part of the pointer class.
450/// For ``CubeCL`` semantics, pointers classes to different buffer IDs should be treated as entirely
451/// separate types.
452#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
453#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)]
454#[format]
455pub enum AddressSpace {
456    #[format("`<` $0 `>`")]
457    Global(usize),
458    Shared,
459    Local,
460}
461
462typed_vec_attr!(AddressSpace, "cube.address_spaces", AddressSpaceVecAttr);
463
464#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
465#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, PartialOrd, Ord)]
466pub enum Type {
467    /// Scalar type containing a single storage element
468    Scalar(ElemType),
469    /// Opaque types that can be stored but not interacted with normally. i.e. barrier,
470    /// arrival tokens and tensor map descriptor.
471    Opaque(OpaqueType),
472    /// Vector wrapping `n` storage elements
473    Vector(Intern<Type>, VectorSize),
474    /// No defined physical representation, purely semantic. i.e. barrier, pipeline
475    Semantic(SemanticType),
476    /// Atomically accessed version of `Type`
477    Atomic(Intern<Type>),
478}
479
480/// `Intern` hashes the pointer, not the values, leading to unstable hashes across runs.
481/// Fix this by manually hashing the value.
482impl core::hash::Hash for Type {
483    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
484        core::mem::discriminant(self).hash(state);
485        match self {
486            Type::Scalar(storage_type) => storage_type.hash(state),
487            Type::Opaque(opaque) => opaque.hash(state),
488            Type::Vector(intern, size) => {
489                intern.as_ref().hash(state);
490                size.hash(state);
491            }
492            Type::Semantic(semantic_type) => semantic_type.hash(state),
493            Type::Atomic(intern) => intern.as_ref().hash(state),
494        }
495    }
496}
497
498pub type VectorSize = usize;
499
500impl Type {
501    pub fn intern(self) -> Intern<Type> {
502        Intern::new(self)
503    }
504
505    /// Create a new type
506    pub fn new(elem: impl Into<ElemType>) -> Self {
507        Type::Scalar(elem.into())
508    }
509
510    pub fn semantic(ty: SemanticType) -> Self {
511        Self::Semantic(ty)
512    }
513
514    pub fn atomic(ty: impl Into<Type>) -> Self {
515        Self::Atomic(ty.into().intern())
516    }
517
518    pub fn with_vector_size(self, vector_size: VectorSize) -> Self {
519        match self {
520            Type::Scalar(inner) if vector_size > 1 => {
521                Type::Vector(Type::new(inner).intern(), vector_size)
522            }
523            Type::Opaque(opaque) => Type::Opaque(opaque),
524            Type::Vector(inner, _) if vector_size <= 1 => *inner,
525            Type::Vector(inner, _) => Type::Vector(inner, vector_size),
526            Type::Atomic(inner) => Type::Atomic(inner.with_vector_size(vector_size).intern()),
527            this @ (Type::Scalar(_) | Type::Semantic(_)) => this,
528        }
529    }
530
531    pub fn vector_size(&self) -> VectorSize {
532        match self {
533            Type::Scalar(_) => 1,
534            Type::Opaque(_) => 1,
535            Type::Vector(inner, vector_size) => inner.vector_size() * *vector_size,
536            Type::Atomic(inner) => inner.vector_size(),
537            Type::Semantic(_) => 0,
538        }
539    }
540
541    pub fn size(&self) -> usize {
542        match self {
543            Type::Scalar(ty) => ty.size(),
544            Type::Opaque(_) => panic!("Can't get size of opaque type"),
545            Type::Vector(ty, vector_size) => ty.size() * *vector_size,
546            Type::Atomic(inner) => inner.size(),
547            Type::Semantic(_) => 0,
548        }
549    }
550
551    pub fn elem_type(&self) -> ElemType {
552        match self {
553            Type::Scalar(ty) => *ty,
554            Type::Semantic(_) | Type::Opaque(_) => {
555                unimplemented!("Can't get storage for semantic type")
556            }
557            Type::Atomic(inner) | Type::Vector(inner, _) => inner.elem_type(),
558        }
559    }
560}
561
562impl Display for Type {
563    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
564        match self {
565            Type::Semantic(ty) => write!(f, "{ty}"),
566            Type::Opaque(ty) => write!(f, "{ty}"),
567            Type::Scalar(ty) => write!(f, "{ty}"),
568            Type::Vector(ty, vector_size) => write!(f, "vector<{ty}, {vector_size}>"),
569            Type::Atomic(ty) => write!(f, "atomic<{ty}>"),
570        }
571    }
572}
573
574impl Display for ElemType {
575    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
576        match self {
577            Self::Index => f.write_str("usize"),
578            Self::Float(kind) => match kind {
579                FloatKind::E2M1 => f.write_str("e2m1"),
580                FloatKind::E2M1x2 => f.write_str("e2m1x2"),
581                FloatKind::E2M3 => f.write_str("e2m3"),
582                FloatKind::E3M2 => f.write_str("e3m2"),
583                FloatKind::E4M3 => f.write_str("e4m3"),
584                FloatKind::E5M2 => f.write_str("e5m2"),
585                FloatKind::UE8M0 => f.write_str("ue8m0"),
586                FloatKind::F16 => f.write_str("f16"),
587                FloatKind::BF16 => f.write_str("bf16"),
588                FloatKind::Flex32 => f.write_str("flex32"),
589                FloatKind::TF32 => f.write_str("tf32"),
590                FloatKind::F32 => f.write_str("f32"),
591                FloatKind::F64 => f.write_str("f64"),
592            },
593            Self::Int(kind) => match kind {
594                IntKind::I8 => f.write_str("i8"),
595                IntKind::I16 => f.write_str("i16"),
596                IntKind::I32 => f.write_str("i32"),
597                IntKind::I64 => f.write_str("i64"),
598            },
599            Self::UInt(kind) => match kind {
600                UIntKind::U8 => f.write_str("u8"),
601                UIntKind::U16 => f.write_str("u16"),
602                UIntKind::U32 => f.write_str("u32"),
603                UIntKind::U64 => f.write_str("u64"),
604            },
605            Self::Complex(ComplexKind::C32) => f.write_str("c32"),
606            Self::Complex(ComplexKind::C64) => f.write_str("c64"),
607            Self::Bool => f.write_str("bool"),
608        }
609    }
610}
611
612impl Display for SemanticType {
613    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
614        match self {
615            SemanticType::TensorLayout(dims, _) => write!(f, "tensor_layout<{dims}>"),
616            SemanticType::TensorView(dims, has_dims, permutation) => {
617                write!(
618                    f,
619                    "tensor_layout<{:?}, has_dims: {has_dims}>",
620                    &permutation[..*dims]
621                )
622            }
623        }
624    }
625}
626
627impl Display for OpaqueType {
628    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
629        match self {
630            OpaqueType::Barrier => write!(f, "barrier"),
631            OpaqueType::TensorMap => f.write_str("tensor_map"),
632        }
633    }
634}
635
636impl Display for AddressSpace {
637    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
638        match self {
639            AddressSpace::Global(id) => write!(f, "global<{id}>"),
640            AddressSpace::Shared => write!(f, "shared"),
641            AddressSpace::Local => f.write_str("local"),
642        }
643    }
644}
645
646#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
647#[derive(Debug, Clone, Copy, PartialEq, Eq, TypeHash, PartialOrd, Ord, Display)]
648pub enum AggregateKind {
649    #[display("ptr<{meta}, {inner_ty}>")]
650    Ptr {
651        inner_ty: Intern<Type>,
652        meta: MetadataKind,
653    },
654}
655
656/// Hashed by value rather than derived, for the same reason as [`Type`]: an [`Intern`] hashes the
657/// pointer it holds, which moves between runs.
658impl core::hash::Hash for AggregateKind {
659    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
660        core::mem::discriminant(self).hash(state);
661        match self {
662            AggregateKind::Ptr { inner_ty, meta } => {
663                inner_ty.as_ref().hash(state);
664                meta.hash(state);
665            }
666        }
667    }
668}
669
670impl AggregateKind {
671    pub fn ptr(inner_ty: Type, meta: MetadataKind) -> Self {
672        AggregateKind::Ptr {
673            inner_ty: inner_ty.intern(),
674            meta,
675        }
676    }
677}
678
679#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
680#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, TypeHash, PartialOrd, Ord, Display)]
681pub enum MetadataKind {
682    /// Slice metadata (offset and length)
683    #[display("slice")]
684    Slice,
685    /// Bounds check (in bounds)
686    #[display("bounds_checked")]
687    BoundsCheck,
688}
689
690pub struct BoundsCheckMetadata;
691impl BoundsCheckMetadata {
692    pub const POINTER: usize = 0;
693    pub const IS_IN_BOUNDS: usize = 1;
694}
695
696pub struct SliceMetadata;
697impl SliceMetadata {
698    pub const LIST: usize = 0;
699    pub const OFFSET: usize = 1;
700    pub const LENGTH: usize = 2;
701}
702
703impl From<e2m1x2> for ExpandValue {
704    fn from(_value: e2m1x2) -> Self {
705        unimplemented!("Can't currently construct e2m1x2")
706    }
707}
708
709impl From<e2m3> for ExpandValue {
710    fn from(_value: e2m3) -> Self {
711        unimplemented!("Can't currently construct fp6")
712    }
713}
714
715impl From<e3m2> for ExpandValue {
716    fn from(_value: e3m2) -> Self {
717        unimplemented!("Can't currently construct fp6")
718    }
719}
720
721impl From<i8> for ConstantValue {
722    fn from(value: i8) -> Self {
723        ConstantValue::Int(value as i64)
724    }
725}
726
727impl From<i16> for ConstantValue {
728    fn from(value: i16) -> Self {
729        ConstantValue::Int(value as i64)
730    }
731}
732
733impl From<i32> for ConstantValue {
734    fn from(value: i32) -> Self {
735        ConstantValue::Int(value as i64)
736    }
737}
738
739impl From<isize> for ConstantValue {
740    fn from(value: isize) -> Self {
741        ConstantValue::Int(value as i64)
742    }
743}
744
745impl From<u8> for ConstantValue {
746    fn from(value: u8) -> Self {
747        ConstantValue::UInt(value as u64)
748    }
749}
750
751impl From<u16> for ConstantValue {
752    fn from(value: u16) -> Self {
753        ConstantValue::UInt(value as u64)
754    }
755}
756
757impl From<u32> for ConstantValue {
758    fn from(value: u32) -> Self {
759        ConstantValue::UInt(value as u64)
760    }
761}
762
763impl From<usize> for ConstantValue {
764    fn from(value: usize) -> Self {
765        ConstantValue::UInt(value as u64)
766    }
767}
768
769impl From<e2m1> for ConstantValue {
770    fn from(value: e2m1) -> Self {
771        ConstantValue::Float(value.to_f64())
772    }
773}
774
775impl From<e4m3> for ConstantValue {
776    fn from(value: e4m3) -> Self {
777        ConstantValue::Float(value.to_f64())
778    }
779}
780
781impl From<e5m2> for ConstantValue {
782    fn from(value: e5m2) -> Self {
783        ConstantValue::Float(value.to_f64())
784    }
785}
786
787impl From<ue8m0> for ConstantValue {
788    fn from(value: ue8m0) -> Self {
789        ConstantValue::Float(value.to_f64())
790    }
791}
792
793impl From<half::f16> for ConstantValue {
794    fn from(value: half::f16) -> Self {
795        ConstantValue::Float(value.to_f64())
796    }
797}
798
799impl From<half::bf16> for ConstantValue {
800    fn from(value: half::bf16) -> Self {
801        ConstantValue::Float(value.to_f64())
802    }
803}
804
805impl From<flex32> for ConstantValue {
806    fn from(value: flex32) -> Self {
807        ConstantValue::Float(value.to_f64())
808    }
809}
810
811impl From<tf32> for ConstantValue {
812    fn from(value: tf32) -> Self {
813        ConstantValue::Float(value.to_f64())
814    }
815}
816
817impl From<f32> for ConstantValue {
818    fn from(value: f32) -> Self {
819        ConstantValue::Float(value as f64)
820    }
821}
822
823macro_rules! impl_into_value {
824    ($($ty: ty => $kind: path,)*) => {
825        $(
826            impl From<$ty> for ExpandValue {
827                fn from(value: $ty) -> Self {
828                    ExpandValue::Constant { value: value.into(), ty: $kind.into() }
829                }
830            }
831        )*
832    };
833}
834
835impl From<num_complex::Complex<f32>> for ExpandValue {
836    fn from(value: num_complex::Complex<f32>) -> Self {
837        ExpandValue::Constant {
838            value: ConstantValue::Complex(value.re as f64, value.im as f64),
839            ty: ElemType::Complex(ComplexKind::C32),
840        }
841    }
842}
843
844impl From<num_complex::Complex<f64>> for ExpandValue {
845    fn from(value: num_complex::Complex<f64>) -> Self {
846        ExpandValue::Constant {
847            value: ConstantValue::Complex(value.re, value.im),
848            ty: ElemType::Complex(ComplexKind::C64),
849        }
850    }
851}
852
853impl_into_value!(
854    bool => ElemType::Bool,
855
856    i8 => IntKind::I8,
857    i16 => IntKind::I16,
858    i32 => IntKind::I32,
859    i64 => IntKind::I64,
860
861    u8 => UIntKind::U8,
862    u16 => UIntKind::U16,
863    u32 => UIntKind::U32,
864    u64 => UIntKind::U64,
865
866    e2m1 => FloatKind::E2M1,
867    e4m3 => FloatKind::E4M3,
868    e5m2 => FloatKind::E5M2,
869    ue8m0 => FloatKind::UE8M0,
870    f16 => FloatKind::F16,
871    bf16 => FloatKind::BF16,
872    f32 => FloatKind::F32,
873    flex32 => FloatKind::Flex32,
874    tf32 => FloatKind::TF32,
875    f64 => FloatKind::F64,
876
877    usize => ElemType::Index,
878    isize => IntKind::I32,
879);
880
881#[cfg(test)]
882mod tests {
883    use super::*;
884    use core::hash::Hash;
885    use cubecl_common::hash::{StableHash, StableHasher};
886
887    fn hash(ty: Type) -> StableHash {
888        let mut hasher = StableHasher::new();
889        ty.hash(&mut hasher);
890        hasher.finalize()
891    }
892
893    #[test]
894    fn vector_size_is_part_of_the_hash() {
895        let f32_ty = Type::Scalar(ElemType::Float(FloatKind::F32));
896
897        assert_ne!(
898            hash(Type::Vector(f32_ty.intern(), 2)),
899            hash(Type::Vector(f32_ty.intern(), 4))
900        );
901    }
902}
cubecl_ir/type.rs

cubecl_ir/
type.rs
