cubecl_cpp/shared/
variable.rs

1use cubecl_core::{
2    e2m1, e2m1x2, e4m3, e5m2,
3    ir::{BarrierLevel, ConstantValue, Id},
4    ue8m0,
5};
6use std::fmt::{Display, Formatter};
7
8use crate::shared::{FP4Kind, FP8Kind};
9
10use super::{COUNTER_TMP_VAR, Dialect, Elem, Fragment, FragmentIdent, Item};
11
12pub trait Component<D: Dialect>: Display + FmtLeft {
13    fn item(&self) -> Item<D>;
14    fn is_const(&self) -> bool;
15    fn index(&self, index: usize) -> IndexedVariable<D>;
16    fn elem(&self) -> Elem<D> {
17        *self.item().elem()
18    }
19}
20
21pub trait FmtLeft: Display {
22    fn fmt_left(&self) -> String;
23}
24
25#[derive(new)]
26pub struct OptimizedArgs<const N: usize, D: Dialect> {
27    pub args: [Variable<D>; N],
28    pub optimization_factor: Option<usize>,
29}
30
31#[derive(Debug, Clone, Copy, PartialEq)]
32pub enum Variable<D: Dialect> {
33    AbsolutePos(Elem<D>),
34    AbsolutePosBaseName, // base name for XYZ
35    AbsolutePosX,
36    AbsolutePosY,
37    AbsolutePosZ,
38    UnitPos,
39    UnitPosBaseName, // base name for XYZ
40    UnitPosX,
41    UnitPosY,
42    UnitPosZ,
43    CubePos(Elem<D>),
44    CubePosBaseName, // base name for XYZ
45    CubePosX,
46    CubePosY,
47    CubePosZ,
48    CubeDim,
49    CubeDimBaseName, // base name for XYZ
50    CubeDimX,
51    CubeDimY,
52    CubeDimZ,
53    CubeCount(Elem<D>),
54    CubeCountBaseName, // base name for XYZ
55    CubeCountX,
56    CubeCountY,
57    CubeCountZ,
58    PlaneDim,
59    PlaneDimChecked,
60    PlanePos,
61    UnitPosPlane,
62    ClusterRank,
63    ClusterIndexX,
64    ClusterIndexY,
65    ClusterIndexZ,
66    GlobalInputArray(Id, Item<D>),
67    GlobalOutputArray(Id, Item<D>),
68    GlobalScalar {
69        id: Id,
70        elem: Elem<D>,
71    },
72    ConstantArray(Id, Item<D>, usize),
73    Constant(ConstantValue, Item<D>),
74    TensorMap(Id),
75    LocalMut {
76        id: Id,
77        item: Item<D>,
78    },
79    LocalConst {
80        id: Id,
81        item: Item<D>,
82    },
83    Named {
84        name: &'static str,
85        item: Item<D>,
86    },
87    Slice {
88        id: Id,
89        item: Item<D>,
90    },
91    SharedArray(Id, Item<D>, usize),
92    Shared(Id, Item<D>),
93    LocalArray(Id, Item<D>, usize),
94    WmmaFragment {
95        id: Id,
96        frag: Fragment<D>,
97    },
98    Pipeline {
99        id: Id,
100    },
101    Barrier {
102        id: Id,
103        level: BarrierLevel,
104    },
105    BarrierToken {
106        id: Id,
107        level: BarrierLevel,
108    },
109    Tmp {
110        id: Id,
111        item: Item<D>,
112        is_declared: bool,
113        is_ptr: bool,
114        is_const: bool,
115    },
116}
117
118impl<D: Dialect> Component<D> for Variable<D> {
119    fn index(&self, index: usize) -> IndexedVariable<D> {
120        self.index(index)
121    }
122
123    fn item(&self) -> Item<D> {
124        match self {
125            Variable::AbsolutePos(elem) => Item::scalar(*elem, true),
126            Variable::AbsolutePosBaseName => Item {
127                elem: Elem::U32,
128                vectorization: 3,
129                native: true,
130            },
131            Variable::AbsolutePosX => Item::scalar(Elem::U32, true),
132            Variable::AbsolutePosY => Item::scalar(Elem::U32, true),
133            Variable::AbsolutePosZ => Item::scalar(Elem::U32, true),
134            Variable::CubeCount(elem) => Item::scalar(*elem, true),
135            Variable::CubeCountBaseName => Item {
136                elem: Elem::U32,
137                vectorization: 3,
138                native: true,
139            },
140            Variable::CubeCountX => Item::scalar(Elem::U32, true),
141            Variable::CubeCountY => Item::scalar(Elem::U32, true),
142            Variable::CubeCountZ => Item::scalar(Elem::U32, true),
143            Variable::CubeDimBaseName => Item {
144                elem: Elem::U32,
145                vectorization: 3,
146                native: true,
147            },
148            Variable::CubeDim => Item::scalar(Elem::U32, true),
149            Variable::CubeDimX => Item::scalar(Elem::U32, true),
150            Variable::CubeDimY => Item::scalar(Elem::U32, true),
151            Variable::CubeDimZ => Item::scalar(Elem::U32, true),
152            Variable::CubePos(elem) => Item::scalar(*elem, true),
153            Variable::CubePosBaseName => Item {
154                elem: Elem::U32,
155                vectorization: 3,
156                native: true,
157            },
158            Variable::CubePosX => Item::scalar(Elem::U32, true),
159            Variable::CubePosY => Item::scalar(Elem::U32, true),
160            Variable::CubePosZ => Item::scalar(Elem::U32, true),
161            Variable::UnitPos => Item::scalar(Elem::U32, true),
162            Variable::UnitPosBaseName => Item {
163                elem: Elem::U32,
164                vectorization: 3,
165                native: true,
166            },
167            Variable::UnitPosX => Item::scalar(Elem::U32, true),
168            Variable::UnitPosY => Item::scalar(Elem::U32, true),
169            Variable::UnitPosZ => Item::scalar(Elem::U32, true),
170            Variable::PlaneDim => Item::scalar(Elem::U32, true),
171            Variable::PlaneDimChecked => Item::scalar(Elem::U32, true),
172            Variable::PlanePos => Item::scalar(Elem::U32, true),
173            Variable::UnitPosPlane => Item::scalar(Elem::U32, true),
174            Variable::ClusterRank => Item::scalar(Elem::U32, true),
175            Variable::ClusterIndexX => Item::scalar(Elem::U32, true),
176            Variable::ClusterIndexY => Item::scalar(Elem::U32, true),
177            Variable::ClusterIndexZ => Item::scalar(Elem::U32, true),
178            Variable::GlobalInputArray(_, e) => *e,
179            Variable::GlobalOutputArray(_, e) => *e,
180            Variable::LocalArray(_, e, _) => *e,
181            Variable::SharedArray(_, e, _) => *e,
182            Variable::Shared(_, e) => *e,
183            Variable::ConstantArray(_, e, _) => *e,
184            Variable::LocalMut { item, .. } => *item,
185            Variable::LocalConst { item, .. } => *item,
186            Variable::Named { item, .. } => *item,
187            Variable::Slice { item, .. } => *item,
188            Variable::Constant(_, e) => *e,
189            Variable::GlobalScalar { elem, .. } => Item::scalar(*elem, false),
190            Variable::WmmaFragment { frag, .. } => Item::scalar(frag.elem, false),
191            Variable::Tmp { item, .. } => *item,
192            Variable::Pipeline { .. }
193            | Variable::Barrier { .. }
194            | Variable::BarrierToken { .. } => Item::new(Elem::Bool, 1, false),
195            Variable::TensorMap(_) => unreachable!(),
196        }
197    }
198
199    fn is_const(&self) -> bool {
200        if let Variable::Tmp { is_const, .. } = self {
201            return *is_const;
202        }
203
204        matches!(
205            self,
206            Variable::LocalConst { .. } | Variable::GlobalInputArray { .. }
207        )
208    }
209}
210
211pub(crate) fn format_const<D: Dialect>(number: &ConstantValue, item: &Item<D>) -> String {
212    // minifloats are represented as raw bits, so use special handling
213    let number = match item.elem() {
214        Elem::FP4(FP4Kind::E2M1) => e2m1::from_f64(number.as_f64()).to_bits(),
215        Elem::FP4x2(FP4Kind::E2M1) => {
216            let v = number.as_f64() as f32;
217            let value = [v, v];
218            e2m1x2::from_f32_slice(&value).remove(0).to_bits()
219        }
220        Elem::FP6(_) | Elem::FP6x2(_) => {
221            todo!("FP6 constants are not yet supported")
222        }
223        Elem::FP8(FP8Kind::E4M3) => e4m3::from_f64(number.as_f64()).to_bits(),
224        Elem::FP8(FP8Kind::E5M2) => e5m2::from_f64(number.as_f64()).to_bits(),
225        Elem::FP8(FP8Kind::UE8M0) => ue8m0::from_f64(number.as_f64()).to_bits(),
226        _ => {
227            return format!("{number}");
228        }
229    };
230    format!("{number}")
231}
232
233impl<D: Dialect> Display for Variable<D> {
234    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
235        match self {
236            Variable::GlobalInputArray(id, _) => f.write_fmt(format_args!("buffer_{id}")),
237            Variable::GlobalOutputArray(id, _) => write!(f, "buffer_{id}"),
238            Variable::TensorMap(id) => write!(f, "tensor_map_{id}"),
239            Variable::LocalMut { id, .. } => f.write_fmt(format_args!("l_mut_{id}")),
240            Variable::LocalConst { id, .. } => f.write_fmt(format_args!("l_{id}")),
241            Variable::Named { name, .. } => f.write_fmt(format_args!("{name}")),
242            Variable::Slice { id, .. } => {
243                write!(f, "slice_{id}")
244            }
245            Variable::GlobalScalar { id, elem } => write!(f, "info.scalars_{elem}[{id}]"),
246            Variable::Constant(number, item) if item.vectorization <= 1 => {
247                let value = format_const(number, item);
248                write!(f, "{item}({value})")
249            }
250            Variable::Constant(number, item) => {
251                let number = format_const(number, item);
252                let values = (0..item.vectorization)
253                    .map(|_| format!("{}({number})", item.elem()))
254                    .collect::<Vec<_>>();
255                write!(f, "{item} {{ {} }}", values.join(","))
256            }
257            Variable::SharedArray(number, _, _) | Variable::Shared(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 is_optimized = args_after.iter().all(|var| var.is_optimized());
400
401        if is_optimized {
402            let vectorization_before = args
403                .iter()
404                .map(|var| var.item().vectorization)
405                .max()
406                .unwrap();
407            let vectorization_after = args_after
408                .iter()
409                .map(|var| var.item().vectorization)
410                .max()
411                .unwrap();
412
413            OptimizedArgs::new(args_after, Some(vectorization_before / vectorization_after))
414        } else {
415            OptimizedArgs::new(args, None)
416        }
417    }
418
419    pub fn optimized(&self) -> Self {
420        match self {
421            Variable::GlobalInputArray(id, item) => {
422                Variable::GlobalInputArray(*id, item.optimized())
423            }
424            Variable::GlobalOutputArray(id, item) => {
425                Variable::GlobalOutputArray(*id, item.optimized())
426            }
427            Variable::LocalMut { id, item } => Variable::LocalMut {
428                id: *id,
429                item: item.optimized(),
430            },
431            Variable::LocalConst { id, item } => Variable::LocalConst {
432                id: *id,
433                item: item.optimized(),
434            },
435            Variable::Slice { id, item } => Variable::Slice {
436                id: *id,
437                item: item.optimized(),
438            },
439            Variable::Tmp {
440                id,
441                item,
442                is_declared,
443                is_ptr,
444                is_const,
445            } => Variable::Tmp {
446                id: *id,
447                item: item.optimized(),
448                is_declared: *is_declared,
449                is_ptr: *is_ptr,
450                is_const: *is_const,
451            },
452            Variable::SharedArray(id, item, size) => {
453                let before = item.vectorization;
454                let item = item.optimized();
455                let after = item.vectorization;
456                let scaling = before / after;
457
458                Variable::SharedArray(*id, item, size / scaling)
459            }
460            Variable::LocalArray(id, item, size) => {
461                let before = item.vectorization;
462                let item = item.optimized();
463                let after = item.vectorization;
464                let scaling = before / after;
465
466                Variable::LocalArray(*id, item.optimized(), size / scaling)
467            }
468            _ => *self,
469        }
470    }
471
472    pub fn is_always_scalar(&self) -> bool {
473        match self {
474            Variable::AbsolutePos(_) => true,
475            Variable::AbsolutePosBaseName => false,
476            Variable::AbsolutePosX => true,
477            Variable::AbsolutePosY => true,
478            Variable::AbsolutePosZ => true,
479            Variable::CubeCount(_) => true,
480            Variable::CubeCountBaseName => false,
481            Variable::CubeCountX => true,
482            Variable::CubeCountY => true,
483            Variable::CubeCountZ => true,
484            Variable::CubeDim => true,
485            Variable::CubeDimBaseName => false,
486            Variable::CubeDimX => true,
487            Variable::CubeDimY => true,
488            Variable::CubeDimZ => true,
489            Variable::CubePos(_) => true,
490            Variable::CubePosBaseName => true,
491            Variable::CubePosX => true,
492            Variable::CubePosY => true,
493            Variable::CubePosZ => true,
494            Variable::UnitPos => true,
495            Variable::UnitPosBaseName => true,
496            Variable::UnitPosPlane => true,
497            Variable::UnitPosX => true,
498            Variable::UnitPosY => true,
499            Variable::UnitPosZ => true,
500            Variable::PlaneDim => true,
501            Variable::PlaneDimChecked => true,
502            Variable::PlanePos => true,
503            Variable::ClusterRank => true,
504            Variable::ClusterIndexX => true,
505            Variable::ClusterIndexY => true,
506            Variable::ClusterIndexZ => true,
507
508            Variable::Barrier { .. } => false,
509            Variable::BarrierToken { .. } => false,
510            Variable::ConstantArray(_, _, _) => false,
511            Variable::Constant(_, _) => true,
512            Variable::GlobalInputArray(_, _) => false,
513            Variable::GlobalOutputArray(_, _) => false,
514            Variable::GlobalScalar { .. } => true,
515            Variable::LocalArray(_, _, _) => false,
516            Variable::LocalConst { .. } => false,
517            Variable::LocalMut { .. } => false,
518            Variable::Named { .. } => false,
519            Variable::Pipeline { .. } => false,
520            Variable::SharedArray(_, _, _) => false,
521            Variable::Shared(_, _) => false,
522            Variable::Slice { .. } => false,
523            Variable::Tmp { .. } => false,
524            Variable::WmmaFragment { .. } => false,
525            Variable::TensorMap { .. } => false,
526        }
527    }
528
529    pub fn index(&self, index: usize) -> IndexedVariable<D> {
530        IndexedVariable {
531            var: *self,
532            index,
533            optimized: self.is_optimized(),
534        }
535    }
536
537    pub fn const_qualifier(&self) -> &str {
538        if self.is_const() { " const" } else { "" }
539    }
540
541    pub fn id(&self) -> Option<Id> {
542        match self {
543            Variable::GlobalInputArray(id, ..) => Some(*id),
544            Variable::GlobalOutputArray(id, ..) => Some(*id),
545            Variable::GlobalScalar { id, .. } => Some(*id),
546            Variable::ConstantArray(id, ..) => Some(*id),
547            Variable::LocalMut { id, .. } => Some(*id),
548            Variable::LocalConst { id, .. } => Some(*id),
549            Variable::Slice { id, .. } => Some(*id),
550            Variable::Shared(id, ..) => Some(*id),
551            Variable::SharedArray(id, ..) => Some(*id),
552            Variable::LocalArray(id, ..) => Some(*id),
553            Variable::WmmaFragment { id, .. } => Some(*id),
554            Variable::Pipeline { id, .. } => Some(*id),
555            Variable::Barrier { id, .. } => Some(*id),
556            Variable::Tmp { id, .. } => Some(*id),
557            _ => None,
558        }
559    }
560
561    /// Format variable for a pointer argument. Slices and buffers are already pointers, so we
562    /// just leave them as is to avoid accidental double pointers
563    pub fn fmt_ptr(&self) -> String {
564        match self {
565            Variable::Slice { .. }
566            | Variable::SharedArray(_, _, _)
567            | Variable::GlobalInputArray(_, _)
568            | Variable::GlobalOutputArray(_, _) => format!("{self}"),
569            _ => format!("&{self}"),
570        }
571    }
572
573    /// Format an item with a specific type, casting if necessary
574    pub fn fmt_cast_to(&self, item: Item<D>) -> String {
575        if self.item() == item {
576            self.to_string()
577        } else {
578            format!("{item}({self})")
579        }
580    }
581}
582
583impl<D: Dialect> FmtLeft for Variable<D> {
584    fn fmt_left(&self) -> String {
585        match self {
586            Self::LocalConst { item, .. } => match item.elem {
587                Elem::Atomic(_) => {
588                    format!("{item}* {self}")
589                }
590                _ => {
591                    format!("const {item} {self}")
592                }
593            },
594            Variable::Tmp {
595                item,
596                is_declared,
597                is_ptr,
598                is_const,
599                ..
600            } => {
601                if *is_declared {
602                    return format!("{self}");
603                }
604                if *is_ptr {
605                    if *is_const {
606                        return format!("const {item} *{self}");
607                    }
608                    return format!("{item} *{self}");
609                }
610
611                format!("{item} {self}")
612            }
613            var => format!("{var}"),
614        }
615    }
616}
617
618#[derive(Debug, Clone)]
619pub struct IndexedVariable<D: Dialect> {
620    var: Variable<D>,
621    optimized: bool,
622    index: usize,
623}
624
625impl<D: Dialect> Component<D> for IndexedVariable<D> {
626    fn item(&self) -> Item<D> {
627        self.var.item()
628    }
629
630    fn index(&self, index: usize) -> IndexedVariable<D> {
631        self.var.index(index)
632    }
633
634    fn is_const(&self) -> bool {
635        self.var.is_const()
636    }
637}
638
639impl<D: Dialect> Display for IndexedVariable<D> {
640    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
641        let var = &self.var;
642
643        if let Variable::Constant(value, item) = var {
644            let value = format_const(value, item);
645            return write!(f, "{}({value})", item.elem());
646        }
647
648        let ref_ = matches!(var, Variable::LocalConst { .. })
649            .then_some("const&")
650            .unwrap_or("&");
651
652        if self.var.item().vectorization > 1 {
653            if self.optimized {
654                let item = self.var.item();
655                let addr_space = D::address_space_for_variable(&self.var);
656                write!(
657                    f,
658                    "(reinterpret_cast<{addr_space}{item} {ref_}>({var})).i_{}",
659                    self.index
660                )
661            } else {
662                write!(f, "{var}.i_{}", self.index)
663            }
664        } else if self.optimized {
665            let item = self.var.item();
666            let addr_space = D::address_space_for_variable(&self.var);
667            write!(f, "reinterpret_cast<{addr_space}{item} {ref_}>({var})")
668        } else {
669            write!(f, "{var}")
670        }
671    }
672}
673
674impl<D: Dialect> FmtLeft for IndexedVariable<D> {
675    fn fmt_left(&self) -> String {
676        let var = &self.var;
677        let ref_ = matches!(var, Variable::LocalConst { .. })
678            .then_some("const&")
679            .unwrap_or("&");
680
681        let name = if self.var.item().vectorization > 1 {
682            if self.optimized {
683                let item = self.var.item();
684                let addr_space = D::address_space_for_variable(&self.var);
685                format!(
686                    "(reinterpret_cast<{addr_space}{item} {ref_}>({var})).i_{}",
687                    self.index
688                )
689            } else {
690                format!("{var}.i_{}", self.index)
691            }
692        } else {
693            format!("{var}")
694        };
695        match var {
696            Variable::LocalConst { item, .. } => format!("const {item} {name}"),
697            Variable::Tmp { item, is_ptr, .. } => {
698                if *is_ptr {
699                    format!("{item} *{name}")
700                } else {
701                    format!("{item} {name}")
702                }
703            }
704            _ => name,
705        }
706    }
707}
708
709impl FmtLeft for &String {
710    fn fmt_left(&self) -> String {
711        self.to_string()
712    }
713}
cubecl_cpp/shared/variable.rs

cubecl_cpp/shared/
variable.rs
