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 => f32::EPSILON.into(),
419                // Ten bits of mantissa: what an `f32` rounded for a tf32 instruction keeps.
420                FloatKind::TF32 => 2f64.powi(-10),
421                FloatKind::F64 => f64::EPSILON,
422            },
423            ElemType::Index | ElemType::Int(_) | ElemType::UInt(_) => 1.0, // step of 1
424            ElemType::Complex(ComplexKind::C32) => f32::EPSILON.into(),
425            ElemType::Complex(ComplexKind::C64) => f64::EPSILON,
426            ElemType::Bool => 1.0,
427        }
428    }
429}
430
431impl From<OpaqueType> for Type {
432    fn from(val: OpaqueType) -> Self {
433        Type::Opaque(val)
434    }
435}
436
437impl<T: Into<ElemType>> From<T> for Type {
438    fn from(val: T) -> Self {
439        Type::new(val.into())
440    }
441}
442
443impl From<SemanticType> for Type {
444    fn from(val: SemanticType) -> Self {
445        Type::semantic(val)
446    }
447}
448
449/// Class of a pointer. For `Global`, the ID contains the underlying buffer ID.
450/// The ID can be used to determine more detailed buffer properties, i.e. for Metal where readability
451/// is part of the pointer class.
452/// For ``CubeCL`` semantics, pointers classes to different buffer IDs should be treated as entirely
453/// separate types.
454#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
455#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)]
456#[format]
457pub enum AddressSpace {
458    #[format("`<` $0 `>`")]
459    Global(usize),
460    Shared,
461    Local,
462}
463
464typed_vec_attr!(AddressSpace, "cube.address_spaces", AddressSpaceVecAttr);
465
466#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
467#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, PartialOrd, Ord)]
468pub enum Type {
469    /// Scalar type containing a single storage element
470    Scalar(ElemType),
471    /// Opaque types that can be stored but not interacted with normally. i.e. barrier,
472    /// arrival tokens and tensor map descriptor.
473    Opaque(OpaqueType),
474    /// Vector wrapping `n` storage elements
475    Vector(Intern<Type>, VectorSize),
476    /// No defined physical representation, purely semantic. i.e. barrier, pipeline
477    Semantic(SemanticType),
478    /// Atomically accessed version of `Type`
479    Atomic(Intern<Type>),
480}
481
482/// `Intern` hashes the pointer, not the values, leading to unstable hashes across runs.
483/// Fix this by manually hashing the value.
484impl core::hash::Hash for Type {
485    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
486        core::mem::discriminant(self).hash(state);
487        match self {
488            Type::Scalar(storage_type) => storage_type.hash(state),
489            Type::Opaque(opaque) => opaque.hash(state),
490            Type::Vector(intern, size) => {
491                intern.as_ref().hash(state);
492                size.hash(state);
493            }
494            Type::Semantic(semantic_type) => semantic_type.hash(state),
495            Type::Atomic(intern) => intern.as_ref().hash(state),
496        }
497    }
498}
499
500pub type VectorSize = usize;
501
502impl Type {
503    pub fn intern(self) -> Intern<Type> {
504        Intern::new(self)
505    }
506
507    /// Create a new type
508    pub fn new(elem: impl Into<ElemType>) -> Self {
509        Type::Scalar(elem.into())
510    }
511
512    pub fn semantic(ty: SemanticType) -> Self {
513        Self::Semantic(ty)
514    }
515
516    pub fn atomic(ty: impl Into<Type>) -> Self {
517        Self::Atomic(ty.into().intern())
518    }
519
520    pub fn with_vector_size(self, vector_size: VectorSize) -> Self {
521        match self {
522            Type::Scalar(inner) if vector_size > 1 => {
523                Type::Vector(Type::new(inner).intern(), vector_size)
524            }
525            Type::Opaque(opaque) => Type::Opaque(opaque),
526            Type::Vector(inner, _) if vector_size <= 1 => *inner,
527            Type::Vector(inner, _) => Type::Vector(inner, vector_size),
528            Type::Atomic(inner) => Type::Atomic(inner.with_vector_size(vector_size).intern()),
529            this @ (Type::Scalar(_) | Type::Semantic(_)) => this,
530        }
531    }
532
533    pub fn vector_size(&self) -> VectorSize {
534        match self {
535            Type::Scalar(_) => 1,
536            Type::Opaque(_) => 1,
537            Type::Vector(inner, vector_size) => inner.vector_size() * *vector_size,
538            Type::Atomic(inner) => inner.vector_size(),
539            Type::Semantic(_) => 0,
540        }
541    }
542
543    pub fn size(&self) -> usize {
544        match self {
545            Type::Scalar(ty) => ty.size(),
546            Type::Opaque(_) => panic!("Can't get size of opaque type"),
547            Type::Vector(ty, vector_size) => ty.size() * *vector_size,
548            Type::Atomic(inner) => inner.size(),
549            Type::Semantic(_) => 0,
550        }
551    }
552
553    pub fn elem_type(&self) -> ElemType {
554        match self {
555            Type::Scalar(ty) => *ty,
556            Type::Semantic(_) | Type::Opaque(_) => {
557                unimplemented!("Can't get storage for semantic type")
558            }
559            Type::Atomic(inner) | Type::Vector(inner, _) => inner.elem_type(),
560        }
561    }
562}
563
564impl Display for Type {
565    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
566        match self {
567            Type::Semantic(ty) => write!(f, "{ty}"),
568            Type::Opaque(ty) => write!(f, "{ty}"),
569            Type::Scalar(ty) => write!(f, "{ty}"),
570            Type::Vector(ty, vector_size) => write!(f, "vector<{ty}, {vector_size}>"),
571            Type::Atomic(ty) => write!(f, "atomic<{ty}>"),
572        }
573    }
574}
575
576impl Display for ElemType {
577    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
578        match self {
579            Self::Index => f.write_str("usize"),
580            Self::Float(kind) => match kind {
581                FloatKind::E2M1 => f.write_str("e2m1"),
582                FloatKind::E2M1x2 => f.write_str("e2m1x2"),
583                FloatKind::E2M3 => f.write_str("e2m3"),
584                FloatKind::E3M2 => f.write_str("e3m2"),
585                FloatKind::E4M3 => f.write_str("e4m3"),
586                FloatKind::E5M2 => f.write_str("e5m2"),
587                FloatKind::UE8M0 => f.write_str("ue8m0"),
588                FloatKind::F16 => f.write_str("f16"),
589                FloatKind::BF16 => f.write_str("bf16"),
590                FloatKind::Flex32 => f.write_str("flex32"),
591                FloatKind::TF32 => f.write_str("tf32"),
592                FloatKind::F32 => f.write_str("f32"),
593                FloatKind::F64 => f.write_str("f64"),
594            },
595            Self::Int(kind) => match kind {
596                IntKind::I8 => f.write_str("i8"),
597                IntKind::I16 => f.write_str("i16"),
598                IntKind::I32 => f.write_str("i32"),
599                IntKind::I64 => f.write_str("i64"),
600            },
601            Self::UInt(kind) => match kind {
602                UIntKind::U8 => f.write_str("u8"),
603                UIntKind::U16 => f.write_str("u16"),
604                UIntKind::U32 => f.write_str("u32"),
605                UIntKind::U64 => f.write_str("u64"),
606            },
607            Self::Complex(ComplexKind::C32) => f.write_str("c32"),
608            Self::Complex(ComplexKind::C64) => f.write_str("c64"),
609            Self::Bool => f.write_str("bool"),
610        }
611    }
612}
613
614impl Display for SemanticType {
615    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
616        match self {
617            SemanticType::TensorLayout(dims, _) => write!(f, "tensor_layout<{dims}>"),
618            SemanticType::TensorView(dims, has_dims, permutation) => {
619                write!(
620                    f,
621                    "tensor_layout<{:?}, has_dims: {has_dims}>",
622                    &permutation[..*dims]
623                )
624            }
625        }
626    }
627}
628
629impl Display for OpaqueType {
630    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
631        match self {
632            OpaqueType::Barrier => write!(f, "barrier"),
633            OpaqueType::TensorMap => f.write_str("tensor_map"),
634        }
635    }
636}
637
638impl Display for AddressSpace {
639    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
640        match self {
641            AddressSpace::Global(id) => write!(f, "global<{id}>"),
642            AddressSpace::Shared => write!(f, "shared"),
643            AddressSpace::Local => f.write_str("local"),
644        }
645    }
646}
647
648#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
649#[derive(Debug, Clone, Copy, PartialEq, Eq, TypeHash, PartialOrd, Ord, Display)]
650pub enum AggregateKind {
651    #[display("ptr<{meta}, {inner_ty}>")]
652    Ptr {
653        inner_ty: Intern<Type>,
654        meta: MetadataKind,
655    },
656}
657
658/// Hashed by value rather than derived, for the same reason as [`Type`]: an [`Intern`] hashes the
659/// pointer it holds, which moves between runs.
660impl core::hash::Hash for AggregateKind {
661    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
662        core::mem::discriminant(self).hash(state);
663        match self {
664            AggregateKind::Ptr { inner_ty, meta } => {
665                inner_ty.as_ref().hash(state);
666                meta.hash(state);
667            }
668        }
669    }
670}
671
672impl AggregateKind {
673    pub fn ptr(inner_ty: Type, meta: MetadataKind) -> Self {
674        AggregateKind::Ptr {
675            inner_ty: inner_ty.intern(),
676            meta,
677        }
678    }
679}
680
681#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
682#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, TypeHash, PartialOrd, Ord, Display)]
683pub enum MetadataKind {
684    /// Slice metadata (offset and length)
685    #[display("slice")]
686    Slice,
687    /// Bounds check (in bounds)
688    #[display("bounds_checked")]
689    BoundsCheck,
690}
691
692pub struct BoundsCheckMetadata;
693impl BoundsCheckMetadata {
694    pub const POINTER: usize = 0;
695    pub const IS_IN_BOUNDS: usize = 1;
696}
697
698pub struct SliceMetadata;
699impl SliceMetadata {
700    pub const LIST: usize = 0;
701    pub const OFFSET: usize = 1;
702    pub const LENGTH: usize = 2;
703}
704
705impl From<e2m1x2> for ExpandValue {
706    fn from(_value: e2m1x2) -> Self {
707        unimplemented!("Can't currently construct e2m1x2")
708    }
709}
710
711impl From<e2m3> for ExpandValue {
712    fn from(_value: e2m3) -> Self {
713        unimplemented!("Can't currently construct fp6")
714    }
715}
716
717impl From<e3m2> for ExpandValue {
718    fn from(_value: e3m2) -> Self {
719        unimplemented!("Can't currently construct fp6")
720    }
721}
722
723impl From<i8> for ConstantValue {
724    fn from(value: i8) -> Self {
725        ConstantValue::Int(value as i64)
726    }
727}
728
729impl From<i16> for ConstantValue {
730    fn from(value: i16) -> Self {
731        ConstantValue::Int(value as i64)
732    }
733}
734
735impl From<i32> for ConstantValue {
736    fn from(value: i32) -> Self {
737        ConstantValue::Int(value as i64)
738    }
739}
740
741impl From<isize> for ConstantValue {
742    fn from(value: isize) -> Self {
743        ConstantValue::Int(value as i64)
744    }
745}
746
747impl From<u8> for ConstantValue {
748    fn from(value: u8) -> Self {
749        ConstantValue::UInt(value as u64)
750    }
751}
752
753impl From<u16> for ConstantValue {
754    fn from(value: u16) -> Self {
755        ConstantValue::UInt(value as u64)
756    }
757}
758
759impl From<u32> for ConstantValue {
760    fn from(value: u32) -> Self {
761        ConstantValue::UInt(value as u64)
762    }
763}
764
765impl From<usize> for ConstantValue {
766    fn from(value: usize) -> Self {
767        ConstantValue::UInt(value as u64)
768    }
769}
770
771impl From<e2m1> for ConstantValue {
772    fn from(value: e2m1) -> Self {
773        ConstantValue::Float(value.to_f64())
774    }
775}
776
777impl From<e4m3> for ConstantValue {
778    fn from(value: e4m3) -> Self {
779        ConstantValue::Float(value.to_f64())
780    }
781}
782
783impl From<e5m2> for ConstantValue {
784    fn from(value: e5m2) -> Self {
785        ConstantValue::Float(value.to_f64())
786    }
787}
788
789impl From<ue8m0> for ConstantValue {
790    fn from(value: ue8m0) -> Self {
791        ConstantValue::Float(value.to_f64())
792    }
793}
794
795impl From<half::f16> for ConstantValue {
796    fn from(value: half::f16) -> Self {
797        ConstantValue::Float(value.to_f64())
798    }
799}
800
801impl From<half::bf16> for ConstantValue {
802    fn from(value: half::bf16) -> Self {
803        ConstantValue::Float(value.to_f64())
804    }
805}
806
807impl From<flex32> for ConstantValue {
808    fn from(value: flex32) -> Self {
809        ConstantValue::Float(value.to_f64())
810    }
811}
812
813impl From<tf32> for ConstantValue {
814    fn from(value: tf32) -> Self {
815        ConstantValue::Float(value.to_f64())
816    }
817}
818
819impl From<f32> for ConstantValue {
820    fn from(value: f32) -> Self {
821        ConstantValue::Float(value as f64)
822    }
823}
824
825macro_rules! impl_into_value {
826    ($($ty: ty => $kind: path,)*) => {
827        $(
828            impl From<$ty> for ExpandValue {
829                fn from(value: $ty) -> Self {
830                    ExpandValue::Constant { value: value.into(), ty: $kind.into() }
831                }
832            }
833        )*
834    };
835}
836
837impl From<num_complex::Complex<f32>> for ExpandValue {
838    fn from(value: num_complex::Complex<f32>) -> Self {
839        ExpandValue::Constant {
840            value: ConstantValue::Complex(value.re as f64, value.im as f64),
841            ty: ElemType::Complex(ComplexKind::C32),
842        }
843    }
844}
845
846impl From<num_complex::Complex<f64>> for ExpandValue {
847    fn from(value: num_complex::Complex<f64>) -> Self {
848        ExpandValue::Constant {
849            value: ConstantValue::Complex(value.re, value.im),
850            ty: ElemType::Complex(ComplexKind::C64),
851        }
852    }
853}
854
855impl_into_value!(
856    bool => ElemType::Bool,
857
858    i8 => IntKind::I8,
859    i16 => IntKind::I16,
860    i32 => IntKind::I32,
861    i64 => IntKind::I64,
862
863    u8 => UIntKind::U8,
864    u16 => UIntKind::U16,
865    u32 => UIntKind::U32,
866    u64 => UIntKind::U64,
867
868    e2m1 => FloatKind::E2M1,
869    e4m3 => FloatKind::E4M3,
870    e5m2 => FloatKind::E5M2,
871    ue8m0 => FloatKind::UE8M0,
872    f16 => FloatKind::F16,
873    bf16 => FloatKind::BF16,
874    f32 => FloatKind::F32,
875    flex32 => FloatKind::Flex32,
876    tf32 => FloatKind::TF32,
877    f64 => FloatKind::F64,
878
879    usize => ElemType::Index,
880    isize => IntKind::I32,
881);
882
883#[cfg(test)]
884mod tests {
885    use super::*;
886    use core::hash::Hash;
887    use cubecl_common::hash::{StableHash, StableHasher};
888
889    fn hash(ty: Type) -> StableHash {
890        let mut hasher = StableHasher::new();
891        ty.hash(&mut hasher);
892        hasher.finalize()
893    }
894
895    #[test]
896    fn vector_size_is_part_of_the_hash() {
897        let f32_ty = Type::Scalar(ElemType::Float(FloatKind::F32));
898
899        assert_ne!(
900            hash(Type::Vector(f32_ty.intern(), 2)),
901            hash(Type::Vector(f32_ty.intern(), 4))
902        );
903    }
904}
cubecl_ir/type.rs

cubecl_ir/
type.rs
