library/portable-simd/crates/core_simd/src/ops.rs

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