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 => bf16_rounded_once(real).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/// The `bf16` nearest `value`, rounded once as a device conversion rounds it.
450/// `half::bf16::from_f64` goes through the nearest `f32` and so rounds twice; rounding to odd at
451/// `f32` instead leaves the last rounding exact. This is the host twin of the conversion
452/// `cubecl_core::post_processing::bf16` lowers on the device, so a folded constant and the same
453/// cast at runtime agree.
454fn bf16_rounded_once(value: ConstantValue) -> half::bf16 {
455    let odd = match value {
456        ConstantValue::Int(value) => {
457            let magnitude = u64_to_f32_round_to_odd(value.unsigned_abs());
458            if value < 0 { -magnitude } else { magnitude }
459        }
460        ConstantValue::UInt(value) => u64_to_f32_round_to_odd(value),
461        value => f64_to_f32_round_to_odd(value.as_f64()),
462    };
463    half::bf16::from_f32(odd)
464}
465
466/// The nearest `f32`, or when inexact the one beside it toward zero with its low bit set.
467fn f64_to_f32_round_to_odd(value: f64) -> f32 {
468    let nearest = value as f32;
469    if nearest as f64 == value || value.is_nan() {
470        return nearest;
471    }
472    let bits = nearest.to_bits();
473    let toward_zero = match (nearest as f64).abs() > value.abs() {
474        true => bits - 1,
475        false => bits,
476    };
477    f32::from_bits(toward_zero | 1)
478}
479
480/// [`f64_to_f32_round_to_odd`] for an integer magnitude: its top 24 significant bits, with a
481/// sticky low bit for whatever was dropped below them.
482fn u64_to_f32_round_to_odd(magnitude: u64) -> f32 {
483    let significant = u64::BITS - magnitude.leading_zeros();
484    let excess = significant.saturating_sub(f32::MANTISSA_DIGITS);
485    let kept = magnitude >> excess;
486    let odd = kept | u64::from(kept << excess != magnitude);
487    let scale = f32::from_bits((excess + (f32::MAX_EXP - 1) as u32) << (f32::MANTISSA_DIGITS - 1));
488    odd as f32 * scale
489}
490
491#[cfg(test)]
492mod tests {
493    use super::*;
494
495    #[test]
496    fn complex_casts_use_the_real_component_except_for_bool() {
497        let value = ConstantValue::Complex(3.5, 2.0);
498        assert_eq!(value.cast_to(FloatKind::F64), ConstantValue::Float(3.5));
499        assert_eq!(value.cast_to(IntKind::I32), ConstantValue::Int(3));
500        assert_eq!(value.cast_to(UIntKind::U32), ConstantValue::UInt(3));
501        assert_eq!(
502            ConstantValue::Complex(0.0, 1.0).cast_to(ElemType::Bool),
503            ConstantValue::Bool(true)
504        );
505        assert_eq!(
506            ConstantValue::Complex(0.0, 0.0).cast_to(ElemType::Bool),
507            ConstantValue::Bool(false)
508        );
509    }
510
511    /// A folded cast to `bf16` rounds once, as the device does: each value sits just past a
512    /// `bf16` tie that rounding to `f32` first would land on, then round to even the wrong way.
513    #[test]
514    fn bf16_constants_round_once() {
515        let bf16 = |value: ConstantValue| match value.cast_to(FloatKind::BF16) {
516            ConstantValue::Float(value) => half::bf16::from_f64(value).to_bits(),
517            other => panic!("a bf16 constant is a float, got {other:?}"),
518        };
519        // 2^31 + 2^23 + 2^7, and its negation shifted down one.
520        assert_eq!(bf16(ConstantValue::UInt(0x8080_0080)), 0x4F01);
521        assert_eq!(bf16(ConstantValue::Int(-0x4040_0040)), 0xCE81);
522        // 1 + 2^-8 + 2^-52.
523        assert_eq!(
524            bf16(ConstantValue::Float(f64::from_bits(0x3FF0_1000_0000_0001))),
525            0x3F81
526        );
527        // 2^63 + 2^55 + 1, past what `f64` holds exactly.
528        assert_eq!(bf16(ConstantValue::UInt((1 << 63) + (1 << 55) + 1)), 0x5F01);
529        assert_eq!(bf16(ConstantValue::Int(i64::MIN)), 0xDF00);
530        assert_eq!(bf16(ConstantValue::Float(f64::INFINITY)), 0x7F80);
531        assert_eq!(bf16(ConstantValue::Float(-0.0)), 0x8000);
532        assert!(half::bf16::from_bits(bf16(ConstantValue::Float(f64::NAN))).is_nan());
533    }
534}
cubecl_ir/value.rs

cubecl_ir/
value.rs
