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

cubecl_core/frontend/operation/
unary.rs
