cubecl_ir/
type.rs

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

cubecl_ir/
type.rs
