cubecl_ir/
value.rs

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

cubecl_ir/
value.rs
