cubecl_cpp/shared/
instruction.rs

1use cubecl_core::ir::Id;
2
3use crate::shared::{Builtin, FmtLeft};
4
5use super::{
6    Component, Dialect, Elem, Item, Value, WarpInstruction, WmmaInstruction, barrier::BarrierOps,
7    binary::*, unary::*,
8};
9use std::{
10    borrow::Cow,
11    fmt::{Display, Formatter, Write},
12    marker::PhantomData,
13};
14
15pub(crate) const INFO_NAME: &str = "info";
16pub(crate) const DYNAMIC_META_NAME: &str = "dynamic_meta";
17pub(crate) const STATIC_META_NAME: &str = "info.static_meta";
18
19#[derive(Debug, Clone, Copy)]
20pub struct BinaryInstruction<D: Dialect> {
21    pub lhs: Value<D>,
22    pub rhs: Value<D>,
23    pub out: Value<D>,
24}
25
26#[derive(Debug, Clone)]
27pub struct IndexInstruction<D: Dialect> {
28    pub list: Value<D>,
29    pub index: Value<D>,
30    pub out: Value<D>,
31}
32
33#[derive(Debug, Clone, Copy)]
34pub struct UnaryInstruction<D: Dialect> {
35    pub input: Value<D>,
36    pub out: Value<D>,
37}
38
39#[derive(Debug, Clone)]
40pub enum Instruction<D: Dialect> {
41    Metadata {
42        info_offset: Value<D>,
43        out: Value<D>,
44    },
45    ExtendedMetadata {
46        info_offset: Value<D>,
47        dim: Value<D>,
48        out: Value<D>,
49    },
50    ConstLength {
51        length: usize,
52        out: Value<D>,
53    },
54    SliceLength {
55        input: Value<D>,
56        out: Value<D>,
57    },
58    DeclareVariable {
59        val: Value<D>,
60        value_ty: Item<D>,
61    },
62    Add(BinaryInstruction<D>),
63    SaturatingAdd(BinaryInstruction<D>),
64    Fma {
65        a: Value<D>,
66        b: Value<D>,
67        c: Value<D>,
68        out: Value<D>,
69    },
70    Div(BinaryInstruction<D>),
71    Rem(BinaryInstruction<D>),
72    ModFloor(BinaryInstruction<D>),
73    FastDiv(BinaryInstruction<D>),
74    FastRecip(UnaryInstruction<D>),
75    Mul(BinaryInstruction<D>),
76    Sub(BinaryInstruction<D>),
77    SaturatingSub(BinaryInstruction<D>),
78    HiMul(BinaryInstruction<D>),
79    Index(IndexInstruction<D>),
80    Assign(UnaryInstruction<D>),
81    ReadBuiltin {
82        builtin: Builtin<D>,
83        out: Value<D>,
84    },
85    ReadScalar {
86        id: Id,
87        out: Value<D>,
88    },
89    Store(UnaryInstruction<D>),
90    Load(UnaryInstruction<D>),
91    SpecialCast(UnaryInstruction<D>),
92    RangeLoop {
93        i: Value<D>,
94        start: Value<D>,
95        end: Value<D>,
96        step: Option<Value<D>>,
97        inclusive: bool,
98        instructions: Vec<Self>,
99    },
100    VecInit {
101        inputs: Vec<Value<D>>,
102        out: Value<D>,
103    },
104    InsertComponent {
105        vector: Value<D>,
106        index: Value<D>,
107        value: Value<D>,
108        out: Value<D>,
109    },
110    ExtractComponent(BinaryInstruction<D>),
111    Loop {
112        instructions: Vec<Self>,
113    },
114    If {
115        cond: Value<D>,
116        instructions: Vec<Self>,
117    },
118    IfElse {
119        cond: Value<D>,
120        instructions_if: Vec<Self>,
121        instructions_else: Vec<Self>,
122    },
123    Select {
124        cond: Value<D>,
125        then: Value<D>,
126        or_else: Value<D>,
127        out: Value<D>,
128    },
129    Switch {
130        value: Value<D>,
131        instructions_default: Vec<Self>,
132        instructions_cases: Vec<(Value<D>, Vec<Self>)>,
133    },
134    Slice {
135        input: Value<D>,
136        start: Value<D>,
137        end: Value<D>,
138        out: Value<D>,
139    },
140    CheckedSlice {
141        input: Value<D>,
142        start: Value<D>,
143        end: Value<D>,
144        out: Value<D>,
145        len: Value<D>,
146    },
147    ReinterpretSlice {
148        input: Value<D>,
149        vector_size: u32,
150        out: Value<D>,
151    },
152    Return,
153    Break,
154    Unreachable,
155    Equal(BinaryInstruction<D>),
156    NotEqual(BinaryInstruction<D>),
157    Lower(BinaryInstruction<D>),
158    Greater(BinaryInstruction<D>),
159    LowerEqual(BinaryInstruction<D>),
160    GreaterEqual(BinaryInstruction<D>),
161    Erf(UnaryInstruction<D>),
162    BitwiseOr(BinaryInstruction<D>),
163    BitwiseAnd(BinaryInstruction<D>),
164    BitwiseXor(BinaryInstruction<D>),
165    CountBits(UnaryInstruction<D>),
166    ReverseBits(UnaryInstruction<D>),
167    ShiftLeft(BinaryInstruction<D>),
168    ShiftRight(BinaryInstruction<D>),
169    BitwiseNot(UnaryInstruction<D>),
170    LeadingZeros(UnaryInstruction<D>),
171    TrailingZeros(UnaryInstruction<D>),
172    FindFirstSet(UnaryInstruction<D>),
173    Abs(UnaryInstruction<D>),
174    Exp(UnaryInstruction<D>),
175    FastExp(UnaryInstruction<D>),
176    Log(UnaryInstruction<D>),
177    FastLog(UnaryInstruction<D>),
178    Log1p(UnaryInstruction<D>),
179    Expm1(UnaryInstruction<D>),
180    Cos(UnaryInstruction<D>),
181    Sin(UnaryInstruction<D>),
182    Tan(UnaryInstruction<D>),
183    Tanh(UnaryInstruction<D>),
184    Sinh(UnaryInstruction<D>),
185    Cosh(UnaryInstruction<D>),
186    ArcCos(UnaryInstruction<D>),
187    ArcSin(UnaryInstruction<D>),
188    ArcTan(UnaryInstruction<D>),
189    ArcSinh(UnaryInstruction<D>),
190    ArcCosh(UnaryInstruction<D>),
191    ArcTanh(UnaryInstruction<D>),
192    Degrees(UnaryInstruction<D>),
193    Radians(UnaryInstruction<D>),
194    ArcTan2(BinaryInstruction<D>),
195    FastSin(UnaryInstruction<D>),
196    FastCos(UnaryInstruction<D>),
197    FastTanh(UnaryInstruction<D>),
198    Powf(BinaryInstruction<D>),
199    FastPowf(BinaryInstruction<D>),
200    Powi(BinaryInstruction<D>),
201    Hypot(BinaryInstruction<D>),
202    Rhypot(BinaryInstruction<D>),
203    Sqrt(UnaryInstruction<D>),
204    FastSqrt(UnaryInstruction<D>),
205    InverseSqrt(UnaryInstruction<D>),
206    FastInverseSqrt(UnaryInstruction<D>),
207    Min(BinaryInstruction<D>),
208    Max(BinaryInstruction<D>),
209    Not(UnaryInstruction<D>),
210    Or(BinaryInstruction<D>),
211    And(BinaryInstruction<D>),
212    Clamp {
213        input: Value<D>,
214        min_value: Value<D>,
215        max_value: Value<D>,
216        out: Value<D>,
217    },
218    IsNan(UnaryInstruction<D>),
219    IsInf(UnaryInstruction<D>),
220    SyncThreads,
221    SyncWarp,
222    ThreadFence,
223    ProxyAsyncToSharedFence,
224    BulkCommitGroup,
225    BulkWaitGroup {
226        max_pending: u32,
227    },
228    BulkWaitGroupRead {
229        max_pending: u32,
230    },
231    TmaReplacePointer {
232        buffer: Value<D>,
233        offset: Value<D>,
234        tensor_map: Value<D>,
235        out: Value<D>,
236    },
237    Round(UnaryInstruction<D>),
238    Ceil(UnaryInstruction<D>),
239    Trunc(UnaryInstruction<D>),
240    Floor(UnaryInstruction<D>),
241    Warp(WarpInstruction<D>),
242    Wmma(WmmaInstruction<D>),
243    Bitcast(UnaryInstruction<D>),
244    AtomicLoad(UnaryInstruction<D>),
245    AtomicStore(UnaryInstruction<D>),
246    AtomicSwap(BinaryInstruction<D>),
247    AtomicAdd(BinaryInstruction<D>),
248    AtomicSub(BinaryInstruction<D>),
249    AtomicMax(BinaryInstruction<D>),
250    AtomicMin(BinaryInstruction<D>),
251    AtomicAnd(BinaryInstruction<D>),
252    AtomicOr(BinaryInstruction<D>),
253    AtomicXor(BinaryInstruction<D>),
254    AtomicCAS {
255        input: Value<D>,
256        cmp: Value<D>,
257        val: Value<D>,
258        out: Value<D>,
259    },
260    Neg(UnaryInstruction<D>),
261    Magnitude(UnaryInstruction<D>),
262    FastMagnitude(UnaryInstruction<D>),
263    Normalize(UnaryInstruction<D>),
264    FastNormalize(UnaryInstruction<D>),
265    Dot(BinaryInstruction<D>),
266    VectorSum(UnaryInstruction<D>),
267    Copy {
268        source: Value<D>,
269        dest: Value<D>,
270        len: u32,
271    },
272    Printf {
273        format_string: String,
274        args: Vec<Value<D>>,
275    },
276    Comment {
277        content: String,
278    },
279    Barrier(BarrierOps<D>),
280    MemCopyAsyncTensorSharedToGlobal {
281        smem_buffer: Value<D>,
282        tensor_map: Value<D>,
283        indices: Vec<Value<D>>,
284    },
285    Line {
286        file: Cow<'static, str>,
287        line: u32,
288    },
289}
290
291impl<D: Dialect> Display for Instruction<D> {
292    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
293        match self {
294            Instruction::Return => f.write_str("return;"),
295            Instruction::Break => f.write_str("break;"),
296            Instruction::Unreachable => D::compile_unreachable(f),
297            Instruction::DeclareVariable { val, value_ty } => {
298                match value_ty {
299                    Item::Fragment(_) => {
300                        D::compile_wmma_fragment_declaration(f, val, value_ty)?;
301                    }
302                    item => {
303                        writeln!(f, "{item} {val}_store;")?;
304                    }
305                };
306                writeln!(f, "{} {val} = &{val}_store;", val.item())
307            }
308            Instruction::Add(it) => Add::format(f, &it.lhs, &it.rhs, &it.out),
309            Instruction::SaturatingAdd(it) => SaturatingAdd::format(f, &it.lhs, &it.rhs, &it.out),
310            Instruction::Slice {
311                input,
312                start,
313                end,
314                out,
315            } => {
316                let item = out.item();
317                let addr_space = D::address_space_for_value(input);
318                writeln!(f, "const uint {out}_length = {end} - {start};")?;
319                writeln!(f, "{addr_space}{item} *{out} = {input} + {start};")
320            }
321            Instruction::CheckedSlice {
322                input,
323                start,
324                end,
325                out,
326                len,
327            } => {
328                let item = out.item();
329                let addr_space = D::address_space_for_value(input);
330                writeln!(f, "const uint {out}_length = min({len}, {end}) - {start};")?;
331                writeln!(f, "{addr_space}{item} *{out} = {input} + {start};")
332            }
333            Instruction::ReinterpretSlice {
334                input,
335                vector_size,
336                out,
337            } => {
338                let item = Item::new(out.elem(), *vector_size as usize);
339                let addr_space = D::address_space_for_value(input);
340
341                writeln!(
342                    f,
343                    "{addr_space}{item} *{out} = reinterpret_cast<{item}*>({input});"
344                )
345            }
346            Instruction::Mul(it) => Mul::format(f, &it.lhs, &it.rhs, &it.out),
347            Instruction::Div(it) => Div::format(f, &it.lhs, &it.rhs, &it.out),
348            Instruction::FastDiv(it) => FastDiv::format(f, &it.lhs, &it.rhs, &it.out),
349            Instruction::FastRecip(it) => FastRecip::format(f, &it.input, &it.out),
350            Instruction::Sub(it) => Sub::format(f, &it.lhs, &it.rhs, &it.out),
351            Instruction::SaturatingSub(it) => SaturatingSub::format(f, &it.lhs, &it.rhs, &it.out),
352            Instruction::HiMul(it) => HiMul::format(f, &it.lhs, &it.rhs, &it.out),
353            Instruction::ModFloor(inst) => ModFloor::format(f, &inst.lhs, &inst.rhs, &inst.out),
354            Instruction::BitwiseOr(it) => BitwiseOr::format(f, &it.lhs, &it.rhs, &it.out),
355            Instruction::BitwiseAnd(it) => BitwiseAnd::format(f, &it.lhs, &it.rhs, &it.out),
356            Instruction::BitwiseXor(it) => BitwiseXor::format(f, &it.lhs, &it.rhs, &it.out),
357            Instruction::CountBits(it) => CountBits::format(f, &it.input, &it.out),
358            Instruction::ReverseBits(it) => ReverseBits::format(f, &it.input, &it.out),
359            Instruction::LeadingZeros(it) => LeadingZeros::format(f, &it.input, &it.out),
360            Instruction::TrailingZeros(it) => TrailingZeros::format(f, &it.input, &it.out),
361            Instruction::FindFirstSet(it) => FindFirstSet::format(f, &it.input, &it.out),
362            Instruction::ShiftLeft(it) => ShiftLeft::format(f, &it.lhs, &it.rhs, &it.out),
363            Instruction::ShiftRight(it) => ShiftRight::format(f, &it.lhs, &it.rhs, &it.out),
364            Instruction::Index(it) => Index::format(f, &it.list, &it.index, &it.out),
365            Instruction::Copy { source, dest, len } => {
366                for i in 0..*len {
367                    writeln!(f, "*({dest} + {i}) = *({source} + {i});")?;
368                }
369                Ok(())
370            }
371            Instruction::Assign(it) => Assign::format(f, &it.input, &it.out),
372            Instruction::Store(it) => {
373                writeln!(f, "*{} = {};", it.out, it.input)
374            }
375            Instruction::Load(it) => {
376                let out = it.out.fmt_left();
377                writeln!(f, "{out} = *{};", it.input)
378            }
379            Instruction::RangeLoop {
380                i,
381                start,
382                end,
383                step,
384                inclusive,
385                instructions,
386            } => {
387                let increment = step
388                    .map(|step| format!("*{i} += {step}"))
389                    .unwrap_or_else(|| format!("++*{i}"));
390                let cmp = if *inclusive { "<=" } else { "<" };
391                write!(
392                    f,
393                    "
394for (*{i} = {start}; *{i} {cmp} {end}; {increment}) {{
395"
396                )?;
397                for instruction in instructions {
398                    write!(f, "{instruction}")?;
399                }
400
401                f.write_str("}\n")
402            }
403            Instruction::Loop { instructions } => {
404                writeln!(f, "while (true) {{")?;
405                for i in instructions {
406                    write!(f, "{i}")?;
407                }
408                f.write_str("}\n")
409            }
410            Instruction::If { cond, instructions } => {
411                writeln!(f, "if ({cond}) {{")?;
412                for i in instructions {
413                    write!(f, "{i}")?;
414                }
415                f.write_str("}\n")
416            }
417            Instruction::IfElse {
418                cond,
419                instructions_if,
420                instructions_else,
421            } => {
422                writeln!(f, "if ({cond}) {{")?;
423                for i in instructions_if {
424                    write!(f, "{i}")?;
425                }
426                f.write_str("} else {\n")?;
427                for i in instructions_else {
428                    write!(f, "{i}")?;
429                }
430                f.write_str("}\n")
431            }
432            Instruction::Select {
433                cond,
434                then,
435                or_else,
436                out,
437            } => {
438                let item_or_else = or_else.item();
439                let item_then = then.item();
440                let item_out = out.item();
441
442                let vf_then = item_then.vectorization();
443                let vf_or_else = item_or_else.vectorization();
444                let vf_out = item_out.vectorization();
445                let vf_cond = cond.item().vectorization();
446
447                let item_out = out.item();
448                let cond_elem = cond.elem();
449                let out = out.fmt_left();
450
451                // It seems to always be faster to broadcast the select, because the compiler is
452                // able to output branchless instructions when the ternary is done on native types
453                // rather than cubecl defined types.
454
455                let vf = usize::max(vf_cond, vf_out);
456                let vf = usize::max(vf, vf_then);
457                let vf = usize::max(vf, vf_or_else);
458                let should_broadcast = vf > 1;
459
460                // Keep the condition here for future testing.
461                //
462                // let should_broadcast =
463                //     vf_cond > 1 || item_out != item_or_else || item_out != item_then;
464
465                if should_broadcast {
466                    writeln!(f, "{out} = {item_out} {{")?;
467                    for i in 0..vf {
468                        let theni = then.index(i);
469                        let or_elsei = or_else.index(i);
470                        let condi = cond.index(i);
471                        let condi = EnsureBoolArg {
472                            val: &condi,
473                            elem: &cond_elem,
474                        };
475
476                        writeln!(f, "({condi}) ? {theni} : {or_elsei},")?;
477                    }
478
479                    writeln!(f, "}};")
480                } else {
481                    let cond = EnsureBoolArg {
482                        val: &cond,
483                        elem: &cond_elem,
484                    };
485                    writeln!(f, "{out} = ({cond}) ? {then} : {or_else};")
486                }
487            }
488            Instruction::Switch {
489                value,
490                instructions_default,
491                instructions_cases,
492            } => {
493                writeln!(f, "switch({value}) {{")?;
494                for (value, block) in instructions_cases {
495                    write!(f, "case {value}:\n{{\n")?;
496                    for i in block {
497                        i.fmt(f)?;
498                    }
499                    f.write_str("break;\n}\n")?;
500                }
501                f.write_str("default:\n{")?;
502                for i in instructions_default {
503                    i.fmt(f)?;
504                }
505                f.write_str("break;\n}\n}\n")
506            }
507            Instruction::Metadata { info_offset, out } => {
508                let out = out.fmt_left();
509                writeln!(f, "{out} = {STATIC_META_NAME}[{info_offset}];")
510            }
511            Instruction::ExtendedMetadata {
512                info_offset,
513                dim,
514                out,
515            } => {
516                let out = out.fmt_left();
517                writeln!(
518                    f,
519                    "{out} = {DYNAMIC_META_NAME}[{STATIC_META_NAME}[{info_offset}] + {dim}];"
520                )
521            }
522            Instruction::Equal(it) => Equal::format(f, &it.lhs, &it.rhs, &it.out),
523            Instruction::NotEqual(it) => NotEqual::format(f, &it.lhs, &it.rhs, &it.out),
524            Instruction::Lower(it) => Lower::format(f, &it.lhs, &it.rhs, &it.out),
525            Instruction::Greater(it) => Greater::format(f, &it.lhs, &it.rhs, &it.out),
526            Instruction::LowerEqual(it) => LowerEqual::format(f, &it.lhs, &it.rhs, &it.out),
527            Instruction::GreaterEqual(it) => GreaterEqual::format(f, &it.lhs, &it.rhs, &it.out),
528            Instruction::Erf(it) => Erf::format(f, &it.input, &it.out),
529            Instruction::Abs(it) => Abs::format(f, &it.input, &it.out),
530            Instruction::Exp(it) => Exp::format(f, &it.input, &it.out),
531            Instruction::FastExp(it) => FastExp::format(f, &it.input, &it.out),
532            Instruction::Log(it) => Log::format(f, &it.input, &it.out),
533            Instruction::FastLog(it) => FastLog::format(f, &it.input, &it.out),
534            Instruction::Log1p(it) => Log1p::format(f, &it.input, &it.out),
535            Instruction::Expm1(it) => Expm1::format(f, &it.input, &it.out),
536            Instruction::Cos(it) => Cos::format(f, &it.input, &it.out),
537            Instruction::FastCos(it) => FastCos::format(f, &it.input, &it.out),
538            Instruction::Sin(it) => Sin::format(f, &it.input, &it.out),
539            Instruction::Tan(it) => Tan::format(f, &it.input, &it.out),
540            Instruction::Tanh(it) => Tanh::format(f, &it.input, &it.out),
541            Instruction::Sinh(it) => Sinh::format(f, &it.input, &it.out),
542            Instruction::Cosh(it) => Cosh::format(f, &it.input, &it.out),
543            Instruction::ArcCos(it) => ArcCos::format(f, &it.input, &it.out),
544            Instruction::ArcSin(it) => ArcSin::format(f, &it.input, &it.out),
545            Instruction::ArcTan(it) => ArcTan::format(f, &it.input, &it.out),
546            Instruction::ArcSinh(it) => ArcSinh::format(f, &it.input, &it.out),
547            Instruction::ArcCosh(it) => ArcCosh::format(f, &it.input, &it.out),
548            Instruction::ArcTanh(it) => ArcTanh::format(f, &it.input, &it.out),
549            Instruction::Degrees(it) => Degrees::format(f, &it.input, &it.out),
550            Instruction::Radians(it) => Radians::format(f, &it.input, &it.out),
551            Instruction::ArcTan2(it) => ArcTan2::format(f, &it.lhs, &it.rhs, &it.out),
552            Instruction::FastSin(it) => FastSin::format(f, &it.input, &it.out),
553            Instruction::FastTanh(it) => FastTanh::format(f, &it.input, &it.out),
554            Instruction::Powf(it) => Powf::format(f, &it.lhs, &it.rhs, &it.out),
555            Instruction::FastPowf(it) => FastPowf::format(f, &it.lhs, &it.rhs, &it.out),
556            Instruction::Powi(it) => Powi::format(f, &it.lhs, &it.rhs, &it.out),
557            Instruction::Hypot(it) => Hypot::format(f, &it.lhs, &it.rhs, &it.out),
558            Instruction::Rhypot(it) => Rhypot::format(f, &it.lhs, &it.rhs, &it.out),
559            Instruction::Sqrt(it) => Sqrt::format(f, &it.input, &it.out),
560            Instruction::FastSqrt(it) => FastSqrt::format(f, &it.input, &it.out),
561            Instruction::InverseSqrt(it) => InverseSqrt::format(f, &it.input, &it.out),
562            Instruction::FastInverseSqrt(it) => FastInverseSqrt::format(f, &it.input, &it.out),
563            Instruction::Max(it) => Max::format(f, &it.lhs, &it.rhs, &it.out),
564            Instruction::Min(it) => Min::format(f, &it.lhs, &it.rhs, &it.out),
565            Instruction::Not(it) => Not::format(f, &it.input, &it.out),
566            Instruction::BitwiseNot(it) => BitwiseNot::format(f, &it.input, &it.out),
567            Instruction::Or(it) => Or::format(f, &it.lhs, &it.rhs, &it.out),
568            Instruction::And(it) => And::format(f, &it.lhs, &it.rhs, &it.out),
569            Instruction::Clamp {
570                input,
571                min_value,
572                max_value,
573                out,
574            } => Clamp::format(f, input, min_value, max_value, out),
575            Instruction::IsNan(it) => IsNan::format(f, &it.input, &it.out),
576            Instruction::IsInf(it) => IsInf::format(f, &it.input, &it.out),
577            Instruction::SyncThreads => D::compile_instruction_sync_threads(f),
578            Instruction::SyncWarp => D::compile_instruction_sync_warp(f),
579            Instruction::ThreadFence => f.write_str("__threadfence();\n"),
580            Instruction::Round(it) => Round::format(f, &it.input, &it.out),
581            Instruction::Ceil(it) => Ceil::format(f, &it.input, &it.out),
582            Instruction::Trunc(it) => Trunc::format(f, &it.input, &it.out),
583            Instruction::Floor(it) => Floor::format(f, &it.input, &it.out),
584            Instruction::SliceLength { input, out } => {
585                let out = out.fmt_left();
586                writeln!(f, "{out} = {input}_length;")
587            }
588            Instruction::ConstLength { length, out } => {
589                let out = out.fmt_left();
590                writeln!(f, "{out} = {length};")
591            }
592            Instruction::Warp(it) => write!(f, "{it}"),
593            Instruction::Fma { a, b, c, out } => Fma::format(f, a, b, c, out),
594            Instruction::Wmma(it) => write!(f, "{it}"),
595            Instruction::Bitcast(UnaryInstruction { input, out }) => {
596                let qualifier = out.const_qualifier();
597                let input_item = input.item();
598                let out_item = out.item();
599
600                if out_item.size() != input_item.size() {
601                    panic!("Unsupported type for bitcasting {out_item:?} from {input_item:?}");
602                } else {
603                    let out = out.fmt_left();
604                    let addr_space = D::address_space_for_value(input);
605                    let input = match input {
606                        Value::Constant(..) => input.ensure_lvalue(f)?,
607                        _ => *input,
608                    };
609                    writeln!(
610                        f,
611                        "{out} = reinterpret_cast<{addr_space}{out_item}{qualifier}&>({input});"
612                    )
613                }
614            }
615            Instruction::AtomicAdd(BinaryInstruction { lhs, rhs, out }) => {
616                D::compile_atomic_add(f, lhs, rhs, out)
617            }
618            Instruction::AtomicAnd(BinaryInstruction { lhs, rhs, out }) => {
619                D::compile_atomic_and(f, lhs, rhs, out)
620            }
621            Instruction::AtomicCAS {
622                input,
623                cmp,
624                val,
625                out,
626            } => D::compile_atomic_cas(f, input, cmp, val, out),
627            Instruction::AtomicLoad(UnaryInstruction { input, out }) => {
628                D::compile_atomic_load(f, input, out)
629            }
630            Instruction::AtomicMax(BinaryInstruction { lhs, rhs, out }) => {
631                D::compile_atomic_max(f, lhs, rhs, out)
632            }
633            Instruction::AtomicMin(BinaryInstruction { lhs, rhs, out }) => {
634                D::compile_atomic_min(f, lhs, rhs, out)
635            }
636            Instruction::AtomicOr(BinaryInstruction { lhs, rhs, out }) => {
637                D::compile_atomic_or(f, lhs, rhs, out)
638            }
639            Instruction::AtomicStore(UnaryInstruction { input, out }) => {
640                D::compile_atomic_store(f, input, out)
641            }
642            Instruction::AtomicSub(BinaryInstruction { lhs, rhs, out }) => {
643                D::compile_atomic_sub(f, lhs, rhs, out)
644            }
645            Instruction::AtomicSwap(BinaryInstruction { lhs, rhs, out }) => {
646                D::compile_atomic_swap(f, lhs, rhs, out)
647            }
648            Instruction::AtomicXor(BinaryInstruction { lhs, rhs, out }) => {
649                D::compile_atomic_xor(f, lhs, rhs, out)
650            }
651            Instruction::Rem(inst) => Remainder::format(f, &inst.lhs, &inst.rhs, &inst.out),
652            Instruction::Neg(UnaryInstruction { input, out }) => Neg::format(f, input, out),
653            Instruction::Normalize(inst) => {
654                Normalize::<D, InverseSqrt>::format(f, &inst.input, &inst.out)
655            }
656            Instruction::FastNormalize(inst) => {
657                Normalize::<D, FastInverseSqrt>::format(f, &inst.input, &inst.out)
658            }
659            Instruction::Magnitude(inst) => Magnitude::<D, Sqrt>::format(f, &inst.input, &inst.out),
660            Instruction::FastMagnitude(inst) => {
661                Magnitude::<D, FastSqrt>::format(f, &inst.input, &inst.out)
662            }
663            Instruction::Dot(inst) => Dot::format(f, &inst.lhs, &inst.rhs, &inst.out),
664            Instruction::VectorSum(inst) => VectorSumFmt::<D>::format(f, &inst.input, &inst.out),
665            Instruction::VecInit { inputs, out } => {
666                let item = out.item();
667                let inputs = inputs
668                    .iter()
669                    .map(|input| format!("{input}"))
670                    .collect::<Vec<_>>();
671                let out = out.fmt_left();
672                writeln!(f, "{out} = {item}{{{}}};", inputs.join(","))
673            }
674            Instruction::InsertComponent {
675                vector,
676                index,
677                value,
678                out,
679            } => InsertComponent::format(f, vector, index, value, out),
680            Instruction::ExtractComponent(inst) => {
681                ExtractComponent::format(f, &inst.lhs, &inst.rhs, &inst.out)
682            }
683            Instruction::Printf {
684                format_string,
685                args,
686            } => D::compile_instruction_printf(f, format_string, args),
687            Instruction::Comment { content } => {
688                if content.contains('\n') {
689                    writeln!(f, "/* {content} */")
690                } else {
691                    writeln!(f, "// {content}")
692                }
693            }
694            Instruction::Barrier(barrier_ops) => write!(f, "{barrier_ops}"),
695            Instruction::Line { file, line } => writeln!(f, "#line {line} \"{file}\""),
696            Instruction::ProxyAsyncToSharedFence => {
697                writeln!(
698                    f,
699                    "cuda::device::experimental::fence_proxy_async_shared_cta();"
700                )
701            }
702            Instruction::BulkCommitGroup => writeln!(
703                f,
704                "cuda::device::experimental::cp_async_bulk_commit_group();"
705            ),
706            Instruction::BulkWaitGroup { max_pending } => writeln!(
707                f,
708                "cuda::device::experimental::cp_async_bulk_wait_group<{max_pending}>();"
709            ),
710            Instruction::BulkWaitGroupRead { max_pending } => writeln!(
711                f,
712                "cuda::device::experimental::cp_async_bulk_wait_group_read<{max_pending}>();"
713            ),
714            Instruction::TmaReplacePointer {
715                buffer,
716                offset,
717                tensor_map,
718                out,
719            } => {
720                let pos = Builtin::<D>::UnitPos;
721                writeln!(f, "__shared__ alignas(128) CUtensorMap {out};")?;
722                writeln!(
723                    f,
724                    "
725if({pos} == 0) {{
726    {out} = {tensor_map};
727    tensormap_replace_global_address({out}, &{buffer}[{offset}]);
728}}"
729                )?;
730                writeln!(f, "__syncthreads();")
731            }
732            Instruction::MemCopyAsyncTensorSharedToGlobal {
733                smem_buffer,
734                tensor_map,
735                indices,
736            } => {
737                let rank = indices.len();
738                let smem_ptr = smem_buffer.fmt_ptr();
739                let indices = indices.iter().rev().fold(String::new(), |mut s, it| {
740                    let _ = write!(s, "{it}, ");
741                    s
742                });
743                writeln!(
744                    f,
745                    "cuda::device::experimental::cp_async_bulk_tensor_{rank}d_shared_to_global(&{tensor_map}, {indices} {smem_ptr});"
746                )
747            }
748            Instruction::SpecialCast(UnaryInstruction { input, out }) => {
749                // Only supported in CUDA so I'm putting it here. Move to dialect if necessary.
750                #[cfg(not(feature = "cuda"))]
751                {
752                    let _ = (input, out);
753                    writeln!(
754                        f,
755                        "#error FP8/FP6/FP4 casting isn't supported outside of CUDA"
756                    )
757                }
758                #[cfg(feature = "cuda")]
759                crate::cuda::convert::special_cast::<D>(f, input, out)
760            }
761            Instruction::ReadBuiltin { builtin, out } => {
762                writeln!(f, "{} = {builtin};", out.fmt_left())
763            }
764            Instruction::ReadScalar { id, out } => {
765                let elem = *out.item().elem();
766                writeln!(f, "{} = info.scalars_{elem}[{id}];", out.fmt_left())
767            }
768        }
769    }
770}
771
772struct Fma<D: Dialect> {
773    _dialect: PhantomData<D>,
774}
775
776impl<D: Dialect> Fma<D> {
777    fn format(
778        f: &mut core::fmt::Formatter<'_>,
779        a: &Value<D>,
780        b: &Value<D>,
781        c: &Value<D>,
782        out: &Value<D>,
783    ) -> core::fmt::Result {
784        let out_item = out.item();
785
786        let out = out.fmt_left();
787        if let Item::Vector(_, num) = out_item {
788            writeln!(f, "{out} = {out_item}{{")?;
789
790            for i in 0..num {
791                let ai = a.index(i);
792                let bi = b.index(i);
793                let ci = c.index(i);
794
795                writeln!(f, "fma({ai}, {bi}, {ci}),")?;
796            }
797            f.write_str("};\n")
798        } else {
799            writeln!(f, "{out} = fma({a}, {b}, {c});")
800        }
801    }
802}
803
804struct Clamp<D: Dialect> {
805    _dialect: PhantomData<D>,
806}
807
808impl<D: Dialect> Clamp<D> {
809    fn format(
810        f: &mut core::fmt::Formatter<'_>,
811        input: &Value<D>,
812        min_value: &Value<D>,
813        max_value: &Value<D>,
814        out: &Value<D>,
815    ) -> core::fmt::Result {
816        let out_item = out.item();
817        if let Item::Vector(..) = out_item {
818            Self::unroll_vec(f, input, min_value, max_value, out)
819        } else {
820            let out = out.fmt_left();
821            write!(f, "{out} = ")?;
822            Self::format_scalar(f, *input, *min_value, *max_value, out_item)?;
823            f.write_str(";\n")
824        }
825    }
826
827    fn format_scalar(
828        f: &mut Formatter<'_>,
829        input: impl Component<D>,
830        min_value: impl Component<D>,
831        max_value: impl Component<D>,
832        item: Item<D>,
833    ) -> std::fmt::Result {
834        D::compile_instruction_max_function_name(f, item)?;
835        write!(f, "({min_value}, ")?;
836        D::compile_instruction_min_function_name(f, item)?;
837        write!(f, "({max_value}, {input}))")
838    }
839
840    fn unroll_vec(
841        f: &mut core::fmt::Formatter<'_>,
842        input: &Value<D>,
843        min_value: &Value<D>,
844        max_value: &Value<D>,
845        out: &Value<D>,
846    ) -> std::fmt::Result {
847        let optimized = Value::optimized_args([*input, *min_value, *max_value, *out]);
848        let [input, min_value, max_value, out_optimized] = optimized.args;
849
850        let item_out_original = out.item();
851        let item_out_optimized = out_optimized.item();
852
853        let index = match item_out_optimized {
854            Item::Vector(_, index) => index,
855            _ => 1,
856        };
857
858        let mut write_op = |input: &Value<D>,
859                            min_value: &Value<D>,
860                            max_value: &Value<D>,
861                            out: &Value<D>,
862                            item_out: Item<D>| {
863            let out = out.fmt_left();
864            writeln!(f, "{out} = {item_out}{{")?;
865            for i in 0..index {
866                let inputi = input.index(i);
867                let min_valuei = min_value.index(i);
868                let max_valuei = max_value.index(i);
869
870                Self::format_scalar(f, inputi, min_valuei, max_valuei, item_out)?;
871                f.write_str(", ")?;
872            }
873
874            f.write_str("};\n")
875        };
876
877        if item_out_original == item_out_optimized {
878            write_op(&input, &min_value, &max_value, out, item_out_optimized)
879        } else {
880            let out_tmp = Value::tmp(item_out_optimized);
881            write_op(&input, &min_value, &max_value, &out_tmp, item_out_optimized)?;
882            let addr_space = D::address_space_for_value(out);
883            let out = out.fmt_left();
884
885            writeln!(
886                f,
887                "{out} = reinterpret_cast<{addr_space}{item_out_original}&>({out_tmp});\n"
888            )?;
889
890            Ok(())
891        }
892    }
893}
894
895struct Magnitude<D: Dialect, S: FunctionFmt<D>> {
896    _dialect: PhantomData<D>,
897    _sqrt: PhantomData<S>,
898}
899
900impl<D: Dialect, S: FunctionFmt<D>> Magnitude<D, S> {
901    fn format(
902        f: &mut core::fmt::Formatter<'_>,
903        input: &Value<D>,
904        out: &Value<D>,
905    ) -> core::fmt::Result {
906        let num = match input.item() {
907            Item::Vector(_, vectorization) => vectorization,
908            _ => 1,
909        };
910        let elem = input.elem();
911
912        let mag = format!("{out}_mag");
913
914        // Use elem cast for the literal to support bfloat
915        writeln!(f, "{} {mag} = {}(0.0);", out.item(), out.item())?;
916
917        for i in 0..num {
918            let input_i = input.index(i);
919            writeln!(f, "{mag} += {input_i} * {input_i};")?;
920        }
921
922        let out = out.fmt_left();
923        write!(f, "{out} = ")?;
924        S::format_unary(f, &mag, elem)?;
925        f.write_str(";\n")
926    }
927}
928
929struct Normalize<D: Dialect, InvS: FunctionFmt<D>> {
930    _dialect: PhantomData<D>,
931    _rsqrt: PhantomData<InvS>,
932}
933
934impl<D: Dialect, InvS: FunctionFmt<D>> Normalize<D, InvS> {
935    fn format(
936        f: &mut core::fmt::Formatter<'_>,
937        input: &Value<D>,
938        out: &Value<D>,
939    ) -> core::fmt::Result {
940        let num = match input.item() {
941            Item::Vector(_, vectorization) => vectorization,
942            _ => 1,
943        };
944        let elem = input.elem();
945        let norm = format!("{out}_norm");
946
947        let out_item = out.item();
948        let out = out.fmt_left();
949        // Use elem cast for the literal to support bfloat
950        writeln!(f, "{elem} {norm} = {elem}(0.0);")?;
951
952        for i in 0..num {
953            let input_i = input.index(i);
954            writeln!(f, "{norm} += {input_i} * {input_i};")?;
955        }
956
957        write!(f, "{norm} = ")?;
958        InvS::format_unary(f, &norm, elem)?;
959        f.write_str(";\n")?;
960
961        if num == 1 {
962            writeln!(f, "{out} = {input} * {norm};")
963        } else {
964            write!(f, "{out} = {out_item}{{")?;
965            for i in 0..num {
966                let input_i = input.index(i);
967
968                writeln!(f, "{input_i} * {norm},")?;
969            }
970
971            f.write_str("};\n")
972        }
973    }
974}
975
976struct Dot<D: Dialect> {
977    _dialect: PhantomData<D>,
978}
979
980impl<D: Dialect> Dot<D> {
981    fn format(
982        f: &mut core::fmt::Formatter<'_>,
983        lhs: &Value<D>,
984        rhs: &Value<D>,
985        out: &Value<D>,
986    ) -> core::fmt::Result {
987        let num = match lhs.item() {
988            Item::Vector(_, vectorization) => vectorization,
989            _ => 1,
990        };
991
992        let muls = (0..num)
993            .map(|i| {
994                let lhs_i = lhs.index(i);
995                let rhs_i = rhs.index(i);
996                format!("{lhs_i} * {rhs_i}")
997            })
998            .collect::<Vec<_>>();
999
1000        let value = muls.join(" + ");
1001        if out.declare_local_ptr_backing(f)? {
1002            writeln!(f, "*{out} = {value};")
1003        } else {
1004            writeln!(f, "{} = {value};", out.fmt_left())
1005        }
1006    }
1007}
1008
1009struct VectorSumFmt<D: Dialect> {
1010    _dialect: PhantomData<D>,
1011}
1012
1013impl<D: Dialect> VectorSumFmt<D> {
1014    fn format(
1015        f: &mut core::fmt::Formatter<'_>,
1016        input: &Value<D>,
1017        out: &Value<D>,
1018    ) -> core::fmt::Result {
1019        let num = input.item().vectorization();
1020
1021        let elems = (0..num)
1022            .map(|i| format!("{}", input.index(i)))
1023            .collect::<Vec<_>>();
1024
1025        let value = elems.join(" + ");
1026        if out.declare_local_ptr_backing(f)? {
1027            writeln!(f, "*{out} = {value};")
1028        } else {
1029            writeln!(f, "{} = {value};", out.fmt_left())
1030        }
1031    }
1032}
1033
1034struct EnsureBoolArg<'a, V: Display, D: Dialect> {
1035    val: &'a V,
1036    elem: &'a Elem<D>,
1037}
1038
1039impl<V: Display, D: Dialect> Display for EnsureBoolArg<'_, V, D> {
1040    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1041        if self.elem != &Elem::Bool {
1042            write!(f, "bool({})", self.val)
1043        } else {
1044            write!(f, "{}", self.val)
1045        }
1046    }
1047}
1048
1049#[cfg(all(test, feature = "metal"))]
1050mod tests {
1051    use cubecl_core::ir::ConstantValue;
1052
1053    use crate::metal::MslDialect;
1054
1055    use super::*;
1056
1057    #[test]
1058    fn bitcast_constant_uses_lvalue() {
1059        let instruction = Instruction::<MslDialect>::Bitcast(UnaryInstruction {
1060            input: Value::Constant(ConstantValue::UInt(0x7f80_0000), Item::Scalar(Elem::U32)),
1061            out: Value::Value {
1062                id: 0,
1063                item: Item::Scalar(Elem::F32),
1064            },
1065        });
1066
1067        let source = instruction.to_string();
1068
1069        assert!(source.contains("uint _tmp_"), "{source}");
1070        assert!(
1071            source.contains("reinterpret_cast<thread float const&>(_tmp_"),
1072            "{source}"
1073        );
1074    }
1075}
cubecl_cpp/shared/instruction.rs

cubecl_cpp/shared/
instruction.rs
