cubecl_cpp/shared/
value.rs

1use cubecl_core::{
2    e2m1, e2m1x2, e4m3, e5m2,
3    ir::{ConstantValue, Id},
4    ue8m0,
5};
6use cubecl_runtime::kernel::Visibility;
7use std::fmt::{Display, Formatter};
8
9use crate::shared::{FP4Kind, FP8Kind, PointerClass, binary::fmt_index};
10
11use super::{COUNTER_TMP_VAR, Dialect, Elem, Item};
12
13pub trait Component<D: Dialect>: Display + FmtLeft {
14    fn item(&self) -> Item<D>;
15    fn is_const(&self) -> bool;
16    fn index(&self, index: usize) -> IndexedValue<D>;
17    fn elem(&self) -> Elem<D> {
18        *self.item().elem()
19    }
20}
21
22pub trait FmtLeft: Display {
23    fn fmt_left(&self) -> String;
24}
25
26#[derive(new, Debug)]
27pub struct OptimizedArgs<const N: usize, D: Dialect> {
28    pub args: [Value<D>; N],
29    pub optimization_factor: Option<usize>,
30}
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
33pub enum Value<D: Dialect> {
34    Constant(ConstantValue, Item<D>),
35    Value {
36        id: Id,
37        item: Item<D>,
38    },
39    Tmp {
40        id: Id,
41        item: Item<D>,
42        is_declared: bool,
43        is_ptr: bool,
44        is_const: bool,
45    },
46}
47
48#[derive(Debug, Clone, Copy)]
49pub enum Builtin<D: Dialect> {
50    AbsolutePos(Elem<D>),
51    AbsolutePosBaseName, // base name for XYZ
52    AbsolutePosX,
53    AbsolutePosY,
54    AbsolutePosZ,
55    UnitPos,
56    UnitPosBaseName, // base name for XYZ
57    UnitPosX,
58    UnitPosY,
59    UnitPosZ,
60    CubePos(Elem<D>),
61    CubePosBaseName, // base name for XYZ
62    CubePosX,
63    CubePosY,
64    CubePosZ,
65    CubeDim,
66    CubeDimBaseName, // base name for XYZ
67    CubeDimX,
68    CubeDimY,
69    CubeDimZ,
70    CubeCount(Elem<D>),
71    CubeCountBaseName, // base name for XYZ
72    CubeCountX,
73    CubeCountY,
74    CubeCountZ,
75    PlaneDim,
76    PlaneDimChecked,
77    PlanePos,
78    UnitPosPlane,
79    ClusterRank,
80    ClusterIndexX,
81    ClusterIndexY,
82    ClusterIndexZ,
83}
84
85impl<D: Dialect> Builtin<D> {
86    /// Format an item with a specific type, casting if necessary
87    pub fn fmt_cast_to(&self, item: Item<D>) -> String {
88        if self.item() == item {
89            self.to_string()
90        } else {
91            format!("{item}({self})")
92        }
93    }
94
95    pub fn item(&self) -> Item<D> {
96        match self {
97            Builtin::AbsolutePos(elem) => Item::Scalar(*elem),
98            Builtin::AbsolutePosBaseName => Item::NativeVector(Elem::U32, 3),
99            Builtin::AbsolutePosX => Item::Scalar(Elem::U32),
100            Builtin::AbsolutePosY => Item::Scalar(Elem::U32),
101            Builtin::AbsolutePosZ => Item::Scalar(Elem::U32),
102            Builtin::CubeCount(elem) => Item::Scalar(*elem),
103            Builtin::CubeCountBaseName => Item::NativeVector(Elem::U32, 3),
104            Builtin::CubeCountX => Item::Scalar(Elem::U32),
105            Builtin::CubeCountY => Item::Scalar(Elem::U32),
106            Builtin::CubeCountZ => Item::Scalar(Elem::U32),
107            Builtin::CubeDimBaseName => Item::NativeVector(Elem::U32, 3),
108            Builtin::CubeDim => Item::Scalar(Elem::U32),
109            Builtin::CubeDimX => Item::Scalar(Elem::U32),
110            Builtin::CubeDimY => Item::Scalar(Elem::U32),
111            Builtin::CubeDimZ => Item::Scalar(Elem::U32),
112            Builtin::CubePos(elem) => Item::Scalar(*elem),
113            Builtin::CubePosBaseName => Item::NativeVector(Elem::U32, 3),
114            Builtin::CubePosX => Item::Scalar(Elem::U32),
115            Builtin::CubePosY => Item::Scalar(Elem::U32),
116            Builtin::CubePosZ => Item::Scalar(Elem::U32),
117            Builtin::UnitPos => Item::Scalar(Elem::U32),
118            Builtin::UnitPosBaseName => Item::NativeVector(Elem::U32, 3),
119            Builtin::UnitPosX => Item::Scalar(Elem::U32),
120            Builtin::UnitPosY => Item::Scalar(Elem::U32),
121            Builtin::UnitPosZ => Item::Scalar(Elem::U32),
122            Builtin::PlaneDim => Item::Scalar(Elem::U32),
123            Builtin::PlaneDimChecked => Item::Scalar(Elem::U32),
124            Builtin::PlanePos => Item::Scalar(Elem::U32),
125            Builtin::UnitPosPlane => Item::Scalar(Elem::U32),
126            Builtin::ClusterRank => Item::Scalar(Elem::U32),
127            Builtin::ClusterIndexX => Item::Scalar(Elem::U32),
128            Builtin::ClusterIndexY => Item::Scalar(Elem::U32),
129            Builtin::ClusterIndexZ => Item::Scalar(Elem::U32),
130        }
131    }
132}
133
134impl<D: Dialect> Component<D> for Value<D> {
135    fn index(&self, index: usize) -> IndexedValue<D> {
136        self.index(index)
137    }
138
139    fn item(&self) -> Item<D> {
140        match self {
141            Value::Value { item, .. } => *item,
142            Value::Constant(_, e) => *e,
143            Value::Tmp { item, .. } => *item,
144        }
145    }
146
147    fn is_const(&self) -> bool {
148        if let Value::Tmp { is_const, .. } = self {
149            return *is_const;
150        }
151        if let Item::Pointer(_, PointerClass::Global(Visibility::Read)) = self.item() {
152            return true;
153        }
154
155        !self.item().is_ptr()
156    }
157}
158
159pub(crate) fn format_const<D: Dialect>(number: &ConstantValue, item: &Item<D>) -> String {
160    // minifloats are represented as raw bits, so use special handling
161    let number = match item.elem() {
162        Elem::FP4(FP4Kind::E2M1) => e2m1::from_f64(number.as_f64()).to_bits(),
163        Elem::FP4x2(FP4Kind::E2M1) => {
164            let v = number.as_f64() as f32;
165            let value = [v, v];
166            e2m1x2::from_f32_slice(&value).remove(0).to_bits()
167        }
168        Elem::FP6(_) | Elem::FP6x2(_) => {
169            todo!("FP6 constants are not yet supported")
170        }
171        Elem::FP8(FP8Kind::E4M3) => e4m3::from_f64(number.as_f64()).to_bits(),
172        Elem::FP8(FP8Kind::E5M2) => e5m2::from_f64(number.as_f64()).to_bits(),
173        Elem::FP8(FP8Kind::UE8M0) => ue8m0::from_f64(number.as_f64()).to_bits(),
174        _ => {
175            // Non-finite floats have no C++ literal, and the math.h macros
176            // (INFINITY/NAN) are not declared in the headerless HIP/nvrtc
177            // sources — the IEEE constant expressions work everywhere.
178            if let ConstantValue::Float(value) = number {
179                if value.is_infinite() {
180                    return if *value < 0.0 {
181                        "(-1.0f/0.0f)"
182                    } else {
183                        "(1.0f/0.0f)"
184                    }
185                    .to_string();
186                }
187                if value.is_nan() {
188                    return "(0.0f/0.0f)".to_string();
189                }
190            }
191            return format!("{number}");
192        }
193    };
194    format!("{number}")
195}
196
197impl<D: Dialect> Display for Value<D> {
198    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
199        match self {
200            Value::Value { id, .. } => write!(f, "val_{id}"),
201            Value::Constant(number, item) if item.vectorization() <= 1 => {
202                let value = format_const(number, item);
203                write!(f, "{item}({value})")
204            }
205            Value::Constant(number, item) => {
206                let number = format_const(number, item);
207                let values = (0..item.vectorization())
208                    .map(|_| format!("{}({number})", item.elem()))
209                    .collect::<Vec<_>>();
210                write!(f, "{item} {{ {} }}", values.join(","))
211            }
212            Value::Tmp { id, .. } => write!(f, "_tmp_{id}"),
213        }
214    }
215}
216
217impl<D: Dialect> Display for Builtin<D> {
218    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
219        match self {
220            Builtin::AbsolutePos(_) => D::compile_absolute_pos(f),
221            Builtin::AbsolutePosBaseName => D::compile_absolute_pos_base_name(f),
222            Builtin::AbsolutePosX => D::compile_absolute_pos_x(f),
223            Builtin::AbsolutePosY => D::compile_absolute_pos_y(f),
224            Builtin::AbsolutePosZ => D::compile_absolute_pos_z(f),
225            Builtin::CubeCount(_) => D::compile_cube_count(f),
226            Builtin::CubeCountBaseName => D::compile_cube_count_base_name(f),
227            Builtin::CubeCountX => D::compile_cube_count_x(f),
228            Builtin::CubeCountY => D::compile_cube_count_y(f),
229            Builtin::CubeCountZ => D::compile_cube_count_z(f),
230            Builtin::CubeDim => D::compile_cube_dim(f),
231            Builtin::CubeDimBaseName => D::compile_cube_dim_base_name(f),
232            Builtin::CubeDimX => D::compile_cube_dim_x(f),
233            Builtin::CubeDimY => D::compile_cube_dim_y(f),
234            Builtin::CubeDimZ => D::compile_cube_dim_z(f),
235            Builtin::CubePos(_) => D::compile_cube_pos(f),
236            Builtin::CubePosBaseName => D::compile_cube_pos_base_name(f),
237            Builtin::CubePosX => D::compile_cube_pos_x(f),
238            Builtin::CubePosY => D::compile_cube_pos_y(f),
239            Builtin::CubePosZ => D::compile_cube_pos_z(f),
240            Builtin::UnitPos => D::compile_unit_pos(f),
241            Builtin::UnitPosBaseName => D::compile_unit_pos_base_name(f),
242            Builtin::UnitPosX => D::compile_unit_pos_x(f),
243            Builtin::UnitPosY => D::compile_unit_pos_y(f),
244            Builtin::UnitPosZ => D::compile_unit_pos_z(f),
245            Builtin::PlaneDim => D::compile_plane_dim(f),
246            Builtin::PlaneDimChecked => D::compile_plane_dim_checked(f),
247            Builtin::PlanePos => D::compile_plane_pos(f),
248            Builtin::UnitPosPlane => D::compile_unit_pos_plane(f),
249            Builtin::ClusterRank => D::compile_cluster_pos(f),
250            Builtin::ClusterIndexX => D::compile_cluster_pos_x(f),
251            Builtin::ClusterIndexY => D::compile_cluster_pos_y(f),
252            Builtin::ClusterIndexZ => D::compile_cluster_pos_z(f),
253        }
254    }
255}
256
257impl<D: Dialect> Value<D> {
258    pub fn is_optimized(&self) -> bool {
259        self.item().is_optimized()
260    }
261
262    /// Create a temporary variable.
263    ///
264    /// Also see [`Self::tmp_declared`] for a version that needs custom declaration.
265    pub fn tmp(item: Item<D>) -> Self {
266        let inc = COUNTER_TMP_VAR.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
267
268        Value::Tmp {
269            id: inc as Id,
270            item,
271            is_declared: false,
272            is_ptr: false,
273            is_const: false,
274        }
275    }
276
277    pub fn to_const(&mut self) {
278        if let Value::Tmp { is_const, .. } = self {
279            *is_const = true;
280        }
281    }
282
283    /// Create a temporary variable with a `reinterpret_cast`.
284    pub fn reinterpret_ptr(&self, f: &mut Formatter<'_>, item: Item<D>) -> Self {
285        let mut out = Self::tmp_ptr(item);
286
287        if self.is_const() {
288            out.to_const();
289        }
290
291        let elem = out.elem();
292        let qualifier = out.const_qualifier();
293        let addr_space = D::address_space_for_value(self);
294        let out_fmt = out.fmt_left();
295
296        writeln!(
297            f,
298            "{out_fmt} = reinterpret_cast<{addr_space}{elem}{qualifier}*>({self});"
299        )
300        .unwrap();
301
302        out
303    }
304
305    /// Create a temporary pointer variable.
306    ///
307    /// Also see [`Self::tmp_declared`] for a version that needs custom declaration.
308    pub fn tmp_ptr(item: Item<D>) -> Self {
309        let inc = COUNTER_TMP_VAR.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
310
311        Value::Tmp {
312            id: inc as Id,
313            item,
314            is_declared: false,
315            is_ptr: true,
316            is_const: false,
317        }
318    }
319
320    /// Create a temporary variable with a custom declaration.
321    ///
322    /// # Notes
323    ///
324    /// Calling `val.fmt_left()` will assume the variable already exist.
325    pub fn tmp_declared(item: Item<D>) -> Self {
326        let inc = COUNTER_TMP_VAR.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
327
328        Value::Tmp {
329            id: inc as Id,
330            item,
331            is_declared: true,
332            is_ptr: false,
333            is_const: false,
334        }
335    }
336
337    pub fn optimized_args<const N: usize>(args: [Self; N]) -> OptimizedArgs<N, D> {
338        let args_after = args.map(|a| a.optimized());
339
340        let is_optimized = args_after.iter().all(|val| val.is_optimized());
341
342        if is_optimized {
343            let vectorization_before = args
344                .iter()
345                .map(|val| val.item().vectorization())
346                .max()
347                .unwrap();
348            let vectorization_after = args_after
349                .iter()
350                .map(|val| val.item().vectorization())
351                .max()
352                .unwrap();
353
354            OptimizedArgs::new(args_after, Some(vectorization_before / vectorization_after))
355        } else {
356            OptimizedArgs::new(args, None)
357        }
358    }
359
360    pub fn optimized(&self) -> Self {
361        match self {
362            Value::Value { id, item } => Value::Value {
363                id: *id,
364                item: item.optimized(),
365            },
366            Value::Tmp {
367                id,
368                item,
369                is_declared,
370                is_ptr,
371                is_const,
372            } => Value::Tmp {
373                id: *id,
374                item: item.optimized(),
375                is_declared: *is_declared,
376                is_ptr: *is_ptr,
377                is_const: *is_const,
378            },
379            _ => *self,
380        }
381    }
382
383    pub fn index(&self, index: usize) -> IndexedValue<D> {
384        IndexedValue {
385            val: *self,
386            index,
387            optimized: self.is_optimized(),
388        }
389    }
390
391    pub fn const_qualifier(&self) -> &str {
392        if self.is_const() { " const" } else { "" }
393    }
394
395    pub fn id(&self) -> Option<Id> {
396        match self {
397            Value::Value { id, .. } => Some(*id),
398            Value::Tmp { id, .. } => Some(*id),
399            _ => None,
400        }
401    }
402
403    /// A value-producing op (e.g. `Dot`/`VectorSum`, or any arithmetic) whose
404    /// output was allocated as a fresh mutable local ends up typed as a local
405    /// pointer. Such an output can't be declared inline via [`FmtLeft::fmt_left`]
406    /// — that yields `T* out = <scalar value>;`, which is a type error. Instead it
407    /// needs backing storage, exactly like [`super::Instruction::DeclareVariable`].
408    ///
409    /// When `self` is such a local pointer, this emits the backing declaration
410    /// (`T out_store; T* out = &out_store;`) and returns `true`, so the caller
411    /// writes the result through the pointer (`*out = value;`). Otherwise it emits
412    /// nothing and returns `false`, and the caller declares the output inline.
413    ///
414    /// This is only reached for freshly-created outputs (the op defines `out`),
415    /// so it never double-declares an already-declared local.
416    pub fn declare_local_ptr_backing(
417        &self,
418        f: &mut Formatter<'_>,
419    ) -> Result<bool, std::fmt::Error> {
420        if let Value::Value {
421            item: Item::Pointer(_, PointerClass::Local),
422            ..
423        } = self
424        {
425            writeln!(f, "{} {self}_store;", self.item().value_ty())?;
426            writeln!(f, "{} {self} = &{self}_store;", self.item())?;
427            Ok(true)
428        } else {
429            Ok(false)
430        }
431    }
432
433    /// Format variable for a pointer argument. Slices and buffers are already pointers, so we
434    /// just leave them as is to avoid accidental double pointers
435    pub fn fmt_ptr(&self) -> String {
436        match self.item() {
437            Item::Pointer(inner, _) if inner.is_array() => {
438                format!("{self}->data")
439            }
440            Item::Array(..) => format!("{self}.data"),
441            Item::DynamicArray(..) | Item::Pointer(..) => format!("{self}"),
442            _ => format!("&{self}"),
443        }
444    }
445
446    /// Format variable for a reference argument. Dereferences pointers while keeping locals as is.
447    pub fn fmt_ref(&self) -> String {
448        match self.item() {
449            Item::DynamicArray(..) | Item::Pointer(..) => format!("*{self}"),
450            _ => format!("{self}"),
451        }
452    }
453
454    /// Format an item with a specific type, casting if necessary
455    pub fn fmt_cast_to(&self, item: Item<D>) -> String {
456        if self.item() == item {
457            self.to_string()
458        } else {
459            format!("{item}({self})")
460        }
461    }
462
463    /// Ensure a variable is a named lvalue, reassigning to a temporary if necessary.
464    /// This is required for reinterpreting constants.
465    pub fn ensure_lvalue(&self, f: &mut Formatter<'_>) -> Result<Value<D>, core::fmt::Error> {
466        if matches!(self, Value::Constant(..)) {
467            let tmp = Value::tmp(self.item());
468            writeln!(f, "{} = {self};", tmp.fmt_left())?;
469            Ok(tmp)
470        } else if matches!(self.item(), Item::Pointer(..)) {
471            let tmp = Value::tmp(*self.item().value_ty());
472            writeln!(f, "{}& {tmp} = *{self};", tmp.item())?;
473            Ok(tmp)
474        } else {
475            Ok(*self)
476        }
477    }
478}
479
480impl<D: Dialect> FmtLeft for Value<D> {
481    fn fmt_left(&self) -> String {
482        match self {
483            Self::Value { item, .. } => match item {
484                // Pointer constness is determined by the type, not variable kind
485                Item::Pointer(..) => {
486                    format!("{item} {self}")
487                }
488                // Barrier is a memory object so can only exist behind a reference
489                Item::Barrier(..) => {
490                    format!("{item}& {self}")
491                }
492                // C++ has weird semantics so this needs to be mutable for use with `std::move`.
493                // `std::move` preserves constness for the moved value, and the API requires
494                // a non-const `BarrierToken&&`.
495                Item::BarrierToken(..) => {
496                    format!("{item} {self}")
497                }
498                _ => {
499                    format!("const {item} {self}")
500                }
501            },
502            Value::Tmp {
503                item,
504                is_declared,
505                is_ptr,
506                is_const,
507                ..
508            } => {
509                if *is_declared {
510                    return format!("{self}");
511                }
512                if *is_const && !*is_ptr {
513                    format!("const {item} {self}")
514                } else {
515                    format!("{item} {self}")
516                }
517            }
518            var => format!("{var}"),
519        }
520    }
521}
522
523#[derive(Debug, Clone)]
524pub struct IndexedValue<D: Dialect> {
525    val: Value<D>,
526    optimized: bool,
527    index: usize,
528}
529
530impl<D: Dialect> Component<D> for IndexedValue<D> {
531    fn item(&self) -> Item<D> {
532        self.val.item()
533    }
534
535    fn index(&self, index: usize) -> IndexedValue<D> {
536        self.val.index(index)
537    }
538
539    fn is_const(&self) -> bool {
540        self.val.is_const()
541    }
542}
543
544impl<D: Dialect> Display for IndexedValue<D> {
545    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
546        let var = &self.val;
547
548        if let Value::Constant(value, item) = var {
549            let value = format_const(value, item);
550            return write!(f, "{}({value})", item.elem());
551        }
552
553        if var.item().unwrap_ptr().is_array_like() {
554            return write!(f, "{}", fmt_index(var, &self.index, &var.item()));
555        }
556
557        let item = var.item();
558        let addr_space = D::address_space_for_value(&self.val);
559        let ty = match var {
560            _ if item.is_ptr() => {
561                format!("{item}")
562            }
563            Value::Value { item, .. } => format!("{addr_space}{item} const&"),
564            _ => format!("{addr_space}{item}&"),
565        };
566        let accessor = match var.item().is_ptr() {
567            true => "->",
568            false => ".",
569        };
570
571        if self.val.item().vectorization() > 1 {
572            if self.optimized {
573                write!(
574                    f,
575                    "(reinterpret_cast<{ty}>({var})){accessor}i_{}",
576                    self.index
577                )
578            } else {
579                write!(f, "{var}{accessor}i_{}", self.index)
580            }
581        } else if self.optimized {
582            write!(f, "reinterpret_cast<{ty}>({var})")
583        } else {
584            write!(f, "{var}")
585        }
586    }
587}
588
589impl<D: Dialect> FmtLeft for IndexedValue<D> {
590    fn fmt_left(&self) -> String {
591        let var = &self.val;
592        let ref_ = matches!(var, Value::Value { .. })
593            .then_some("const&")
594            .unwrap_or("&");
595
596        let name = if self.val.item().vectorization() > 1 {
597            if self.optimized {
598                let item = self.val.item();
599                let addr_space = D::address_space_for_value(&self.val);
600                format!(
601                    "(reinterpret_cast<{addr_space}{item} {ref_}>({var})).i_{}",
602                    self.index
603                )
604            } else {
605                format!("{var}.i_{}", self.index)
606            }
607        } else {
608            format!("{var}")
609        };
610        match var {
611            Value::Value { item, .. } => format!("const {item} {name}"),
612            Value::Tmp { item, is_ptr, .. } => {
613                if *is_ptr {
614                    format!("{item} *{name}")
615                } else {
616                    format!("{item} {name}")
617                }
618            }
619            _ => name,
620        }
621    }
622}
623
624impl FmtLeft for &String {
625    fn fmt_left(&self) -> String {
626        self.to_string()
627    }
628}
cubecl_cpp/shared/value.rs

cubecl_cpp/shared/
value.rs
