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

cubecl_cpp/shared/
variable.rs
