cubecl_ir/
value.rs

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

cubecl_ir/
value.rs
