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        binary_expand(scope, lhs.into(), rhs.into(), Arithmetic::Min).into()
236    }
237}
238
239/// The maximum of two values, not requiring `Ord`. Provided for clarity in certain cases, though
240/// `clamp_min` may sometimes be more clear.
241pub fn max<T: PartialOrd + RudaPrimitive>(lhs: T, rhs: T) -> T {
242    clamp_min(lhs, rhs)
243}
244
245pub mod max {
246    use super::*;
247
248    pub fn expand<C: PartialOrd + RudaPrimitive>(
249        scope: &mut Scope,
250        lhs: NativeExpand<C>,
251        rhs: NativeExpand<C>,
252    ) -> NativeExpand<C> {
253        binary_expand(scope, lhs.into(), rhs.into(), Arithmetic::Max).into()
254    }
255}
256
257/// For binary functions without special syntax
258macro_rules! impl_binary_func {
259    ($trait_name:ident, $method_name:ident, $operator:expr, $($type:ty),*) => {
260        paste::paste! {
261            pub trait $trait_name: RudaPrimitive + RudaType<ExpandType: [<$trait_name Expand>]> + Sized {
262                fn $method_name(self, _rhs: Self) -> Self {
263                    unexpanded!()
264                }
265
266                fn [<__expand_ $method_name>](
267                    scope: &mut Scope,
268                    lhs: NativeExpand<Self>,
269                    rhs: NativeExpand<Self>,
270                ) -> NativeExpand<Self> {
271                    lhs.[<__expand_ $method_name _method>](scope, rhs)
272                }
273            }
274
275            pub trait [<$trait_name Expand>] {
276                fn [<__expand_ $method_name _method>](self, scope: &mut Scope, rhs: Self) -> Self;
277            }
278
279            $(impl $trait_name for $type {})*
280            impl<T: RudaPrimitive + $trait_name> [<$trait_name Expand>] for NativeExpand<T> {
281                fn [<__expand_ $method_name _method>](self, scope: &mut Scope, rhs: Self) -> Self {
282                    binary_expand(scope, self.into(), rhs.into(), $operator).into()
283                }
284            }
285        }
286    }
287}
288
289macro_rules! impl_binary_func_scalar_out {
290    ($trait_name:ident, $method_name:ident, $operator:expr, $($type:ty),*) => {
291        paste::paste! {
292            pub trait $trait_name: RudaPrimitive
293                + RudaType<ExpandType: [<$trait_name Expand>]
294                + RudaPrimitiveExpand<Scalar = NativeExpand<Self::Scalar>>>
295                + Sized {
296                fn $method_name(self, _rhs: Self) -> Self::Scalar {
297                    unexpanded!()
298                }
299
300                fn [<__expand_ $method_name>](
301                    scope: &mut Scope,
302                    lhs: NativeExpand<Self>,
303                    rhs: NativeExpand<Self>,
304                ) -> NativeExpand<Self::Scalar> {
305                    lhs.[<__expand_ $method_name _method>](scope, rhs)
306                }
307            }
308
309            pub trait [<$trait_name Expand>]: RudaPrimitiveExpand {
310                fn [<__expand_ $method_name _method>](self, scope: &mut Scope, rhs: Self) -> Self::Scalar;
311            }
312
313            $(impl $trait_name for $type {})*
314            impl<T: RudaPrimitive + $trait_name> [<$trait_name Expand>] for NativeExpand<T> {
315                fn [<__expand_ $method_name _method>](self, scope: &mut Scope, rhs: Self) -> Self::Scalar {
316                    let lhs: ManagedVariable = self.into();
317                    let item = lhs.ty.with_vector_size(0);
318                    binary_expand_fixed_output(scope, lhs, rhs.into(), item, $operator).into()
319                }
320            }
321        }
322    }
323}
324
325macro_rules! impl_binary_func_mixed_types {
326    ($trait_name:ident, $method_name:ident, $rhs_ty: ident, $operator:expr, $($type:ty),*) => {
327        paste::paste! {
328            pub trait $trait_name<Rhs: RudaPrimitive + RudaType<ExpandType: Into<ManagedVariable>> + Sized>:
329                RudaPrimitive + RudaType<ExpandType: [<$trait_name Expand>]<Rhs>> + Sized {
330                fn $method_name(self, _rhs: Rhs) -> Self {
331                    unexpanded!()
332                }
333
334                fn [<__expand_ $method_name>](
335                    scope: &mut Scope,
336                    lhs: NativeExpand<Self>,
337                    rhs: NativeExpand<Rhs>,
338                ) -> NativeExpand<Self> {
339                    binary_expand(scope, lhs.into(), rhs.into(), $operator).into()
340                }
341            }
342
343            pub trait [<$trait_name Expand>]<Rhs: RudaType>{
344                fn [<__expand_ $method_name _method>](self, scope: &mut Scope, rhs: Rhs::ExpandType) -> Self;
345            }
346
347            $(impl $trait_name<$rhs_ty> for $type {})*
348            impl<Rhs: RudaPrimitive, T: RudaPrimitive + $trait_name<Rhs>> [<$trait_name Expand>]<Rhs> for NativeExpand<T> {
349                fn [<__expand_ $method_name _method>](self, scope: &mut Scope, rhs: NativeExpand<Rhs>) -> Self {
350                    binary_expand(scope, self.into(), rhs.into(), $operator).into()
351                }
352            }
353        }
354    }
355}
356
357macro_rules! impl_core_binop {
358    ($trait: ident, $method: ident, $op: expr) => {
359        paste::paste! {
360            pub trait [<Ruda $trait>]: $trait<Output = Self> + RudaPrimitive + RudaType<ExpandType: [<$trait Expand>]> + Sized {
361                fn [<__expand_ $method>](
362                    scope: &mut Scope,
363                    lhs: NativeExpand<Self>,
364                    rhs: NativeExpand<Self>,
365                ) -> NativeExpand<Self> {
366                    lhs.[<__expand_ $method _method>](scope, rhs)
367                }
368            }
369
370            pub trait [<$trait Expand>] {
371                fn [<__expand_ $method _method>](self, scope: &mut Scope, rhs: Self) -> Self;
372            }
373
374            impl<T: $trait<Output = T> + RudaPrimitive> [<Ruda $trait>] for T {}
375            impl<T: $trait<Output = T> + RudaPrimitive> [<$trait Expand>] for NativeExpand<T> {
376                fn [<__expand_ $method _method>](self, scope: &mut Scope, rhs: Self) -> Self {
377                    binary_expand(scope, self.into(), rhs.into(), $op).into()
378                }
379            }
380        }
381    };
382}
383
384macro_rules! impl_core_assign_binop {
385    ($trait: ident, $method: ident, $op: expr) => {
386        paste::paste! {
387            pub trait [<Ruda $trait>]: $trait + RudaPrimitive + RudaType<ExpandType: [<$trait Expand>]> + Sized {
388                fn [<__expand_ $method>](
389                    scope: &mut Scope,
390                    lhs: NativeExpand<Self>,
391                    rhs: NativeExpand<Self>,
392                ) {
393                    lhs.[<__expand_ $method _method>](scope, rhs)
394                }
395            }
396
397            pub trait [<$trait Expand>] {
398                fn [<__expand_ $method _method>](self, scope: &mut Scope, rhs: Self);
399            }
400
401            impl<T: $trait + RudaPrimitive> [<Ruda $trait>] for T {}
402            impl<T: $trait + RudaPrimitive> [<$trait Expand>] for NativeExpand<T> {
403                fn [<__expand_ $method _method>](self, scope: &mut Scope, rhs: Self) {
404                    assign_op_expand(scope, self.into(), rhs.into(), $op);
405                }
406            }
407        }
408    };
409}
410
411impl_core_binop!(Add, add, Arithmetic::Add);
412impl_core_binop!(Sub, sub, Arithmetic::Sub);
413impl_core_binop!(Mul, mul, Arithmetic::Mul);
414impl_core_binop!(Div, mul, Arithmetic::Div);
415impl_core_binop!(Rem, rem, Arithmetic::Modulo);
416
417impl_core_assign_binop!(AddAssign, add_assign, Arithmetic::Add);
418impl_core_assign_binop!(SubAssign, sub_assign, Arithmetic::Sub);
419impl_core_assign_binop!(MulAssign, mul_assign, Arithmetic::Mul);
420impl_core_assign_binop!(DivAssign, div_assign, Arithmetic::Div);
421impl_core_assign_binop!(RemAssign, rem_assign, Arithmetic::Modulo);
422
423#[derive_expand(RudaType, RudaTypeMut, IntoRuntime)]
424#[ruda(runtime_variants, no_constructors)]
425pub enum Ordering {
426    Less = -1,
427    Equal = 0,
428    Greater = 1,
429}
430
431fn ordering_disc(name: &'static str) -> NativeExpand<i32> {
432    OrderingExpand::discriminant_of(name).into()
433}
434
435#[allow(non_snake_case)]
436pub trait RudaOrdering {
437    fn Less() -> Ordering {
438        Ordering::Less
439    }
440    fn Equal() -> Ordering {
441        Ordering::Equal
442    }
443    fn Greater() -> Ordering {
444        Ordering::Greater
445    }
446    fn __expand_Less(_scope: &mut Scope) -> OrderingExpand {
447        OrderingExpand {
448            discriminant: ordering_disc("Less"),
449            value: (),
450        }
451    }
452    fn __expand_Equal(_scope: &mut Scope) -> OrderingExpand {
453        OrderingExpand {
454            discriminant: ordering_disc("Equal"),
455            value: (),
456        }
457    }
458    fn __expand_Greater(_scope: &mut Scope) -> OrderingExpand {
459        OrderingExpand {
460            discriminant: ordering_disc("Greater"),
461            value: (),
462        }
463    }
464}
465
466impl RudaOrdering for Ordering {}
467
468pub trait RudaOrd: Ord + RudaType<ExpandType: OrdExpand> + Sized {
469    fn __expand_cmp(
470        scope: &mut Scope,
471        lhs: Self::ExpandType,
472        rhs: Self::ExpandType,
473    ) -> OrderingExpand {
474        lhs.__expand_cmp_method(scope, rhs)
475    }
476
477    fn __expand_min(
478        scope: &mut Scope,
479        lhs: Self::ExpandType,
480        rhs: Self::ExpandType,
481    ) -> Self::ExpandType {
482        lhs.__expand_min_method(scope, rhs)
483    }
484
485    fn __expand_max(
486        scope: &mut Scope,
487        lhs: Self::ExpandType,
488        rhs: Self::ExpandType,
489    ) -> Self::ExpandType {
490        lhs.__expand_max_method(scope, rhs)
491    }
492
493    fn __expand_clamp(
494        scope: &mut Scope,
495        lhs: Self::ExpandType,
496        min: Self::ExpandType,
497        max: Self::ExpandType,
498    ) -> Self::ExpandType {
499        lhs.__expand_clamp_method(scope, min, max)
500    }
501}
502pub trait OrdExpand {
503    fn __expand_cmp_method(self, scope: &mut Scope, rhs: Self) -> OrderingExpand;
504    fn __expand_min_method(self, scope: &mut Scope, rhs: Self) -> Self;
505    fn __expand_max_method(self, scope: &mut Scope, rhs: Self) -> Self;
506    fn __expand_clamp_method(self, scope: &mut Scope, min: Self, max: Self) -> Self;
507}
508
509impl<T: Ord + RudaPrimitive> RudaOrd for T {}
510impl<T: Ord + RudaPrimitive> OrdExpand for NativeExpand<T> {
511    fn __expand_cmp_method(self, scope: &mut Scope, rhs: Self) -> OrderingExpand {
512        let lhs_lt_rhs = lt::expand(scope, self.clone(), rhs.clone());
513        let lhs_gt_rhs = gt::expand(scope, self, rhs);
514        let less = ordering_disc("Less");
515        let equal = ordering_disc("Equal");
516        let greater = ordering_disc("Greater");
517        let eq_or_gt = select::expand(scope, lhs_gt_rhs, greater, equal);
518        let discriminant = select::expand(scope, lhs_lt_rhs, less, eq_or_gt);
519        OrderingExpand {
520            discriminant,
521            value: (),
522        }
523    }
524    fn __expand_min_method(self, scope: &mut Scope, rhs: Self) -> Self {
525        binary_expand(scope, self.into(), rhs.into(), Arithmetic::Min).into()
526    }
527    fn __expand_max_method(self, scope: &mut Scope, rhs: Self) -> Self {
528        binary_expand(scope, self.into(), rhs.into(), Arithmetic::Max).into()
529    }
530    fn __expand_clamp_method(self, scope: &mut Scope, min: Self, max: Self) -> Self {
531        unary_expand(scope, self.into(), |op| {
532            Arithmetic::Clamp(ClampOperator {
533                input: op.input,
534                min_value: *min.expand,
535                max_value: *max.expand,
536            })
537        })
538        .into()
539    }
540}
541
542impl_binary_func!(
543    Powf,
544    powf,
545    Arithmetic::Powf,
546    f16,
547    bf16,
548    flex32,
549    tf32,
550    f32,
551    f64
552);
553
554impl_binary_func!(
555    Hypot,
556    hypot,
557    Arithmetic::Hypot,
558    f16,
559    bf16,
560    flex32,
561    tf32,
562    f32,
563    f64
564);
565
566impl_binary_func!(
567    Rhypot,
568    rhypot,
569    Arithmetic::Rhypot,
570    f16,
571    bf16,
572    flex32,
573    tf32,
574    f32,
575    f64
576);
577
578impl_binary_func!(
579    ArcTan2,
580    atan2,
581    Arithmetic::ArcTan2,
582    f16,
583    bf16,
584    flex32,
585    tf32,
586    f32,
587    f64
588);
589impl_binary_func!(
590    Remainder,
591    rem,
592    Arithmetic::Remainder,
593    e2m1,
594    e4m3,
595    e5m2,
596    ue8m0,
597    f16,
598    bf16,
599    flex32,
600    tf32,
601    f32,
602    f64,
603    i8,
604    i16,
605    i32,
606    i64,
607    u8,
608    u16,
609    u32,
610    u64,
611    usize,
612    isize
613);
614impl_binary_func!(MulHi, mul_hi, Arithmetic::MulHi, i32, u32, usize, isize);
615impl_binary_func!(
616    SaturatingAdd,
617    saturating_add,
618    Arithmetic::SaturatingAdd,
619    i8,
620    i16,
621    i32,
622    i64,
623    u8,
624    u16,
625    u32,
626    u64,
627    usize,
628    isize
629);
630impl_binary_func!(
631    SaturatingSub,
632    saturating_sub,
633    Arithmetic::SaturatingSub,
634    i8,
635    i16,
636    i32,
637    i64,
638    u8,
639    u16,
640    u32,
641    u64,
642    usize,
643    isize
644);
645impl_binary_func_scalar_out!(
646    Dot,
647    dot,
648    Arithmetic::Dot,
649    f16,
650    bf16,
651    flex32,
652    tf32,
653    f32,
654    f64,
655    i8,
656    i16,
657    i32,
658    i64,
659    u8,
660    u16,
661    u32,
662    u64,
663    usize,
664    isize
665);
666
667impl_binary_func_mixed_types!(
668    Powi,
669    powi,
670    i32,
671    Arithmetic::Powi,
672    f16,
673    bf16,
674    flex32,
675    tf32,
676    f32,
677    f64,
678    i8,
679    i16,
680    i32,
681    i64,
682    u8,
683    u16,
684    u32,
685    u64,
686    usize,
687    isize
688);
ruda_kernel/dsl/frontend/operation/binary.rs

ruda_kernel/dsl/frontend/operation/
binary.rs
