cubecl_cpp/shared/
dialect.rs

1use std::{collections::HashSet, fmt::Debug};
2use std::{fmt::Display, hash::Hash};
3
4use cubecl_core::ir::Processor;
5
6use crate::shared::{
7    Builtin, FmtLeft, IndexedValue, MmaShape, SupportedMmaCombinations,
8    SupportedScaledMmaCombinations, reduce_comparison, reduce_exclusive, reduce_inclusive,
9    reduce_operator, reduce_quantifier,
10    unary::{Neg, Unary},
11};
12
13use super::{
14    Architecture, Body, Component, CubeIndexFlags, Elem, Flags, FragmentIdent, FragmentLayout,
15    FragmentType, Instruction, Item, KernelArg, SharedMemory, Value, WarpInstruction,
16    WmmaInstruction,
17};
18
19// Base dialect
20
21pub trait Dialect:
22    DialectIncludes<Self>
23    + DialectTypes<Self>
24    + DialectBindings<Self>
25    + DialectWarpReduceCompiler<Self>
26    + DialectCubeBuiltins<Self>
27    + DialectInstructions<Self>
28    + DialectWmmaCompiler<Self>
29    + DialectProcessors<Self>
30    + Default
31    + Clone
32    + Copy
33    + Debug
34    + Send
35    + Sync
36    + Eq
37    + Hash
38    + 'static
39{
40    type Architecture: Architecture;
41}
42
43// Includes
44
45pub trait DialectIncludes<D: Dialect> {
46    type Extension: Debug + Clone + Sync + Send;
47
48    fn compile_includes(f: &mut std::fmt::Formatter<'_>, flags: &Flags<D>) -> std::fmt::Result;
49    fn compile_extensions(
50        f: &mut std::fmt::Formatter<'_>,
51        extensions: &[Self::Extension],
52    ) -> std::fmt::Result;
53    fn register_instruction_extension(
54        extensions: &mut Vec<Self::Extension>,
55        instruction: &Instruction<D>,
56    );
57    fn register_warp_instruction_extension(
58        extensions: &mut Vec<Self::Extension>,
59        instruction: &WarpInstruction<D>,
60    );
61    #[allow(unused_variables)]
62    fn register_wmma_instruction_extension(
63        extensions: &mut Vec<Self::Extension>,
64        instruction: &WmmaInstruction<D>,
65    ) {
66    }
67}
68
69// Types
70
71pub trait DialectTypes<D: Dialect> {
72    fn item_can_be_optimized() -> bool;
73    fn compile_elem(
74        f: &mut std::fmt::Formatter<'_>,
75        elem: &Elem<D>,
76        word: bool,
77    ) -> std::fmt::Result;
78
79    fn compile_item(f: &mut std::fmt::Formatter<'_>, item: &Item<D>) -> std::fmt::Result;
80    fn compile_type_definitions(
81        f: &mut std::fmt::Formatter<'_>,
82        items: &HashSet<Item<D>>,
83        scalars: &[(Elem<D>, usize)],
84        info: &cubecl_core::Info,
85        flags: &Flags<D>,
86    ) -> std::fmt::Result;
87    fn compile_local_memory_qualifier(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result;
88    fn compile_shared_memory_declaration(
89        f: &mut std::fmt::Formatter<'_>,
90        shared: &SharedMemory<D>,
91    ) -> std::fmt::Result {
92        let SharedMemory { ptr, offset, .. } = shared;
93        let ptr_ty = ptr.item();
94        let size_bytes = shared.size();
95        writeln!(f, "// Shared value size: {size_bytes} bytes")?;
96        writeln!(
97            f,
98            "{ptr_ty} {ptr} = reinterpret_cast<{ptr_ty}>(&dynamic_shared_mem[{offset}]);"
99        )
100    }
101    fn compile_polyfills(_f: &mut std::fmt::Formatter<'_>, _flags: &Flags<D>) -> std::fmt::Result {
102        Ok(())
103    }
104    /// Address space (for Metal dialect only).
105    fn address_space_for_value(_value: &Value<D>) -> String {
106        "".to_string()
107    }
108}
109
110// Kernel argument bindings
111
112pub trait DialectBindings<D: Dialect> {
113    fn compile_kernel_signature(
114        f: &mut std::fmt::Formatter<'_>,
115        kernel_name: &str,
116        tensor_maps: &[KernelArg<D>],
117        buffers: &[KernelArg<D>],
118        flags: &Flags<D>,
119    ) -> std::fmt::Result;
120    fn compile_bindings_body(
121        _f: &mut std::fmt::Formatter<'_>,
122        _body: &Body<D>,
123    ) -> std::fmt::Result {
124        Ok(())
125    }
126}
127
128// Cube builtins dialect
129
130pub trait DialectCubeBuiltins<D: Dialect> {
131    /// Depending on the dialect available built-ins the
132    /// inclusion rules might change.
133    /// For instance in metal we have a built-in for the Unit plane position
134    /// but in other dialects there is none so we have to compute it using
135    /// other built-ins.
136    fn builtin_rules(flags: &CubeIndexFlags) -> CubeIndexFlags {
137        let unit_pos_plane = flags.unit_pos_plane;
138        let plane_dim_checked = flags.plane_dim_checked;
139        let plane_dim = flags.plane_dim || plane_dim_checked || unit_pos_plane;
140        let plane_pos = flags.plane_pos;
141        let absolute_pos = flags.absolute_pos || unit_pos_plane;
142        let absolute_pos_tuple = flags.absolute_pos_tuple || absolute_pos;
143        let cube_dim = flags.cube_dim;
144        let cube_dim_tuple = flags.cube_dim_tuple || cube_dim || absolute_pos || plane_dim_checked;
145        let unit_pos = flags.unit_pos;
146        let unit_pos_tuple = flags.unit_pos_tuple || unit_pos;
147        let cube_count = flags.cube_count;
148        let cube_count_tuple = flags.cube_count_tuple || absolute_pos;
149        let cube_pos = flags.cube_pos;
150        let cube_pos_tuple = flags.cube_pos_tuple || cube_pos;
151        let cluster_group = flags.cluster_pos;
152
153        CubeIndexFlags {
154            absolute_pos,
155            absolute_pos_tuple,
156            cube_count,
157            cube_count_tuple,
158            cube_dim,
159            cube_dim_tuple,
160            cube_pos,
161            cube_pos_tuple,
162            plane_dim,
163            plane_dim_checked,
164            plane_pos,
165            unit_pos_tuple,
166            unit_pos,
167            unit_pos_plane,
168            cluster_pos: cluster_group,
169        }
170    }
171
172    fn compile_absolute_pos_tuple_computation(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
173        let value = Builtin::<D>::AbsolutePosBaseName;
174        let ty = value.item();
175        let cube_pos_x = Builtin::<D>::CubePosX;
176        let cube_pos_y = Builtin::<D>::CubePosY;
177        let cube_pos_z = Builtin::<D>::CubePosZ;
178        let cube_dim_x = Builtin::<D>::CubeDimX;
179        let cube_dim_y = Builtin::<D>::CubeDimY;
180        let cube_dim_z = Builtin::<D>::CubeDimZ;
181        let unit_pos_x = Builtin::<D>::UnitPosX;
182        let unit_pos_y = Builtin::<D>::UnitPosY;
183        let unit_pos_z = Builtin::<D>::UnitPosZ;
184        writeln!(
185            f,
186            "{ty} {value} = make_{ty}(
187    {cube_pos_x} * {cube_dim_x} + {unit_pos_x},
188    {cube_pos_y} * {cube_dim_y} + {unit_pos_y},
189    {cube_pos_z} * {cube_dim_z} + {unit_pos_z}
190);"
191        )
192    }
193
194    fn compile_absolute_pos_base_name(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
195        f.write_str("absoluteIdx")
196    }
197
198    fn compile_absolute_pos(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
199        f.write_str("idxGlobal")
200    }
201
202    fn compile_absolute_pos_x(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
203        Self::compile_absolute_pos_base_name(f)?;
204        write!(f, ".x")
205    }
206
207    fn compile_absolute_pos_y(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
208        Self::compile_absolute_pos_base_name(f)?;
209        write!(f, ".y")
210    }
211
212    fn compile_absolute_pos_z(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
213        Self::compile_absolute_pos_base_name(f)?;
214        write!(f, ".z")
215    }
216
217    fn compile_cube_count_base_name(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
218        f.write_str("gridDim")
219    }
220
221    fn compile_cube_count(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
222        f.write_str("gridDimGlobal")
223    }
224
225    fn compile_cube_count_x(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
226        Self::compile_cube_count_base_name(f)?;
227        write!(f, ".x")
228    }
229
230    fn compile_cube_count_y(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
231        Self::compile_cube_count_base_name(f)?;
232        write!(f, ".y")
233    }
234
235    fn compile_cube_count_z(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
236        Self::compile_cube_count_base_name(f)?;
237        write!(f, ".z")
238    }
239
240    fn compile_cube_dim_base_name(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
241        f.write_str("blockDim")
242    }
243
244    fn compile_cube_dim(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
245        f.write_str("blockDimGlobal")
246    }
247
248    fn compile_cube_dim_x(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
249        Self::compile_cube_dim_base_name(f)?;
250        write!(f, ".x")
251    }
252
253    fn compile_cube_dim_y(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
254        Self::compile_cube_dim_base_name(f)?;
255        write!(f, ".y")
256    }
257
258    fn compile_cube_dim_z(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
259        Self::compile_cube_dim_base_name(f)?;
260        write!(f, ".z")
261    }
262
263    fn compile_cube_pos_base_name(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
264        f.write_str("blockIdx")
265    }
266
267    fn compile_cube_pos(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
268        f.write_str("blockIdxGlobal")
269    }
270
271    fn compile_cube_pos_x(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
272        Self::compile_cube_pos_base_name(f)?;
273        write!(f, ".x")
274    }
275
276    fn compile_cube_pos_y(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
277        Self::compile_cube_pos_base_name(f)?;
278        write!(f, ".y")
279    }
280
281    fn compile_cube_pos_z(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
282        Self::compile_cube_pos_base_name(f)?;
283        write!(f, ".z")
284    }
285
286    fn compile_unit_pos_computation(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
287        let value = Builtin::<D>::UnitPos;
288        let ty = value.item();
289        let cube_dim_x = Builtin::<D>::CubeDimX;
290        let cube_dim_y = Builtin::<D>::CubeDimY;
291        let unit_pos_x = Builtin::<D>::UnitPosX;
292        let unit_pos_y = Builtin::<D>::UnitPosY;
293        let unit_pos_z = Builtin::<D>::UnitPosZ;
294        writeln!(
295            f,
296            "{ty} {value} = {unit_pos_x} + {unit_pos_y} * {cube_dim_x} + {unit_pos_z} * ({cube_dim_x} * {cube_dim_y});"
297        )
298    }
299
300    fn compile_unit_pos(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
301        f.write_str("threadIdxGlobal")
302    }
303
304    fn compile_unit_pos_base_name(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
305        f.write_str("threadIdx")
306    }
307
308    fn compile_unit_pos_x(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
309        Self::compile_unit_pos_base_name(f)?;
310        write!(f, ".x")
311    }
312
313    fn compile_unit_pos_y(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
314        Self::compile_unit_pos_base_name(f)?;
315        write!(f, ".y")
316    }
317
318    fn compile_unit_pos_z(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
319        Self::compile_unit_pos_base_name(f)?;
320        write!(f, ".z")
321    }
322
323    fn compile_plane_dim(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
324        f.write_str("warpSize")
325    }
326
327    fn compile_plane_dim_checked(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
328        f.write_str("warpSizeChecked")
329    }
330
331    fn compile_plane_pos(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
332        let unit_pos_x = Builtin::<D>::UnitPosX;
333        let plane_dim = Builtin::<D>::PlaneDim;
334        write!(f, "{unit_pos_x} / {plane_dim}")
335    }
336
337    fn compile_unit_pos_plane(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
338        let absolute_pos = Builtin::<D>::AbsolutePos(Elem::U32);
339        let plane_dim = Builtin::<D>::PlaneDim;
340        let ty = Item::<D>::Scalar(Elem::U32);
341        write!(f, "{ty}({absolute_pos}) % {plane_dim}")
342    }
343
344    fn compile_cluster_pos(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
345        write!(f, "0")
346    }
347    fn compile_cluster_pos_x(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
348        write!(f, "0")
349    }
350    fn compile_cluster_pos_y(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
351        write!(f, "0")
352    }
353    fn compile_cluster_pos_z(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
354        write!(f, "0")
355    }
356}
357
358// Instructions
359
360pub trait DialectInstructions<D: Dialect> {
361    // atomics
362    fn compile_atomic_add(
363        f: &mut std::fmt::Formatter<'_>,
364        lhs: &Value<D>,
365        rhs: &Value<D>,
366        out: &Value<D>,
367    ) -> std::fmt::Result {
368        let rhs = rhs.ensure_lvalue(f)?;
369
370        let optimized = Value::optimized_args([*lhs, rhs, *out]);
371        let [lhs, rhs, out_optimized] = optimized.args;
372
373        let addr_space = D::address_space_for_value(out);
374        let out_item = out.item();
375        let out = out.fmt_left();
376
377        match out_optimized.item() {
378            Item::Scalar(Elem::I64) => writeln!(
379                f,
380                "{out} = atomicAdd(reinterpret_cast<{uint}*>({lhs}), {uint}({rhs}));",
381                uint = Elem::<D>::U64
382            ),
383            Item::Vector(inner, vectorization) if matches!(inner.elem(), Elem::F32) => {
384                let vec_ty = Item::NativeVector(*inner.elem(), vectorization);
385                let out_tmp = Value::tmp(out_optimized.item());
386                writeln!(
387                    f,
388                    "{vec_ty} {out_tmp} = atomicAdd(
389                    reinterpret_cast<{addr_space}{vec_ty}*>({lhs}),
390                    reinterpret_cast<const {addr_space}{vec_ty}&>({rhs}));",
391                )?;
392                writeln!(
393                    f,
394                    "{out} = reinterpret_cast<{addr_space}{out_item}&>({out_tmp});"
395                )
396            }
397            Item::Scalar(Elem::F16x2) | Item::Scalar(Elem::BF16x2) => {
398                let out_tmp = Value::tmp(out_optimized.item());
399                writeln!(
400                    f,
401                    "{} = atomicAdd(
402                    reinterpret_cast<{}>({lhs}),
403                    reinterpret_cast<const {addr_space}{}&>({rhs}));",
404                    out_tmp.fmt_left(),
405                    lhs.item(),
406                    rhs.item()
407                )?;
408                writeln!(
409                    f,
410                    "{out} = reinterpret_cast<{addr_space}{out_item}&>({out_tmp});"
411                )
412            }
413            _ => writeln!(f, "{out} = atomicAdd({lhs}, {rhs});"),
414        }
415    }
416
417    fn compile_atomic_and(
418        f: &mut std::fmt::Formatter<'_>,
419        lhs: &Value<D>,
420        rhs: &Value<D>,
421        out: &Value<D>,
422    ) -> std::fmt::Result {
423        let out = out.fmt_left();
424        writeln!(f, "{out} = atomicAnd({lhs}, {rhs});")
425    }
426
427    fn compile_atomic_cas(
428        f: &mut std::fmt::Formatter<'_>,
429        input: &Value<D>,
430        cmp: &Value<D>,
431        val: &Value<D>,
432        out: &Value<D>,
433    ) -> std::fmt::Result {
434        let addr_space = D::address_space_for_value(out);
435        let out_item = out.item();
436        let out = out.fmt_left();
437
438        match val.item() {
439            // vec4 is automatically supported by the new 128-bit template version
440            Item::Vector(inner, 2) if matches!(inner.elem(), Elem::F32) => {
441                let cmp = cmp.ensure_lvalue(f)?;
442                let val = val.ensure_lvalue(f)?;
443                let u64 = Item::Scalar(Elem::<D>::U64);
444                let out_tmp = Value::tmp(u64);
445                writeln!(
446                    f,
447                    "{} = atomicCAS(
448                reinterpret_cast<{addr_space}{u64}*>({input}),
449                reinterpret_cast<{u64}&>({cmp}),
450                reinterpret_cast<{u64}&>({val}));",
451                    out_tmp.fmt_left()
452                )?;
453                writeln!(f, "{out} = reinterpret_cast<{out_item}&>({out_tmp});")
454            }
455            Item::Vector(inner, 2) if matches!(inner.elem(), Elem::F16 | Elem::BF16) => {
456                let cmp = cmp.ensure_lvalue(f)?;
457                let val = val.ensure_lvalue(f)?;
458                let u32 = Item::Scalar(Elem::<D>::U32);
459                let out_tmp = Value::tmp(u32);
460                writeln!(
461                    f,
462                    "{} = atomicCAS(
463                reinterpret_cast<{addr_space}{u32}*>({input}),
464                reinterpret_cast<{u32}&>({cmp}),
465                reinterpret_cast<{u32}&>({val}));",
466                    out_tmp.fmt_left()
467                )?;
468                writeln!(f, "{out} = reinterpret_cast<{out_item}&>({out_tmp});")
469            }
470            _ => writeln!(f, "{out} = atomicCAS({input}, {cmp}, {val});"),
471        }
472    }
473
474    fn compile_atomic_load(
475        f: &mut std::fmt::Formatter<'_>,
476        input: &Value<D>,
477        out: &Value<D>,
478    ) -> std::fmt::Result {
479        let out_item = out.item();
480        let out = out.fmt_left();
481
482        let Item::Pointer(_, class) = input.item() else {
483            unreachable!()
484        };
485
486        let unsigned_ty = match out_item.size() {
487            1 => Item::Scalar(Elem::<D>::U8),
488            2 => Item::Scalar(Elem::<D>::U16),
489            4 => Item::Scalar(Elem::<D>::U32),
490            8 => Item::Scalar(Elem::<D>::U64),
491            // Hacky, but it's CUDA only for now. We should really migrate to a more modern API in
492            // general
493            16 => {
494                let out_tmp = Value::tmp(out_item);
495                writeln!(f, "{};", out_tmp.fmt_left())?;
496                writeln!(
497                    f,
498                    "__nv_atomic_load({input}, &{out_tmp}, __NV_ATOMIC_RELAXED);"
499                )?;
500                return writeln!(f, "{out} = {out_tmp};");
501            }
502            _ => unreachable!(),
503        };
504        let unsigned_ptr_ty = Item::Pointer(unsigned_ty.intern(), class);
505
506        let ptr_tmp = Value::tmp(unsigned_ptr_ty);
507        let out_tmp = Value::tmp(unsigned_ty);
508        writeln!(
509            f,
510            "volatile {} = reinterpret_cast<volatile {unsigned_ptr_ty}>({input});",
511            ptr_tmp.fmt_left()
512        )?;
513        writeln!(f, "{} = *{ptr_tmp};", out_tmp.fmt_left())?;
514        writeln!(f, "{out} = reinterpret_cast<const {out_item}&>({out_tmp});")
515    }
516
517    fn compile_atomic_max(
518        f: &mut std::fmt::Formatter<'_>,
519        lhs: &Value<D>,
520        rhs: &Value<D>,
521        out: &Value<D>,
522    ) -> std::fmt::Result {
523        let out = out.fmt_left();
524        writeln!(f, "{out} = atomicMax({lhs}, {rhs});")
525    }
526
527    fn compile_atomic_min(
528        f: &mut std::fmt::Formatter<'_>,
529        lhs: &Value<D>,
530        rhs: &Value<D>,
531        out: &Value<D>,
532    ) -> std::fmt::Result {
533        let out = out.fmt_left();
534        writeln!(f, "{out} = atomicMin({lhs}, {rhs});")
535    }
536
537    fn compile_atomic_or(
538        f: &mut std::fmt::Formatter<'_>,
539        lhs: &Value<D>,
540        rhs: &Value<D>,
541        out: &Value<D>,
542    ) -> std::fmt::Result {
543        let out = out.fmt_left();
544        writeln!(f, "{out} = atomicOr({lhs}, {rhs});")
545    }
546
547    fn compile_atomic_store(
548        f: &mut std::fmt::Formatter<'_>,
549        input: &Value<D>,
550        out: &Value<D>,
551    ) -> std::fmt::Result {
552        let tmp = Value::tmp(input.item());
553        Self::compile_atomic_swap(f, out, input, &tmp)
554    }
555
556    fn compile_atomic_sub(
557        f: &mut std::fmt::Formatter<'_>,
558        lhs: &Value<D>,
559        rhs: &Value<D>,
560        out: &Value<D>,
561    ) -> std::fmt::Result {
562        let tmp = Value::tmp(rhs.item());
563        Neg::format(f, rhs, &tmp)?;
564        Self::compile_atomic_add(f, lhs, &tmp, out)
565    }
566
567    fn compile_atomic_swap(
568        f: &mut std::fmt::Formatter<'_>,
569        lhs: &Value<D>,
570        rhs: &Value<D>,
571        out: &Value<D>,
572    ) -> std::fmt::Result {
573        let out_item = out.item();
574        let out = out.fmt_left();
575
576        let unsigned_elem = match rhs.item().size() {
577            1 => Elem::<D>::U8,
578            2 => Elem::<D>::U16,
579            4 => Elem::<D>::U32,
580            8 => Elem::<D>::U64,
581            // 128-bit wide uses a generic template that accepts arbitrary types
582            _ => return writeln!(f, "{out} = atomicExch({lhs}, {rhs});"),
583        };
584
585        let rhs = rhs.ensure_lvalue(f)?;
586        let Item::Pointer(_, class) = lhs.item() else {
587            unreachable!()
588        };
589        let unsigned_ty = Item::Scalar(unsigned_elem);
590        let unsigned_ptr_ty = Item::Pointer(unsigned_ty.intern(), class);
591
592        let out_tmp = Value::tmp(unsigned_ty);
593        writeln!(
594            f,
595            "{} = atomicExch(
596                    reinterpret_cast<{unsigned_ptr_ty}>({lhs}),
597                    reinterpret_cast<const {unsigned_ty}&>({rhs}));",
598            out_tmp.fmt_left()
599        )?;
600        writeln!(f, "{out} = reinterpret_cast<const {out_item}&>({out_tmp});")
601    }
602
603    fn compile_atomic_xor(
604        f: &mut std::fmt::Formatter<'_>,
605        lhs: &Value<D>,
606        rhs: &Value<D>,
607        out: &Value<D>,
608    ) -> std::fmt::Result {
609        let out = out.fmt_left();
610        writeln!(f, "{out} = atomicXor({lhs}, {rhs});")
611    }
612
613    fn compile_saturating_add(
614        f: &mut std::fmt::Formatter<'_>,
615        lhs: impl Display,
616        rhs: impl Display,
617        item: Item<D>,
618    ) -> std::fmt::Result;
619
620    fn compile_saturating_sub(
621        f: &mut std::fmt::Formatter<'_>,
622        lhs: impl Display,
623        rhs: impl Display,
624        item: Item<D>,
625    ) -> std::fmt::Result;
626
627    // debug
628    fn compile_instruction_printf(
629        f: &mut std::fmt::Formatter<'_>,
630        format_string: &str,
631        args: &[Value<D>],
632    ) -> std::fmt::Result {
633        let args = args.iter().map(|arg| format!("{arg}")).collect::<Vec<_>>();
634        let args = match args.is_empty() {
635            true => "".to_string(),
636            false => format!(", {}", args.join(",")),
637        };
638        writeln!(f, "printf({format_string:?}{args});")
639    }
640
641    // logs
642    fn compile_instruction_log1p_scalar<T: Component<D>>(
643        f: &mut std::fmt::Formatter<'_>,
644        input: T,
645    ) -> std::fmt::Result {
646        let elem = input.elem();
647        match elem {
648            Elem::F16 | Elem::F16x2 | Elem::BF16 | Elem::BF16x2 => {
649                write!(f, "{elem}(log1p(float({input})))")
650            }
651            _ => write!(f, "log1p({input})"),
652        }
653    }
654
655    // exp
656    fn compile_instruction_expm1_scalar<T: Component<D>>(
657        f: &mut std::fmt::Formatter<'_>,
658        input: T,
659    ) -> std::fmt::Result {
660        let elem = input.elem();
661        match elem {
662            Elem::F16 | Elem::F16x2 | Elem::BF16 | Elem::BF16x2 => {
663                write!(f, "{elem}(expm1(float({input})))")
664            }
665            _ => write!(f, "expm1({input})"),
666        }
667    }
668
669    // sync
670    fn compile_instruction_sync_threads(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result;
671    fn compile_instruction_sync_warp(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result;
672    fn compile_instruction_thread_fence(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result;
673
674    // trigo
675    fn compile_instruction_tanh_scalar<T: Component<D>>(
676        f: &mut std::fmt::Formatter<'_>,
677        input: T,
678    ) -> std::fmt::Result {
679        let elem = input.elem();
680        match elem {
681            Elem::F16 | Elem::F16x2 | Elem::BF16 | Elem::BF16x2 => {
682                write!(f, "{elem}(tanh(float({input})))")
683            }
684            _ => write!(f, "tanh({input})"),
685        }
686    }
687
688    // unary
689    fn compile_instruction_find_first_set<T: Component<D>>(
690        f: &mut std::fmt::Formatter<'_>,
691        input: T,
692        out_elem: Elem<D>,
693    ) -> std::fmt::Result;
694    fn compile_instruction_leading_zeros_scalar<T: Component<D>>(
695        f: &mut std::fmt::Formatter<'_>,
696        input: T,
697        out_elem: Elem<D>,
698    ) -> std::fmt::Result;
699
700    fn compile_instruction_trailing_zeros_scalar<T: Component<D>>(
701        f: &mut std::fmt::Formatter<'_>,
702        input: T,
703        out_elem: Elem<D>,
704    ) -> std::fmt::Result;
705
706    fn compile_instruction_popcount_scalar<T: Component<D>>(
707        f: &mut std::fmt::Formatter<'_>,
708        input: T,
709        out_elem: Elem<D>,
710    ) -> std::fmt::Result {
711        write!(f, "{out_elem}(")?;
712        match input.elem() {
713            Elem::I32 => write!(f, "__popc({}({input}))", Elem::<D>::U32),
714            Elem::U32 => write!(f, "__popc({input})"),
715            Elem::I64 => write!(f, "__popcll({}({input}))", Elem::<D>::U64),
716            Elem::U64 => write!(f, "__popcll({input})"),
717            _ => write!(f, "__popc({})", super::unary::zero_extend(input)),
718        }?;
719        write!(f, ")")
720    }
721
722    fn compile_instruction_reverse_bits_scalar<T: Component<D>>(
723        f: &mut std::fmt::Formatter<'_>,
724        input: T,
725        out_elem: Elem<D>,
726    ) -> std::fmt::Result {
727        write!(f, "{out_elem}(")?;
728        match out_elem {
729            Elem::I32 => write!(f, "__brev({}({input}))", Elem::<D>::U32),
730            Elem::U32 => write!(f, "__brev({input})"),
731            Elem::I64 => write!(f, "__brevll({}({input}))", Elem::<D>::U64),
732            Elem::U64 => write!(f, "__brevll({input})"),
733            _ => write!(
734                f,
735                "__brev({}) >> {}",
736                super::unary::zero_extend(input),
737                (size_of::<u32>() - out_elem.size()) * 8
738            ),
739        }?;
740        write!(f, ")")
741    }
742
743    // others
744    fn compile_instruction_max_function_name(
745        f: &mut std::fmt::Formatter<'_>,
746        item: Item<D>,
747    ) -> std::fmt::Result;
748
749    fn compile_instruction_min_function_name(
750        f: &mut std::fmt::Formatter<'_>,
751        item: Item<D>,
752    ) -> std::fmt::Result;
753
754    fn compile_instruction_powf(
755        f: &mut std::fmt::Formatter<'_>,
756        lhs: &str,
757        rhs: &str,
758        elem: Elem<D>,
759    ) -> std::fmt::Result {
760        match elem {
761            Elem::F32 => write!(f, "powf({lhs}, {rhs})"),
762            Elem::F64 => write!(f, "pow({lhs}, {rhs})"),
763            _ => write!(f, "#error Unsupported type for powf: {elem}"),
764        }
765    }
766
767    fn compile_instruction_hypot(
768        f: &mut std::fmt::Formatter<'_>,
769        lhs: &str,
770        rhs: &str,
771        elem: Elem<D>,
772    ) -> std::fmt::Result {
773        match elem {
774            Elem::F32 => write!(f, "hypotf({lhs}, {rhs})"),
775            Elem::F64 => write!(f, "hypot({lhs}, {rhs})"),
776            _ => write!(f, "#error Unsupported type for hypot: {elem}"),
777        }
778    }
779
780    fn compile_instruction_rhypot(
781        f: &mut std::fmt::Formatter<'_>,
782        lhs: &str,
783        rhs: &str,
784        elem: Elem<D>,
785    ) -> std::fmt::Result {
786        match elem {
787            Elem::F32 => write!(f, "rhypotf({lhs}, {rhs})"),
788            Elem::F64 => write!(f, "rhypot({lhs}, {rhs})"),
789            _ => write!(f, "#error Unsupported type for rhypot: {elem}"),
790        }
791    }
792
793    fn compile_instruction_half_function_name_prefix() -> &'static str {
794        "h"
795    }
796
797    fn compile_instruction_half2_function_name_prefix() -> &'static str {
798        "h2"
799    }
800
801    /// Remaps a math-function name for the dialect (default: unchanged), e.g. fast-math
802    /// intrinsics to each dialect's spelling (`__expf` -> `fast::exp`).
803    fn compile_fast_math_function_name(name: &'static str) -> &'static str {
804        name
805    }
806
807    // warp
808    fn compile_warp_shuffle(
809        f: &mut std::fmt::Formatter<'_>,
810        val: &str,
811        elem: &Elem<D>,
812        source: &str,
813    ) -> std::fmt::Result;
814    fn compile_warp_shuffle_xor(
815        f: &mut std::fmt::Formatter<'_>,
816        val: &str,
817        elem: &Elem<D>,
818        offset: &str,
819    ) -> std::fmt::Result;
820    fn compile_warp_shuffle_up(
821        f: &mut std::fmt::Formatter<'_>,
822        val: &str,
823        elem: &Elem<D>,
824        offset: &str,
825    ) -> std::fmt::Result;
826    fn compile_warp_shuffle_down(
827        f: &mut std::fmt::Formatter<'_>,
828        val: &str,
829        elem: &Elem<D>,
830        offset: &str,
831    ) -> std::fmt::Result;
832    fn compile_warp_all<T: Component<D>>(
833        f: &mut std::fmt::Formatter<'_>,
834        input: &T,
835    ) -> std::fmt::Result;
836    fn compile_warp_any<T: Component<D>>(
837        f: &mut std::fmt::Formatter<'_>,
838        input: &T,
839    ) -> std::fmt::Result;
840    fn compile_warp_ballot(
841        f: &mut std::fmt::Formatter<'_>,
842        input: &Value<D>,
843        out_elem: &Elem<D>,
844    ) -> std::fmt::Result;
845    fn compile_warp_elect(f: &mut std::fmt::Formatter<'_>, out: &str) -> std::fmt::Result {
846        write!(
847            f,
848            "
849unsigned int mask = __activemask();
850unsigned int leader = __ffs(mask) - 1;
851{out} = threadIdx.x % warpSize == leader;
852            "
853        )
854    }
855    fn compile_unreachable(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result;
856}
857
858#[derive(Debug, Clone, Copy, new)]
859pub struct ManualMma<'a, D: Dialect> {
860    pub shape: MmaShape<D>,
861    pub frag_a: &'a Value<D>,
862    pub frag_b: &'a Value<D>,
863    pub frag_c: &'a Value<D>,
864    pub frag_d: &'a Value<D>,
865}
866
867pub trait DialectWarpReduceCompiler<D: Dialect>:
868    Default + Clone + Copy + Debug + Send + Sync + Eq + Hash + 'static
869{
870    fn warp_reduce_sum(
871        f: &mut core::fmt::Formatter<'_>,
872        input: &Value<D>,
873        out: &Value<D>,
874    ) -> core::fmt::Result {
875        reduce_operator(f, input, out, "+=")
876    }
877    fn warp_reduce_prod(
878        f: &mut core::fmt::Formatter<'_>,
879        input: &Value<D>,
880        out: &Value<D>,
881    ) -> core::fmt::Result {
882        reduce_operator(f, input, out, "*=")
883    }
884    fn warp_reduce_max(
885        f: &mut core::fmt::Formatter<'_>,
886        input: &Value<D>,
887        out: &Value<D>,
888    ) -> core::fmt::Result {
889        reduce_comparison(f, input, out, D::compile_instruction_max_function_name)
890    }
891    fn warp_reduce_min(
892        f: &mut core::fmt::Formatter<'_>,
893        input: &Value<D>,
894        out: &Value<D>,
895    ) -> core::fmt::Result {
896        reduce_comparison(f, input, out, D::compile_instruction_min_function_name)
897    }
898    fn warp_reduce_all(
899        f: &mut core::fmt::Formatter<'_>,
900        input: &Value<D>,
901        out: &Value<D>,
902    ) -> core::fmt::Result {
903        reduce_quantifier(f, input, out, D::compile_warp_all::<IndexedValue<D>>)
904    }
905    fn warp_reduce_any(
906        f: &mut core::fmt::Formatter<'_>,
907        input: &Value<D>,
908        out: &Value<D>,
909    ) -> core::fmt::Result {
910        reduce_quantifier(f, input, out, D::compile_warp_any::<IndexedValue<D>>)
911    }
912    fn warp_reduce_sum_inclusive(
913        f: &mut core::fmt::Formatter<'_>,
914        input: &Value<D>,
915        out: &Value<D>,
916    ) -> core::fmt::Result {
917        reduce_inclusive(f, input, out, "+=")
918    }
919    fn warp_reduce_prod_inclusive(
920        f: &mut core::fmt::Formatter<'_>,
921        input: &Value<D>,
922        out: &Value<D>,
923    ) -> core::fmt::Result {
924        reduce_inclusive(f, input, out, "*=")
925    }
926    fn warp_reduce_sum_exclusive(
927        f: &mut core::fmt::Formatter<'_>,
928        input: &Value<D>,
929        out: &Value<D>,
930    ) -> core::fmt::Result {
931        reduce_exclusive(f, input, out, "+=", "0")
932    }
933    fn warp_reduce_prod_exclusive(
934        f: &mut core::fmt::Formatter<'_>,
935        input: &Value<D>,
936        out: &Value<D>,
937    ) -> core::fmt::Result {
938        reduce_exclusive(f, input, out, "*=", "1")
939    }
940}
941
942pub trait DialectWmmaCompiler<D: Dialect>:
943    Default + Clone + Copy + Debug + Send + Sync + Eq + Hash + 'static
944{
945    #[allow(unused_variables)]
946    fn compile_wmma_includes(
947        f: &mut std::fmt::Formatter<'_>,
948        flags: &Flags<D>,
949    ) -> std::fmt::Result {
950        Ok(())
951    }
952    #[allow(unused_variables)]
953    fn compile_wmma_type_definitions(
954        f: &mut std::fmt::Formatter<'_>,
955        flags: &Flags<D>,
956    ) -> std::fmt::Result {
957        Ok(())
958    }
959    #[allow(unused_variables)]
960    fn compile_wmma_local_variables(f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
961        Ok(())
962    }
963    #[allow(unused_variables)]
964    fn compile_wwma_fragment_ident(
965        f: &mut std::fmt::Formatter<'_>,
966        ident: &FragmentIdent<D>,
967    ) -> std::fmt::Result {
968        Ok(())
969    }
970    #[allow(unused_variables)]
971    fn compile_wmma_fragment_layout(
972        f: &mut std::fmt::Formatter<'_>,
973        layout: &FragmentLayout<D>,
974    ) -> std::fmt::Result {
975        Ok(())
976    }
977    #[allow(unused_variables)]
978    fn compile_wmma_fragment(
979        f: &mut std::fmt::Formatter<'_>,
980        fragment: &FragmentType<D>,
981    ) -> std::fmt::Result {
982        Ok(())
983    }
984
985    fn compile_wmma_fragment_declaration(
986        f: &mut std::fmt::Formatter<'_>,
987        val: &Value<D>,
988        value_ty: &Item<D>,
989    ) -> std::fmt::Result;
990
991    fn compile_wmma_instruction(
992        f: &mut std::fmt::Formatter<'_>,
993        instruction: &WmmaInstruction<D>,
994    ) -> std::fmt::Result;
995    fn compile_manual_mma(f: &mut std::fmt::Formatter<'_>, mma: ManualMma<D>) -> std::fmt::Result;
996    fn compile_scaled_mma(
997        f: &mut std::fmt::Formatter<'_>,
998        mma: ManualMma<D>,
999        scales_a: Value<D>,
1000        scales_b: Value<D>,
1001        scales_factor: u32,
1002    ) -> std::fmt::Result;
1003    fn supported_wmma_combinations(arch: &D::Architecture) -> SupportedMmaCombinations;
1004    fn supported_mma_combinations(arch: &D::Architecture) -> SupportedMmaCombinations;
1005    fn supported_scaled_mma_combinations(
1006        _arch: &D::Architecture,
1007    ) -> SupportedScaledMmaCombinations {
1008        Vec::new()
1009    }
1010}
1011
1012/// IR Processors to be applied to the scopes during processing. ``CheckedIO`` is always applied
1013/// by default, so these are only for target specific processors like MMA index processors.
1014pub trait DialectProcessors<D: Dialect> {
1015    fn processors() -> Vec<Box<dyn Processor>>;
1016}
cubecl_cpp/shared/dialect.rs

cubecl_cpp/shared/
dialect.rs
