cubecl_cpp/shared/
binary.rs

1use cubecl_core::{
2    self as cubecl, define_scalar, define_size,
3    frontend::polyfills::{expand_dp4a_polyfill, powi_int},
4    ir::{
5        dialect::{
6            bitwise::*,
7            cmp::*,
8            general::{BoolAndOp, BoolOrOp},
9            math::*,
10            memory::{self, IndexOp},
11        },
12        interfaces::TypedExt,
13        prelude::*,
14        types::{ArrayType, PointerType},
15    },
16    prelude::*,
17};
18use itertools::Itertools;
19
20use crate::{
21    cuda::packed_ops::packable,
22    shared::{
23        CppValue,
24        convert::{no_half, promotes_int},
25        lowering::LowerOp,
26        shared_op, shared_op_with_out,
27        ty::{TypeExtCPP, TypedExtCPP},
28        unroll::unrolling,
29    },
30    target::{CtxTarget, Hip, Target},
31};
32
33#[op_interface_impl]
34impl LowerOp for Dp4aOp {
35    fn should_lower(&self, ctx: &Context) -> bool {
36        ctx.target() != Target::Cuda
37    }
38
39    fn lower(&self, scope: &Scope) -> Vec<Value> {
40        let ctx = scope.ctx();
41        vec![expand_dp4a_polyfill(
42            scope,
43            self.a(ctx),
44            self.b(ctx),
45            self.c(ctx),
46        )]
47    }
48}
49
50macro_rules! operator {
51    ($name:ident, $op:expr) => {
52        shared_op_with_out!($name, |op, ctx| {
53            let lhs = op.lhs(ctx).name(ctx);
54            let rhs = op.rhs(ctx).name(ctx);
55            format!("{lhs} {} {rhs}", $op)
56        });
57        unrolling!($name);
58        promotes_int!($name);
59    };
60}
61
62operator!(IAddOp, "+");
63operator!(FAddOp, "+");
64operator!(ISubOp, "-");
65operator!(FSubOp, "-");
66operator!(SDivOp, "/");
67operator!(UDivOp, "/");
68operator!(FDivOp, "/");
69operator!(IMulOp, "*");
70operator!(FMulOp, "*");
71operator!(IEqualOp, "==");
72operator!(FEqualOp, "==");
73operator!(INotEqualOp, "!=");
74operator!(FNotEqualOp, "!=");
75operator!(SLessThanOp, "<");
76operator!(ULessThanOp, "<");
77operator!(FLessThanOp, "<");
78operator!(SLessThanOrEqualOp, "<=");
79operator!(ULessThanOrEqualOp, "<=");
80operator!(FLessThanOrEqualOp, "<=");
81operator!(SGreaterThanOp, ">");
82operator!(UGreaterThanOp, ">");
83operator!(FGreaterThanOp, ">");
84operator!(SGreaterThanOrEqualOp, ">=");
85operator!(UGreaterThanOrEqualOp, ">=");
86operator!(FGreaterThanOrEqualOp, ">=");
87operator!(ShiftLeftOp, "<<");
88operator!(ShiftRightOp, ">>");
89operator!(BitwiseOrOp, "|");
90operator!(BitwiseAndOp, "&");
91operator!(BitwiseXorOp, "^");
92operator!(BoolOrOp, "||");
93operator!(BoolAndOp, "&&");
94
95shared_op_with_out!(SRemOp, |op, ctx| {
96    let lhs = op.lhs(ctx).name(ctx);
97    let rhs = op.rhs(ctx).name(ctx);
98    format!("{lhs} % {rhs}")
99});
100unrolling!(SRemOp);
101promotes_int!(SRemOp);
102
103shared_op_with_out!(URemOp, |op, ctx| {
104    let lhs = op.lhs(ctx).name(ctx);
105    let rhs = op.rhs(ctx).name(ctx);
106    format!("{lhs} % {rhs}")
107});
108unrolling!(URemOp);
109promotes_int!(URemOp);
110
111shared_op_with_out!(FRemOp, |op, ctx| {
112    let lhs = op.lhs(ctx).name(ctx);
113    let rhs = op.rhs(ctx).name(ctx);
114    format!("fmod({lhs}, {rhs})")
115});
116unrolling!(FRemOp);
117no_half!(FRemOp);
118
119shared_op_with_out!(SModFloorOp, |op, ctx| {
120    let lhs = op.lhs(ctx).name(ctx);
121    let rhs = op.rhs(ctx).name(ctx);
122    let out_elem = op.get_result(ctx).get_type(ctx).to_cpp(ctx);
123    format!("{lhs} - {rhs} * ({out_elem})floor((float){lhs} / (float){rhs})")
124});
125unrolling!(SModFloorOp);
126
127shared_op_with_out!(FModFloorOp, |op, ctx| {
128    let lhs = op.lhs(ctx).name(ctx);
129    let rhs = op.rhs(ctx).name(ctx);
130    let prefix = ctx.target().ty_prefix(ctx, op.get_result(ctx));
131    let floor = format!("{prefix}floor");
132    format!("{lhs} - {rhs} * {floor}({lhs} / {rhs})")
133});
134unrolling!(FModFloorOp);
135packable!(FModFloorOp);
136
137// pub struct FastDiv;
138
139// impl<D: Dialect> Binary<D> for FastDiv {
140//     fn format_scalar<Lhs: Display, Rhs: Display>(
141//         f: &mut std::fmt::Formatter<'_>,
142//         lhs: Lhs,
143//         rhs: Rhs,
144//         _out_item: Item<D>,
145//     ) -> std::fmt::Result {
146//         // f32 only
147//         write!(f, "__fdividef({lhs}, {rhs})")
148//     }
149// }
150
151shared_op_with_out!(SMulHiOp, |op, ctx| {
152    let lhs = op.lhs(ctx);
153    let rhs = op.rhs(ctx).name(ctx);
154    match lhs.size(ctx) {
155        4 => format!("__mulhi({}, {rhs})", lhs.name(ctx)),
156        8 => format!("__mul64hi({}, {rhs})", lhs.name(ctx)),
157        _ => unreachable!("HiMul only supports 32 and 64 bit ints"),
158    }
159});
160unrolling!(SMulHiOp);
161
162shared_op_with_out!(UMulHiOp, |op, ctx| {
163    let lhs = op.lhs(ctx);
164    let rhs = op.rhs(ctx).name(ctx);
165    match lhs.size(ctx) {
166        4 => format!("__umulhi({}, {rhs})", lhs.name(ctx)),
167        8 => format!("__umul64hi({}, {rhs})", lhs.name(ctx)),
168        _ => unreachable!("HiMul only supports 32 and 64 bit ints"),
169    }
170});
171unrolling!(UMulHiOp);
172
173macro_rules! lower_binop {
174    ($ty: ty, $name: ident, $pred: expr) => {
175        $crate::shared::binary::lower_target_binop!($ty, $name, $crate::target::Shared, $pred);
176    };
177    ($ty: ty, $name: ident) => {
178        $crate::shared::binary::lower_binop!($ty, $name, |_, _| true);
179    };
180}
181pub(crate) use lower_binop;
182
183macro_rules! lower_target_binop {
184    ($ty: ty, $name: ident, $target: ty, $pred: expr) => {
185        #[op_interface_impl]
186        impl $crate::shared::lowering::LowerOp<$target> for $ty {
187            fn should_lower(&self, ctx: &Context) -> bool {
188                $crate::shared::closure_inference_hack::<$ty, bool>(self, ctx, $pred)
189            }
190
191            fn lower(&self, scope: &Scope) -> Vec<Value> {
192                use cubecl_core::ir::dialect::OperationPtrExt;
193                define_scalar!(T);
194                define_size!(S);
195                let lhs = self.get_operation().operand(scope.ctx(), 0);
196                let rhs = self.get_operation().operand(scope.ctx(), 1);
197                scope.register_value_type::<T, S>(rhs);
198                vec![$name::expand::<T, S>(scope, lhs.into(), rhs.into()).read_value(scope)]
199            }
200        }
201    };
202    ($ty: ty, $name: ident, $target: ty) => {
203        lower_target_binop!($ty, $name, $target, |_, _| true);
204    };
205}
206pub(crate) use lower_target_binop;
207
208#[cube]
209fn min_bf16<T: Numeric, N: Size>(lhs: Vector<T, N>, rhs: Vector<T, N>) -> Vector<T, N> {
210    let lhs = Vector::<f32, N>::cast_from(lhs);
211    let rhs = Vector::<f32, N>::cast_from(rhs);
212    Vector::cast_from(lhs.min(rhs))
213}
214
215#[cube]
216fn max_bf16<T: Numeric, N: Size>(lhs: Vector<T, N>, rhs: Vector<T, N>) -> Vector<T, N> {
217    let lhs = Vector::<f32, N>::cast_from(lhs);
218    let rhs = Vector::<f32, N>::cast_from(rhs);
219    Vector::cast_from(lhs.max(rhs))
220}
221
222lower_target_binop!(FMinOp, min_bf16, Hip, |op, ctx| {
223    op.lhs(ctx).is_bfloat16(ctx)
224});
225lower_target_binop!(FMaxOp, max_bf16, Hip, |op, ctx| {
226    op.lhs(ctx).is_bfloat16(ctx)
227});
228
229shared_op_with_out!(SMinOp, |op, ctx| {
230    let lhs = op.lhs(ctx).name(ctx);
231    let rhs = op.rhs(ctx).name(ctx);
232    format!("min({lhs}, {rhs})")
233});
234unrolling!(SMinOp);
235promotes_int!(SMinOp);
236
237shared_op_with_out!(UMinOp, |op, ctx| {
238    let lhs = op.lhs(ctx).name(ctx);
239    let rhs = op.rhs(ctx).name(ctx);
240    format!("min({lhs}, {rhs})")
241});
242unrolling!(UMinOp);
243promotes_int!(UMinOp);
244
245shared_op_with_out!(FMinOp, |op, ctx| {
246    let lhs = op.lhs(ctx);
247    let rhs = op.rhs(ctx).name(ctx);
248    if lhs.is_half(ctx) {
249        format!("__hmin({}, {rhs})", lhs.name(ctx))
250    } else if lhs.is_half2(ctx) {
251        format!("__hmin2({}, {rhs})", lhs.name(ctx))
252    } else {
253        format!("min({}, {rhs})", lhs.name(ctx))
254    }
255});
256unrolling!(FMinOp);
257packable!(FMinOp);
258
259shared_op_with_out!(SMaxOp, |op, ctx| {
260    let lhs = op.lhs(ctx).name(ctx);
261    let rhs = op.rhs(ctx).name(ctx);
262    format!("max({lhs}, {rhs})")
263});
264unrolling!(SMaxOp);
265promotes_int!(SMaxOp);
266
267shared_op_with_out!(UMaxOp, |op, ctx| {
268    let lhs = op.lhs(ctx).name(ctx);
269    let rhs = op.rhs(ctx).name(ctx);
270    format!("max({lhs}, {rhs})")
271});
272unrolling!(UMaxOp);
273promotes_int!(UMaxOp);
274
275shared_op_with_out!(FMaxOp, |op, ctx| {
276    let lhs = op.lhs(ctx);
277    let rhs = op.rhs(ctx).name(ctx);
278    if lhs.is_half(ctx) {
279        format!("__hmax({}, {rhs})", lhs.name(ctx))
280    } else if lhs.is_half2(ctx) {
281        format!("__hmax2({}, {rhs})", lhs.name(ctx))
282    } else {
283        format!("max({}, {rhs})", lhs.name(ctx))
284    }
285});
286unrolling!(FMaxOp);
287packable!(FMaxOp);
288
289shared_op_with_out!(SClampOp, |op, ctx| {
290    let input = op.input(ctx).name(ctx);
291    let min = op.min(ctx).name(ctx);
292    let max = op.max(ctx).name(ctx);
293    format!("max(min({input}, {max}), {min})")
294});
295unrolling!(SClampOp);
296promotes_int!(SClampOp);
297
298shared_op_with_out!(UClampOp, |op, ctx| {
299    let input = op.input(ctx).name(ctx);
300    let min = op.min(ctx).name(ctx);
301    let max = op.max(ctx).name(ctx);
302    format!("max(min({input}, {max}), {min})")
303});
304unrolling!(UClampOp);
305promotes_int!(UClampOp);
306
307shared_op_with_out!(FClampOp, |op, ctx| {
308    let input = op.input(ctx);
309    let min = op.min(ctx).name(ctx);
310    let max = op.max(ctx).name(ctx);
311    if input.is_half(ctx) {
312        format!("__hmax(__hmin({}, {max}), {min})", input.name(ctx))
313    } else if input.is_half2(ctx) {
314        format!("__hmax2(__hmin2({}, {max}), {min})", input.name(ctx))
315    } else {
316        format!("max(min({}, {max}), {min})", input.name(ctx))
317    }
318});
319unrolling!(FClampOp);
320packable!(FClampOp);
321
322shared_op_with_out!(PowfOp, |op, ctx| {
323    format!("pow({}, {})", op.lhs(ctx).name(ctx), op.rhs(ctx).name(ctx))
324});
325unrolling!(PowfOp);
326no_half!(PowfOp);
327
328#[op_interface_impl]
329impl LowerOp for PowiOp {
330    fn should_lower(&self, ctx: &Context) -> bool {
331        let ty = self.get_result(ctx).scalar_ty(ctx);
332        ty.is_int(ctx) || ty.is_index(ctx)
333    }
334
335    fn lower(&self, scope: &Scope) -> Vec<Value> {
336        use cubecl_core::ir::dialect::OperationPtrExt;
337
338        define_scalar!(T);
339        define_size!(S);
340        let lhs = self.get_operation().operand(scope.ctx(), 0);
341        let rhs = self.get_operation().operand(scope.ctx(), 1);
342        scope.register_value_type::<T, S>(lhs);
343        vec![powi_int::expand::<T, S>(scope, lhs.into(), rhs.into()).read_value(scope)]
344    }
345}
346
347shared_op_with_out!(PowiOp, |op, ctx| {
348    let lhs = op.lhs(ctx);
349    let rhs = op.rhs(ctx).name(ctx);
350    format!("pow({}, {rhs})", lhs.name(ctx))
351});
352unrolling!(PowiOp);
353no_half!(PowiOp);
354
355// pub struct FastPowf;
356
357// impl<D: Dialect> Binary<D> for FastPowf {
358//     // Only executed for f32
359//     fn format_scalar<Lhs: Display, Rhs: Display>(
360//         f: &mut std::fmt::Formatter<'_>,
361//         lhs: Lhs,
362//         rhs: Rhs,
363//         _item: Item<D>,
364//     ) -> std::fmt::Result {
365//         write!(f, "__powf({lhs}, {rhs})")
366//     }
367// }
368
369shared_op_with_out!(ArcTan2Op, |op, ctx| {
370    let lhs = op.lhs(ctx).name(ctx);
371    let rhs = op.rhs(ctx).name(ctx);
372    format!("atan2({lhs}, {rhs})")
373});
374unrolling!(ArcTan2Op);
375no_half!(ArcTan2Op);
376
377shared_op_with_out!(HypotOp, |op, ctx| {
378    let lhs = op.lhs(ctx);
379    let rhs = op.rhs(ctx).name(ctx);
380    format!("hypot({}, {rhs})", lhs.name(ctx))
381});
382unrolling!(HypotOp);
383no_half!(HypotOp);
384
385shared_op_with_out!(RhypotOp, |op, ctx| {
386    let lhs = op.lhs(ctx);
387    let rhs = op.rhs(ctx).name(ctx);
388    if lhs.is_float32(ctx) {
389        format!("rhypotf({}, {rhs})", lhs.name(ctx))
390    } else {
391        format!("rhypot({}, {rhs})", lhs.name(ctx))
392    }
393});
394unrolling!(RhypotOp);
395no_half!(RhypotOp);
396
397shared_op_with_out!(IndexOp, |op, ctx| {
398    format!("&{}", fmt_index(ctx, op.base(ctx), op.index(ctx)))
399});
400
401pub fn fmt_index(ctx: &Context, list: Value, index: Value) -> String {
402    let list_ty = list.get_type(ctx).deref(ctx);
403    let list = list.name(ctx);
404    let index = index.name(ctx);
405    // Array nested in pointer, deref first
406    if let Some(PointerType { inner, .. }) = list_ty.downcast_ref()
407        && inner.deref(ctx).is::<ArrayType>()
408    {
409        format!("(*{list})[{index}]")
410    } else {
411        format!("{list}[{index}]")
412    }
413}
414
415shared_op!(memory::CopyOp, |op, ctx| {
416    let source = op.source(ctx).name(ctx);
417    let dest = op.destination(ctx).name(ctx);
418    (0..op.len(ctx).0)
419        .map(|i| format!("*({dest} + {i}) = *({source} + {i});\n"))
420        .join("")
421});
cubecl_cpp/shared/binary.rs

cubecl_cpp/shared/
binary.rs
