cubecl_core/frontend/operation/
unary.rs

1use core::ops::{Neg, Not};
2use cubecl_common::{e2m1, e2m1x2, e4m3, e5m2, ue8m0};
3use cubecl_ir::{Bitwise, Comparison, Operator};
4use half::{bf16, f16};
5
6use crate::{
7    flex32,
8    ir::{Arithmetic, Scope, Value},
9    prelude::{
10        CubePrimitive, CubePrimitiveExpand, CubeType, IntoExpand, NativeExpand, Reinterpret,
11    },
12    tf32, unexpanded,
13};
14
15use super::base::{unary_expand, unary_expand_fixed_output};
16
17pub mod not {
18    use super::*;
19
20    pub fn expand<T: CubeNot>(scope: &Scope, x: NativeExpand<T>) -> NativeExpand<T> {
21        if x.expand.ty.is_bool() {
22            unary_expand(scope, x.into(), Operator::Not).into()
23        } else {
24            unary_expand(scope, x.into(), Bitwise::BitwiseNot).into()
25        }
26    }
27}
28
29macro_rules! impl_unary_func {
30    ($trait_name:ident, $method_name:ident, $operator:expr, $($type:ty),*) => {
31        paste::paste! {
32            pub trait $trait_name: CubePrimitive + CubeType<ExpandType: [<$trait_name Expand>]> + Sized {
33                #[allow(unused_variables)]
34                fn $method_name(self) -> Self {
35                    unexpanded!()
36                }
37
38                fn [<__expand_ $method_name>](scope: &Scope, x: NativeExpand<Self>) -> NativeExpand<Self> {
39                    x.[<__expand_ $method_name _method>](scope)
40                }
41            }
42
43            pub trait [<$trait_name Expand>] {
44                fn [<__expand_ $method_name _method>](self, scope: &Scope) -> Self;
45            }
46
47            $(impl $trait_name for $type {})*
48            impl<T: $trait_name + CubePrimitive> [<$trait_name Expand>] for NativeExpand<T> {
49                fn [<__expand_ $method_name _method>](self, scope: &Scope) -> Self {
50                    unary_expand(scope, self.into(), $operator).into()
51                }
52            }
53        }
54    }
55}
56
57impl Exp for f32 {
58    fn exp(self) -> Self {
59        self.exp()
60    }
61}
62
63macro_rules! impl_unary_func_scalar_out {
64    ($trait_name:ident, $method_name:ident, $operator:expr, $($type:ty),*) => {
65        paste::paste! {
66            pub trait $trait_name: CubePrimitive
67                + CubeType<ExpandType: [<$trait_name Expand>]
68                + CubePrimitiveExpand<Scalar = NativeExpand<Self::Scalar>>>
69                + Sized {
70                #[allow(unused_variables)]
71                fn $method_name(self) -> Self::Scalar {
72                    unexpanded!()
73                }
74
75                fn [<__expand_ $method_name>](scope: &Scope, x: NativeExpand<Self>) -> NativeExpand<Self::Scalar> {
76                    x.[<__expand_ $method_name _method>](scope)
77                }
78            }
79
80            pub trait [<$trait_name Expand>]: CubePrimitiveExpand {
81                fn [<__expand_ $method_name _method>](self, scope: &Scope) -> Self::Scalar;
82            }
83
84            $(impl $trait_name for $type {})*
85            impl<T: $trait_name + CubePrimitive> [<$trait_name Expand>] for NativeExpand<T> {
86                fn [<__expand_ $method_name _method>](self, scope: &Scope) -> Self::Scalar {
87                    let expand_element = self.expand;
88                    let item = expand_element.value_type().with_vector_size(0);
89                    unary_expand_fixed_output(scope, expand_element, item, $operator).into()
90                }
91            }
92        }
93    }
94}
95
96macro_rules! impl_unary_func_fixed_out_ty {
97    ($trait_name:ident, $method_name:ident, $out_ty: ty, $operator:expr, $($type:ty),*) => {
98        paste::paste! {
99            pub trait $trait_name: CubePrimitive + CubeType<ExpandType: [<$trait_name Expand>]
100            + CubePrimitiveExpand<WithScalar<$out_ty> = NativeExpand<Self::WithScalar<$out_ty>>>> + Sized {
101                #[allow(unused_variables, clippy::wrong_self_convention)]
102                fn $method_name(self) -> Self::WithScalar<$out_ty> {
103                    unexpanded!()
104                }
105
106                fn [<__expand_ $method_name>](scope: &Scope, x: NativeExpand<Self>) -> NativeExpand<Self::WithScalar<$out_ty>> {
107                    x.[<__expand_ $method_name _method>](scope)
108                }
109            }
110
111            pub trait [<$trait_name Expand>]: CubePrimitiveExpand {
112                fn [<__expand_ $method_name _method>](self, scope: &Scope) -> Self::WithScalar<$out_ty>;
113            }
114
115            $(impl $trait_name for $type {})*
116            impl<T: $trait_name + CubePrimitive> [<$trait_name Expand>] for NativeExpand<T> {
117                fn [<__expand_ $method_name _method>](self, scope: &Scope) -> Self::WithScalar<$out_ty> {
118                    let expand_element: Value = self.into();
119                    let item = <$out_ty as CubePrimitive>::__expand_as_type(scope).with_vector_size(expand_element.ty.vector_size());
120                    unary_expand_fixed_output(scope, expand_element, item, $operator).into()
121                }
122            }
123        }
124    }
125}
126
127// Needs special handling because Rust combines bitwise and logical or into one trait
128macro_rules! impl_not {
129    ($trait_name:ident, $method_name:ident, $($type:ty),*) => {
130        paste::paste! {
131            pub trait [<Cube $trait_name>]:
132                $trait_name<Output = Self>
133                + CubePrimitive
134                + CubeType<ExpandType: [<$trait_name Expand>]>
135                + IntoExpand<Expand = <Self as CubeType>::ExpandType> {
136                fn [<__expand_ $method_name _method>](self, scope: &Scope) -> NativeExpand<Self> {
137                    let this = self.into_expand(scope);
138                    this.[<__expand_ $method_name _method>](scope)
139                }
140
141                fn [<__expand_ $method_name>](scope: &Scope, x: NativeExpand<Self>) -> NativeExpand<Self> {
142                    x.[<__expand_ $method_name _method>](scope)
143                }
144            }
145
146            pub trait [<$trait_name Expand>] {
147                fn [<__expand_ $method_name _method>](self, scope: &Scope) -> Self;
148            }
149
150            $(impl [<Cube $trait_name>] for $type {})*
151            impl<T: [<Cube $trait_name>] + CubePrimitive> [<$trait_name Expand>] for NativeExpand<T> {
152                fn [<__expand_ $method_name _method>](self, scope: &Scope) -> Self {
153                    not::expand(scope, self.into())
154                }
155            }
156        }
157    }
158}
159
160macro_rules! impl_core_unop {
161    ($trait: ident, $method: ident, $op: expr) => {
162        paste::paste! {
163            pub trait [<Cube $trait>]: $trait<Output = Self> + CubePrimitive + Into<Value> + CubeType<ExpandType: [<$trait Expand>]> + Sized {
164                fn [<__expand_ $method _method>](self, scope: &Scope) -> NativeExpand<Self> {
165                    let this: Value = self.into();
166                    let this: NativeExpand<Self> = this.into();
167                    this.[<__expand_ $method _method>](scope)
168                }
169
170                fn [<__expand_ $method>](
171                    scope: &Scope,
172                    lhs: NativeExpand<Self>,
173                ) -> NativeExpand<Self> {
174                    lhs.[<__expand_ $method _method>](scope)
175                }
176            }
177
178            pub trait [<$trait Expand>] {
179                fn [<__expand_ $method _method>](self, scope: &Scope) -> Self;
180            }
181
182            impl<T: $trait<Output = T> + CubePrimitive + Into<Value>> [<Cube $trait>] for T {}
183            impl<T: $trait<Output = T> + CubePrimitive> [<$trait Expand>] for NativeExpand<T> {
184                fn [<__expand_ $method _method>](self, scope: &Scope) -> Self {
185                    unary_expand(scope, self.into(), $op).into()
186                }
187            }
188        }
189    };
190}
191
192impl_not!(
193    Not, not, bool, u8, u16, u32, u64, i8, i16, i32, i64, isize, usize
194);
195
196impl_core_unop!(Neg, neg, Arithmetic::Neg);
197
198impl_unary_func!(
199    Abs,
200    abs,
201    Arithmetic::Abs,
202    e2m1,
203    e4m3,
204    e5m2,
205    ue8m0,
206    f16,
207    bf16,
208    flex32,
209    tf32,
210    f32,
211    f64,
212    i8,
213    i16,
214    i32,
215    i64,
216    u8,
217    u16,
218    u32,
219    u64,
220    usize,
221    isize
222);
223impl_unary_func!(
224    Exp,
225    exp,
226    Arithmetic::Exp,
227    f16,
228    bf16,
229    flex32,
230    tf32,
231    // f32,
232    f64
233);
234impl_unary_func!(Log, ln, Arithmetic::Log, f16, bf16, flex32, tf32, f32, f64);
235impl_unary_func!(
236    Log1p,
237    log1p,
238    Arithmetic::Log1p,
239    f16,
240    bf16,
241    flex32,
242    tf32,
243    f32,
244    f64
245);
246impl_unary_func!(
247    Expm1,
248    exp_m1,
249    Arithmetic::Expm1,
250    f16,
251    bf16,
252    flex32,
253    tf32,
254    f32,
255    f64
256);
257impl_unary_func!(Cos, cos, Arithmetic::Cos, f16, bf16, flex32, tf32, f32, f64);
258impl_unary_func!(Sin, sin, Arithmetic::Sin, f16, bf16, flex32, tf32, f32, f64);
259impl_unary_func!(Tan, tan, Arithmetic::Tan, f16, bf16, flex32, tf32, f32, f64);
260impl_unary_func!(
261    Tanh,
262    tanh,
263    Arithmetic::Tanh,
264    f16,
265    bf16,
266    flex32,
267    tf32,
268    f32,
269    f64
270);
271impl_unary_func!(
272    Sinh,
273    sinh,
274    Arithmetic::Sinh,
275    f16,
276    bf16,
277    flex32,
278    tf32,
279    f32,
280    f64
281);
282impl_unary_func!(
283    Cosh,
284    cosh,
285    Arithmetic::Cosh,
286    f16,
287    bf16,
288    flex32,
289    tf32,
290    f32,
291    f64
292);
293impl_unary_func!(
294    ArcCos,
295    acos,
296    Arithmetic::ArcCos,
297    f16,
298    bf16,
299    flex32,
300    tf32,
301    f32,
302    f64
303);
304impl_unary_func!(
305    ArcSin,
306    asin,
307    Arithmetic::ArcSin,
308    f16,
309    bf16,
310    flex32,
311    tf32,
312    f32,
313    f64
314);
315impl_unary_func!(
316    ArcTan,
317    atan,
318    Arithmetic::ArcTan,
319    f16,
320    bf16,
321    flex32,
322    tf32,
323    f32,
324    f64
325);
326impl_unary_func!(
327    ArcSinh,
328    asinh,
329    Arithmetic::ArcSinh,
330    f16,
331    bf16,
332    flex32,
333    tf32,
334    f32,
335    f64
336);
337impl_unary_func!(
338    ArcCosh,
339    acosh,
340    Arithmetic::ArcCosh,
341    f16,
342    bf16,
343    flex32,
344    tf32,
345    f32,
346    f64
347);
348impl_unary_func!(
349    ArcTanh,
350    atanh,
351    Arithmetic::ArcTanh,
352    f16,
353    bf16,
354    flex32,
355    tf32,
356    f32,
357    f64
358);
359impl_unary_func!(
360    Degrees,
361    to_degrees,
362    Arithmetic::Degrees,
363    f16,
364    bf16,
365    flex32,
366    tf32,
367    f32,
368    f64
369);
370impl_unary_func!(
371    Radians,
372    to_radians,
373    Arithmetic::Radians,
374    f16,
375    bf16,
376    flex32,
377    tf32,
378    f32,
379    f64
380);
381impl_unary_func!(
382    Sqrt,
383    sqrt,
384    Arithmetic::Sqrt,
385    f16,
386    bf16,
387    flex32,
388    tf32,
389    f32,
390    f64
391);
392impl_unary_func!(
393    InverseSqrt,
394    inverse_sqrt,
395    Arithmetic::InverseSqrt,
396    f16,
397    bf16,
398    flex32,
399    tf32,
400    f32,
401    f64
402);
403impl_unary_func!(
404    Round,
405    round,
406    Arithmetic::Round,
407    f16,
408    bf16,
409    flex32,
410    tf32,
411    f32,
412    f64
413);
414impl_unary_func!(
415    Floor,
416    floor,
417    Arithmetic::Floor,
418    f16,
419    bf16,
420    flex32,
421    tf32,
422    f32,
423    f64
424);
425impl_unary_func!(
426    Ceil,
427    ceil,
428    Arithmetic::Ceil,
429    f16,
430    bf16,
431    flex32,
432    tf32,
433    f32,
434    f64
435);
436impl_unary_func!(
437    Trunc,
438    trunc,
439    Arithmetic::Trunc,
440    f16,
441    bf16,
442    flex32,
443    tf32,
444    f32,
445    f64
446);
447impl_unary_func!(Erf, erf, Arithmetic::Erf, f16, bf16, flex32, tf32, f32, f64);
448impl_unary_func!(
449    Recip,
450    recip,
451    Arithmetic::Recip,
452    f16,
453    bf16,
454    flex32,
455    tf32,
456    f32,
457    f64
458);
459impl_unary_func_scalar_out!(
460    Magnitude,
461    magnitude,
462    Arithmetic::Magnitude,
463    f16,
464    bf16,
465    flex32,
466    tf32,
467    f32,
468    f64
469);
470impl_unary_func_scalar_out!(
471    VectorSum,
472    vector_sum,
473    Arithmetic::VectorSum,
474    e2m1,
475    e4m3,
476    e5m2,
477    ue8m0,
478    f16,
479    bf16,
480    flex32,
481    tf32,
482    f32,
483    f64,
484    i8,
485    i16,
486    i32,
487    i64,
488    u8,
489    u16,
490    u32,
491    u64,
492    usize,
493    isize
494);
495impl_unary_func!(
496    Normalize,
497    normalize,
498    Arithmetic::Normalize,
499    f16,
500    bf16,
501    flex32,
502    tf32,
503    f32,
504    f64
505);
506impl_unary_func_fixed_out_ty!(
507    CountOnes,
508    count_ones,
509    u32,
510    Bitwise::CountOnes,
511    u8,
512    i8,
513    u16,
514    i16,
515    u32,
516    i32,
517    u64,
518    i64,
519    usize,
520    isize
521);
522impl_unary_func!(
523    ReverseBits,
524    reverse_bits,
525    Bitwise::ReverseBits,
526    u8,
527    i8,
528    u16,
529    i16,
530    u32,
531    i32,
532    u64,
533    i64,
534    usize,
535    isize
536);
537
538impl_unary_func_fixed_out_ty!(
539    LeadingZeros,
540    leading_zeros,
541    u32,
542    Bitwise::LeadingZeros,
543    u8,
544    i8,
545    u16,
546    i16,
547    u32,
548    i32,
549    u64,
550    i64,
551    usize,
552    isize
553);
554impl_unary_func_fixed_out_ty!(
555    TrailingZeros,
556    trailing_zeros,
557    u32,
558    Bitwise::TrailingZeros,
559    u8,
560    i8,
561    u16,
562    i16,
563    u32,
564    i32,
565    u64,
566    i64,
567    usize,
568    isize
569);
570impl_unary_func_fixed_out_ty!(
571    FindFirstSet,
572    find_first_set,
573    u32,
574    Bitwise::FindFirstSet,
575    u8,
576    i8,
577    u16,
578    i16,
579    u32,
580    i32,
581    u64,
582    i64,
583    usize,
584    isize
585);
586impl_unary_func_fixed_out_ty!(
587    IsNan,
588    is_nan,
589    bool,
590    Comparison::IsNan,
591    f16,
592    bf16,
593    flex32,
594    tf32,
595    f32,
596    f64
597);
598impl_unary_func_fixed_out_ty!(
599    IsInf,
600    is_inf,
601    bool,
602    Comparison::IsInf,
603    f16,
604    bf16,
605    flex32,
606    tf32,
607    f32,
608    f64
609);
610
611pub trait FloatBits:
612    CubePrimitive + CubeType<ExpandType: FloatBitsExpand<Bits = Self::Bits>>
613{
614    type Bits: CubePrimitive;
615
616    fn __expand_from_bits(scope: &Scope, bits: NativeExpand<Self::Bits>) -> NativeExpand<Self> {
617        Self::__expand_reinterpret(scope, bits)
618    }
619
620    fn __expand_to_bits(scope: &Scope, this: NativeExpand<Self>) -> NativeExpand<Self::Bits> {
621        <Self::Bits as Reinterpret>::__expand_reinterpret(scope, this)
622    }
623}
624
625pub trait FloatBitsExpand: Sized {
626    type Bits: CubePrimitive;
627
628    fn __expand_to_bits_method(self, scope: &Scope) -> NativeExpand<Self::Bits>;
629}
630
631impl<F: FloatBits> FloatBitsExpand for NativeExpand<F> {
632    type Bits = F::Bits;
633
634    fn __expand_to_bits_method(self, scope: &Scope) -> NativeExpand<Self::Bits> {
635        <Self::Bits as Reinterpret>::__expand_reinterpret(scope, self)
636    }
637}
638
639impl FloatBits for e2m1x2 {
640    type Bits = u8;
641}
642
643impl FloatBits for e5m2 {
644    type Bits = u8;
645}
646
647impl FloatBits for e4m3 {
648    type Bits = u8;
649}
650
651impl FloatBits for f16 {
652    type Bits = u16;
653}
654
655impl FloatBits for bf16 {
656    type Bits = u16;
657}
658
659impl FloatBits for f32 {
660    type Bits = u32;
661}
662
663impl FloatBits for f64 {
664    type Bits = u64;
665}
cubecl_core/frontend/operation/unary.rs

cubecl_core/frontend/operation/
unary.rs
