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!(Cbrt, cbrt, CbrtOp, f16, bf16, flex32, tf32, f32, f64);
393define_unary_func!(
394    InverseSqrt,
395    inverse_sqrt,
396    RsqrtOp,
397    f16,
398    bf16,
399    flex32,
400    tf32,
401    f32,
402    f64
403);
404define_unary_func!(
405    InverseCbrt,
406    inverse_cbrt,
407    RcbrtOp,
408    f16,
409    bf16,
410    flex32,
411    tf32,
412    f32,
413    f64
414);
415define_unary_func!(Round, round, RoundOp, f16, bf16, flex32, tf32, f32, f64);
416define_unary_func!(Floor, floor, FloorOp, f16, bf16, flex32, tf32, f32, f64);
417define_unary_func!(Ceil, ceil, CeilOp, f16, bf16, flex32, tf32, f32, f64);
418define_unary_func!(Trunc, trunc, TruncOp, f16, bf16, flex32, tf32, f32, f64);
419define_unary_func!(Erf, erf, ErfOp, f16, bf16, flex32, tf32, f32, f64);
420define_unary_func!(Recip, recip, RecipOp, f16, bf16, flex32, tf32, f32, f64);
421
422define_unary_func_scalar_out!(Magnitude, magnitude);
423impl_unary_func_scalar_out!(f16, bf16, flex32, tf32, f32, f64; Magnitude, magnitude, MagnitudeOp, |scope, input| unary_expand(scope, input, FAbsOp::new));
424
425define_unary_func_scalar_out!(VectorSum, vector_sum);
426impl_unary_func_scalar_out!(i8, i16, i32, i64, isize, u8, u16, u32, u64, usize; VectorSum, vector_sum, ISumOp, |_, input| input);
427impl_unary_func_scalar_out!(f16, bf16, f32, flex32, tf32, f64; VectorSum, vector_sum, FSumOp, |_, input| input);
428
429impl_normalize!(
430    Normalize,
431    normalize,
432    NormalizeOp,
433    f16,
434    bf16,
435    flex32,
436    tf32,
437    f32,
438    f64
439);
440impl_unary_func_fixed_out_ty!(
441    CountOnes,
442    count_ones,
443    u32,
444    CountOnesOp,
445    u8,
446    i8,
447    u16,
448    i16,
449    u32,
450    i32,
451    u64,
452    i64,
453    usize,
454    isize
455);
456define_unary_func!(
457    ReverseBits,
458    reverse_bits,
459    ReverseBitsOp,
460    u8,
461    i8,
462    u16,
463    i16,
464    u32,
465    i32,
466    u64,
467    i64,
468    usize,
469    isize
470);
471
472impl_unary_func_fixed_out_ty!(
473    LeadingZeros,
474    leading_zeros,
475    u32,
476    LeadingZerosBitsOp,
477    u8,
478    i8,
479    u16,
480    i16,
481    u32,
482    i32,
483    u64,
484    i64,
485    usize,
486    isize
487);
488impl_unary_func_fixed_out_ty!(
489    TrailingZeros,
490    trailing_zeros,
491    u32,
492    TrailingZerosBitsOp,
493    u8,
494    i8,
495    u16,
496    i16,
497    u32,
498    i32,
499    u64,
500    i64,
501    usize,
502    isize
503);
504impl_unary_func_fixed_out_ty!(
505    FindFirstSet,
506    find_first_set,
507    u32,
508    FindFirstSetOp,
509    u8,
510    i8,
511    u16,
512    i16,
513    u32,
514    i32,
515    u64,
516    i64,
517    usize,
518    isize
519);
520impl_unary_func_fixed_out_ty!(
521    IsNan, is_nan, bool, IsNanOp, f16, bf16, flex32, tf32, f32, f64
522);
523impl_unary_func_fixed_out_ty!(
524    IsInf, is_inf, bool, IsInfOp, f16, bf16, flex32, tf32, f32, f64
525);
526
527pub trait FloatBits:
528    CubePrimitive + CubeType<ExpandType: FloatBitsExpand<Bits = Self::Bits>>
529{
530    type Bits: CubePrimitive;
531
532    fn __expand_from_bits(scope: &Scope, bits: NativeExpand<Self::Bits>) -> NativeExpand<Self> {
533        Self::__expand_reinterpret(scope, bits)
534    }
535
536    fn __expand_to_bits(scope: &Scope, this: NativeExpand<Self>) -> NativeExpand<Self::Bits> {
537        <Self::Bits as Reinterpret>::__expand_reinterpret(scope, this)
538    }
539}
540
541pub trait FloatBitsExpand: Sized {
542    type Bits: CubePrimitive;
543
544    fn __expand_to_bits_method(self, scope: &Scope) -> NativeExpand<Self::Bits>;
545}
546
547impl<F: FloatBits> FloatBitsExpand for NativeExpand<F> {
548    type Bits = F::Bits;
549
550    fn __expand_to_bits_method(self, scope: &Scope) -> NativeExpand<Self::Bits> {
551        <Self::Bits as Reinterpret>::__expand_reinterpret(scope, self)
552    }
553}
554
555impl FloatBits for e2m1x2 {
556    type Bits = u8;
557}
558
559impl FloatBits for e5m2 {
560    type Bits = u8;
561}
562
563impl FloatBits for e4m3 {
564    type Bits = u8;
565}
566
567impl FloatBits for f16 {
568    type Bits = u16;
569}
570
571impl FloatBits for bf16 {
572    type Bits = u16;
573}
574
575impl FloatBits for f32 {
576    type Bits = u32;
577}
578
579impl FloatBits for f64 {
580    type Bits = u64;
581}
cubecl_core/frontend/operation/unary.rs

cubecl_core/frontend/operation/
unary.rs
