ruda_kernel/dsl/frontend/operation/
binary.rs

1use crate::dsl::ir::{Arithmetic, Bitwise, ManagedVariable, Operator, Scope};
2use crate::dsl::{
3    flex32,
4    frontend::{RudaPrimitive, NativeExpand},
5    prelude::*,
6};
7use crate::dsl::{frontend::RudaType, tf32};
8use crate::dsl::{
9    frontend::operation::base::{binary_expand, binary_expand_fixed_output},
10    unexpanded,
11};
12use core::{cmp::Ordering, ops::*};
13use ruda_core::{e2m1, e4m3, e5m2, ue8m0};
14use ruda_core::ir::ClampOperator;
15use ruda_kernel_macros::derive_expand;
16use half::{bf16, f16};
17
18pub mod add {
19    use super::*;
20
21    pub fn expand<C: RudaPrimitive>(
22        scope: &mut Scope,
23        lhs: NativeExpand<C>,
24        rhs: NativeExpand<C>,
25    ) -> NativeExpand<C> {
26        binary_expand(scope, lhs.into(), rhs.into(), Arithmetic::Add).into()
27    }
28}
29
30pub mod sub {
31    use ruda_core::ir::{ConstantValue, Variable};
32
33    use super::*;
34
35    pub fn expand<C: RudaPrimitive>(
36        scope: &mut Scope,
37        lhs: NativeExpand<C>,
38        rhs: NativeExpand<C>,
39    ) -> NativeExpand<C> {
40        // Dirty hack to enable slice destructuring with trailing patterns on `Sequence`
41        match (lhs.expand.as_const(), rhs.expand.as_const()) {
42            (Some(ConstantValue::UInt(lhs_val)), Some(ConstantValue::UInt(rhs_val))) => {
43                let item_lhs = lhs.expand.ty;
44                let item_rhs = rhs.expand.ty;
45
46                let vector_size = find_vectorization(item_lhs, item_rhs);
47
48                let item = item_lhs.with_vector_size(vector_size);
49                let value = (lhs_val - rhs_val).into();
50                ManagedVariable::Plain(Variable::constant(value, item)).into()
51            }
52            _ => binary_expand(scope, lhs.into(), rhs.into(), Arithmetic::Sub).into(),
53        }
54    }
55}
56
57pub mod mul {
58    use super::*;
59
60    pub fn expand<C: RudaPrimitive>(
61        scope: &mut Scope,
62        lhs: NativeExpand<C>,
63        rhs: NativeExpand<C>,
64    ) -> NativeExpand<C> {
65        binary_expand(scope, lhs.into(), rhs.into(), Arithmetic::Mul).into()
66    }
67}
68
69pub mod div {
70    use super::*;
71
72    pub fn expand<C: RudaPrimitive>(
73        scope: &mut Scope,
74        lhs: NativeExpand<C>,
75        rhs: NativeExpand<C>,
76    ) -> NativeExpand<C> {
77        binary_expand(scope, lhs.into(), rhs.into(), Arithmetic::Div).into()
78    }
79}
80
81pub mod rem {
82    use super::*;
83
84    pub fn expand<C: RudaPrimitive>(
85        scope: &mut Scope,
86        lhs: NativeExpand<C>,
87        rhs: NativeExpand<C>,
88    ) -> NativeExpand<C> {
89        binary_expand(scope, lhs.into(), rhs.into(), Arithmetic::Modulo).into()
90    }
91}
92
93pub mod and {
94    use super::*;
95
96    pub fn expand<C: RudaPrimitive>(
97        scope: &mut Scope,
98        lhs: NativeExpand<C>,
99        rhs: NativeExpand<C>,
100    ) -> NativeExpand<bool> {
101        binary_expand(scope, lhs.into(), rhs.into(), Operator::And).into()
102    }
103}
104
105pub mod bitand {
106    use super::*;
107
108    pub fn expand<C: RudaPrimitive>(
109        scope: &mut Scope,
110        lhs: NativeExpand<C>,
111        rhs: NativeExpand<C>,
112    ) -> NativeExpand<C> {
113        binary_expand(scope, lhs.into(), rhs.into(), Bitwise::BitwiseAnd).into()
114    }
115}
116
117pub mod bitor {
118    use super::*;
119
120    pub fn expand<C: RudaPrimitive>(
121        scope: &mut Scope,
122        lhs: NativeExpand<C>,
123        rhs: NativeExpand<C>,
124    ) -> NativeExpand<C> {
125        binary_expand(scope, lhs.into(), rhs.into(), Bitwise::BitwiseOr).into()
126    }
127}
128
129pub mod or {
130    use super::*;
131
132    pub fn expand<C: RudaPrimitive>(
133        scope: &mut Scope,
134        lhs: NativeExpand<C>,
135        rhs: NativeExpand<C>,
136    ) -> NativeExpand<bool> {
137        binary_expand(scope, lhs.into(), rhs.into(), Operator::Or).into()
138    }
139}
140
141pub mod bitxor {
142    use super::*;
143
144    pub fn expand<C: RudaPrimitive>(
145        scope: &mut Scope,
146        lhs: NativeExpand<C>,
147        rhs: NativeExpand<C>,
148    ) -> NativeExpand<C> {
149        binary_expand(scope, lhs.into(), rhs.into(), Bitwise::BitwiseXor).into()
150    }
151}
152
153pub mod shl {
154    use super::*;
155
156    pub fn expand<C: RudaPrimitive>(
157        scope: &mut Scope,
158        lhs: NativeExpand<C>,
159        rhs: NativeExpand<C>,
160    ) -> NativeExpand<C> {
161        binary_expand(scope, lhs.into(), rhs.into(), Bitwise::ShiftLeft).into()
162    }
163}
164
165pub mod shr {
166    use super::*;
167
168    pub fn expand<C: RudaPrimitive>(
169        scope: &mut Scope,
170        lhs: NativeExpand<C>,
171        rhs: NativeExpand<C>,
172    ) -> NativeExpand<C> {
173        binary_expand(scope, lhs.into(), rhs.into(), Bitwise::ShiftRight).into()
174    }
175}
176
177pub mod clamp {
178    use super::*;
179
180    pub fn expand<C: PartialOrd + RudaPrimitive>(
181        scope: &mut Scope,
182        input: NativeExpand<C>,
183        min: NativeExpand<C>,
184        max: NativeExpand<C>,
185    ) -> NativeExpand<C> {
186        unary_expand(scope, input.into(), |op| {
187            Arithmetic::Clamp(ClampOperator {
188                input: op.input,
189                min_value: *min.expand,
190                max_value: *max.expand,
191            })
192        })
193        .into()
194    }
195}
196
197pub mod clamp_max {
198    use super::*;
199
200    pub fn expand<C: PartialOrd + RudaPrimitive>(
201        scope: &mut Scope,
202        lhs: NativeExpand<C>,
203        rhs: NativeExpand<C>,
204    ) -> NativeExpand<C> {
205        binary_expand(scope, lhs.into(), rhs.into(), Arithmetic::Min).into()
206    }
207}
208
209pub mod clamp_min {
210    use super::*;
211
212    pub fn expand<C: PartialOrd + RudaPrimitive>(
213        scope: &mut Scope,
214        lhs: NativeExpand<C>,
215        rhs: NativeExpand<C>,
216    ) -> NativeExpand<C> {
217        binary_expand(scope, lhs.into(), rhs.into(), Arithmetic::Max).into()
218    }
219}
220
221/// The minimum of two values, not requiring `Ord`. Provided for clarity in certain cases, though
222/// `clamp_max` may sometimes be more clear.
223pub fn min<T: PartialOrd + RudaPrimitive>(lhs: T, rhs: T) -> T {
224    clamp_max(lhs, rhs)
225}
226
227pub mod min {
228    use super::*;
229
230    pub fn expand<C: PartialOrd + RudaPrimitive>(
231        scope: &mut Scope,
232        lhs: NativeExpand<C>,
233        rhs: NativeExpand<C>,
234    ) -> NativeExpand<C> {
235        let lhs = lhs.constant().map(|value| NativeExpand::from_lit(scope, value)).unwrap_or(lhs);
236        let rhs = rhs.constant().map(|value| NativeExpand::from_lit(scope, value)).unwrap_or(rhs);
237        binary_expand(scope, lhs.into(), rhs.into(), Arithmetic::Min).into()
238    }
239}
240
241/// The maximum of two values, not requiring `Ord`. Provided for clarity in certain cases, though
242/// `clamp_min` may sometimes be more clear.
243pub fn max<T: PartialOrd + RudaPrimitive>(lhs: T, rhs: T) -> T {
244    clamp_min(lhs, rhs)
245}
246
247pub mod max {
248    use super::*;
249
250    pub fn expand<C: PartialOrd + RudaPrimitive>(
251        scope: &mut Scope,
252        lhs: NativeExpand<C>,
253        rhs: NativeExpand<C>,
254    ) -> NativeExpand<C> {
255        binary_expand(scope, lhs.into(), rhs.into(), Arithmetic::Max).into()
256    }
257}
258
259/// For binary functions without special syntax
260macro_rules! impl_binary_func {
261    ($trait_name:ident, $method_name:ident, $operator:expr, $($type:ty),*) => {
262        paste::paste! {
263            pub trait $trait_name: RudaPrimitive + RudaType<ExpandType: [<$trait_name Expand>]> + Sized {
264                fn $method_name(self, _rhs: Self) -> Self {
265                    unexpanded!()
266                }
267
268                fn [<__expand_ $method_name>](
269                    scope: &mut Scope,
270                    lhs: NativeExpand<Self>,
271                    rhs: NativeExpand<Self>,
272                ) -> NativeExpand<Self> {
273                    lhs.[<__expand_ $method_name _method>](scope, rhs)
274                }
275            }
276
277            pub trait [<$trait_name Expand>] {
278                fn [<__expand_ $method_name _method>](self, scope: &mut Scope, rhs: Self) -> Self;
279            }
280
281            $(impl $trait_name for $type {})*
282            impl<T: RudaPrimitive + $trait_name> [<$trait_name Expand>] for NativeExpand<T> {
283                fn [<__expand_ $method_name _method>](self, scope: &mut Scope, rhs: Self) -> Self {
284                    binary_expand(scope, self.into(), rhs.into(), $operator).into()
285                }
286            }
287        }
288    }
289}
290
291macro_rules! impl_binary_func_scalar_out {
292    ($trait_name:ident, $method_name:ident, $operator:expr, $($type:ty),*) => {
293        paste::paste! {
294            pub trait $trait_name: RudaPrimitive
295                + RudaType<ExpandType: [<$trait_name Expand>]
296                + RudaPrimitiveExpand<Scalar = NativeExpand<Self::Scalar>>>
297                + Sized {
298                fn $method_name(self, _rhs: Self) -> Self::Scalar {
299                    unexpanded!()
300                }
301
302                fn [<__expand_ $method_name>](
303                    scope: &mut Scope,
304                    lhs: NativeExpand<Self>,
305                    rhs: NativeExpand<Self>,
306                ) -> NativeExpand<Self::Scalar> {
307                    lhs.[<__expand_ $method_name _method>](scope, rhs)
308                }
309            }
310
311            pub trait [<$trait_name Expand>]: RudaPrimitiveExpand {
312                fn [<__expand_ $method_name _method>](self, scope: &mut Scope, rhs: Self) -> Self::Scalar;
313            }
314
315            $(impl $trait_name for $type {})*
316            impl<T: RudaPrimitive + $trait_name> [<$trait_name Expand>] for NativeExpand<T> {
317                fn [<__expand_ $method_name _method>](self, scope: &mut Scope, rhs: Self) -> Self::Scalar {
318                    let lhs: ManagedVariable = self.into();
319                    let item = lhs.ty.with_vector_size(0);
320                    binary_expand_fixed_output(scope, lhs, rhs.into(), item, $operator).into()
321                }
322            }
323        }
324    }
325}
326
327macro_rules! impl_binary_func_mixed_types {
328    ($trait_name:ident, $method_name:ident, $rhs_ty: ident, $operator:expr, $($type:ty),*) => {
329        paste::paste! {
330            pub trait $trait_name<Rhs: RudaPrimitive + RudaType<ExpandType: Into<ManagedVariable>> + Sized>:
331                RudaPrimitive + RudaType<ExpandType: [<$trait_name Expand>]<Rhs>> + Sized {
332                fn $method_name(self, _rhs: Rhs) -> Self {
333                    unexpanded!()
334                }
335
336                fn [<__expand_ $method_name>](
337                    scope: &mut Scope,
338                    lhs: NativeExpand<Self>,
339                    rhs: NativeExpand<Rhs>,
340                ) -> NativeExpand<Self> {
341                    binary_expand(scope, lhs.into(), rhs.into(), $operator).into()
342                }
343            }
344
345            pub trait [<$trait_name Expand>]<Rhs: RudaType>{
346                fn [<__expand_ $method_name _method>](self, scope: &mut Scope, rhs: Rhs::ExpandType) -> Self;
347            }
348
349            $(impl $trait_name<$rhs_ty> for $type {})*
350            impl<Rhs: RudaPrimitive, T: RudaPrimitive + $trait_name<Rhs>> [<$trait_name Expand>]<Rhs> for NativeExpand<T> {
351                fn [<__expand_ $method_name _method>](self, scope: &mut Scope, rhs: NativeExpand<Rhs>) -> Self {
352                    binary_expand(scope, self.into(), rhs.into(), $operator).into()
353                }
354            }
355        }
356    }
357}
358
359macro_rules! impl_core_binop {
360    ($trait: ident, $method: ident, $op: expr) => {
361        paste::paste! {
362            pub trait [<Ruda $trait>]: $trait<Output = Self> + RudaPrimitive + RudaType<ExpandType: [<$trait Expand>]> + Sized {
363                fn [<__expand_ $method>](
364                    scope: &mut Scope,
365                    lhs: NativeExpand<Self>,
366                    rhs: NativeExpand<Self>,
367                ) -> NativeExpand<Self> {
368                    lhs.[<__expand_ $method _method>](scope, rhs)
369                }
370            }
371
372            pub trait [<$trait Expand>] {
373                fn [<__expand_ $method _method>](self, scope: &mut Scope, rhs: Self) -> Self;
374            }
375
376            impl<T: $trait<Output = T> + RudaPrimitive> [<Ruda $trait>] for T {}
377            impl<T: $trait<Output = T> + RudaPrimitive> [<$trait Expand>] for NativeExpand<T> {
378                fn [<__expand_ $method _method>](self, scope: &mut Scope, rhs: Self) -> Self {
379                    binary_expand(scope, self.into(), rhs.into(), $op).into()
380                }
381            }
382        }
383    };
384}
385
386macro_rules! impl_core_assign_binop {
387    ($trait: ident, $method: ident, $op: expr) => {
388        paste::paste! {
389            pub trait [<Ruda $trait>]: $trait + RudaPrimitive + RudaType<ExpandType: [<$trait Expand>]> + Sized {
390                fn [<__expand_ $method>](
391                    scope: &mut Scope,
392                    lhs: NativeExpand<Self>,
393                    rhs: NativeExpand<Self>,
394                ) {
395                    lhs.[<__expand_ $method _method>](scope, rhs)
396                }
397            }
398
399            pub trait [<$trait Expand>] {
400                fn [<__expand_ $method _method>](self, scope: &mut Scope, rhs: Self);
401            }
402
403            impl<T: $trait + RudaPrimitive> [<Ruda $trait>] for T {}
404            impl<T: $trait + RudaPrimitive> [<$trait Expand>] for NativeExpand<T> {
405                fn [<__expand_ $method _method>](self, scope: &mut Scope, rhs: Self) {
406                    assign_op_expand(scope, self.into(), rhs.into(), $op);
407                }
408            }
409        }
410    };
411}
412
413impl_core_binop!(Add, add, Arithmetic::Add);
414impl_core_binop!(Sub, sub, Arithmetic::Sub);
415impl_core_binop!(Mul, mul, Arithmetic::Mul);
416impl_core_binop!(Div, mul, Arithmetic::Div);
417impl_core_binop!(Rem, rem, Arithmetic::Modulo);
418
419impl_core_assign_binop!(AddAssign, add_assign, Arithmetic::Add);
420impl_core_assign_binop!(SubAssign, sub_assign, Arithmetic::Sub);
421impl_core_assign_binop!(MulAssign, mul_assign, Arithmetic::Mul);
422impl_core_assign_binop!(DivAssign, div_assign, Arithmetic::Div);
423impl_core_assign_binop!(RemAssign, rem_assign, Arithmetic::Modulo);
424
425#[derive_expand(RudaType, RudaTypeMut, IntoRuntime)]
426#[ruda(runtime_variants, no_constructors)]
427pub enum Ordering {
428    Less = -1,
429    Equal = 0,
430    Greater = 1,
431}
432
433fn ordering_disc(name: &'static str) -> NativeExpand<i32> {
434    OrderingExpand::discriminant_of(name).into()
435}
436
437#[allow(non_snake_case)]
438pub trait RudaOrdering {
439    fn Less() -> Ordering {
440        Ordering::Less
441    }
442    fn Equal() -> Ordering {
443        Ordering::Equal
444    }
445    fn Greater() -> Ordering {
446        Ordering::Greater
447    }
448    fn __expand_Less(_scope: &mut Scope) -> OrderingExpand {
449        OrderingExpand {
450            discriminant: ordering_disc("Less"),
451            value: (),
452        }
453    }
454    fn __expand_Equal(_scope: &mut Scope) -> OrderingExpand {
455        OrderingExpand {
456            discriminant: ordering_disc("Equal"),
457            value: (),
458        }
459    }
460    fn __expand_Greater(_scope: &mut Scope) -> OrderingExpand {
461        OrderingExpand {
462            discriminant: ordering_disc("Greater"),
463            value: (),
464        }
465    }
466}
467
468impl RudaOrdering for Ordering {}
469
470pub trait RudaOrd: Ord + RudaType<ExpandType: OrdExpand> + Sized {
471    fn __expand_cmp(
472        scope: &mut Scope,
473        lhs: Self::ExpandType,
474        rhs: Self::ExpandType,
475    ) -> OrderingExpand {
476        lhs.__expand_cmp_method(scope, rhs)
477    }
478
479    fn __expand_min(
480        scope: &mut Scope,
481        lhs: Self::ExpandType,
482        rhs: Self::ExpandType,
483    ) -> Self::ExpandType {
484        lhs.__expand_min_method(scope, rhs)
485    }
486
487    fn __expand_max(
488        scope: &mut Scope,
489        lhs: Self::ExpandType,
490        rhs: Self::ExpandType,
491    ) -> Self::ExpandType {
492        lhs.__expand_max_method(scope, rhs)
493    }
494
495    fn __expand_clamp(
496        scope: &mut Scope,
497        lhs: Self::ExpandType,
498        min: Self::ExpandType,
499        max: Self::ExpandType,
500    ) -> Self::ExpandType {
501        lhs.__expand_clamp_method(scope, min, max)
502    }
503}
504pub trait OrdExpand {
505    fn __expand_cmp_method(self, scope: &mut Scope, rhs: Self) -> OrderingExpand;
506    fn __expand_min_method(self, scope: &mut Scope, rhs: Self) -> Self;
507    fn __expand_max_method(self, scope: &mut Scope, rhs: Self) -> Self;
508    fn __expand_clamp_method(self, scope: &mut Scope, min: Self, max: Self) -> Self;
509}
510
511impl<T: Ord + RudaPrimitive> RudaOrd for T {}
512impl<T: Ord + RudaPrimitive> OrdExpand for NativeExpand<T> {
513    fn __expand_cmp_method(self, scope: &mut Scope, rhs: Self) -> OrderingExpand {
514        let lhs_lt_rhs = lt::expand(scope, self.clone(), rhs.clone());
515        let lhs_gt_rhs = gt::expand(scope, self, rhs);
516        let less = ordering_disc("Less");
517        let equal = ordering_disc("Equal");
518        let greater = ordering_disc("Greater");
519        let eq_or_gt = select::expand(scope, lhs_gt_rhs, greater, equal);
520        let discriminant = select::expand(scope, lhs_lt_rhs, less, eq_or_gt);
521        OrderingExpand {
522            discriminant,
523            value: (),
524        }
525    }
526    fn __expand_min_method(self, scope: &mut Scope, rhs: Self) -> Self {
527        binary_expand(scope, self.into(), rhs.into(), Arithmetic::Min).into()
528    }
529    fn __expand_max_method(self, scope: &mut Scope, rhs: Self) -> Self {
530        binary_expand(scope, self.into(), rhs.into(), Arithmetic::Max).into()
531    }
532    fn __expand_clamp_method(self, scope: &mut Scope, min: Self, max: Self) -> Self {
533        unary_expand(scope, self.into(), |op| {
534            Arithmetic::Clamp(ClampOperator {
535                input: op.input,
536                min_value: *min.expand,
537                max_value: *max.expand,
538            })
539        })
540        .into()
541    }
542}
543
544impl_binary_func!(
545    Powf,
546    powf,
547    Arithmetic::Powf,
548    f16,
549    bf16,
550    flex32,
551    tf32,
552    f32,
553    f64
554);
555
556impl_binary_func!(
557    Hypot,
558    hypot,
559    Arithmetic::Hypot,
560    f16,
561    bf16,
562    flex32,
563    tf32,
564    f32,
565    f64
566);
567
568impl_binary_func!(
569    Rhypot,
570    rhypot,
571    Arithmetic::Rhypot,
572    f16,
573    bf16,
574    flex32,
575    tf32,
576    f32,
577    f64
578);
579
580impl_binary_func!(
581    ArcTan2,
582    atan2,
583    Arithmetic::ArcTan2,
584    f16,
585    bf16,
586    flex32,
587    tf32,
588    f32,
589    f64
590);
591impl_binary_func!(
592    Remainder,
593    rem,
594    Arithmetic::Remainder,
595    e2m1,
596    e4m3,
597    e5m2,
598    ue8m0,
599    f16,
600    bf16,
601    flex32,
602    tf32,
603    f32,
604    f64,
605    i8,
606    i16,
607    i32,
608    i64,
609    u8,
610    u16,
611    u32,
612    u64,
613    usize,
614    isize
615);
616impl_binary_func!(MulHi, mul_hi, Arithmetic::MulHi, i32, u32, usize, isize);
617impl_binary_func!(
618    SaturatingAdd,
619    saturating_add,
620    Arithmetic::SaturatingAdd,
621    i8,
622    i16,
623    i32,
624    i64,
625    u8,
626    u16,
627    u32,
628    u64,
629    usize,
630    isize
631);
632impl_binary_func!(
633    SaturatingSub,
634    saturating_sub,
635    Arithmetic::SaturatingSub,
636    i8,
637    i16,
638    i32,
639    i64,
640    u8,
641    u16,
642    u32,
643    u64,
644    usize,
645    isize
646);
647impl_binary_func_scalar_out!(
648    Dot,
649    dot,
650    Arithmetic::Dot,
651    f16,
652    bf16,
653    flex32,
654    tf32,
655    f32,
656    f64,
657    i8,
658    i16,
659    i32,
660    i64,
661    u8,
662    u16,
663    u32,
664    u64,
665    usize,
666    isize
667);
668
669impl_binary_func_mixed_types!(
670    Powi,
671    powi,
672    i32,
673    Arithmetic::Powi,
674    f16,
675    bf16,
676    flex32,
677    tf32,
678    f32,
679    f64,
680    i8,
681    i16,
682    i32,
683    i64,
684    u8,
685    u16,
686    u32,
687    u64,
688    usize,
689    isize
690);
ruda_kernel/dsl/frontend/operation/binary.rs

ruda_kernel/dsl/frontend/operation/
binary.rs
