cubecl_ir/
variable.rs

1use core::{fmt::Display, hash::Hash};
2
3use crate::{AddressSpace, FloatKind, IntKind, StorageType, TypeHash};
4
5use super::{ElemType, Type, UIntKind};
6use cubecl_common::{e2m1, e4m3, e5m2, ue8m0};
7use derive_more::From;
8use float_ord::FloatOrd;
9
10#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
11#[derive(Debug, Clone, Copy, TypeHash, Hash, PartialEq, Eq, PartialOrd, Ord)]
12#[allow(missing_docs)]
13pub struct Value {
14    pub kind: ValueKind,
15    pub ty: Type,
16}
17
18impl Value {
19    pub fn new(id: Id, ty: Type) -> Self {
20        Self {
21            kind: ValueKind::Value { id },
22            ty,
23        }
24    }
25
26    pub fn constant(value: ConstantValue, ty: impl Into<Type>) -> Self {
27        let ty = ty.into();
28        let value = value.cast_to(ty);
29        Self {
30            kind: ValueKind::Constant(value),
31            ty,
32        }
33    }
34
35    pub fn elem_type(&self) -> ElemType {
36        self.ty.elem_type()
37    }
38
39    pub fn storage_type(&self) -> StorageType {
40        self.ty.storage_type()
41    }
42
43    pub fn can_mutate(&self) -> bool {
44        self.ty.is_ptr()
45    }
46
47    pub fn address_space(&self) -> AddressSpace {
48        match self.ty {
49            Type::Pointer(_, addr_space) => addr_space,
50            _ => match self.kind {
51                ValueKind::Value { .. } | ValueKind::Constant(..) => AddressSpace::Local,
52            },
53        }
54    }
55
56    pub fn value_type(&self) -> Type {
57        self.ty.value_type()
58    }
59}
60
61pub type Id = u32;
62
63#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
64#[derive(Debug, Clone, Copy, TypeHash, PartialEq, Eq, Hash, PartialOrd, Ord)]
65pub enum ValueKind {
66    Value { id: Id },
67    Constant(ConstantValue),
68}
69
70#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
71#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, TypeHash, PartialOrd, Ord)]
72#[repr(u32)]
73pub enum Builtin {
74    UnitPos,
75    UnitPosX,
76    UnitPosY,
77    UnitPosZ,
78    CubePosCluster,
79    CubePosClusterX,
80    CubePosClusterY,
81    CubePosClusterZ,
82    CubePos,
83    CubePosX,
84    CubePosY,
85    CubePosZ,
86    CubeDim,
87    CubeDimX,
88    CubeDimY,
89    CubeDimZ,
90    CubeClusterDim,
91    CubeClusterDimX,
92    CubeClusterDimY,
93    CubeClusterDimZ,
94    CubeCount,
95    CubeCountX,
96    CubeCountY,
97    CubeCountZ,
98    PlaneDim,
99    PlanePos,
100    UnitPosPlane,
101    AbsolutePos,
102    AbsolutePosX,
103    AbsolutePosY,
104    AbsolutePosZ,
105}
106
107impl Value {
108    /// Whether a value is always immutable. Used for optimizations to determine whether it's
109    /// safe to inline/merge
110    pub fn is_immutable(&self) -> bool {
111        !self.can_mutate()
112    }
113
114    /// Is this an array type that yields items when indexed,
115    /// or a scalar/vector that yields elems/slices when indexed?
116    pub fn is_array_like(&self) -> bool {
117        self.ty.is_array_like()
118    }
119
120    pub fn is_value(&self) -> bool {
121        self.ty.is_value()
122    }
123
124    /// Is this an value type that is contained in concrete memory,
125    /// or a local array/scalar/vector?
126    pub fn is_memory(&self) -> bool {
127        matches!(
128            self.address_space(),
129            AddressSpace::Global(_) | AddressSpace::Shared
130        )
131    }
132
133    pub fn has_buffer_length(&self) -> bool {
134        matches!(self.address_space(), AddressSpace::Global(_))
135    }
136
137    /// Determines if the value is a constant with the specified value (converted if necessary)
138    pub fn is_constant(&self, value: i64) -> bool {
139        match self.kind {
140            ValueKind::Constant(ConstantValue::Int(val)) => val == value,
141            ValueKind::Constant(ConstantValue::UInt(val)) => val as i64 == value,
142            ValueKind::Constant(ConstantValue::Float(val)) => val == value as f64,
143            _ => false,
144        }
145    }
146
147    /// Determines if the value is a boolean constant with the `true` value
148    pub fn is_true(&self) -> bool {
149        match self.kind {
150            ValueKind::Constant(ConstantValue::Bool(val)) => val,
151            _ => false,
152        }
153    }
154
155    /// Determines if the value is a boolean constant with the `false` value
156    pub fn is_false(&self) -> bool {
157        match self.kind {
158            ValueKind::Constant(ConstantValue::Bool(val)) => !val,
159            _ => false,
160        }
161    }
162}
163
164/// The scalars are stored with the highest precision possible, but they might get reduced during
165/// compilation. For constant propagation, casts are always executed before converting back to the
166/// larger type to ensure deterministic output.
167#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
168#[derive(Debug, Clone, Copy, TypeHash, PartialEq, PartialOrd, From)]
169#[allow(missing_docs, clippy::derive_ord_xor_partial_ord)]
170pub enum ConstantValue {
171    Int(i64),
172    Float(f64),
173    UInt(u64),
174    Bool(bool),
175}
176
177impl Ord for ConstantValue {
178    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
179        // Override float-float comparison with `FloatOrd` since `f64` isn't `Ord`. All other
180        // comparisons are safe to unwrap since they're either `Ord` or only compare discriminants.
181        match (self, other) {
182            (ConstantValue::Float(this), ConstantValue::Float(other)) => {
183                FloatOrd(*this).cmp(&FloatOrd(*other))
184            }
185            _ => self.partial_cmp(other).unwrap(),
186        }
187    }
188}
189
190impl Eq for ConstantValue {}
191impl Hash for ConstantValue {
192    fn hash<H: core::hash::Hasher>(&self, ra_expand_state: &mut H) {
193        core::mem::discriminant(self).hash(ra_expand_state);
194        match self {
195            ConstantValue::Int(f0) => {
196                f0.hash(ra_expand_state);
197            }
198            ConstantValue::Float(f0) => {
199                FloatOrd(*f0).hash(ra_expand_state);
200            }
201            ConstantValue::UInt(f0) => {
202                f0.hash(ra_expand_state);
203            }
204            ConstantValue::Bool(f0) => {
205                f0.hash(ra_expand_state);
206            }
207        }
208    }
209}
210
211impl ConstantValue {
212    /// Returns the value of the constant as a usize.
213    ///
214    /// It will return [None] if the constant type is a float or a bool.
215    pub fn try_as_usize(&self) -> Option<usize> {
216        match self {
217            ConstantValue::UInt(val) => Some(*val as usize),
218            ConstantValue::Int(val) => Some(*val as usize),
219            ConstantValue::Float(_) => None,
220            ConstantValue::Bool(_) => None,
221        }
222    }
223
224    /// Returns the value of the constant as a usize.
225    pub fn as_usize(&self) -> usize {
226        match self {
227            ConstantValue::UInt(val) => *val as usize,
228            ConstantValue::Int(val) => *val as usize,
229            ConstantValue::Float(val) => *val as usize,
230            ConstantValue::Bool(val) => *val as usize,
231        }
232    }
233
234    /// Returns the value of the scalar as a u32.
235    ///
236    /// It will return [None] if the scalar type is a float or a bool.
237    pub fn try_as_u32(&self) -> Option<u32> {
238        self.try_as_u64().map(|it| it as u32)
239    }
240
241    /// Returns the value of the scalar as a u32.
242    ///
243    /// It will panic if the scalar type is a float or a bool.
244    pub fn as_u32(&self) -> u32 {
245        self.as_u64() as u32
246    }
247
248    /// Returns the value of the scalar as a u64.
249    ///
250    /// It will return [None] if the scalar type is a float or a bool.
251    pub fn try_as_u64(&self) -> Option<u64> {
252        match self {
253            ConstantValue::UInt(val) => Some(*val),
254            ConstantValue::Int(val) => Some(*val as u64),
255            ConstantValue::Float(_) => None,
256            ConstantValue::Bool(_) => None,
257        }
258    }
259
260    /// Returns the value of the scalar as a u64.
261    pub fn as_u64(&self) -> u64 {
262        match self {
263            ConstantValue::UInt(val) => *val,
264            ConstantValue::Int(val) => *val as u64,
265            ConstantValue::Float(val) => *val as u64,
266            ConstantValue::Bool(val) => *val as u64,
267        }
268    }
269
270    /// Returns the value of the scalar as a i64.
271    ///
272    /// It will return [None] if the scalar type is a float or a bool.
273    pub fn try_as_i64(&self) -> Option<i64> {
274        match self {
275            ConstantValue::UInt(val) => Some(*val as i64),
276            ConstantValue::Int(val) => Some(*val),
277            ConstantValue::Float(_) => None,
278            ConstantValue::Bool(_) => None,
279        }
280    }
281
282    /// Returns the value of the scalar as a i128.
283    pub fn as_i128(&self) -> i128 {
284        match self {
285            ConstantValue::UInt(val) => *val as i128,
286            ConstantValue::Int(val) => *val as i128,
287            ConstantValue::Float(val) => *val as i128,
288            ConstantValue::Bool(val) => *val as i128,
289        }
290    }
291
292    /// Returns the value of the scalar as a i64.
293    pub fn as_i64(&self) -> i64 {
294        match self {
295            ConstantValue::UInt(val) => *val as i64,
296            ConstantValue::Int(val) => *val,
297            ConstantValue::Float(val) => *val as i64,
298            ConstantValue::Bool(val) => *val as i64,
299        }
300    }
301
302    /// Returns the value of the scalar as a i64.
303    pub fn as_i32(&self) -> i32 {
304        match self {
305            ConstantValue::UInt(val) => *val as i32,
306            ConstantValue::Int(val) => *val as i32,
307            ConstantValue::Float(val) => *val as i32,
308            ConstantValue::Bool(val) => *val as i32,
309        }
310    }
311
312    /// Returns the value of the scalar as a f64.
313    ///
314    /// It will return [None] if the scalar type is an int or a bool.
315    pub fn try_as_f64(&self) -> Option<f64> {
316        match self {
317            ConstantValue::Float(val) => Some(*val),
318            _ => None,
319        }
320    }
321
322    /// Returns the value of the scalar as a f64.
323    pub fn as_f64(&self) -> f64 {
324        match self {
325            ConstantValue::UInt(val) => *val as f64,
326            ConstantValue::Int(val) => *val as f64,
327            ConstantValue::Float(val) => *val,
328            ConstantValue::Bool(val) => *val as u8 as f64,
329        }
330    }
331
332    /// Returns the value of the variable as a bool if it actually is a bool.
333    pub fn try_as_bool(&self) -> Option<bool> {
334        match self {
335            ConstantValue::Bool(val) => Some(*val),
336            _ => None,
337        }
338    }
339
340    /// Returns the value of the variable as a bool.
341    ///
342    /// It will panic if the scalar isn't a bool.
343    pub fn as_bool(&self) -> bool {
344        match self {
345            ConstantValue::UInt(val) => *val != 0,
346            ConstantValue::Int(val) => *val != 0,
347            ConstantValue::Float(val) => *val != 0.,
348            ConstantValue::Bool(val) => *val,
349        }
350    }
351
352    pub fn is_zero(&self) -> bool {
353        match self {
354            ConstantValue::Int(val) => *val == 0,
355            ConstantValue::Float(val) => *val == 0.0,
356            ConstantValue::UInt(val) => *val == 0,
357            ConstantValue::Bool(val) => !*val,
358        }
359    }
360
361    pub fn is_one(&self) -> bool {
362        match self {
363            ConstantValue::Int(val) => *val == 1,
364            ConstantValue::Float(val) => *val == 1.0,
365            ConstantValue::UInt(val) => *val == 1,
366            ConstantValue::Bool(val) => *val,
367        }
368    }
369
370    pub fn cast_to(&self, other: impl Into<Type>) -> ConstantValue {
371        match other.into().storage_type() {
372            StorageType::Scalar(elem_type) => match elem_type {
373                ElemType::Float(kind) => match kind {
374                    FloatKind::E2M1 => e2m1::from_f64(self.as_f64()).to_f64(),
375                    FloatKind::E2M3 | FloatKind::E3M2 => {
376                        unimplemented!("FP6 constants not yet supported")
377                    }
378                    FloatKind::E4M3 => e4m3::from_f64(self.as_f64()).to_f64(),
379                    FloatKind::E5M2 => e5m2::from_f64(self.as_f64()).to_f64(),
380                    FloatKind::UE8M0 => ue8m0::from_f64(self.as_f64()).to_f64(),
381                    FloatKind::F16 => half::f16::from_f64(self.as_f64()).to_f64(),
382                    FloatKind::BF16 => half::bf16::from_f64(self.as_f64()).to_f64(),
383                    FloatKind::Flex32 | FloatKind::TF32 | FloatKind::F32 => {
384                        self.as_f64() as f32 as f64
385                    }
386                    FloatKind::F64 => self.as_f64(),
387                }
388                .into(),
389                ElemType::Int(kind) => match kind {
390                    IntKind::I8 => self.as_i64() as i8 as i64,
391                    IntKind::I16 => self.as_i64() as i16 as i64,
392                    IntKind::I32 => self.as_i64() as i32 as i64,
393                    IntKind::I64 => self.as_i64(),
394                }
395                .into(),
396                ElemType::UInt(kind) => match kind {
397                    UIntKind::U8 => self.as_u64() as u8 as u64,
398                    UIntKind::U16 => self.as_u64() as u16 as u64,
399                    UIntKind::U32 => self.as_u64() as u32 as u64,
400                    UIntKind::U64 => self.as_u64(),
401                }
402                .into(),
403                ElemType::Bool => self.as_bool().into(),
404            },
405            StorageType::Packed(ElemType::Float(FloatKind::E2M1), 2) => {
406                e2m1::from_f64(self.as_f64()).to_f64().into()
407            }
408            StorageType::Packed(..) => unimplemented!("Unsupported packed type"),
409        }
410    }
411}
412
413impl Display for ConstantValue {
414    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
415        match self {
416            ConstantValue::Int(val) => write!(f, "{val}"),
417            ConstantValue::Float(val) => write!(f, "{val:?}"),
418            ConstantValue::UInt(val) => write!(f, "{val}"),
419            ConstantValue::Bool(val) => write!(f, "{val}"),
420        }
421    }
422}
423
424impl Value {
425    pub fn vector_size(&self) -> usize {
426        self.ty.vector_size()
427    }
428
429    pub fn id(&self) -> Id {
430        match self.kind {
431            ValueKind::Value { id, .. } => id,
432            _ => panic!("Can't get ID of constant"),
433        }
434    }
435
436    pub fn as_const(&self) -> Option<ConstantValue> {
437        match self.kind {
438            ValueKind::Constant(constant) => Some(constant),
439            _ => None,
440        }
441    }
442}
443
444impl Display for Value {
445    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
446        match self.kind {
447            ValueKind::Constant(constant) => write!(f, "{}({constant})", self.ty),
448            other => write!(f, "{other}"),
449        }
450    }
451}
452
453impl Display for ValueKind {
454    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
455        match self {
456            ValueKind::Constant(constant) => write!(f, "{constant}"),
457            ValueKind::Value { id } => write!(f, "%{id}"),
458        }
459    }
460}
461
462// Useful with the cube_inline macro.
463impl From<&Value> for Value {
464    fn from(value: &Value) -> Self {
465        *value
466    }
467}
cubecl_ir/variable.rs

cubecl_ir/
variable.rs
