cubecl_core/frontend/operation/
unary.rs

1use core::ops::{Div, Neg, Not};
2use cubecl_common::{e2m1, e2m1x2, e4m3, e5m2, ue8m0};
3use cubecl_ir::dialect::{bitwise::*, general::BoolNotOp, math::*, vector::*};
4use half::{bf16, f16};
5
6use crate::{
7    flex32,
8    frontend::{DivExpand, Scalar},
9    ir::{ExpandValue, Scope},
10    prelude::{
11        CubePrimitive, CubePrimitiveExpand, CubeType, IntoExpand, NativeExpand, Reinterpret,
12    },
13    tf32, unexpanded,
14};
15
16use super::base::unary_expand;
17
18pub trait Abs:
19    CubePrimitive<Scalar: AbsNativeExpand<AbsElem = Self::AbsElem>>
20    + CubeType<
21        ExpandType: AbsExpand<
22            AbsElem = Self::AbsElem,
23            AbsOut = NativeExpand<Self::WithScalar<Self::AbsElem>>,
24        >,
25    > + Sized
26{
27    type AbsElem: Scalar;
28
29    fn abs(self) -> Self::WithScalar<Self::AbsElem> {
30        unexpanded!()
31    }
32
33    fn __expand_abs(
34        scope: &Scope,
35        x: NativeExpand<Self>,
36    ) -> NativeExpand<Self::WithScalar<Self::AbsElem>> {
37        x.__expand_abs_method(scope)
38    }
39}
40
41pub trait AbsExpand {
42    type AbsElem: Scalar;
43    type AbsOut;
44    fn __expand_abs_method(self, scope: &Scope) -> Self::AbsOut;
45}
46
47pub trait AbsNativeExpand {
48    type AbsElem: Scalar;
49    fn __expand_native_abs(scope: &Scope, input: ExpandValue) -> ExpandValue;
50}
51
52pub trait ScalarAbs: Abs<AbsElem = Self> + AbsNativeExpand<AbsElem = Self> {}
53impl<T: Abs<AbsElem = Self> + AbsNativeExpand<AbsElem = Self>> ScalarAbs for T {}
54
55impl<T: Abs> AbsExpand for NativeExpand<T> {
56    type AbsElem = T::AbsElem;
57    type AbsOut = NativeExpand<T::WithScalar<T::AbsElem>>;
58
59    fn __expand_abs_method(self, scope: &Scope) -> Self::AbsOut {
60        T::Scalar::__expand_native_abs(scope, self.into()).into()
61    }
62}
63
64macro_rules! impl_abs {
65    ($($type:ty),*; $operator:expr) => {
66        $(
67            impl Abs for $type { type AbsElem = $type; }
68            impl AbsNativeExpand for $type {
69                type AbsElem = $type;
70                fn __expand_native_abs(scope: &Scope, input: ExpandValue) -> ExpandValue {
71                    unary_expand(scope, input, $operator)
72                }
73            }
74        )*
75    };
76}
77
78macro_rules! impl_abs_nop {
79    ($($type:ty),*) => {
80        $(
81            impl Abs for $type { type AbsElem = $type; }
82            impl AbsNativeExpand for $type {
83                type AbsElem = $type;
84                fn __expand_native_abs(_scope: &Scope, input: ExpandValue) -> ExpandValue {
85                    input
86                }
87            }
88        )*
89    };
90}
91
92pub mod not {
93    use super::*;
94
95    pub fn expand<T: CubeNot>(scope: &Scope, x: NativeExpand<T>) -> NativeExpand<T> {
96        if T::Scalar::elem_type(scope).is_bool() {
97            unary_expand(scope, x.into(), BoolNotOp::new).into()
98        } else {
99            unary_expand(scope, x.into(), BitwiseNotOp::new).into()
100        }
101    }
102}
103
104macro_rules! define_unary_func {
105    ($trait_name:ident, $method_name:ident, $operator:expr, $($type:ty),*) => {
106        paste::paste! {
107            pub trait [<Scalar $trait_name>]: $trait_name + [<$trait_name NativeExpand>] {}
108            impl<T: $trait_name + [<$trait_name NativeExpand>]> [<Scalar $trait_name>] for T {}
109
110            pub trait $trait_name: CubePrimitive<Scalar: [<$trait_name NativeExpand>]>
111                + CubeType<ExpandType: [<$trait_name Expand>]> + Sized {
112                #[allow(unused_variables)]
113                fn $method_name(self) -> Self {
114                    unexpanded!()
115                }
116
117                fn [<__expand_ $method_name>](scope: &Scope, x: NativeExpand<Self>) -> NativeExpand<Self> {
118                    x.[<__expand_ $method_name _method>](scope)
119                }
120            }
121
122            pub trait [<$trait_name Expand>] {
123                fn [<__expand_ $method_name _method>](self, scope: &Scope) -> Self;
124            }
125
126            pub trait [<$trait_name NativeExpand>] {
127                fn [<__expand_native_ $method_name>](scope: &Scope, input: ExpandValue) -> ExpandValue;
128            }
129
130            $(impl $trait_name for $type {})*
131            $(impl [<$trait_name NativeExpand>] for $type {
132                fn [<__expand_native_ $method_name>](scope: &Scope, input: ExpandValue) -> ExpandValue {
133                    unary_expand(scope, input, $operator::new)
134                }
135            })*
136
137            impl<T: $trait_name + CubePrimitive> [<$trait_name Expand>] for NativeExpand<T> {
138                fn [<__expand_ $method_name _method>](self, scope: &Scope) -> Self {
139                    T::Scalar::[<__expand_native_ $method_name>](scope, self.into()).into()
140                }
141            }
142        }
143    }
144}
145
146// Special handling for scalars
147macro_rules! impl_normalize {
148    ($trait_name:ident, $method_name:ident, $operator:expr, $($type:ty),*) => {
149        paste::paste! {
150            pub trait $trait_name: CubePrimitive + CubeType<ExpandType: [<$trait_name Expand>]> + Sized + Abs + Div<Output = Self> {
151                #[allow(unused_variables)]
152                fn $method_name(self) -> Self {
153                    unexpanded!()
154                }
155
156                fn [<__expand_ $method_name>](scope: &Scope, x: NativeExpand<Self>) -> NativeExpand<Self> {
157                    x.[<__expand_ $method_name _method>](scope)
158                }
159            }
160
161            pub trait [<$trait_name Expand>] {
162                fn [<__expand_ $method_name _method>](self, scope: &Scope) -> Self;
163            }
164
165            $(impl $trait_name for $type {})*
166            impl<T: $trait_name + CubePrimitive<WithScalar<<T as Abs>::AbsElem> = T>> [<$trait_name Expand>] for NativeExpand<T> where NativeExpand<T>: DivExpand {
167                fn [<__expand_ $method_name _method>](self, scope: &Scope) -> Self {
168                    if self.__expand_vector_size_method(scope) == 1 {
169                        // Sign might work, but dividing by `abs` preserves the NaN when normalizing 0.0
170                        let abs = self.__expand_abs_method(scope);
171                        self.__expand_div_method(scope, abs)
172                    } else {
173                        unary_expand(scope, self.into(), $operator::new).into()
174                    }
175                }
176            }
177        }
178    }
179}
180
181macro_rules! define_unary_func_scalar_out {
182    ($trait_name:ident, $method_name:ident) => {
183        paste::paste! {
184            pub trait [<Scalar $trait_name>]: $trait_name + [<$trait_name NativeExpand>] {}
185            impl<T: $trait_name + [<$trait_name NativeExpand>]> [<Scalar $trait_name>] for T {}
186
187            pub trait $trait_name: CubePrimitive<Scalar: [<$trait_name NativeExpand>]>
188                + CubeType<ExpandType: [<$trait_name Expand>]
189                + CubePrimitiveExpand<Scalar = NativeExpand<Self::Scalar>>>
190                + Sized {
191                #[allow(unused_variables)]
192                fn $method_name(self) -> Self::Scalar {
193                    unexpanded!()
194                }
195
196                fn [<__expand_ $method_name>](scope: &Scope, x: NativeExpand<Self>) -> NativeExpand<Self::Scalar> {
197                    x.[<__expand_ $method_name _method>](scope)
198                }
199            }
200
201            pub trait [<$trait_name Expand>]: CubePrimitiveExpand {
202                fn [<__expand_ $method_name _method>](self, scope: &Scope) -> Self::Scalar;
203            }
204
205            pub trait [<$trait_name NativeExpand>] {
206                fn [<__expand_native_ $method_name _scalar>](scope: &Scope, input: ExpandValue) -> ExpandValue;
207                fn [<__expand_native_ $method_name>](scope: &Scope, input: ExpandValue) -> ExpandValue;
208            }
209
210            impl<T: $trait_name + CubePrimitive> [<$trait_name Expand>] for NativeExpand<T> {
211                fn [<__expand_ $method_name _method>](self, scope: &Scope) -> Self::Scalar {
212                    // A lot of backends can't deal with 1-sized vectors, and we don't want to deal
213                    // with the scalar edge case.
214                    if self.__expand_vector_size_method(scope) == 1 {
215                        T::Scalar::[<__expand_native_ $method_name _scalar>](scope, self.into()).into()
216                    } else {
217                        T::Scalar::[<__expand_native_ $method_name>](scope, self.into()).into()
218                    }
219                }
220            }
221        }
222    }
223}
224
225macro_rules! impl_unary_func_scalar_out {
226    ($($type:ty),*; $trait_name:ident, $method_name:ident, $operator:expr, $scalar_op: expr) => {
227        paste::paste! {
228            $(impl $trait_name for $type {})*
229            $(impl [<$trait_name NativeExpand>] for $type {
230                fn [<__expand_native_ $method_name _scalar>](scope: &Scope, input: ExpandValue) -> ExpandValue {
231                    ($scalar_op)(scope, input)
232                }
233                fn [<__expand_native_ $method_name>](scope: &Scope, input: ExpandValue) -> ExpandValue {
234                    unary_expand(scope, input, $operator::new)
235                }
236            })*
237        }
238    }
239}
240
241macro_rules! impl_unary_func_fixed_out_ty {
242    ($trait_name:ident, $method_name:ident, $out_ty: ty, $operator:expr, $($type:ty),*) => {
243        paste::paste! {
244            pub trait $trait_name: CubePrimitive + CubeType<ExpandType: [<$trait_name Expand>]
245            + CubePrimitiveExpand<WithScalar<$out_ty> = NativeExpand<Self::WithScalar<$out_ty>>>> + Sized {
246                #[allow(unused_variables, clippy::wrong_self_convention)]
247                fn $method_name(self) -> Self::WithScalar<$out_ty> {
248                    unexpanded!()
249                }
250
251                fn [<__expand_ $method_name>](scope: &Scope, x: NativeExpand<Self>) -> NativeExpand<Self::WithScalar<$out_ty>> {
252                    x.[<__expand_ $method_name _method>](scope)
253                }
254            }
255
256            pub trait [<$trait_name Expand>]: CubePrimitiveExpand {
257                fn [<__expand_ $method_name _method>](self, scope: &Scope) -> Self::WithScalar<$out_ty>;
258            }
259
260            $(impl $trait_name for $type {})*
261            impl<T: $trait_name + CubePrimitive> [<$trait_name Expand>] for NativeExpand<T> {
262                fn [<__expand_ $method_name _method>](self, scope: &Scope) -> Self::WithScalar<$out_ty> {
263                    unary_expand(scope, self.into(), $operator::new).into()
264                }
265            }
266        }
267    }
268}
269
270// Needs special handling because Rust combines bitwise and logical or into one trait
271macro_rules! impl_not {
272    ($trait:ident, $method_name:ident, $($type:ty),*) => {
273        paste::paste! {
274            pub trait [<Cube $trait>]:
275                $trait<Output = Self>
276                + CubePrimitive
277                + CubeType<ExpandType: [<$trait Expand>]>
278                + IntoExpand<Expand = <Self as CubeType>::ExpandType> {
279                fn [<__expand_ $method_name _method>](self, scope: &Scope) -> NativeExpand<Self> {
280                    let this = self.into_expand(scope);
281                    this.[<__expand_ $method_name _method>](scope)
282                }
283
284                fn [<__expand_ $method_name>](scope: &Scope, x: NativeExpand<Self>) -> NativeExpand<Self> {
285                    x.[<__expand_ $method_name _method>](scope)
286                }
287            }
288
289            pub trait [<$trait Expand>] {
290                fn [<__expand_ $method_name _method>](self, scope: &Scope) -> Self;
291            }
292
293            $(impl [<Cube $trait>] for $type {})*
294            impl<T: [<Cube $trait>] + CubePrimitive> [<$trait Expand>] for NativeExpand<T> {
295                fn [<__expand_ $method_name _method>](self, scope: &Scope) -> Self {
296                    not::expand(scope, self.into())
297                }
298            }
299        }
300    }
301}
302
303macro_rules! define_core_unop {
304    ($trait: ident, $method: ident) => {
305        paste::paste! {
306            pub trait [<Scalar $trait>]: [<Cube $trait>] + [<$trait NativeExpand>] {}
307            impl<T: [<Cube $trait>] + [<$trait NativeExpand>]> [<Scalar $trait>] for T {}
308
309            pub trait [<Cube $trait>]:
310                $trait<Output = Self> + CubePrimitive<Scalar: [<$trait NativeExpand>]>
311                + IntoExpand<Expand = <Self as CubeType>::ExpandType>
312                + CubeType<ExpandType: [<$trait Expand>]> + Sized {
313                fn [<__expand_ $method _method>](self, scope: &Scope) -> NativeExpand<Self> {
314                    let this: NativeExpand<Self> = self.into_expand(scope);
315                    this.[<__expand_ $method _method>](scope)
316                }
317
318                fn [<__expand_ $method>](
319                    scope: &Scope,
320                    lhs: NativeExpand<Self>,
321                ) -> NativeExpand<Self> {
322                    lhs.[<__expand_ $method _method>](scope)
323                }
324            }
325
326            pub trait [<$trait Expand>] {
327                fn [<__expand_ $method _method>](self, scope: &Scope) -> Self;
328            }
329
330            pub trait [<$trait NativeExpand>] {
331                fn [<__expand_native_ $method>](scope: &Scope, this: ExpandValue) -> ExpandValue;
332            }
333
334            impl<T: $trait<Output = Self> + CubePrimitive<Scalar: [<$trait NativeExpand>]>
335                + IntoExpand<Expand = <Self as CubeType>::ExpandType>> [<Cube $trait>] for T {}
336            impl<T: [<Cube $trait>]> [<$trait Expand>] for NativeExpand<T> {
337                fn [<__expand_ $method _method>](self, scope: &Scope) -> Self {
338                    T::Scalar::[<__expand_native_ $method>](scope, self.expand).into()
339                }
340            }
341        }
342    };
343}
344
345macro_rules! impl_core_unop {
346    ($($ty: ty),*; $trait: ident, $method: ident, $op: expr) => {
347        paste::paste! {
348            $(impl [<$trait NativeExpand>] for $ty {
349                fn [<__expand_native_ $method>](scope: &Scope, this: ExpandValue) -> ExpandValue {
350                    unary_expand(scope, this, $op::new).into()
351                }
352            })*
353        }
354    };
355}
356
357impl_not!(
358    Not, not, bool, u8, u16, u32, u64, i8, i16, i32, i64, isize, usize
359);
360
361define_core_unop!(Neg, neg);
362impl_core_unop!(i8, i16, i32, i64, isize; Neg, neg, SNegOp);
363impl_core_unop!(f16, bf16, f32, flex32, tf32, f64; Neg, neg, FNegOp);
364
365impl_abs!(i8, i16, i32, i64, isize; SAbsOp::new);
366impl_abs!(e2m1, e4m3, e5m2, ue8m0, f16, bf16, flex32, tf32, f32, f64; FAbsOp::new);
367impl_abs_nop!(u8, u16, u32, u64, usize);
368
369define_unary_func!(Exp, exp, ExpOp, f16, bf16, flex32, tf32, f32, f64);
370define_unary_func!(Log, ln, LogOp, f16, bf16, flex32, tf32, f32, f64);
371define_unary_func!(Log1p, log1p, Log1pOp, f16, bf16, flex32, tf32, f32, f64);
372define_unary_func!(Expm1, exp_m1, Expm1Op, f16, bf16, flex32, tf32, f32, f64);
373define_unary_func!(Cos, cos, CosOp, f16, bf16, flex32, tf32, f32, f64);
374define_unary_func!(Sin, sin, SinOp, f16, bf16, flex32, tf32, f32, f64);
375define_unary_func!(Tan, tan, TanOp, f16, bf16, flex32, tf32, f32, f64);
376define_unary_func!(Tanh, tanh, TanhOp, f16, bf16, flex32, tf32, f32, f64);
377define_unary_func!(Sinh, sinh, SinhOp, f16, bf16, flex32, tf32, f32, f64);
378define_unary_func!(Cosh, cosh, CoshOp, f16, bf16, flex32, tf32, f32, f64);
379define_unary_func!(ArcCos, acos, ArcCosOp, f16, bf16, flex32, tf32, f32, f64);
380define_unary_func!(ArcSin, asin, ArcSinOp, f16, bf16, flex32, tf32, f32, f64);
381define_unary_func!(ArcTan, atan, ArcTanOp, f16, bf16, flex32, tf32, f32, f64);
382define_unary_func!(ArcSinh, asinh, ArcSinhOp, f16, bf16, flex32, tf32, f32, f64);
383define_unary_func!(ArcCosh, acosh, ArcCoshOp, f16, bf16, flex32, tf32, f32, f64);
384define_unary_func!(ArcTanh, atanh, ArcTanhOp, f16, bf16, flex32, tf32, f32, f64);
385define_unary_func!(
386    Degrees, to_degrees, DegreesOp, f16, bf16, flex32, tf32, f32, f64
387);
388define_unary_func!(
389    Radians, to_radians, RadiansOp, f16, bf16, flex32, tf32, f32, f64
390);
391define_unary_func!(Sqrt, sqrt, SqrtOp, f16, bf16, flex32, tf32, f32, f64);
392define_unary_func!(
393    InverseSqrt,
394    inverse_sqrt,
395    RsqrtOp,
396    f16,
397    bf16,
398    flex32,
399    tf32,
400    f32,
401    f64
402);
403define_unary_func!(Round, round, RoundOp, f16, bf16, flex32, tf32, f32, f64);
404define_unary_func!(Floor, floor, FloorOp, f16, bf16, flex32, tf32, f32, f64);
405define_unary_func!(Ceil, ceil, CeilOp, f16, bf16, flex32, tf32, f32, f64);
406define_unary_func!(Trunc, trunc, TruncOp, f16, bf16, flex32, tf32, f32, f64);
407define_unary_func!(Erf, erf, ErfOp, f16, bf16, flex32, tf32, f32, f64);
408define_unary_func!(Recip, recip, RecipOp, f16, bf16, flex32, tf32, f32, f64);
409
410define_unary_func_scalar_out!(Magnitude, magnitude);
411impl_unary_func_scalar_out!(f16, bf16, flex32, tf32, f32, f64; Magnitude, magnitude, MagnitudeOp, |scope, input| unary_expand(scope, input, FAbsOp::new));
412
413define_unary_func_scalar_out!(VectorSum, vector_sum);
414impl_unary_func_scalar_out!(i8, i16, i32, i64, isize, u8, u16, u32, u64, usize; VectorSum, vector_sum, ISumOp, |_, input| input);
415impl_unary_func_scalar_out!(f16, bf16, f32, flex32, tf32, f64; VectorSum, vector_sum, FSumOp, |_, input| input);
416
417impl_normalize!(
418    Normalize,
419    normalize,
420    NormalizeOp,
421    f16,
422    bf16,
423    flex32,
424    tf32,
425    f32,
426    f64
427);
428impl_unary_func_fixed_out_ty!(
429    CountOnes,
430    count_ones,
431    u32,
432    CountOnesOp,
433    u8,
434    i8,
435    u16,
436    i16,
437    u32,
438    i32,
439    u64,
440    i64,
441    usize,
442    isize
443);
444define_unary_func!(
445    ReverseBits,
446    reverse_bits,
447    ReverseBitsOp,
448    u8,
449    i8,
450    u16,
451    i16,
452    u32,
453    i32,
454    u64,
455    i64,
456    usize,
457    isize
458);
459
460impl_unary_func_fixed_out_ty!(
461    LeadingZeros,
462    leading_zeros,
463    u32,
464    LeadingZerosBitsOp,
465    u8,
466    i8,
467    u16,
468    i16,
469    u32,
470    i32,
471    u64,
472    i64,
473    usize,
474    isize
475);
476impl_unary_func_fixed_out_ty!(
477    TrailingZeros,
478    trailing_zeros,
479    u32,
480    TrailingZerosBitsOp,
481    u8,
482    i8,
483    u16,
484    i16,
485    u32,
486    i32,
487    u64,
488    i64,
489    usize,
490    isize
491);
492impl_unary_func_fixed_out_ty!(
493    FindFirstSet,
494    find_first_set,
495    u32,
496    FindFirstSetOp,
497    u8,
498    i8,
499    u16,
500    i16,
501    u32,
502    i32,
503    u64,
504    i64,
505    usize,
506    isize
507);
508impl_unary_func_fixed_out_ty!(
509    IsNan, is_nan, bool, IsNanOp, f16, bf16, flex32, tf32, f32, f64
510);
511impl_unary_func_fixed_out_ty!(
512    IsInf, is_inf, bool, IsInfOp, f16, bf16, flex32, tf32, f32, f64
513);
514
515pub trait FloatBits:
516    CubePrimitive + CubeType<ExpandType: FloatBitsExpand<Bits = Self::Bits>>
517{
518    type Bits: CubePrimitive;
519
520    fn __expand_from_bits(scope: &Scope, bits: NativeExpand<Self::Bits>) -> NativeExpand<Self> {
521        Self::__expand_reinterpret(scope, bits)
522    }
523
524    fn __expand_to_bits(scope: &Scope, this: NativeExpand<Self>) -> NativeExpand<Self::Bits> {
525        <Self::Bits as Reinterpret>::__expand_reinterpret(scope, this)
526    }
527}
528
529pub trait FloatBitsExpand: Sized {
530    type Bits: CubePrimitive;
531
532    fn __expand_to_bits_method(self, scope: &Scope) -> NativeExpand<Self::Bits>;
533}
534
535impl<F: FloatBits> FloatBitsExpand for NativeExpand<F> {
536    type Bits = F::Bits;
537
538    fn __expand_to_bits_method(self, scope: &Scope) -> NativeExpand<Self::Bits> {
539        <Self::Bits as Reinterpret>::__expand_reinterpret(scope, self)
540    }
541}
542
543impl FloatBits for e2m1x2 {
544    type Bits = u8;
545}
546
547impl FloatBits for e5m2 {
548    type Bits = u8;
549}
550
551impl FloatBits for e4m3 {
552    type Bits = u8;
553}
554
555impl FloatBits for f16 {
556    type Bits = u16;
557}
558
559impl FloatBits for bf16 {
560    type Bits = u16;
561}
562
563impl FloatBits for f32 {
564    type Bits = u32;
565}
566
567impl FloatBits for f64 {
568    type Bits = u64;
569}
cubecl_core/frontend/operation/unary.rs

cubecl_core/frontend/operation/
unary.rs
