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

cubecl_core/frontend/operation/
unary.rs
