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                    writeln!(
606                        f,
607                        "{out} = reinterpret_cast<{addr_space}{out_item}{qualifier}&>({input});"
608                    )
609                }
610            }
611            Instruction::AtomicAdd(BinaryInstruction { lhs, rhs, out }) => {
612                D::compile_atomic_add(f, lhs, rhs, out)
613            }
614            Instruction::AtomicAnd(BinaryInstruction { lhs, rhs, out }) => {
615                D::compile_atomic_and(f, lhs, rhs, out)
616            }
617            Instruction::AtomicCAS {
618                input,
619                cmp,
620                val,
621                out,
622            } => D::compile_atomic_cas(f, input, cmp, val, out),
623            Instruction::AtomicLoad(UnaryInstruction { input, out }) => {
624                D::compile_atomic_load(f, input, out)
625            }
626            Instruction::AtomicMax(BinaryInstruction { lhs, rhs, out }) => {
627                D::compile_atomic_max(f, lhs, rhs, out)
628            }
629            Instruction::AtomicMin(BinaryInstruction { lhs, rhs, out }) => {
630                D::compile_atomic_min(f, lhs, rhs, out)
631            }
632            Instruction::AtomicOr(BinaryInstruction { lhs, rhs, out }) => {
633                D::compile_atomic_or(f, lhs, rhs, out)
634            }
635            Instruction::AtomicStore(UnaryInstruction { input, out }) => {
636                D::compile_atomic_store(f, input, out)
637            }
638            Instruction::AtomicSub(BinaryInstruction { lhs, rhs, out }) => {
639                D::compile_atomic_sub(f, lhs, rhs, out)
640            }
641            Instruction::AtomicSwap(BinaryInstruction { lhs, rhs, out }) => {
642                D::compile_atomic_swap(f, lhs, rhs, out)
643            }
644            Instruction::AtomicXor(BinaryInstruction { lhs, rhs, out }) => {
645                D::compile_atomic_xor(f, lhs, rhs, out)
646            }
647            Instruction::Rem(inst) => Remainder::format(f, &inst.lhs, &inst.rhs, &inst.out),
648            Instruction::Neg(UnaryInstruction { input, out }) => Neg::format(f, input, out),
649            Instruction::Normalize(inst) => {
650                Normalize::<D, InverseSqrt>::format(f, &inst.input, &inst.out)
651            }
652            Instruction::FastNormalize(inst) => {
653                Normalize::<D, FastInverseSqrt>::format(f, &inst.input, &inst.out)
654            }
655            Instruction::Magnitude(inst) => Magnitude::<D, Sqrt>::format(f, &inst.input, &inst.out),
656            Instruction::FastMagnitude(inst) => {
657                Magnitude::<D, FastSqrt>::format(f, &inst.input, &inst.out)
658            }
659            Instruction::Dot(inst) => Dot::format(f, &inst.lhs, &inst.rhs, &inst.out),
660            Instruction::VectorSum(inst) => VectorSumFmt::<D>::format(f, &inst.input, &inst.out),
661            Instruction::VecInit { inputs, out } => {
662                let item = out.item();
663                let inputs = inputs
664                    .iter()
665                    .map(|input| format!("{input}"))
666                    .collect::<Vec<_>>();
667                let out = out.fmt_left();
668                writeln!(f, "{out} = {item}{{{}}};", inputs.join(","))
669            }
670            Instruction::InsertComponent {
671                vector,
672                index,
673                value,
674                out,
675            } => InsertComponent::format(f, vector, index, value, out),
676            Instruction::ExtractComponent(inst) => {
677                ExtractComponent::format(f, &inst.lhs, &inst.rhs, &inst.out)
678            }
679            Instruction::Printf {
680                format_string,
681                args,
682            } => D::compile_instruction_printf(f, format_string, args),
683            Instruction::Comment { content } => {
684                if content.contains('\n') {
685                    writeln!(f, "/* {content} */")
686                } else {
687                    writeln!(f, "// {content}")
688                }
689            }
690            Instruction::Barrier(barrier_ops) => write!(f, "{barrier_ops}"),
691            Instruction::Line { file, line } => writeln!(f, "#line {line} \"{file}\""),
692            Instruction::ProxyAsyncToSharedFence => {
693                writeln!(
694                    f,
695                    "cuda::device::experimental::fence_proxy_async_shared_cta();"
696                )
697            }
698            Instruction::BulkCommitGroup => writeln!(
699                f,
700                "cuda::device::experimental::cp_async_bulk_commit_group();"
701            ),
702            Instruction::BulkWaitGroup { max_pending } => writeln!(
703                f,
704                "cuda::device::experimental::cp_async_bulk_wait_group<{max_pending}>();"
705            ),
706            Instruction::BulkWaitGroupRead { max_pending } => writeln!(
707                f,
708                "cuda::device::experimental::cp_async_bulk_wait_group_read<{max_pending}>();"
709            ),
710            Instruction::TmaReplacePointer {
711                buffer,
712                offset,
713                tensor_map,
714                out,
715            } => {
716                let pos = Builtin::<D>::UnitPos;
717                writeln!(f, "__shared__ alignas(128) CUtensorMap {out};")?;
718                writeln!(
719                    f,
720                    "
721if({pos} == 0) {{
722    {out} = {tensor_map};
723    tensormap_replace_global_address({out}, &{buffer}[{offset}]);
724}}"
725                )?;
726                writeln!(f, "__syncthreads();")
727            }
728            Instruction::MemCopyAsyncTensorSharedToGlobal {
729                smem_buffer,
730                tensor_map,
731                indices,
732            } => {
733                let rank = indices.len();
734                let smem_ptr = smem_buffer.fmt_ptr();
735                let indices = indices.iter().rev().fold(String::new(), |mut s, it| {
736                    let _ = write!(s, "{it}, ");
737                    s
738                });
739                writeln!(
740                    f,
741                    "cuda::device::experimental::cp_async_bulk_tensor_{rank}d_shared_to_global(&{tensor_map}, {indices} {smem_ptr});"
742                )
743            }
744            Instruction::SpecialCast(UnaryInstruction { input, out }) => {
745                // Only supported in CUDA so I'm putting it here. Move to dialect if necessary.
746                #[cfg(not(feature = "cuda"))]
747                {
748                    let _ = (input, out);
749                    writeln!(
750                        f,
751                        "#error FP8/FP6/FP4 casting isn't supported outside of CUDA"
752                    )
753                }
754                #[cfg(feature = "cuda")]
755                crate::cuda::convert::special_cast::<D>(f, input, out)
756            }
757            Instruction::ReadBuiltin { builtin, out } => {
758                writeln!(f, "{} = {builtin};", out.fmt_left())
759            }
760            Instruction::ReadScalar { id, out } => {
761                let elem = *out.item().elem();
762                writeln!(f, "{} = info.scalars_{elem}[{id}];", out.fmt_left())
763            }
764        }
765    }
766}
767
768struct Fma<D: Dialect> {
769    _dialect: PhantomData<D>,
770}
771
772impl<D: Dialect> Fma<D> {
773    fn format(
774        f: &mut core::fmt::Formatter<'_>,
775        a: &Value<D>,
776        b: &Value<D>,
777        c: &Value<D>,
778        out: &Value<D>,
779    ) -> core::fmt::Result {
780        let out_item = out.item();
781
782        let out = out.fmt_left();
783        if let Item::Vector(_, num) = out_item {
784            writeln!(f, "{out} = {out_item}{{")?;
785
786            for i in 0..num {
787                let ai = a.index(i);
788                let bi = b.index(i);
789                let ci = c.index(i);
790
791                writeln!(f, "fma({ai}, {bi}, {ci}),")?;
792            }
793            f.write_str("};\n")
794        } else {
795            writeln!(f, "{out} = fma({a}, {b}, {c});")
796        }
797    }
798}
799
800struct Clamp<D: Dialect> {
801    _dialect: PhantomData<D>,
802}
803
804impl<D: Dialect> Clamp<D> {
805    fn format(
806        f: &mut core::fmt::Formatter<'_>,
807        input: &Value<D>,
808        min_value: &Value<D>,
809        max_value: &Value<D>,
810        out: &Value<D>,
811    ) -> core::fmt::Result {
812        let out_item = out.item();
813        if let Item::Vector(..) = out_item {
814            Self::unroll_vec(f, input, min_value, max_value, out)
815        } else {
816            let out = out.fmt_left();
817            write!(f, "{out} = ")?;
818            Self::format_scalar(f, *input, *min_value, *max_value, out_item)?;
819            f.write_str(";\n")
820        }
821    }
822
823    fn format_scalar(
824        f: &mut Formatter<'_>,
825        input: impl Component<D>,
826        min_value: impl Component<D>,
827        max_value: impl Component<D>,
828        item: Item<D>,
829    ) -> std::fmt::Result {
830        D::compile_instruction_max_function_name(f, item)?;
831        write!(f, "({min_value}, ")?;
832        D::compile_instruction_min_function_name(f, item)?;
833        write!(f, "({max_value}, {input}))")
834    }
835
836    fn unroll_vec(
837        f: &mut core::fmt::Formatter<'_>,
838        input: &Value<D>,
839        min_value: &Value<D>,
840        max_value: &Value<D>,
841        out: &Value<D>,
842    ) -> std::fmt::Result {
843        let optimized = Value::optimized_args([*input, *min_value, *max_value, *out]);
844        let [input, min_value, max_value, out_optimized] = optimized.args;
845
846        let item_out_original = out.item();
847        let item_out_optimized = out_optimized.item();
848
849        let index = match item_out_optimized {
850            Item::Vector(_, index) => index,
851            _ => 1,
852        };
853
854        let mut write_op = |input: &Value<D>,
855                            min_value: &Value<D>,
856                            max_value: &Value<D>,
857                            out: &Value<D>,
858                            item_out: Item<D>| {
859            let out = out.fmt_left();
860            writeln!(f, "{out} = {item_out}{{")?;
861            for i in 0..index {
862                let inputi = input.index(i);
863                let min_valuei = min_value.index(i);
864                let max_valuei = max_value.index(i);
865
866                Self::format_scalar(f, inputi, min_valuei, max_valuei, item_out)?;
867                f.write_str(", ")?;
868            }
869
870            f.write_str("};\n")
871        };
872
873        if item_out_original == item_out_optimized {
874            write_op(&input, &min_value, &max_value, out, item_out_optimized)
875        } else {
876            let out_tmp = Value::tmp(item_out_optimized);
877            write_op(&input, &min_value, &max_value, &out_tmp, item_out_optimized)?;
878            let addr_space = D::address_space_for_value(out);
879            let out = out.fmt_left();
880
881            writeln!(
882                f,
883                "{out} = reinterpret_cast<{addr_space}{item_out_original}&>({out_tmp});\n"
884            )?;
885
886            Ok(())
887        }
888    }
889}
890
891struct Magnitude<D: Dialect, S: FunctionFmt<D>> {
892    _dialect: PhantomData<D>,
893    _sqrt: PhantomData<S>,
894}
895
896impl<D: Dialect, S: FunctionFmt<D>> Magnitude<D, S> {
897    fn format(
898        f: &mut core::fmt::Formatter<'_>,
899        input: &Value<D>,
900        out: &Value<D>,
901    ) -> core::fmt::Result {
902        let num = match input.item() {
903            Item::Vector(_, vectorization) => vectorization,
904            _ => 1,
905        };
906        let elem = input.elem();
907
908        let mag = format!("{out}_mag");
909
910        // Use elem cast for the literal to support bfloat
911        writeln!(f, "{} {mag} = {}(0.0);", out.item(), out.item())?;
912
913        for i in 0..num {
914            let input_i = input.index(i);
915            writeln!(f, "{mag} += {input_i} * {input_i};")?;
916        }
917
918        let out = out.fmt_left();
919        write!(f, "{out} = ")?;
920        S::format_unary(f, &mag, elem)?;
921        f.write_str(";\n")
922    }
923}
924
925struct Normalize<D: Dialect, InvS: FunctionFmt<D>> {
926    _dialect: PhantomData<D>,
927    _rsqrt: PhantomData<InvS>,
928}
929
930impl<D: Dialect, InvS: FunctionFmt<D>> Normalize<D, InvS> {
931    fn format(
932        f: &mut core::fmt::Formatter<'_>,
933        input: &Value<D>,
934        out: &Value<D>,
935    ) -> core::fmt::Result {
936        let num = match input.item() {
937            Item::Vector(_, vectorization) => vectorization,
938            _ => 1,
939        };
940        let elem = input.elem();
941        let norm = format!("{out}_norm");
942
943        let out_item = out.item();
944        let out = out.fmt_left();
945        // Use elem cast for the literal to support bfloat
946        writeln!(f, "{elem} {norm} = {elem}(0.0);")?;
947
948        for i in 0..num {
949            let input_i = input.index(i);
950            writeln!(f, "{norm} += {input_i} * {input_i};")?;
951        }
952
953        write!(f, "{norm} = ")?;
954        InvS::format_unary(f, &norm, elem)?;
955        f.write_str(";\n")?;
956
957        if num == 1 {
958            writeln!(f, "{out} = {input} * {norm};")
959        } else {
960            write!(f, "{out} = {out_item}{{")?;
961            for i in 0..num {
962                let input_i = input.index(i);
963
964                writeln!(f, "{input_i} * {norm},")?;
965            }
966
967            f.write_str("};\n")
968        }
969    }
970}
971
972struct Dot<D: Dialect> {
973    _dialect: PhantomData<D>,
974}
975
976impl<D: Dialect> Dot<D> {
977    fn format(
978        f: &mut core::fmt::Formatter<'_>,
979        lhs: &Value<D>,
980        rhs: &Value<D>,
981        out: &Value<D>,
982    ) -> core::fmt::Result {
983        let num = match lhs.item() {
984            Item::Vector(_, vectorization) => vectorization,
985            _ => 1,
986        };
987
988        let muls = (0..num)
989            .map(|i| {
990                let lhs_i = lhs.index(i);
991                let rhs_i = rhs.index(i);
992                format!("{lhs_i} * {rhs_i}")
993            })
994            .collect::<Vec<_>>();
995
996        let value = muls.join(" + ");
997        if out.declare_local_ptr_backing(f)? {
998            writeln!(f, "*{out} = {value};")
999        } else {
1000            writeln!(f, "{} = {value};", out.fmt_left())
1001        }
1002    }
1003}
1004
1005struct VectorSumFmt<D: Dialect> {
1006    _dialect: PhantomData<D>,
1007}
1008
1009impl<D: Dialect> VectorSumFmt<D> {
1010    fn format(
1011        f: &mut core::fmt::Formatter<'_>,
1012        input: &Value<D>,
1013        out: &Value<D>,
1014    ) -> core::fmt::Result {
1015        let num = input.item().vectorization();
1016
1017        let elems = (0..num)
1018            .map(|i| format!("{}", input.index(i)))
1019            .collect::<Vec<_>>();
1020
1021        let value = elems.join(" + ");
1022        if out.declare_local_ptr_backing(f)? {
1023            writeln!(f, "*{out} = {value};")
1024        } else {
1025            writeln!(f, "{} = {value};", out.fmt_left())
1026        }
1027    }
1028}
1029
1030struct EnsureBoolArg<'a, V: Display, D: Dialect> {
1031    val: &'a V,
1032    elem: &'a Elem<D>,
1033}
1034
1035impl<V: Display, D: Dialect> Display for EnsureBoolArg<'_, V, D> {
1036    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1037        if self.elem != &Elem::Bool {
1038            write!(f, "bool({})", self.val)
1039        } else {
1040            write!(f, "{}", self.val)
1041        }
1042    }
1043}
cubecl_cpp/shared/instruction.rs

cubecl_cpp/shared/
instruction.rs
