cubecl_core/frontend/operation/
binary.rs

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

cubecl_core/frontend/operation/
binary.rs
