cubecl_core/frontend/operation/
binary.rs

1use crate::ir::{ExpandValue, Scope};
2use crate::{
3    flex32,
4    frontend::{CubePrimitive, NativeExpand},
5    prelude::*,
6};
7use crate::{frontend::CubeType, tf32};
8use crate::{frontend::operation::base::binary_expand, unexpanded};
9use core::ops::*;
10use cubecl_ir::dialect::{
11    bitwise::*,
12    general::{BoolAndOp, BoolOrOp},
13    math::*,
14    vector::{FDotOp, SDotOp, UDotOp},
15};
16use half::{bf16, f16};
17
18pub mod sub {
19    use cubecl_ir::{ConstantValue, ExpandValue};
20
21    use super::*;
22
23    pub fn expand<C: CubeSub>(
24        scope: &Scope,
25        lhs: NativeExpand<C>,
26        rhs: NativeExpand<C>,
27    ) -> NativeExpand<C> {
28        // Dirty hack to enable slice destructuring with trailing patterns on `Sequence`
29        match (lhs.expand, rhs.expand.as_const()) {
30            (
31                ExpandValue::Constant {
32                    value: ConstantValue::UInt(lhs_val),
33                    ty,
34                },
35                Some(ConstantValue::UInt(rhs_val)),
36            ) => {
37                let value = (lhs_val - rhs_val).into();
38                ExpandValue::constant(value, ty).into()
39            }
40            _ => C::Scalar::__expand_native_sub(scope, lhs.into(), rhs.into()).into(),
41        }
42    }
43}
44
45pub mod clamp {
46    use super::*;
47
48    pub fn expand<C: CubePartialOrd>(
49        scope: &Scope,
50        input: NativeExpand<C>,
51        min: NativeExpand<C>,
52        max: NativeExpand<C>,
53    ) -> NativeExpand<C> {
54        C::Scalar::__expand_native_clamp(scope, input.into(), min.into(), max.into()).into()
55    }
56}
57
58pub mod clamp_max {
59    use super::*;
60
61    pub fn expand<C: CubePartialOrd>(
62        scope: &Scope,
63        lhs: NativeExpand<C>,
64        rhs: NativeExpand<C>,
65    ) -> NativeExpand<C> {
66        C::Scalar::__expand_native_min(scope, lhs.into(), rhs.into()).into()
67    }
68}
69
70pub mod clamp_min {
71    use super::*;
72
73    pub fn expand<C: CubePartialOrd>(
74        scope: &Scope,
75        lhs: NativeExpand<C>,
76        rhs: NativeExpand<C>,
77    ) -> NativeExpand<C> {
78        C::Scalar::__expand_native_max(scope, lhs.into(), rhs.into()).into()
79    }
80}
81
82/// The minimum of two values, not requiring `Ord`. Provided for clarity in certain cases, though
83/// `clamp_max` may sometimes be more clear.
84pub fn min<T: CubePartialOrd>(lhs: T, rhs: T) -> T {
85    clamp_max(lhs, rhs)
86}
87
88pub mod min {
89    use super::*;
90
91    pub fn expand<C: CubePartialOrd>(
92        scope: &Scope,
93        lhs: NativeExpand<C>,
94        rhs: NativeExpand<C>,
95    ) -> NativeExpand<C> {
96        C::Scalar::__expand_native_min(scope, lhs.into(), rhs.into()).into()
97    }
98}
99
100/// The maximum of two values, not requiring `Ord`. Provided for clarity in certain cases, though
101/// `clamp_min` may sometimes be more clear.
102pub fn max<T: CubePartialOrd>(lhs: T, rhs: T) -> T {
103    clamp_min(lhs, rhs)
104}
105
106pub mod max {
107    use super::*;
108
109    pub fn expand<C: CubePartialOrd>(
110        scope: &Scope,
111        lhs: NativeExpand<C>,
112        rhs: NativeExpand<C>,
113    ) -> NativeExpand<C> {
114        C::Scalar::__expand_native_max(scope, lhs.into(), rhs.into()).into()
115    }
116}
117
118/// For binary functions without special syntax
119macro_rules! define_binary_func {
120    ($trait_name:ident, $method_name:ident) => {
121        paste::paste! {
122            pub trait [<Scalar $trait_name>]: $trait_name + [<$trait_name NativeExpand>] {}
123            impl<T: $trait_name + [<$trait_name NativeExpand>]> [<Scalar $trait_name>] for T {}
124
125            pub trait $trait_name: CubePrimitive<Scalar: [<$trait_name NativeExpand>]>
126                + CubeType<ExpandType: [<$trait_name Expand>]> + Sized {
127                fn $method_name(self, _rhs: Self) -> Self {
128                    unexpanded!()
129                }
130
131                fn [<__expand_ $method_name>](
132                    scope: &Scope,
133                    lhs: NativeExpand<Self>,
134                    rhs: NativeExpand<Self>,
135                ) -> NativeExpand<Self> {
136                    lhs.[<__expand_ $method_name _method>](scope, rhs)
137                }
138            }
139
140            pub trait [<$trait_name Expand>] {
141                fn [<__expand_ $method_name _method>](self, scope: &Scope, rhs: Self) -> Self;
142            }
143
144            pub trait [<$trait_name NativeExpand>] {
145                fn [<__expand_native_ $method_name>](scope: &Scope, lhs: ExpandValue, rhs: ExpandValue) -> ExpandValue;
146            }
147
148            impl<T: $trait_name> [<$trait_name Expand>] for NativeExpand<T> {
149                fn [<__expand_ $method_name _method>](self, scope: &Scope, rhs: Self) -> Self {
150                    T::Scalar::[<__expand_native_ $method_name>](scope, self.into(), rhs.into()).into()
151                }
152            }
153        }
154    }
155}
156
157macro_rules! impl_binary_func {
158    ($($type:ty),*; $trait_name:ident, $method_name:ident, $operator:expr) => {
159        paste::paste! {
160            $(impl $trait_name for $type {})*
161            $(impl [<$trait_name NativeExpand>] for $type {
162                fn [<__expand_native_ $method_name>](scope: &Scope, lhs: ExpandValue, rhs: ExpandValue) -> ExpandValue {
163                    binary_expand(scope, lhs, rhs, $operator::new)
164                }
165            })*
166        }
167    }
168}
169
170macro_rules! define_binary_func_scalar_out {
171    ($trait_name:ident, $method_name:ident) => {
172        paste::paste! {
173            pub trait [<Scalar $trait_name>]: $trait_name + [<$trait_name NativeExpand>] {}
174            impl<T: $trait_name + [<$trait_name NativeExpand>]> [<Scalar $trait_name>] for T {}
175
176            pub trait $trait_name: CubePrimitive<Scalar: [<$trait_name NativeExpand>]>
177                + CubeType<ExpandType: [<$trait_name Expand>]
178                + CubePrimitiveExpand<Scalar = NativeExpand<Self::Scalar>>>
179                + Sized {
180                fn $method_name(self, _rhs: Self) -> Self::Scalar {
181                    unexpanded!()
182                }
183
184                fn [<__expand_ $method_name>](
185                    scope: &Scope,
186                    lhs: NativeExpand<Self>,
187                    rhs: NativeExpand<Self>,
188                ) -> NativeExpand<Self::Scalar> {
189                    lhs.[<__expand_ $method_name _method>](scope, rhs)
190                }
191            }
192
193            pub trait [<$trait_name Expand>]: CubePrimitiveExpand {
194                fn [<__expand_ $method_name _method>](self, scope: &Scope, rhs: Self) -> Self::Scalar;
195            }
196
197            pub trait [<$trait_name NativeExpand>] {
198                fn [<__expand_native_ $method_name _scalar>](scope: &Scope, lhs: ExpandValue, rhs: ExpandValue) -> ExpandValue;
199                fn [<__expand_native_ $method_name>](scope: &Scope, lhs: ExpandValue, rhs: ExpandValue) -> ExpandValue;
200            }
201
202            impl<T: CubePrimitive + $trait_name> [<$trait_name Expand>] for NativeExpand<T> {
203                fn [<__expand_ $method_name _method>](self, scope: &Scope, rhs: Self) -> Self::Scalar {
204                    // A lot of backends can't deal with 1-sized vectors, and we want to validate
205                    // that the input is a vector
206                    if self.__expand_vector_size_method(scope) == 1 {
207                        T::Scalar::[<__expand_native_ $method_name _scalar>](scope, self.into(), rhs.into()).into()
208                    } else {
209                        T::Scalar::[<__expand_native_ $method_name>](scope, self.into(), rhs.into()).into()
210                    }
211                }
212            }
213        }
214    }
215}
216
217macro_rules! impl_binary_func_scalar_out {
218    ($($type:ty),*; $trait_name:ident, $method_name:ident, $operator:expr, $scalar_op:expr) => {
219        paste::paste! {
220            $(impl $trait_name for $type {})*
221            $(impl [<$trait_name NativeExpand>] for $type {
222                fn [<__expand_native_ $method_name _scalar>](scope: &Scope, lhs: ExpandValue, rhs: ExpandValue) -> ExpandValue {
223                    binary_expand(scope, lhs, rhs, $scalar_op::new)
224                }
225                fn [<__expand_native_ $method_name>](scope: &Scope, lhs: ExpandValue, rhs: ExpandValue) -> ExpandValue {
226                    binary_expand(scope, lhs, rhs, $operator::new)
227                }
228            })*
229        }
230    }
231}
232
233macro_rules! impl_binary_func_mixed_types {
234    ($trait_name:ident, $method_name:ident, $rhs_ty: ident, $operator:expr, $($type:ty),*) => {
235        paste::paste! {
236            pub trait $trait_name<Rhs: CubePrimitive + CubeType<ExpandType: Into<ExpandValue>> + Sized>:
237                CubePrimitive + CubeType<ExpandType: [<$trait_name Expand>]<Rhs>> + Sized {
238                fn $method_name(self, _rhs: Rhs) -> Self {
239                    unexpanded!()
240                }
241
242                fn [<__expand_ $method_name>](
243                    scope: &Scope,
244                    lhs: NativeExpand<Self>,
245                    rhs: NativeExpand<Rhs>,
246                ) -> NativeExpand<Self> {
247                    binary_expand(scope, lhs.into(), rhs.into(), $operator::new).into()
248                }
249            }
250
251            pub trait [<$trait_name Expand>]<Rhs: CubeType>{
252                fn [<__expand_ $method_name _method>](self, scope: &Scope, rhs: Rhs::ExpandType) -> Self;
253            }
254
255            $(impl $trait_name<$rhs_ty> for $type {})*
256            impl<Rhs: CubePrimitive, T: CubePrimitive + $trait_name<Rhs>> [<$trait_name Expand>]<Rhs> for NativeExpand<T> {
257                fn [<__expand_ $method_name _method>](self, scope: &Scope, rhs: NativeExpand<Rhs>) -> Self {
258                    binary_expand(scope, self.into(), rhs.into(), $operator::new).into()
259                }
260            }
261        }
262    }
263}
264
265macro_rules! define_core_binop {
266    ($trait: ident, $method: ident) => {
267        paste::paste! {
268            pub trait [<Scalar $trait>]: [<Cube $trait>] + [<$trait NativeExpand>] {}
269            impl<T: [<Cube $trait>] + [<$trait NativeExpand>]> [<Scalar $trait>] for T {}
270
271            pub trait [<Cube $trait>]:
272                $trait<Output = Self> + CubePrimitive<Scalar: [<$trait NativeExpand>]> + IntoRuntime
273                + CubeType<ExpandType: [<$trait Expand>]> + Sized {
274                fn [<__expand_ $method _method>](self, scope: &Scope, rhs: NativeExpand<Self>) -> NativeExpand<Self> {
275                    let this = self.__expand_runtime_method(scope);
276                    this.[<__expand_ $method _method>](scope, rhs)
277                }
278
279                fn [<__expand_ $method>](
280                    scope: &Scope,
281                    lhs: NativeExpand<Self>,
282                    rhs: NativeExpand<Self>,
283                ) -> NativeExpand<Self> {
284                    lhs.[<__expand_ $method _method>](scope, rhs)
285                }
286            }
287
288            pub trait [<$trait Expand>] {
289                fn [<__expand_ $method _method>](self, scope: &Scope, rhs: Self) -> Self;
290            }
291
292            pub trait [<$trait NativeExpand>] {
293                fn [<__expand_native_ $method>](scope: &Scope, this: ExpandValue, rhs: ExpandValue) -> ExpandValue;
294            }
295
296            impl<T: $trait<Output = Self> + CubePrimitive<Scalar: [<$trait NativeExpand>]> + IntoRuntime> [<Cube $trait>] for T {}
297            impl<T: [<Cube $trait>]> [<$trait Expand>] for NativeExpand<T> {
298                fn [<__expand_ $method _method>](self, scope: &Scope, rhs: Self) -> Self {
299                    T::Scalar::[<__expand_native_ $method>](scope, self.into(), rhs.into()).into()
300                }
301            }
302        }
303    };
304}
305
306macro_rules! impl_core_binop {
307    ($($ty: ty),*; $trait: ident, $method: ident, $op: expr) => {
308        paste::paste! {
309            $(impl [<$trait NativeExpand>] for $ty {
310                fn [<__expand_native_ $method>](scope: &Scope, this: ExpandValue, rhs: ExpandValue) -> ExpandValue {
311                    binary_expand(scope, this, rhs, $op::new)
312                }
313            })*
314        }
315    };
316}
317macro_rules! define_core_assign_binop {
318    ($trait: ident, $base_trait: ident, $method: ident, $base_method: ident) => {
319        paste::paste! {
320            pub trait [<Cube $trait>]: $trait + CubePrimitive<Scalar: [<$base_trait NativeExpand>]>
321                + CubeType<ExpandType: [<$trait Expand>]> + Sized {
322                fn [<__expand_ $method>](
323                    scope: &Scope,
324                    lhs: &mut NativeExpand<Self>,
325                    rhs: NativeExpand<Self>,
326                ) {
327                    lhs.[<__expand_ $method _method>](scope, rhs)
328                }
329            }
330
331            pub trait [<$trait Expand>] {
332                fn [<__expand_ $method _method>](&mut self, scope: &Scope, rhs: Self);
333            }
334
335            impl<T: $trait + [<Cube $base_trait>]> [<Cube $trait>] for T {}
336            impl<T: $trait + [<Cube $base_trait>]> [<$trait Expand>] for NativeExpand<T> {
337                fn [<__expand_ $method _method>](&mut self, scope: &Scope, rhs: Self) {
338                    assign_binop_expand(scope, self, rhs, T::Scalar::[<__expand_native_ $base_method>]);
339                }
340            }
341        }
342    };
343}
344
345define_core_binop!(Add, add);
346define_core_binop!(Sub, sub);
347define_core_binop!(Mul, mul);
348define_core_binop!(Div, div);
349define_core_binop!(Rem, rem);
350
351impl_core_binop!(i8, i16, i32, i64, isize, u8, u16, u32, u64, usize; Add, add, IAddOp);
352impl_core_binop!(f16, bf16, f32, flex32, tf32, f64; Add, add, FAddOp);
353
354impl_core_binop!(i8, i16, i32, i64, isize, u8, u16, u32, u64, usize; Sub, sub, ISubOp);
355impl_core_binop!(f16, bf16, f32, flex32, tf32, f64; Sub, sub, FSubOp);
356
357impl_core_binop!(i8, i16, i32, i64, isize, u8, u16, u32, u64, usize; Mul, mul, IMulOp);
358impl_core_binop!(f16, bf16, f32, flex32, tf32, f64; Mul, mul, FMulOp);
359
360impl_core_binop!(i8, i16, i32, i64, isize; Div, div, SDivOp);
361impl_core_binop!(u8, u16, u32, u64, usize; Div, div, UDivOp);
362impl_core_binop!(f16, bf16, f32, flex32, tf32, f64; Div, div, FDivOp);
363
364impl_core_binop!(i8, i16, i32, i64, isize; Rem, rem, SRemOp);
365impl_core_binop!(u8, u16, u32, u64, usize; Rem, rem, URemOp);
366impl_core_binop!(f16, bf16, f32, flex32, tf32, f64; Rem, rem, FRemOp);
367
368define_core_assign_binop!(AddAssign, Add, add_assign, add);
369define_core_assign_binop!(SubAssign, Sub, sub_assign, sub);
370define_core_assign_binop!(MulAssign, Mul, mul_assign, mul);
371define_core_assign_binop!(DivAssign, Div, div_assign, div);
372define_core_assign_binop!(RemAssign, Rem, rem_assign, rem);
373
374define_core_binop!(BitAnd, bitand);
375define_core_binop!(BitOr, bitor);
376define_core_binop!(BitXor, bitxor);
377define_core_binop!(Shl, shl);
378define_core_binop!(Shr, shr);
379
380impl_core_binop!(i8, i16, i32, i64, isize, u8, u16, u32, u64, usize; BitAnd, bitand, BitwiseAndOp);
381impl_core_binop!(i8, i16, i32, i64, isize, u8, u16, u32, u64, usize; BitOr, bitor, BitwiseOrOp);
382impl_core_binop!(i8, i16, i32, i64, isize, u8, u16, u32, u64, usize; BitXor, bitxor, BitwiseXorOp);
383impl_core_binop!(i8, i16, i32, i64, isize, u8, u16, u32, u64, usize; Shl, shl, ShiftLeftOp);
384impl_core_binop!(i8, i16, i32, i64, isize, u8, u16, u32, u64, usize; Shr, shr, ShiftRightOp);
385
386define_core_assign_binop!(BitAndAssign, BitAnd, bitand_assign, bitand);
387define_core_assign_binop!(BitOrAssign, BitOr, bitor_assign, bitor);
388define_core_assign_binop!(BitXorAssign, BitXor, bitxor_assign, bitxor);
389define_core_assign_binop!(ShlAssign, Shl, shl_assign, shl);
390define_core_assign_binop!(ShrAssign, Shr, shr_assign, shr);
391
392impl BitAndNativeExpand for bool {
393    fn __expand_native_bitand(scope: &Scope, this: ExpandValue, rhs: ExpandValue) -> ExpandValue {
394        binary_expand(scope, this, rhs, BoolAndOp::new)
395    }
396}
397impl BitOrNativeExpand for bool {
398    fn __expand_native_bitor(scope: &Scope, this: ExpandValue, rhs: ExpandValue) -> ExpandValue {
399        binary_expand(scope, this, rhs, BoolOrOp::new)
400    }
401}
402
403pub trait CubeAnd:
404    CubePrimitive + Into<ExpandValue> + CubeType<ExpandType: AndExpand> + Sized
405{
406    fn __expand_and_method(self, scope: &Scope, rhs: NativeExpand<Self>) -> NativeExpand<Self> {
407        let this: ExpandValue = self.into();
408        let this: NativeExpand<Self> = this.into();
409        this.__expand_and_method(scope, rhs)
410    }
411    fn __expand_and(
412        scope: &Scope,
413        lhs: NativeExpand<Self>,
414        rhs: NativeExpand<Self>,
415    ) -> NativeExpand<Self> {
416        lhs.__expand_and_method(scope, rhs)
417    }
418}
419pub trait AndExpand {
420    fn __expand_and_method(self, scope: &Scope, rhs: Self) -> Self;
421}
422
423impl CubeAnd for bool {}
424impl<T: CubeAnd + CubePrimitive> AndExpand for NativeExpand<T> {
425    fn __expand_and_method(self, scope: &Scope, rhs: Self) -> Self {
426        binary_expand(scope, self.into(), rhs.into(), BoolAndOp::new).into()
427    }
428}
429
430pub trait CubeOr:
431    CubePrimitive + Into<ExpandValue> + CubeType<ExpandType: OrExpand> + Sized
432{
433    fn __expand_or_method(self, scope: &Scope, rhs: NativeExpand<Self>) -> NativeExpand<Self> {
434        let this: ExpandValue = self.into();
435        let this: NativeExpand<Self> = this.into();
436        this.__expand_or_method(scope, rhs)
437    }
438    fn __expand_or(
439        scope: &Scope,
440        lhs: NativeExpand<Self>,
441        rhs: NativeExpand<Self>,
442    ) -> NativeExpand<Self> {
443        lhs.__expand_or_method(scope, rhs)
444    }
445}
446pub trait OrExpand {
447    fn __expand_or_method(self, scope: &Scope, rhs: Self) -> Self;
448}
449
450impl CubeOr for bool {}
451impl<T: CubeOr + CubePrimitive> OrExpand for NativeExpand<T> {
452    fn __expand_or_method(self, scope: &Scope, rhs: Self) -> Self {
453        binary_expand(scope, self.into(), rhs.into(), BoolOrOp::new).into()
454    }
455}
456
457define_binary_func!(Powf, powf);
458impl_binary_func!(f16, bf16, flex32, tf32, f32, f64; Powf, powf, PowfOp);
459
460define_binary_func!(Hypot, hypot);
461impl_binary_func!(f16, bf16, flex32, tf32, f32, f64; Hypot, hypot, HypotOp);
462
463define_binary_func!(Rhypot, rhypot);
464impl_binary_func!(f16, bf16, flex32, tf32, f32, f64; Rhypot, rhypot, RhypotOp);
465
466define_binary_func!(ArcTan2, atan2);
467impl_binary_func!(f16, bf16, flex32, tf32, f32, f64; ArcTan2, atan2, ArcTan2Op);
468
469define_binary_func!(ModFloor, mod_floor);
470impl_binary_func!(i8, i16, i32, i64, isize; ModFloor, mod_floor, SModFloorOp);
471impl_binary_func!(u8, u16, u32, u64, usize; ModFloor, mod_floor, URemOp);
472impl_binary_func!(f16, bf16, flex32, tf32, f32, f64; ModFloor, mod_floor, FModFloorOp);
473
474define_binary_func!(MulHi, mul_hi);
475impl_binary_func!(i32, i64, isize; MulHi, mul_hi, SMulHiOp);
476impl_binary_func!(u32, u64, usize; MulHi, mul_hi, UMulHiOp);
477
478define_binary_func!(SaturatingAdd, saturating_add);
479impl_binary_func!(i8, i16, i32, i64, isize; SaturatingAdd, saturating_add, SaturatingSAddOp);
480impl_binary_func!(u8, u16, u32, u64, usize; SaturatingAdd, saturating_add, SaturatingUAddOp);
481
482define_binary_func!(SaturatingSub, saturating_sub);
483impl_binary_func!(i8, i16, i32, i64, isize; SaturatingSub, saturating_sub, SaturatingSSubOp);
484impl_binary_func!(u8, u16, u32, u64, usize; SaturatingSub, saturating_sub, SaturatingUSubOp);
485
486define_binary_func_scalar_out!(Dot, dot);
487impl_binary_func_scalar_out!(i8, i16, i32, i64, isize; Dot, dot, SDotOp, IMulOp);
488impl_binary_func_scalar_out!(u8, u16, u32, u64, usize; Dot, dot, UDotOp, IMulOp);
489impl_binary_func_scalar_out!(f16, bf16, flex32, tf32, f32, f64; Dot, dot, FDotOp, FMulOp);
490
491impl_binary_func_mixed_types!(
492    Powi, powi, i32, PowiOp, f16, bf16, flex32, tf32, f32, f64, i8, i16, i32, i64, u8, u16, u32,
493    u64, usize, isize
494);
cubecl_core/frontend/operation/binary.rs

cubecl_core/frontend/operation/
binary.rs
