cubecl_cpp/shared/
variable.rs

1use cubecl_core::ir::{self as gpu, BarrierLevel, ConstantScalarValue, Id};
2use std::fmt::{Display, Formatter};
3
4use super::{COUNTER_TMP_VAR, Dialect, Elem, Fragment, FragmentIdent, Item};
5
6pub trait Component<D: Dialect>: Display + FmtLeft {
7    fn item(&self) -> Item<D>;
8    fn is_const(&self) -> bool;
9    fn index(&self, index: usize) -> IndexedVariable<D>;
10    fn elem(&self) -> Elem<D> {
11        *self.item().elem()
12    }
13}
14
15pub trait FmtLeft: Display {
16    fn fmt_left(&self) -> String;
17}
18
19#[derive(new)]
20pub struct OptimizedArgs<const N: usize, D: Dialect> {
21    pub args: [Variable<D>; N],
22    pub optimization_factor: Option<usize>,
23}
24
25#[derive(Debug, Clone, Copy, PartialEq)]
26pub enum Variable<D: Dialect> {
27    AbsolutePos,
28    AbsolutePosBaseName, // base name for XYZ
29    AbsolutePosX,
30    AbsolutePosY,
31    AbsolutePosZ,
32    UnitPos,
33    UnitPosBaseName, // base name for XYZ
34    UnitPosX,
35    UnitPosY,
36    UnitPosZ,
37    CubePos,
38    CubePosBaseName, // base name for XYZ
39    CubePosX,
40    CubePosY,
41    CubePosZ,
42    CubeDim,
43    CubeDimBaseName, // base name for XYZ
44    CubeDimX,
45    CubeDimY,
46    CubeDimZ,
47    CubeCount,
48    CubeCountBaseName, // base name for XYZ
49    CubeCountX,
50    CubeCountY,
51    CubeCountZ,
52    PlaneDim,
53    PlaneDimChecked,
54    PlanePos,
55    UnitPosPlane,
56    ClusterRank,
57    ClusterIndexX,
58    ClusterIndexY,
59    ClusterIndexZ,
60    GlobalInputArray(Id, Item<D>),
61    GlobalOutputArray(Id, Item<D>),
62    GlobalScalar {
63        id: Id,
64        elem: Elem<D>,
65        in_struct: bool,
66    },
67    ConstantArray(Id, Item<D>, u32),
68    ConstantScalar(ConstantScalarValue, Elem<D>),
69    TensorMap(Id),
70    LocalMut {
71        id: Id,
72        item: Item<D>,
73    },
74    LocalConst {
75        id: Id,
76        item: Item<D>,
77    },
78    Named {
79        name: &'static str,
80        item: Item<D>,
81    },
82    Slice {
83        id: Id,
84        item: Item<D>,
85    },
86    SharedMemory(Id, Item<D>, u32),
87    LocalArray(Id, Item<D>, u32),
88    WmmaFragment {
89        id: Id,
90        frag: Fragment<D>,
91    },
92    Pipeline {
93        id: Id,
94    },
95    Barrier {
96        id: Id,
97        level: BarrierLevel,
98    },
99    Tmp {
100        id: Id,
101        item: Item<D>,
102        is_declared: bool,
103        is_ptr: bool,
104        is_const: bool,
105    },
106}
107
108impl<D: Dialect> Component<D> for Variable<D> {
109    fn index(&self, index: usize) -> IndexedVariable<D> {
110        self.index(index)
111    }
112
113    fn item(&self) -> Item<D> {
114        match self {
115            Variable::AbsolutePos => Item::scalar(Elem::U32, true),
116            Variable::AbsolutePosBaseName => Item {
117                elem: Elem::U32,
118                vectorization: 3,
119                native: true,
120            },
121            Variable::AbsolutePosX => Item::scalar(Elem::U32, true),
122            Variable::AbsolutePosY => Item::scalar(Elem::U32, true),
123            Variable::AbsolutePosZ => Item::scalar(Elem::U32, true),
124            Variable::CubeCount => Item::scalar(Elem::U32, true),
125            Variable::CubeCountBaseName => Item {
126                elem: Elem::U32,
127                vectorization: 3,
128                native: true,
129            },
130            Variable::CubeCountX => Item::scalar(Elem::U32, true),
131            Variable::CubeCountY => Item::scalar(Elem::U32, true),
132            Variable::CubeCountZ => Item::scalar(Elem::U32, true),
133            Variable::CubeDimBaseName => Item {
134                elem: Elem::U32,
135                vectorization: 3,
136                native: true,
137            },
138            Variable::CubeDim => Item::scalar(Elem::U32, true),
139            Variable::CubeDimX => Item::scalar(Elem::U32, true),
140            Variable::CubeDimY => Item::scalar(Elem::U32, true),
141            Variable::CubeDimZ => Item::scalar(Elem::U32, true),
142            Variable::CubePos => Item::scalar(Elem::U32, true),
143            Variable::CubePosBaseName => Item {
144                elem: Elem::U32,
145                vectorization: 3,
146                native: true,
147            },
148            Variable::CubePosX => Item::scalar(Elem::U32, true),
149            Variable::CubePosY => Item::scalar(Elem::U32, true),
150            Variable::CubePosZ => Item::scalar(Elem::U32, true),
151            Variable::UnitPos => Item::scalar(Elem::U32, true),
152            Variable::UnitPosBaseName => Item {
153                elem: Elem::U32,
154                vectorization: 3,
155                native: true,
156            },
157            Variable::UnitPosX => Item::scalar(Elem::U32, true),
158            Variable::UnitPosY => Item::scalar(Elem::U32, true),
159            Variable::UnitPosZ => Item::scalar(Elem::U32, true),
160            Variable::PlaneDim => Item::scalar(Elem::U32, true),
161            Variable::PlaneDimChecked => Item::scalar(Elem::U32, true),
162            Variable::PlanePos => Item::scalar(Elem::U32, true),
163            Variable::UnitPosPlane => Item::scalar(Elem::U32, true),
164            Variable::ClusterRank => Item::scalar(Elem::U32, true),
165            Variable::ClusterIndexX => Item::scalar(Elem::U32, true),
166            Variable::ClusterIndexY => Item::scalar(Elem::U32, true),
167            Variable::ClusterIndexZ => Item::scalar(Elem::U32, true),
168            Variable::GlobalInputArray(_, e) => *e,
169            Variable::GlobalOutputArray(_, e) => *e,
170            Variable::LocalArray(_, e, _) => *e,
171            Variable::SharedMemory(_, e, _) => *e,
172            Variable::ConstantArray(_, e, _) => *e,
173            Variable::LocalMut { item, .. } => *item,
174            Variable::LocalConst { item, .. } => *item,
175            Variable::Named { item, .. } => *item,
176            Variable::Slice { item, .. } => *item,
177            Variable::ConstantScalar(_, e) => Item::scalar(*e, false),
178            Variable::GlobalScalar { elem, .. } => Item::scalar(*elem, false),
179            Variable::WmmaFragment { frag, .. } => Item::scalar(frag.elem, false),
180            Variable::Tmp { item, .. } => *item,
181            Variable::Pipeline { .. } | Variable::Barrier { .. } => Item::new(Elem::Bool, 1, false),
182            Variable::TensorMap(_) => unreachable!(),
183        }
184    }
185
186    fn is_const(&self) -> bool {
187        if let Variable::Tmp { is_const, .. } = self {
188            return *is_const;
189        }
190
191        matches!(
192            self,
193            Variable::LocalConst { .. } | Variable::GlobalInputArray { .. }
194        )
195    }
196}
197
198impl<D: Dialect> Display for Variable<D> {
199    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
200        match self {
201            Variable::GlobalInputArray(id, _) => f.write_fmt(format_args!("buffer_{id}")),
202            Variable::GlobalOutputArray(id, _) => write!(f, "buffer_{id}"),
203            Variable::TensorMap(id) => write!(f, "tensor_map_{id}"),
204            Variable::LocalMut { id, .. } => f.write_fmt(format_args!("l_mut_{id}")),
205            Variable::LocalConst { id, .. } => f.write_fmt(format_args!("l_{id}")),
206            Variable::Named { name, .. } => f.write_fmt(format_args!("{name}")),
207            Variable::Slice { id, .. } => {
208                write!(f, "slice_{id}")
209            }
210            Variable::GlobalScalar {
211                id,
212                elem,
213                in_struct,
214            } => match *in_struct {
215                true => write!(f, "scalars_{elem}.x[{id}]"),
216                false => write!(f, "scalars_{elem}[{id}]"),
217            },
218            Variable::ConstantScalar(number, elem) => match number {
219                ConstantScalarValue::Int(val, kind) => match kind {
220                    gpu::IntKind::I8 => write!(f, "{elem}({})", *val as i8),
221                    gpu::IntKind::I16 => write!(f, "{elem}({})", *val as i16),
222                    gpu::IntKind::I32 => write!(f, "{elem}({})", *val as i32),
223                    gpu::IntKind::I64 => write!(f, "{elem}({})", *val),
224                },
225                ConstantScalarValue::Float(val, kind) => match kind {
226                    gpu::FloatKind::E2M1
227                    | gpu::FloatKind::E2M3
228                    | gpu::FloatKind::E3M2
229                    | gpu::FloatKind::E4M3
230                    | gpu::FloatKind::E5M2
231                    | gpu::FloatKind::UE8M0 => todo!("Minifloat constants not supported yet"),
232                    gpu::FloatKind::F16 => {
233                        write!(f, "{elem}({:?})", half::f16::from_f64(*val))
234                    }
235                    gpu::FloatKind::BF16 => {
236                        write!(f, "{elem}({:?})", half::bf16::from_f64(*val))
237                    }
238                    gpu::FloatKind::Flex32 => write!(f, "{elem}({:?})", *val as f32),
239                    gpu::FloatKind::TF32 => write!(f, "{elem}({:?})", *val as f32),
240                    gpu::FloatKind::F32 => write!(f, "{elem}({:?})", *val as f32),
241                    gpu::FloatKind::F64 => write!(f, "{elem}({:?})", *val),
242                },
243                ConstantScalarValue::UInt(val, kind) => match kind {
244                    gpu::UIntKind::U8 => write!(f, "{elem}({})", *val as u8),
245                    gpu::UIntKind::U16 => write!(f, "{elem}({})", *val as u16),
246                    gpu::UIntKind::U32 => write!(f, "{elem}({})", *val as u32),
247                    gpu::UIntKind::U64 => write!(f, "{elem}({})", *val),
248                },
249                ConstantScalarValue::Bool(val) => write!(f, "{val}"),
250            },
251            Variable::SharedMemory(number, _, _) => {
252                write!(f, "shared_memory_{number}")
253            }
254
255            Variable::AbsolutePos => D::compile_absolute_pos(f),
256            Variable::AbsolutePosBaseName => D::compile_absolute_pos_base_name(f),
257            Variable::AbsolutePosX => D::compile_absolute_pos_x(f),
258            Variable::AbsolutePosY => D::compile_absolute_pos_y(f),
259            Variable::AbsolutePosZ => D::compile_absolute_pos_z(f),
260            Variable::CubeCount => D::compile_cube_count(f),
261            Variable::CubeCountBaseName => D::compile_cube_count_base_name(f),
262            Variable::CubeCountX => D::compile_cube_count_x(f),
263            Variable::CubeCountY => D::compile_cube_count_y(f),
264            Variable::CubeCountZ => D::compile_cube_count_z(f),
265            Variable::CubeDim => D::compile_cube_dim(f),
266            Variable::CubeDimBaseName => D::compile_cube_dim_base_name(f),
267            Variable::CubeDimX => D::compile_cube_dim_x(f),
268            Variable::CubeDimY => D::compile_cube_dim_y(f),
269            Variable::CubeDimZ => D::compile_cube_dim_z(f),
270            Variable::CubePos => D::compile_cube_pos(f),
271            Variable::CubePosBaseName => D::compile_cube_pos_base_name(f),
272            Variable::CubePosX => D::compile_cube_pos_x(f),
273            Variable::CubePosY => D::compile_cube_pos_y(f),
274            Variable::CubePosZ => D::compile_cube_pos_z(f),
275            Variable::UnitPos => D::compile_unit_pos(f),
276            Variable::UnitPosBaseName => D::compile_unit_pos_base_name(f),
277            Variable::UnitPosX => D::compile_unit_pos_x(f),
278            Variable::UnitPosY => D::compile_unit_pos_y(f),
279            Variable::UnitPosZ => D::compile_unit_pos_z(f),
280            Variable::PlaneDim => D::compile_plane_dim(f),
281            Variable::PlaneDimChecked => D::compile_plane_dim_checked(f),
282            Variable::PlanePos => D::compile_plane_pos(f),
283            Variable::UnitPosPlane => D::compile_unit_pos_plane(f),
284            Variable::ClusterRank => D::compile_cluster_pos(f),
285            Variable::ClusterIndexX => D::compile_cluster_pos_x(f),
286            Variable::ClusterIndexY => D::compile_cluster_pos_y(f),
287            Variable::ClusterIndexZ => D::compile_cluster_pos_z(f),
288
289            Variable::ConstantArray(number, _, _) => f.write_fmt(format_args!("arrays_{number}")),
290            Variable::LocalArray(id, _, _) => {
291                write!(f, "l_arr_{id}")
292            }
293            Variable::WmmaFragment { id: index, frag } => {
294                let name = match frag.ident {
295                    FragmentIdent::A => "a",
296                    FragmentIdent::B => "b",
297                    FragmentIdent::Accumulator => "acc",
298                    FragmentIdent::_Dialect(_) => "",
299                };
300                write!(f, "frag_{name}_{index}")
301            }
302            Variable::Tmp { id, .. } => write!(f, "_tmp_{id}"),
303            Variable::Pipeline { id, .. } => write!(f, "pipeline_{id}"),
304            Variable::Barrier { id, .. } => write!(f, "barrier_{id}"),
305        }
306    }
307}
308
309impl<D: Dialect> Variable<D> {
310    pub fn is_optimized(&self) -> bool {
311        self.item().is_optimized()
312    }
313
314    /// Create a temporary variable.
315    ///
316    /// Also see [Self::tmp_declared] for a version that needs custom declaration.
317    pub fn tmp(item: Item<D>) -> Self {
318        let inc = COUNTER_TMP_VAR.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
319
320        Variable::Tmp {
321            id: inc as Id,
322            item,
323            is_declared: false,
324            is_ptr: false,
325            is_const: false,
326        }
327    }
328
329    pub fn to_const(&mut self) {
330        if let Variable::Tmp { is_const, .. } = self {
331            *is_const = true;
332        }
333    }
334
335    /// Create a temporary variable with a reinterpret_cast.
336    pub fn reinterpret_ptr(&self, f: &mut Formatter<'_>, item: Item<D>) -> Self {
337        let mut out = Self::tmp_ptr(item);
338
339        if self.is_const() {
340            out.to_const();
341        }
342
343        let elem = out.elem();
344        let qualifier = out.const_qualifier();
345        let addr_space = D::address_space_for_variable(self);
346        let out_fmt = out.fmt_left();
347
348        writeln!(
349            f,
350            "{out_fmt} = reinterpret_cast<{addr_space}{elem}{qualifier}*>({self});"
351        )
352        .unwrap();
353
354        out
355    }
356
357    /// Create a temporary pointer variable.
358    ///
359    /// Also see [Self::tmp_declared] for a version that needs custom declaration.
360    pub fn tmp_ptr(item: Item<D>) -> Self {
361        let inc = COUNTER_TMP_VAR.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
362
363        Variable::Tmp {
364            id: inc as Id,
365            item,
366            is_declared: false,
367            is_ptr: true,
368            is_const: false,
369        }
370    }
371
372    /// Create a temporary variable with a custom declaration.
373    ///
374    /// # Notes
375    ///
376    /// Calling `var.fmt_left()` will assume the variable already exist.
377    pub fn tmp_declared(item: Item<D>) -> Self {
378        let inc = COUNTER_TMP_VAR.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
379
380        Variable::Tmp {
381            id: inc as Id,
382            item,
383            is_declared: true,
384            is_ptr: false,
385            is_const: false,
386        }
387    }
388
389    pub fn optimized_args<const N: usize>(args: [Self; N]) -> OptimizedArgs<N, D> {
390        let args_after = args.map(|a| a.optimized());
391
392        let item_reference_after = args_after[0].item();
393
394        let is_optimized = args_after
395            .iter()
396            .all(|var| var.elem() == item_reference_after.elem && var.is_optimized());
397
398        if is_optimized {
399            let vectorization_before = args
400                .iter()
401                .map(|var| var.item().vectorization)
402                .max()
403                .unwrap();
404            let vectorization_after = args_after
405                .iter()
406                .map(|var| var.item().vectorization)
407                .max()
408                .unwrap();
409
410            OptimizedArgs::new(args_after, Some(vectorization_before / vectorization_after))
411        } else {
412            OptimizedArgs::new(args, None)
413        }
414    }
415
416    pub fn optimized(&self) -> Self {
417        match self {
418            Variable::GlobalInputArray(id, item) => {
419                Variable::GlobalInputArray(*id, item.optimized())
420            }
421            Variable::GlobalOutputArray(id, item) => {
422                Variable::GlobalOutputArray(*id, item.optimized())
423            }
424            Variable::LocalMut { id, item } => Variable::LocalMut {
425                id: *id,
426                item: item.optimized(),
427            },
428            Variable::LocalConst { id, item } => Variable::LocalConst {
429                id: *id,
430                item: item.optimized(),
431            },
432            Variable::Slice { id, item } => Variable::Slice {
433                id: *id,
434                item: item.optimized(),
435            },
436            Variable::Tmp {
437                id,
438                item,
439                is_declared,
440                is_ptr,
441                is_const,
442            } => Variable::Tmp {
443                id: *id,
444                item: item.optimized(),
445                is_declared: *is_declared,
446                is_ptr: *is_ptr,
447                is_const: *is_const,
448            },
449            Variable::SharedMemory(id, item, size) => {
450                let before = item.vectorization;
451                let item = item.optimized();
452                let after = item.vectorization;
453                let scaling = (before / after) as u32;
454
455                Variable::SharedMemory(*id, item, size / scaling)
456            }
457            Variable::LocalArray(id, item, size) => {
458                let before = item.vectorization;
459                let item = item.optimized();
460                let after = item.vectorization;
461                let scaling = (before / after) as u32;
462
463                Variable::LocalArray(*id, item.optimized(), size / scaling)
464            }
465            _ => *self,
466        }
467    }
468
469    pub fn is_always_scalar(&self) -> bool {
470        match self {
471            Variable::AbsolutePos => true,
472            Variable::AbsolutePosBaseName => false,
473            Variable::AbsolutePosX => true,
474            Variable::AbsolutePosY => true,
475            Variable::AbsolutePosZ => true,
476            Variable::CubeCount => true,
477            Variable::CubeCountBaseName => false,
478            Variable::CubeCountX => true,
479            Variable::CubeCountY => true,
480            Variable::CubeCountZ => true,
481            Variable::CubeDim => true,
482            Variable::CubeDimBaseName => false,
483            Variable::CubeDimX => true,
484            Variable::CubeDimY => true,
485            Variable::CubeDimZ => true,
486            Variable::CubePos => true,
487            Variable::CubePosBaseName => true,
488            Variable::CubePosX => true,
489            Variable::CubePosY => true,
490            Variable::CubePosZ => true,
491            Variable::UnitPos => true,
492            Variable::UnitPosBaseName => true,
493            Variable::UnitPosPlane => true,
494            Variable::UnitPosX => true,
495            Variable::UnitPosY => true,
496            Variable::UnitPosZ => true,
497            Variable::PlaneDim => true,
498            Variable::PlaneDimChecked => true,
499            Variable::PlanePos => true,
500            Variable::ClusterRank => true,
501            Variable::ClusterIndexX => true,
502            Variable::ClusterIndexY => true,
503            Variable::ClusterIndexZ => true,
504
505            Variable::Barrier { .. } => false,
506            Variable::ConstantArray(_, _, _) => false,
507            Variable::ConstantScalar(_, _) => true,
508            Variable::GlobalInputArray(_, _) => false,
509            Variable::GlobalOutputArray(_, _) => false,
510            Variable::GlobalScalar { .. } => true,
511            Variable::LocalArray(_, _, _) => false,
512            Variable::LocalConst { .. } => false,
513            Variable::LocalMut { .. } => false,
514            Variable::Named { .. } => false,
515            Variable::Pipeline { .. } => false,
516            Variable::SharedMemory(_, _, _) => false,
517            Variable::Slice { .. } => false,
518            Variable::Tmp { .. } => false,
519            Variable::WmmaFragment { .. } => false,
520            Variable::TensorMap { .. } => false,
521        }
522    }
523
524    pub fn index(&self, index: usize) -> IndexedVariable<D> {
525        IndexedVariable {
526            var: *self,
527            index,
528            optimized: self.is_optimized(),
529        }
530    }
531
532    pub fn const_qualifier(&self) -> &str {
533        if self.is_const() { " const" } else { "" }
534    }
535
536    pub fn id(&self) -> Option<Id> {
537        match self {
538            Variable::GlobalInputArray(id, ..) => Some(*id),
539            Variable::GlobalOutputArray(id, ..) => Some(*id),
540            Variable::GlobalScalar { id, .. } => Some(*id),
541            Variable::ConstantArray(id, ..) => Some(*id),
542            Variable::LocalMut { id, .. } => Some(*id),
543            Variable::LocalConst { id, .. } => Some(*id),
544            Variable::Slice { id, .. } => Some(*id),
545            Variable::SharedMemory(id, ..) => Some(*id),
546            Variable::LocalArray(id, ..) => Some(*id),
547            Variable::WmmaFragment { id, .. } => Some(*id),
548            Variable::Pipeline { id, .. } => Some(*id),
549            Variable::Barrier { id, .. } => Some(*id),
550            Variable::Tmp { id, .. } => Some(*id),
551            _ => None,
552        }
553    }
554
555    /// Format variable for a pointer argument. Slices and buffers are already pointers, so we
556    /// just leave them as is to avoid accidental double pointers
557    pub fn fmt_ptr(&self) -> String {
558        match self {
559            Variable::Slice { .. }
560            | Variable::SharedMemory(_, _, _)
561            | Variable::GlobalInputArray(_, _)
562            | Variable::GlobalOutputArray(_, _) => format!("{self}"),
563            _ => format!("&{self}"),
564        }
565    }
566}
567
568impl<D: Dialect> FmtLeft for Variable<D> {
569    fn fmt_left(&self) -> String {
570        match self {
571            Self::LocalConst { item, .. } => match item.elem {
572                Elem::Atomic(_) => {
573                    format!("{item}* {self}")
574                }
575                _ => {
576                    format!("const {item} {self}")
577                }
578            },
579            Variable::Tmp {
580                item,
581                is_declared,
582                is_ptr,
583                is_const,
584                ..
585            } => {
586                if *is_declared {
587                    return format!("{self}");
588                }
589                if *is_ptr {
590                    if *is_const {
591                        return format!("const {item} *{self}");
592                    }
593                    return format!("{item} *{self}");
594                }
595
596                format!("{item} {self}")
597            }
598            var => format!("{var}"),
599        }
600    }
601}
602
603#[derive(Debug, Clone)]
604pub struct IndexedVariable<D: Dialect> {
605    var: Variable<D>,
606    optimized: bool,
607    index: usize,
608}
609
610impl<D: Dialect> Component<D> for IndexedVariable<D> {
611    fn item(&self) -> Item<D> {
612        self.var.item()
613    }
614
615    fn index(&self, index: usize) -> IndexedVariable<D> {
616        self.var.index(index)
617    }
618
619    fn is_const(&self) -> bool {
620        self.var.is_const()
621    }
622}
623
624impl<D: Dialect> Display for IndexedVariable<D> {
625    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
626        let var = &self.var;
627        let ref_ = matches!(var, Variable::LocalConst { .. })
628            .then_some("const&")
629            .unwrap_or("&");
630
631        if self.var.item().vectorization > 1 {
632            if self.optimized {
633                let item = self.var.item();
634                let addr_space = D::address_space_for_variable(&self.var);
635                write!(
636                    f,
637                    "(reinterpret_cast<{addr_space}{item} {ref_}>({var})).i_{}",
638                    self.index
639                )
640            } else {
641                write!(f, "{var}.i_{}", self.index)
642            }
643        } else if self.optimized {
644            let item = self.var.item();
645            let addr_space = D::address_space_for_variable(&self.var);
646            write!(f, "reinterpret_cast<{addr_space}{item} {ref_}>({var})")
647        } else {
648            write!(f, "{var}")
649        }
650    }
651}
652
653impl<D: Dialect> FmtLeft for IndexedVariable<D> {
654    fn fmt_left(&self) -> String {
655        let var = &self.var;
656        let ref_ = matches!(var, Variable::LocalConst { .. })
657            .then_some("const&")
658            .unwrap_or("&");
659
660        let name = if self.var.item().vectorization > 1 {
661            if self.optimized {
662                let item = self.var.item();
663                let addr_space = D::address_space_for_variable(&self.var);
664                format!(
665                    "(reinterpret_cast<{addr_space}{item} {ref_}>({var})).i_{}",
666                    self.index
667                )
668            } else {
669                format!("{var}.i_{}", self.index)
670            }
671        } else {
672            format!("{var}")
673        };
674        match var {
675            Variable::LocalConst { item, .. } => format!("const {item} {name}"),
676            Variable::Tmp { item, is_ptr, .. } => {
677                if *is_ptr {
678                    format!("{item} *{name}")
679                } else {
680                    format!("{item} {name}")
681                }
682            }
683            _ => name,
684        }
685    }
686}
687
688impl FmtLeft for &String {
689    fn fmt_left(&self) -> String {
690        self.to_string()
691    }
692}
cubecl_cpp/shared/variable.rs

cubecl_cpp/shared/
variable.rs
