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

cubecl_cpp/shared/
variable.rs
