]> git.lizzy.rs Git - rust.git/blob - crates/core_simd/src/ops.rs
Begin reducing mask API
[rust.git] / crates / core_simd / src / ops.rs
1 use crate::LanesAtMost32;
2
3 /// Checks if the right-hand side argument of a left- or right-shift would cause overflow.
4 fn invalid_shift_rhs<T>(rhs: T) -> bool
5 where
6     T: Default + PartialOrd + core::convert::TryFrom<usize>,
7     <T as core::convert::TryFrom<usize>>::Error: core::fmt::Debug,
8 {
9     let bits_in_type = T::try_from(8 * core::mem::size_of::<T>()).unwrap();
10     rhs < T::default() || rhs >= bits_in_type
11 }
12
13 /// Automatically implements operators over references in addition to the provided operator.
14 macro_rules! impl_ref_ops {
15     // binary op
16     {
17         impl<const $lanes:ident: usize> core::ops::$trait:ident<$rhs:ty> for $type:ty
18         where
19             $($bound:path: LanesAtMost32,)*
20         {
21             type Output = $output:ty;
22
23             $(#[$attrs:meta])*
24             fn $fn:ident($self_tok:ident, $rhs_arg:ident: $rhs_arg_ty:ty) -> Self::Output $body:tt
25         }
26     } => {
27         impl<const $lanes: usize> core::ops::$trait<$rhs> for $type
28         where
29             $($bound: LanesAtMost32,)*
30         {
31             type Output = $output;
32
33             $(#[$attrs])*
34             fn $fn($self_tok, $rhs_arg: $rhs_arg_ty) -> Self::Output $body
35         }
36
37         impl<const $lanes: usize> core::ops::$trait<&'_ $rhs> for $type
38         where
39             $($bound: LanesAtMost32,)*
40         {
41             type Output = <$type as core::ops::$trait<$rhs>>::Output;
42
43             $(#[$attrs])*
44             fn $fn($self_tok, $rhs_arg: &$rhs) -> Self::Output {
45                 core::ops::$trait::$fn($self_tok, *$rhs_arg)
46             }
47         }
48
49         impl<const $lanes: usize> core::ops::$trait<$rhs> for &'_ $type
50         where
51             $($bound: LanesAtMost32,)*
52         {
53             type Output = <$type as core::ops::$trait<$rhs>>::Output;
54
55             $(#[$attrs])*
56             fn $fn($self_tok, $rhs_arg: $rhs) -> Self::Output {
57                 core::ops::$trait::$fn(*$self_tok, $rhs_arg)
58             }
59         }
60
61         impl<const $lanes: usize> core::ops::$trait<&'_ $rhs> for &'_ $type
62         where
63             $($bound: LanesAtMost32,)*
64         {
65             type Output = <$type as core::ops::$trait<$rhs>>::Output;
66
67             $(#[$attrs])*
68             fn $fn($self_tok, $rhs_arg: &$rhs) -> Self::Output {
69                 core::ops::$trait::$fn(*$self_tok, *$rhs_arg)
70             }
71         }
72     };
73
74     // binary assignment op
75     {
76         impl<const $lanes:ident: usize> core::ops::$trait:ident<$rhs:ty> for $type:ty
77         where
78             $($bound:path: LanesAtMost32,)*
79         {
80             $(#[$attrs:meta])*
81             fn $fn:ident(&mut $self_tok:ident, $rhs_arg:ident: $rhs_arg_ty:ty) $body:tt
82         }
83     } => {
84         impl<const $lanes: usize> core::ops::$trait<$rhs> for $type
85         where
86             $($bound: LanesAtMost32,)*
87         {
88             $(#[$attrs])*
89             fn $fn(&mut $self_tok, $rhs_arg: $rhs_arg_ty) $body
90         }
91
92         impl<const $lanes: usize> core::ops::$trait<&'_ $rhs> for $type
93         where
94             $($bound: LanesAtMost32,)*
95         {
96             $(#[$attrs])*
97             fn $fn(&mut $self_tok, $rhs_arg: &$rhs_arg_ty) {
98                 core::ops::$trait::$fn($self_tok, *$rhs_arg)
99             }
100         }
101     };
102
103     // unary op
104     {
105         impl<const $lanes:ident: usize> core::ops::$trait:ident for $type:ty
106         where
107             $($bound:path: LanesAtMost32,)*
108         {
109             type Output = $output:ty;
110             fn $fn:ident($self_tok:ident) -> Self::Output $body:tt
111         }
112     } => {
113         impl<const $lanes: usize> core::ops::$trait for $type
114         where
115             $($bound: LanesAtMost32,)*
116         {
117             type Output = $output;
118             fn $fn($self_tok) -> Self::Output $body
119         }
120
121         impl<const $lanes: usize> core::ops::$trait for &'_ $type
122         where
123             $($bound: LanesAtMost32,)*
124         {
125             type Output = <$type as core::ops::$trait>::Output;
126             fn $fn($self_tok) -> Self::Output {
127                 core::ops::$trait::$fn(*$self_tok)
128             }
129         }
130     }
131 }
132
133 /// Automatically implements operators over vectors and scalars for a particular vector.
134 macro_rules! impl_op {
135     { impl Add for $type:ident, $scalar:ty } => {
136         impl_op! { @binary $type, $scalar, Add::add, AddAssign::add_assign, simd_add }
137     };
138     { impl Sub for $type:ident, $scalar:ty } => {
139         impl_op! { @binary $type, $scalar, Sub::sub, SubAssign::sub_assign, simd_sub }
140     };
141     { impl Mul for $type:ident, $scalar:ty } => {
142         impl_op! { @binary $type, $scalar, Mul::mul, MulAssign::mul_assign, simd_mul }
143     };
144     { impl Div for $type:ident, $scalar:ty } => {
145         impl_op! { @binary $type, $scalar, Div::div, DivAssign::div_assign, simd_div }
146     };
147     { impl Rem for $type:ident, $scalar:ty } => {
148         impl_op! { @binary $type, $scalar, Rem::rem, RemAssign::rem_assign, simd_rem }
149     };
150     { impl Shl for $type:ident, $scalar:ty } => {
151         impl_op! { @binary $type, $scalar, Shl::shl, ShlAssign::shl_assign, simd_shl }
152     };
153     { impl Shr for $type:ident, $scalar:ty } => {
154         impl_op! { @binary $type, $scalar, Shr::shr, ShrAssign::shr_assign, simd_shr }
155     };
156     { impl BitAnd for $type:ident, $scalar:ty } => {
157         impl_op! { @binary $type, $scalar, BitAnd::bitand, BitAndAssign::bitand_assign, simd_and }
158     };
159     { impl BitOr for $type:ident, $scalar:ty } => {
160         impl_op! { @binary $type, $scalar, BitOr::bitor, BitOrAssign::bitor_assign, simd_or }
161     };
162     { impl BitXor for $type:ident, $scalar:ty } => {
163         impl_op! { @binary $type, $scalar, BitXor::bitxor, BitXorAssign::bitxor_assign, simd_xor }
164     };
165
166     { impl Not for $type:ident, $scalar:ty } => {
167         impl_ref_ops! {
168             impl<const LANES: usize> core::ops::Not for crate::$type<LANES>
169             where
170                 crate::$type<LANES>: LanesAtMost32,
171             {
172                 type Output = Self;
173                 fn not(self) -> Self::Output {
174                     self ^ Self::splat(!<$scalar>::default())
175                 }
176             }
177         }
178     };
179
180     { impl Neg for $type:ident, $scalar:ty } => {
181         impl_ref_ops! {
182             impl<const LANES: usize> core::ops::Neg for crate::$type<LANES>
183             where
184                 crate::$type<LANES>: LanesAtMost32,
185             {
186                 type Output = Self;
187                 fn neg(self) -> Self::Output {
188                     unsafe { crate::intrinsics::simd_neg(self) }
189                 }
190             }
191         }
192     };
193
194     { impl Index for $type:ident, $scalar:ty } => {
195         impl<I, const LANES: usize> core::ops::Index<I> for crate::$type<LANES>
196         where
197             Self: LanesAtMost32,
198             I: core::slice::SliceIndex<[$scalar]>,
199         {
200             type Output = I::Output;
201             fn index(&self, index: I) -> &Self::Output {
202                 let slice: &[_] = self.as_ref();
203                 &slice[index]
204             }
205         }
206
207         impl<I, const LANES: usize> core::ops::IndexMut<I> for crate::$type<LANES>
208         where
209             Self: LanesAtMost32,
210             I: core::slice::SliceIndex<[$scalar]>,
211         {
212             fn index_mut(&mut self, index: I) -> &mut Self::Output {
213                 let slice: &mut [_] = self.as_mut();
214                 &mut slice[index]
215             }
216         }
217     };
218
219     // generic binary op with assignment when output is `Self`
220     { @binary $type:ident, $scalar:ty, $trait:ident :: $trait_fn:ident, $assign_trait:ident :: $assign_trait_fn:ident, $intrinsic:ident } => {
221         impl_ref_ops! {
222             impl<const LANES: usize> core::ops::$trait<Self> for crate::$type<LANES>
223             where
224                 crate::$type<LANES>: LanesAtMost32,
225             {
226                 type Output = Self;
227
228                 #[inline]
229                 fn $trait_fn(self, rhs: Self) -> Self::Output {
230                     unsafe {
231                         crate::intrinsics::$intrinsic(self, rhs)
232                     }
233                 }
234             }
235         }
236
237         impl_ref_ops! {
238             impl<const LANES: usize> core::ops::$trait<$scalar> for crate::$type<LANES>
239             where
240                 crate::$type<LANES>: LanesAtMost32,
241             {
242                 type Output = Self;
243
244                 #[inline]
245                 fn $trait_fn(self, rhs: $scalar) -> Self::Output {
246                     core::ops::$trait::$trait_fn(self, Self::splat(rhs))
247                 }
248             }
249         }
250
251         impl_ref_ops! {
252             impl<const LANES: usize> core::ops::$trait<crate::$type<LANES>> for $scalar
253             where
254                 crate::$type<LANES>: LanesAtMost32,
255             {
256                 type Output = crate::$type<LANES>;
257
258                 #[inline]
259                 fn $trait_fn(self, rhs: crate::$type<LANES>) -> Self::Output {
260                     core::ops::$trait::$trait_fn(crate::$type::splat(self), rhs)
261                 }
262             }
263         }
264
265         impl_ref_ops! {
266             impl<const LANES: usize> core::ops::$assign_trait<Self> for crate::$type<LANES>
267             where
268                 crate::$type<LANES>: LanesAtMost32,
269             {
270                 #[inline]
271                 fn $assign_trait_fn(&mut self, rhs: Self) {
272                     unsafe {
273                         *self = crate::intrinsics::$intrinsic(*self, rhs);
274                     }
275                 }
276             }
277         }
278
279         impl_ref_ops! {
280             impl<const LANES: usize> core::ops::$assign_trait<$scalar> for crate::$type<LANES>
281             where
282                 crate::$type<LANES>: LanesAtMost32,
283             {
284                 #[inline]
285                 fn $assign_trait_fn(&mut self, rhs: $scalar) {
286                     core::ops::$assign_trait::$assign_trait_fn(self, Self::splat(rhs));
287                 }
288             }
289         }
290     };
291 }
292
293 /// Implements floating-point operators for the provided types.
294 macro_rules! impl_float_ops {
295     { $($scalar:ty => $($vector:ident),*;)* } => {
296         $( // scalar
297             $( // vector
298                 impl_op! { impl Add for $vector, $scalar }
299                 impl_op! { impl Sub for $vector, $scalar }
300                 impl_op! { impl Mul for $vector, $scalar }
301                 impl_op! { impl Div for $vector, $scalar }
302                 impl_op! { impl Rem for $vector, $scalar }
303                 impl_op! { impl Neg for $vector, $scalar }
304                 impl_op! { impl Index for $vector, $scalar }
305             )*
306         )*
307     };
308 }
309
310 /// Implements unsigned integer operators for the provided types.
311 macro_rules! impl_unsigned_int_ops {
312     { $($scalar:ty => $($vector:ident),*;)* } => {
313         $( // scalar
314             $( // vector
315                 impl_op! { impl Add for $vector, $scalar }
316                 impl_op! { impl Sub for $vector, $scalar }
317                 impl_op! { impl Mul for $vector, $scalar }
318                 impl_op! { impl BitAnd for $vector, $scalar }
319                 impl_op! { impl BitOr  for $vector, $scalar }
320                 impl_op! { impl BitXor for $vector, $scalar }
321                 impl_op! { impl Not for $vector, $scalar }
322                 impl_op! { impl Index for $vector, $scalar }
323
324                 // Integers panic on divide by 0
325                 impl_ref_ops! {
326                     impl<const LANES: usize> core::ops::Div<Self> for crate::$vector<LANES>
327                     where
328                         crate::$vector<LANES>: LanesAtMost32,
329                     {
330                         type Output = Self;
331
332                         #[inline]
333                         fn div(self, rhs: Self) -> Self::Output {
334                             if rhs.as_slice()
335                                 .iter()
336                                 .any(|x| *x == 0)
337                             {
338                                 panic!("attempt to divide by zero");
339                             }
340
341                             // Guards for div(MIN, -1),
342                             // this check only applies to signed ints
343                             if <$scalar>::MIN != 0 && self.as_slice().iter()
344                                     .zip(rhs.as_slice().iter())
345                                     .any(|(x,y)| *x == <$scalar>::MIN && *y == -1 as _) {
346                                 panic!("attempt to divide with overflow");
347                             }
348                             unsafe { crate::intrinsics::simd_div(self, rhs) }
349                         }
350                     }
351                 }
352
353                 impl_ref_ops! {
354                     impl<const LANES: usize> core::ops::Div<$scalar> for crate::$vector<LANES>
355                     where
356                         crate::$vector<LANES>: LanesAtMost32,
357                     {
358                         type Output = Self;
359
360                         #[inline]
361                         fn div(self, rhs: $scalar) -> Self::Output {
362                             if rhs == 0 {
363                                 panic!("attempt to divide by zero");
364                             }
365                             if <$scalar>::MIN != 0 &&
366                                 self.as_slice().iter().any(|x| *x == <$scalar>::MIN) &&
367                                 rhs == -1 as _ {
368                                     panic!("attempt to divide with overflow");
369                             }
370                             let rhs = Self::splat(rhs);
371                             unsafe { crate::intrinsics::simd_div(self, rhs) }
372                         }
373                     }
374                 }
375
376                 impl_ref_ops! {
377                     impl<const LANES: usize> core::ops::Div<crate::$vector<LANES>> for $scalar
378                     where
379                         crate::$vector<LANES>: LanesAtMost32,
380                     {
381                         type Output = crate::$vector<LANES>;
382
383                         #[inline]
384                         fn div(self, rhs: crate::$vector<LANES>) -> Self::Output {
385                             crate::$vector::splat(self) / rhs
386                         }
387                     }
388                 }
389
390                 impl_ref_ops! {
391                     impl<const LANES: usize> core::ops::DivAssign<Self> for crate::$vector<LANES>
392                     where
393                         crate::$vector<LANES>: LanesAtMost32,
394                     {
395                         #[inline]
396                         fn div_assign(&mut self, rhs: Self) {
397                             *self = *self / rhs;
398                         }
399                     }
400                 }
401
402                 impl_ref_ops! {
403                     impl<const LANES: usize> core::ops::DivAssign<$scalar> for crate::$vector<LANES>
404                     where
405                         crate::$vector<LANES>: LanesAtMost32,
406                     {
407                         #[inline]
408                         fn div_assign(&mut self, rhs: $scalar) {
409                             *self = *self / rhs;
410                         }
411                     }
412                 }
413
414                 // remainder panics on zero divisor
415                 impl_ref_ops! {
416                     impl<const LANES: usize> core::ops::Rem<Self> for crate::$vector<LANES>
417                     where
418                         crate::$vector<LANES>: LanesAtMost32,
419                     {
420                         type Output = Self;
421
422                         #[inline]
423                         fn rem(self, rhs: Self) -> Self::Output {
424                             if rhs.as_slice()
425                                 .iter()
426                                 .any(|x| *x == 0)
427                             {
428                                 panic!("attempt to calculate the remainder with a divisor of zero");
429                             }
430
431                             // Guards for rem(MIN, -1)
432                             // this branch applies the check only to signed ints
433                             if <$scalar>::MIN != 0 && self.as_slice().iter()
434                                     .zip(rhs.as_slice().iter())
435                                     .any(|(x,y)| *x == <$scalar>::MIN && *y == -1 as _) {
436                                 panic!("attempt to calculate the remainder with overflow");
437                             }
438                             unsafe { crate::intrinsics::simd_rem(self, rhs) }
439                         }
440                     }
441                 }
442
443                 impl_ref_ops! {
444                     impl<const LANES: usize> core::ops::Rem<$scalar> for crate::$vector<LANES>
445                     where
446                         crate::$vector<LANES>: LanesAtMost32,
447                     {
448                         type Output = Self;
449
450                         #[inline]
451                         fn rem(self, rhs: $scalar) -> Self::Output {
452                             if rhs == 0 {
453                                 panic!("attempt to calculate the remainder with a divisor of zero");
454                             }
455                             if <$scalar>::MIN != 0 &&
456                                 self.as_slice().iter().any(|x| *x == <$scalar>::MIN) &&
457                                 rhs == -1 as _ {
458                                     panic!("attempt to calculate the remainder with overflow");
459                             }
460                             let rhs = Self::splat(rhs);
461                             unsafe { crate::intrinsics::simd_rem(self, rhs) }
462                         }
463                     }
464                 }
465
466                 impl_ref_ops! {
467                     impl<const LANES: usize> core::ops::Rem<crate::$vector<LANES>> for $scalar
468                     where
469                         crate::$vector<LANES>: LanesAtMost32,
470                     {
471                         type Output = crate::$vector<LANES>;
472
473                         #[inline]
474                         fn rem(self, rhs: crate::$vector<LANES>) -> Self::Output {
475                             crate::$vector::splat(self) % rhs
476                         }
477                     }
478                 }
479
480                 impl_ref_ops! {
481                     impl<const LANES: usize> core::ops::RemAssign<Self> for crate::$vector<LANES>
482                     where
483                         crate::$vector<LANES>: LanesAtMost32,
484                     {
485                         #[inline]
486                         fn rem_assign(&mut self, rhs: Self) {
487                             *self = *self % rhs;
488                         }
489                     }
490                 }
491
492                 impl_ref_ops! {
493                     impl<const LANES: usize> core::ops::RemAssign<$scalar> for crate::$vector<LANES>
494                     where
495                         crate::$vector<LANES>: LanesAtMost32,
496                     {
497                         #[inline]
498                         fn rem_assign(&mut self, rhs: $scalar) {
499                             *self = *self % rhs;
500                         }
501                     }
502                 }
503
504                 // shifts panic on overflow
505                 impl_ref_ops! {
506                     impl<const LANES: usize> core::ops::Shl<Self> for crate::$vector<LANES>
507                     where
508                         crate::$vector<LANES>: LanesAtMost32,
509                     {
510                         type Output = Self;
511
512                         #[inline]
513                         fn shl(self, rhs: Self) -> Self::Output {
514                             // TODO there is probably a better way of doing this
515                             if rhs.as_slice()
516                                 .iter()
517                                 .copied()
518                                 .any(invalid_shift_rhs)
519                             {
520                                 panic!("attempt to shift left with overflow");
521                             }
522                             unsafe { crate::intrinsics::simd_shl(self, rhs) }
523                         }
524                     }
525                 }
526
527                 impl_ref_ops! {
528                     impl<const LANES: usize> core::ops::Shl<$scalar> for crate::$vector<LANES>
529                     where
530                         crate::$vector<LANES>: LanesAtMost32,
531                     {
532                         type Output = Self;
533
534                         #[inline]
535                         fn shl(self, rhs: $scalar) -> Self::Output {
536                             if invalid_shift_rhs(rhs) {
537                                 panic!("attempt to shift left with overflow");
538                             }
539                             let rhs = Self::splat(rhs);
540                             unsafe { crate::intrinsics::simd_shl(self, rhs) }
541                         }
542                     }
543                 }
544
545
546                 impl_ref_ops! {
547                     impl<const LANES: usize> core::ops::ShlAssign<Self> for crate::$vector<LANES>
548                     where
549                         crate::$vector<LANES>: LanesAtMost32,
550                     {
551                         #[inline]
552                         fn shl_assign(&mut self, rhs: Self) {
553                             *self = *self << rhs;
554                         }
555                     }
556                 }
557
558                 impl_ref_ops! {
559                     impl<const LANES: usize> core::ops::ShlAssign<$scalar> for crate::$vector<LANES>
560                     where
561                         crate::$vector<LANES>: LanesAtMost32,
562                     {
563                         #[inline]
564                         fn shl_assign(&mut self, rhs: $scalar) {
565                             *self = *self << rhs;
566                         }
567                     }
568                 }
569
570                 impl_ref_ops! {
571                     impl<const LANES: usize> core::ops::Shr<Self> for crate::$vector<LANES>
572                     where
573                         crate::$vector<LANES>: LanesAtMost32,
574                     {
575                         type Output = Self;
576
577                         #[inline]
578                         fn shr(self, rhs: Self) -> Self::Output {
579                             // TODO there is probably a better way of doing this
580                             if rhs.as_slice()
581                                 .iter()
582                                 .copied()
583                                 .any(invalid_shift_rhs)
584                             {
585                                 panic!("attempt to shift with overflow");
586                             }
587                             unsafe { crate::intrinsics::simd_shr(self, rhs) }
588                         }
589                     }
590                 }
591
592                 impl_ref_ops! {
593                     impl<const LANES: usize> core::ops::Shr<$scalar> for crate::$vector<LANES>
594                     where
595                         crate::$vector<LANES>: LanesAtMost32,
596                     {
597                         type Output = Self;
598
599                         #[inline]
600                         fn shr(self, rhs: $scalar) -> Self::Output {
601                             if invalid_shift_rhs(rhs) {
602                                 panic!("attempt to shift with overflow");
603                             }
604                             let rhs = Self::splat(rhs);
605                             unsafe { crate::intrinsics::simd_shr(self, rhs) }
606                         }
607                     }
608                 }
609
610
611                 impl_ref_ops! {
612                     impl<const LANES: usize> core::ops::ShrAssign<Self> for crate::$vector<LANES>
613                     where
614                         crate::$vector<LANES>: LanesAtMost32,
615                     {
616                         #[inline]
617                         fn shr_assign(&mut self, rhs: Self) {
618                             *self = *self >> rhs;
619                         }
620                     }
621                 }
622
623                 impl_ref_ops! {
624                     impl<const LANES: usize> core::ops::ShrAssign<$scalar> for crate::$vector<LANES>
625                     where
626                         crate::$vector<LANES>: LanesAtMost32,
627                     {
628                         #[inline]
629                         fn shr_assign(&mut self, rhs: $scalar) {
630                             *self = *self >> rhs;
631                         }
632                     }
633                 }
634             )*
635         )*
636     };
637 }
638
639 /// Implements unsigned integer operators for the provided types.
640 macro_rules! impl_signed_int_ops {
641     { $($scalar:ty => $($vector:ident),*;)* } => {
642         impl_unsigned_int_ops! { $($scalar => $($vector),*;)* }
643         $( // scalar
644             $( // vector
645                 impl_op! { impl Neg for $vector, $scalar }
646             )*
647         )*
648     };
649 }
650
651 impl_unsigned_int_ops! {
652     u8 => SimdU8;
653     u16 => SimdU16;
654     u32 => SimdU32;
655     u64 => SimdU64;
656     usize => SimdUsize;
657 }
658
659 impl_signed_int_ops! {
660     i8 => SimdI8;
661     i16 => SimdI16;
662     i32 => SimdI32;
663     i64 => SimdI64;
664     isize => SimdIsize;
665 }
666
667 impl_float_ops! {
668     f32 => SimdF32;
669     f64 => SimdF64;
670 }