cubecl_cpp/shared/
binary.rs

1use itertools::Itertools;
2
3use crate::shared::FmtLeft;
4
5use super::{Component, Dialect, Elem, Item, Value};
6use std::{
7    fmt::{Display, Formatter},
8    marker::PhantomData,
9};
10
11pub trait Binary<D: Dialect> {
12    fn format(
13        f: &mut Formatter<'_>,
14        lhs: &Value<D>,
15        rhs: &Value<D>,
16        out: &Value<D>,
17    ) -> std::fmt::Result {
18        let out_item = *out.item().value_ty();
19        if let Item::Vector(..) = out_item {
20            Self::unroll_vec(f, lhs, rhs, out)
21        } else {
22            if out.declare_local_ptr_backing(f)? {
23                write!(f, "*{out} = ")?;
24            } else {
25                write!(f, "{} = ", out.fmt_left())?;
26            }
27            Self::format_scalar(f, *lhs, *rhs, out_item)?;
28            f.write_str(";\n")
29        }
30    }
31
32    fn format_scalar<Lhs: Component<D>, Rhs: Component<D>>(
33        f: &mut Formatter<'_>,
34        lhs: Lhs,
35        rhs: Rhs,
36        item: Item<D>,
37    ) -> std::fmt::Result;
38
39    fn unroll_vec(
40        f: &mut Formatter<'_>,
41        lhs: &Value<D>,
42        rhs: &Value<D>,
43        out: &Value<D>,
44    ) -> core::fmt::Result {
45        let mut write_op =
46            |index: usize, lhs: &Value<D>, rhs: &Value<D>, out: &Value<D>, item_out: Item<D>| {
47                if out.declare_local_ptr_backing(f)? {
48                    writeln!(f, "*{out} = {item_out}{{")?;
49                } else {
50                    let out = out.fmt_left();
51                    writeln!(f, "{out} = {item_out}{{")?;
52                }
53                for i in 0..index {
54                    let lhsi = lhs.index(i);
55                    let rhsi = rhs.index(i);
56
57                    Self::format_scalar(f, lhsi, rhsi, item_out)?;
58                    f.write_str(", ")?;
59                }
60
61                f.write_str("};\n")
62            };
63
64        if Self::can_optimize() {
65            let optimized = Value::optimized_args([*lhs, *rhs, *out]);
66            let [lhs, rhs, out_optimized] = optimized.args;
67
68            let item_out_original = *out.item().value_ty();
69            let item_out_optimized = *out_optimized.item().value_ty();
70
71            let index = match item_out_optimized {
72                Item::Vector(_, vectorization) => vectorization,
73                _ => 1,
74            };
75
76            if item_out_original == item_out_optimized {
77                write_op(index, &lhs, &rhs, out, item_out_optimized)
78            } else {
79                let out_tmp = Value::tmp(item_out_optimized);
80                write_op(index, &lhs, &rhs, &out_tmp, item_out_optimized)?;
81                let addr_space = D::address_space_for_value(out);
82
83                if out.declare_local_ptr_backing(f)? {
84                    writeln!(
85                        f,
86                        "*{out} = reinterpret_cast<{addr_space}{item_out_original}&>({out_tmp});\n"
87                    )?;
88                } else {
89                    let out = out.fmt_left();
90                    writeln!(
91                        f,
92                        "{out} = reinterpret_cast<{addr_space}{item_out_original}&>({out_tmp});\n"
93                    )?;
94                }
95
96                Ok(())
97            }
98        } else {
99            let index = match out.item() {
100                Item::Vector(_, vectorization) => vectorization,
101                _ => 1,
102            };
103
104            write_op(index, lhs, rhs, out, out.item())
105        }
106    }
107
108    fn can_optimize() -> bool {
109        true
110    }
111}
112
113macro_rules! operator {
114    ($name:ident, $op:expr) => {
115        pub struct $name;
116
117        impl<D: Dialect> Binary<D> for $name {
118            fn format_scalar<Lhs: Display, Rhs: Display>(
119                f: &mut std::fmt::Formatter<'_>,
120                lhs: Lhs,
121                rhs: Rhs,
122                out_item: Item<D>,
123            ) -> std::fmt::Result {
124                let out_elem = out_item.elem();
125                match out_elem {
126                    // prevent auto-promotion rules to kick-in in order to stay in the same type
127                    // this is because of fusion and vectorization that can do elemwise operations on vectorized type,
128                    // the resulting elements need to be of the same type.
129                    Elem::<D>::I16 | Elem::<D>::U16 | Elem::<D>::I8 | Elem::<D>::U8 => {
130                        write!(f, "{out_elem}({lhs} {} {rhs})", $op)
131                    }
132                    _ => write!(f, "{lhs} {} {rhs}", $op),
133                }
134            }
135        }
136    };
137}
138
139operator!(Add, "+");
140operator!(Sub, "-");
141operator!(Div, "/");
142operator!(Mul, "*");
143operator!(Equal, "==");
144operator!(NotEqual, "!=");
145operator!(Lower, "<");
146operator!(LowerEqual, "<=");
147operator!(Greater, ">");
148operator!(GreaterEqual, ">=");
149operator!(ShiftLeft, "<<");
150operator!(ShiftRight, ">>");
151operator!(BitwiseOr, "|");
152operator!(BitwiseAnd, "&");
153operator!(BitwiseXor, "^");
154operator!(Or, "||");
155operator!(And, "&&");
156
157pub struct Remainder;
158
159impl<D: Dialect> Binary<D> for Remainder {
160    fn format_scalar<Lhs: Display, Rhs: Display>(
161        f: &mut std::fmt::Formatter<'_>,
162        lhs: Lhs,
163        rhs: Rhs,
164        out_item: Item<D>,
165    ) -> std::fmt::Result {
166        let out_elem = out_item.elem();
167        match out_elem {
168            Elem::<D>::I16 | Elem::<D>::U16 | Elem::<D>::I8 | Elem::<D>::U8 => {
169                write!(f, "{out_elem}({lhs} % {rhs})")
170            }
171            Elem::<D>::F16 | Elem::<D>::BF16 => {
172                let f32 = Elem::<D>::F32;
173                write!(f, "{out_elem}(fmodf({f32}({lhs}), {f32}({rhs}))))")
174            }
175            Elem::<D>::F32 => {
176                write!(f, "fmodf({lhs}, {rhs})")
177            }
178            Elem::<D>::F64 => {
179                write!(f, "fmod({lhs}, {rhs})")
180            }
181            _ => write!(f, "{lhs} % {rhs}"),
182        }
183    }
184
185    fn can_optimize() -> bool {
186        false
187    }
188}
189
190pub struct ModFloor;
191
192impl<D: Dialect> Binary<D> for ModFloor {
193    fn format_scalar<Lhs: Component<D>, Rhs: Component<D>>(
194        f: &mut Formatter<'_>,
195        lhs: Lhs,
196        rhs: Rhs,
197        item: Item<D>,
198    ) -> std::fmt::Result {
199        let is_uint = matches!(item.elem(), Elem::U8 | Elem::U16 | Elem::U32 | Elem::U64);
200        if is_uint {
201            // Remainder is cheaper and unsigned ints don't have a difference
202            return Remainder::format_scalar(f, lhs, rhs, item);
203        }
204
205        let floor = {
206            let prefix = match item.elem() {
207                Elem::F16 | Elem::BF16 => D::compile_instruction_half_function_name_prefix(),
208                Elem::F16x2 | Elem::BF16x2 => D::compile_instruction_half2_function_name_prefix(),
209                _ => "",
210            };
211            format!("{prefix}floor")
212        };
213
214        let is_int = matches!(item.elem(), Elem::I8 | Elem::I16 | Elem::I32 | Elem::I64);
215        let out_elem = item.elem();
216        if is_int {
217            write!(
218                f,
219                "{lhs} - {rhs} * ({out_elem}){floor}((float){lhs} / (float){rhs})"
220            )
221        } else {
222            write!(f, "{lhs} - {rhs} * {floor}({lhs} / {rhs})")
223        }
224    }
225}
226
227pub struct FastDiv;
228
229impl<D: Dialect> Binary<D> for FastDiv {
230    fn format_scalar<Lhs: Display, Rhs: Display>(
231        f: &mut std::fmt::Formatter<'_>,
232        lhs: Lhs,
233        rhs: Rhs,
234        _out_item: Item<D>,
235    ) -> std::fmt::Result {
236        // f32 only
237        write!(
238            f,
239            "{}({lhs}, {rhs})",
240            D::compile_fast_math_function_name("__fdividef")
241        )
242    }
243}
244
245pub struct HiMul;
246
247impl<D: Dialect> Binary<D> for HiMul {
248    fn format_scalar<Lhs: Display, Rhs: Display>(
249        f: &mut std::fmt::Formatter<'_>,
250        lhs: Lhs,
251        rhs: Rhs,
252        out: Item<D>,
253    ) -> std::fmt::Result {
254        let out_elem = out.elem();
255        match out_elem {
256            Elem::I32 => write!(f, "__mulhi({lhs}, {rhs})"),
257            Elem::U32 => write!(f, "__umulhi({lhs}, {rhs})"),
258            Elem::I64 => write!(f, "__mul64hi({lhs}, {rhs})"),
259            Elem::U64 => write!(f, "__umul64hi({lhs}, {rhs})"),
260            _ => writeln!(f, "#error HiMul only supports 32 and 64 bit ints"),
261        }
262    }
263
264    fn can_optimize() -> bool {
265        false
266    }
267}
268
269pub struct SaturatingAdd;
270
271impl<D: Dialect> Binary<D> for SaturatingAdd {
272    fn format_scalar<Lhs: Display, Rhs: Display>(
273        f: &mut std::fmt::Formatter<'_>,
274        lhs: Lhs,
275        rhs: Rhs,
276        out: Item<D>,
277    ) -> std::fmt::Result {
278        D::compile_saturating_add(f, lhs, rhs, out)
279    }
280}
281
282pub struct SaturatingSub;
283
284impl<D: Dialect> Binary<D> for SaturatingSub {
285    fn format_scalar<Lhs: Display, Rhs: Display>(
286        f: &mut std::fmt::Formatter<'_>,
287        lhs: Lhs,
288        rhs: Rhs,
289        out: Item<D>,
290    ) -> std::fmt::Result {
291        D::compile_saturating_sub(f, lhs, rhs, out)
292    }
293}
294
295pub struct Powf;
296
297impl<D: Dialect> Binary<D> for Powf {
298    // Powf doesn't support half and no half equivalent exists
299    fn format_scalar<Lhs: Display, Rhs: Display>(
300        f: &mut std::fmt::Formatter<'_>,
301        lhs: Lhs,
302        rhs: Rhs,
303        item: Item<D>,
304    ) -> std::fmt::Result {
305        let elem = *item.elem();
306        let lhs = lhs.to_string();
307        let rhs = rhs.to_string();
308        match elem {
309            Elem::F16 | Elem::F16x2 | Elem::BF16 | Elem::BF16x2 => {
310                let lhs = format!("float({lhs})");
311                let rhs = format!("float({rhs})");
312                write!(f, "{elem}(")?;
313                D::compile_instruction_powf(f, &lhs, &rhs, Elem::F32)?;
314                write!(f, ")")
315            }
316            _ => D::compile_instruction_powf(f, &lhs, &rhs, elem),
317        }
318    }
319
320    fn can_optimize() -> bool {
321        false
322    }
323}
324
325pub struct FastPowf;
326
327impl<D: Dialect> Binary<D> for FastPowf {
328    // Only executed for f32
329    fn format_scalar<Lhs: Display, Rhs: Display>(
330        f: &mut std::fmt::Formatter<'_>,
331        lhs: Lhs,
332        rhs: Rhs,
333        _item: Item<D>,
334    ) -> std::fmt::Result {
335        write!(
336            f,
337            "{}({lhs}, {rhs})",
338            D::compile_fast_math_function_name("__powf")
339        )
340    }
341}
342
343pub struct Powi;
344
345impl<D: Dialect> Binary<D> for Powi {
346    // Powi doesn't support half and no half equivalent exists
347    fn format_scalar<Lhs: Display, Rhs: Display>(
348        f: &mut std::fmt::Formatter<'_>,
349        lhs: Lhs,
350        rhs: Rhs,
351        item: Item<D>,
352    ) -> std::fmt::Result {
353        let elem = *item.elem();
354        let lhs = lhs.to_string();
355        let rhs = rhs.to_string();
356        match elem {
357            Elem::F16 | Elem::F16x2 | Elem::BF16 | Elem::BF16x2 => {
358                let lhs = format!("float({lhs})");
359
360                write!(f, "{elem}(")?;
361                D::compile_instruction_powf(f, &lhs, &rhs, Elem::F32)?;
362                write!(f, ")")
363            }
364            Elem::F64 => {
365                // RHS needs to be a double.
366                let rhs = format!("double({rhs})");
367
368                D::compile_instruction_powf(f, &lhs, &rhs, elem)
369            }
370            _ => D::compile_instruction_powf(f, &lhs, &rhs, elem),
371        }
372    }
373}
374pub struct ArcTan2;
375
376impl<D: Dialect> Binary<D> for ArcTan2 {
377    // ArcTan2 doesn't support half and no half equivalent exists
378    fn format_scalar<Lhs: Display, Rhs: Display>(
379        f: &mut std::fmt::Formatter<'_>,
380        lhs: Lhs,
381        rhs: Rhs,
382        item: Item<D>,
383    ) -> std::fmt::Result {
384        let elem = item.elem();
385        match elem {
386            Elem::F16 | Elem::F16x2 | Elem::BF16 | Elem::BF16x2 => {
387                write!(f, "{elem}(atan2(float({lhs}), float({rhs})))")
388            }
389            _ => {
390                write!(f, "atan2({lhs}, {rhs})")
391            }
392        }
393    }
394
395    fn can_optimize() -> bool {
396        false
397    }
398}
399
400pub struct Hypot;
401
402impl<D: Dialect> Binary<D> for Hypot {
403    // Hypot doesn't support half and no half equivalent exists
404    fn format_scalar<Lhs, Rhs>(
405        f: &mut Formatter<'_>,
406        lhs: Lhs,
407        rhs: Rhs,
408        item: Item<D>,
409    ) -> std::fmt::Result
410    where
411        Lhs: Component<D>,
412        Rhs: Component<D>,
413    {
414        let elem = *item.elem();
415        let lhs = lhs.to_string();
416        let rhs = rhs.to_string();
417        match elem {
418            Elem::F16 | Elem::F16x2 | Elem::BF16 | Elem::BF16x2 => {
419                let lhs = format!("float({lhs})");
420                let rhs = format!("float({rhs})");
421                write!(f, "{elem}(")?;
422                D::compile_instruction_hypot(f, &lhs, &rhs, Elem::F32)?;
423                write!(f, ")")
424            }
425            _ => D::compile_instruction_hypot(f, &lhs, &rhs, elem),
426        }
427    }
428
429    fn can_optimize() -> bool {
430        false
431    }
432}
433
434pub struct Rhypot;
435
436impl<D: Dialect> Binary<D> for Rhypot {
437    // Rhypot doesn't support half and no half equivalent exists
438    fn format_scalar<Lhs, Rhs>(
439        f: &mut Formatter<'_>,
440        lhs: Lhs,
441        rhs: Rhs,
442        item: Item<D>,
443    ) -> std::fmt::Result
444    where
445        Lhs: Component<D>,
446        Rhs: Component<D>,
447    {
448        let elem = *item.elem();
449        let lhs = lhs.to_string();
450        let rhs = rhs.to_string();
451        match elem {
452            Elem::F16 | Elem::F16x2 | Elem::BF16 | Elem::BF16x2 => {
453                let lhs = format!("float({lhs})");
454                let rhs = format!("float({rhs})");
455                write!(f, "{elem}(")?;
456                D::compile_instruction_rhypot(f, &lhs, &rhs, Elem::F32)?;
457                write!(f, ")")
458            }
459            _ => D::compile_instruction_rhypot(f, &lhs, &rhs, elem),
460        }
461    }
462
463    fn can_optimize() -> bool {
464        false
465    }
466}
467
468pub struct Max;
469
470impl<D: Dialect> Binary<D> for Max {
471    fn format_scalar<Lhs: Display, Rhs: Display>(
472        f: &mut std::fmt::Formatter<'_>,
473        lhs: Lhs,
474        rhs: Rhs,
475        item: Item<D>,
476    ) -> std::fmt::Result {
477        D::compile_instruction_max_function_name(f, item)?;
478        write!(f, "({lhs}, {rhs})")
479    }
480}
481
482pub struct Min;
483
484impl<D: Dialect> Binary<D> for Min {
485    fn format_scalar<Lhs: Display, Rhs: Display>(
486        f: &mut std::fmt::Formatter<'_>,
487        lhs: Lhs,
488        rhs: Rhs,
489        item: Item<D>,
490    ) -> std::fmt::Result {
491        D::compile_instruction_min_function_name(f, item)?;
492        write!(f, "({lhs}, {rhs})")
493    }
494}
495
496pub struct Index;
497
498impl Index {
499    pub(crate) fn format<D: Dialect>(
500        f: &mut Formatter<'_>,
501        list: &Value<D>,
502        index: &Value<D>,
503        out: &Value<D>,
504    ) -> std::fmt::Result {
505        if list.item().vectorization() != out.item().vectorization() {
506            let item_ptr = out.item();
507            let tmp = Value::tmp_declared(item_ptr);
508
509            writeln!(
510                f,
511                "{item_ptr} {tmp} = reinterpret_cast<{item_ptr}>({list});"
512            )?;
513
514            let index = fmt_index(&tmp, index, &tmp.item());
515            writeln!(f, "{item_ptr} {out} = &{index};")
516        } else {
517            let item_out = out.item();
518            let index = fmt_index(list, index, &list.item());
519            if matches!(item_out, Item::Barrier(_)) {
520                let addr_space = D::address_space_for_value(list);
521                writeln!(f, "{addr_space}{}& {out} = {index};", item_out.elem())
522            } else {
523                writeln!(f, "{item_out} {out} = &{index};")
524            }
525        }
526    }
527}
528
529pub fn fmt_index<D: Dialect>(
530    list: &impl Display,
531    index: &impl Display,
532    list_ty: &Item<D>,
533) -> String {
534    // Array nested in pointer, deref first
535    if let Item::Pointer(list_ty, _) = list_ty
536        && matches!(**list_ty, Item::Array(..))
537    {
538        format!("(*{list})[{index}]")
539    } else {
540        format!("{list}[{index}]")
541    }
542}
543
544/// The goal is to support indexing of vectorized types.
545///
546/// # Examples
547///
548/// ```c
549/// float4 rhs;
550/// float item = val[0]; // We want that.
551/// float item = val.x; // So we compile to that.
552/// ```
553pub struct ExtractComponent<D: Dialect> {
554    _dialect: PhantomData<D>,
555}
556
557/// The goal is to support indexing of vectorized types.
558///
559/// # Examples
560///
561/// ```c
562/// float4 val;
563///
564/// val[0] = 1.0; // We want that.
565/// val.x = 1.0;  // So we compile to that.
566/// ```
567pub struct InsertComponent<D: Dialect> {
568    _dialect: PhantomData<D>,
569}
570
571impl<D: Dialect> ExtractComponent<D> {
572    pub fn format(
573        f: &mut Formatter<'_>,
574        lhs: &Value<D>,
575        rhs: &Value<D>,
576        out: &Value<D>,
577    ) -> std::fmt::Result {
578        match rhs {
579            Value::Constant(value, _elem) => {
580                let index = value.as_usize();
581                let out = out.index(index);
582                let lhs = lhs.index(index);
583                let out = out.fmt_left();
584                writeln!(f, "{out} = {lhs};")
585            }
586            _ => {
587                let elem = out.elem();
588                let qualifier = out.const_qualifier();
589                let addr_space = D::address_space_for_value(out);
590                let lhs = lhs.ensure_lvalue(f)?;
591                let out = out.fmt_left();
592                writeln!(
593                    f,
594                    "{out} = reinterpret_cast<{addr_space}{elem}{qualifier}*>(&{lhs})[{rhs}];"
595                )
596            }
597        }
598    }
599}
600
601impl<D: Dialect> InsertComponent<D> {
602    pub fn format(
603        f: &mut Formatter<'_>,
604        vector: &Value<D>,
605        index: &Value<D>,
606        value: &Value<D>,
607        out: &Value<D>,
608    ) -> std::fmt::Result {
609        let index = match index {
610            Value::Constant(value, _) => value.as_usize(),
611            _ => {
612                let tmp = Value::tmp(out.item());
613                writeln!(f, "{} = {vector};", tmp.fmt_left())?;
614
615                let elem = out.elem();
616                let addr_space = D::address_space_for_value(out);
617                let ptr = tmp.fmt_ptr();
618                writeln!(f, "*(({addr_space}{elem}*){ptr} + {index}) = {value};")?;
619                return writeln!(f, "{} = {tmp};", out.fmt_left());
620            }
621        };
622
623        let elements = (0..out.item().vectorization())
624            .map(|i| {
625                if i == index {
626                    format!("{value}")
627                } else {
628                    format!("{}", vector.index(i))
629                }
630            })
631            .join(", ");
632
633        write!(f, "{} = {{{elements}}};", out.fmt_left())
634    }
635}
cubecl_cpp/shared/binary.rs

cubecl_cpp/shared/
binary.rs
