crates/core_simd/src/masks.rs

   1 //! Types and traits associated with masking lanes of vectors.
   2 //! Types representing
   3 #![allow(non_camel_case_types)]
   4
   5 #[cfg_attr(
   6     not(all(target_arch = "x86_64", target_feature = "avx512f")),
   7     path = "masks/full_masks.rs"
   8 )]
   9 #[cfg_attr(
  10     all(target_arch = "x86_64", target_feature = "avx512f"),
  11     path = "masks/bitmask.rs"
  12 )]
  13 mod mask_impl;
  14
  15 use crate::{SimdI16, SimdI32, SimdI64, SimdI8, SimdIsize};
  16
  17 mod sealed {
  18     pub trait Sealed {}
  19 }
  20
  21 /// Helper trait for mask types.
  22 pub trait Mask: sealed::Sealed {
  23     /// The number of lanes for this mask.
  24     const LANES: usize;
  25
  26     /// Generates a mask with the same value in every lane.
  27     #[must_use]
  28     fn splat(val: bool) -> Self;
  29 }
  30
  31 macro_rules! define_opaque_mask {
  32     {
  33         $(#[$attr:meta])*
  34         struct $name:ident<const $lanes:ident: usize>($inner_ty:ty);
  35         @bits $bits_ty:ident
  36     } => {
  37         $(#[$attr])*
  38         #[allow(non_camel_case_types)]
  39         pub struct $name<const LANES: usize>($inner_ty)
  40         where
  41             crate::LaneCount<LANES>: crate::SupportedLaneCount;
  42
  43         impl<const LANES: usize> sealed::Sealed for $name<LANES>
  44         where
  45             crate::LaneCount<LANES>: crate::SupportedLaneCount,
  46         {}
  47
  48         impl<const LANES: usize> Mask for $name<LANES>
  49         where
  50             crate::LaneCount<LANES>: crate::SupportedLaneCount,
  51         {
  52             const LANES: usize = LANES;
  53
  54             #[inline]
  55             fn splat(value: bool) -> Self {
  56                 Self::splat(value)
  57             }
  58         }
  59
  60         impl_opaque_mask_reductions! { $name, $bits_ty }
  61
  62         impl<const LANES: usize> $name<LANES>
  63         where
  64             crate::LaneCount<LANES>: crate::SupportedLaneCount,
  65         {
  66             /// Construct a mask by setting all lanes to the given value.
  67             pub fn splat(value: bool) -> Self {
  68                 Self(<$inner_ty>::splat(value))
  69             }
  70
  71             /// Converts an array to a SIMD vector.
  72             pub fn from_array(array: [bool; LANES]) -> Self {
  73                 let mut vector = Self::splat(false);
  74                 let mut i = 0;
  75                 while i < $lanes {
  76                     vector.set(i, array[i]);
  77                     i += 1;
  78                 }
  79                 vector
  80             }
  81
  82             /// Converts a SIMD vector to an array.
  83             pub fn to_array(self) -> [bool; LANES] {
  84                 let mut array = [false; LANES];
  85                 let mut i = 0;
  86                 while i < $lanes {
  87                     array[i] = self.test(i);
  88                     i += 1;
  89                 }
  90                 array
  91             }
  92
  93             /// Converts a vector of integers to a mask, where 0 represents `false` and -1
  94             /// represents `true`.
  95             ///
  96             /// # Safety
  97             /// All lanes must be either 0 or -1.
  98             #[inline]
  99             pub unsafe fn from_int_unchecked(value: $bits_ty<LANES>) -> Self {
 100                 Self(<$inner_ty>::from_int_unchecked(value))
 101             }
 102
 103             /// Converts a vector of integers to a mask, where 0 represents `false` and -1
 104             /// represents `true`.
 105             ///
 106             /// # Panics
 107             /// Panics if any lane is not 0 or -1.
 108             #[inline]
 109             pub fn from_int(value: $bits_ty<LANES>) -> Self {
 110                 assert!(
 111                     (value.lanes_eq($bits_ty::splat(0)) | value.lanes_eq($bits_ty::splat(-1))).all(),
 112                     "all values must be either 0 or -1",
 113                 );
 114                 unsafe { Self::from_int_unchecked(value) }
 115             }
 116
 117             /// Converts the mask to a vector of integers, where 0 represents `false` and -1
 118             /// represents `true`.
 119             #[inline]
 120             pub fn to_int(self) -> $bits_ty<LANES> {
 121                 self.0.to_int()
 122             }
 123
 124             /// Tests the value of the specified lane.
 125             ///
 126             /// # Safety
 127             /// `lane` must be less than `LANES`.
 128             #[inline]
 129             pub unsafe fn test_unchecked(&self, lane: usize) -> bool {
 130                 self.0.test_unchecked(lane)
 131             }
 132
 133             /// Tests the value of the specified lane.
 134             ///
 135             /// # Panics
 136             /// Panics if `lane` is greater than or equal to the number of lanes in the vector.
 137             #[inline]
 138             pub fn test(&self, lane: usize) -> bool {
 139                 assert!(lane < LANES, "lane index out of range");
 140                 unsafe { self.test_unchecked(lane) }
 141             }
 142
 143             /// Sets the value of the specified lane.
 144             ///
 145             /// # Safety
 146             /// `lane` must be less than `LANES`.
 147             #[inline]
 148             pub unsafe fn set_unchecked(&mut self, lane: usize, value: bool) {
 149                 self.0.set_unchecked(lane, value);
 150             }
 151
 152             /// Sets the value of the specified lane.
 153             ///
 154             /// # Panics
 155             /// Panics if `lane` is greater than or equal to the number of lanes in the vector.
 156             #[inline]
 157             pub fn set(&mut self, lane: usize, value: bool) {
 158                 assert!(lane < LANES, "lane index out of range");
 159                 unsafe { self.set_unchecked(lane, value); }
 160             }
 161
 162             /// Convert this mask to a bitmask, with one bit set per lane.
 163             pub fn to_bitmask(self) -> [u8; crate::LaneCount::<LANES>::BITMASK_LEN] {
 164                 self.0.to_bitmask()
 165             }
 166
 167             /// Convert a bitmask to a mask.
 168             pub fn from_bitmask(bitmask: [u8; crate::LaneCount::<LANES>::BITMASK_LEN]) -> Self {
 169                 Self(<$inner_ty>::from_bitmask(bitmask))
 170             }
 171         }
 172
 173         // vector/array conversion
 174         impl<const LANES: usize> From<[bool; LANES]> for $name<LANES>
 175         where
 176             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 177         {
 178             fn from(array: [bool; LANES]) -> Self {
 179                 Self::from_array(array)
 180             }
 181         }
 182
 183         impl <const LANES: usize> From<$name<LANES>> for [bool; LANES]
 184         where
 185             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 186         {
 187             fn from(vector: $name<LANES>) -> Self {
 188                 vector.to_array()
 189             }
 190         }
 191
 192         impl<const LANES: usize> Copy for $name<LANES>
 193         where
 194             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 195         {}
 196
 197         impl<const LANES: usize> Clone for $name<LANES>
 198         where
 199             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 200         {
 201             #[inline]
 202             fn clone(&self) -> Self {
 203                 *self
 204             }
 205         }
 206
 207         impl<const LANES: usize> Default for $name<LANES>
 208         where
 209             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 210         {
 211             #[inline]
 212             fn default() -> Self {
 213                 Self::splat(false)
 214             }
 215         }
 216
 217         impl<const LANES: usize> PartialEq for $name<LANES>
 218         where
 219             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 220         {
 221             #[inline]
 222             fn eq(&self, other: &Self) -> bool {
 223                 self.0 == other.0
 224             }
 225         }
 226
 227         impl<const LANES: usize> PartialOrd for $name<LANES>
 228         where
 229             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 230         {
 231             #[inline]
 232             fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
 233                 self.0.partial_cmp(&other.0)
 234             }
 235         }
 236
 237         impl<const LANES: usize> core::fmt::Debug for $name<LANES>
 238         where
 239             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 240         {
 241             fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
 242                 f.debug_list()
 243                     .entries((0..LANES).map(|lane| self.test(lane)))
 244                     .finish()
 245             }
 246         }
 247
 248         impl<const LANES: usize> core::ops::BitAnd for $name<LANES>
 249         where
 250             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 251         {
 252             type Output = Self;
 253             #[inline]
 254             fn bitand(self, rhs: Self) -> Self {
 255                 Self(self.0 & rhs.0)
 256             }
 257         }
 258
 259         impl<const LANES: usize> core::ops::BitAnd<bool> for $name<LANES>
 260         where
 261             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 262         {
 263             type Output = Self;
 264             #[inline]
 265             fn bitand(self, rhs: bool) -> Self {
 266                 self & Self::splat(rhs)
 267             }
 268         }
 269
 270         impl<const LANES: usize> core::ops::BitAnd<$name<LANES>> for bool
 271         where
 272             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 273         {
 274             type Output = $name<LANES>;
 275             #[inline]
 276             fn bitand(self, rhs: $name<LANES>) -> $name<LANES> {
 277                 $name::<LANES>::splat(self) & rhs
 278             }
 279         }
 280
 281         impl<const LANES: usize> core::ops::BitOr for $name<LANES>
 282         where
 283             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 284         {
 285             type Output = Self;
 286             #[inline]
 287             fn bitor(self, rhs: Self) -> Self {
 288                 Self(self.0 | rhs.0)
 289             }
 290         }
 291
 292         impl<const LANES: usize> core::ops::BitOr<bool> for $name<LANES>
 293         where
 294             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 295         {
 296             type Output = Self;
 297             #[inline]
 298             fn bitor(self, rhs: bool) -> Self {
 299                 self | Self::splat(rhs)
 300             }
 301         }
 302
 303         impl<const LANES: usize> core::ops::BitOr<$name<LANES>> for bool
 304         where
 305             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 306         {
 307             type Output = $name<LANES>;
 308             #[inline]
 309             fn bitor(self, rhs: $name<LANES>) -> $name<LANES> {
 310                 $name::<LANES>::splat(self) | rhs
 311             }
 312         }
 313
 314         impl<const LANES: usize> core::ops::BitXor for $name<LANES>
 315         where
 316             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 317         {
 318             type Output = Self;
 319             #[inline]
 320             fn bitxor(self, rhs: Self) -> Self::Output {
 321                 Self(self.0 ^ rhs.0)
 322             }
 323         }
 324
 325         impl<const LANES: usize> core::ops::BitXor<bool> for $name<LANES>
 326         where
 327             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 328         {
 329             type Output = Self;
 330             #[inline]
 331             fn bitxor(self, rhs: bool) -> Self::Output {
 332                 self ^ Self::splat(rhs)
 333             }
 334         }
 335
 336         impl<const LANES: usize> core::ops::BitXor<$name<LANES>> for bool
 337         where
 338             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 339         {
 340             type Output = $name<LANES>;
 341             #[inline]
 342             fn bitxor(self, rhs: $name<LANES>) -> Self::Output {
 343                 $name::<LANES>::splat(self) ^ rhs
 344             }
 345         }
 346
 347         impl<const LANES: usize> core::ops::Not for $name<LANES>
 348         where
 349             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 350         {
 351             type Output = $name<LANES>;
 352             #[inline]
 353             fn not(self) -> Self::Output {
 354                 Self(!self.0)
 355             }
 356         }
 357
 358         impl<const LANES: usize> core::ops::BitAndAssign for $name<LANES>
 359         where
 360             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 361         {
 362             #[inline]
 363             fn bitand_assign(&mut self, rhs: Self) {
 364                 self.0 = self.0 & rhs.0;
 365             }
 366         }
 367
 368         impl<const LANES: usize> core::ops::BitAndAssign<bool> for $name<LANES>
 369         where
 370             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 371         {
 372             #[inline]
 373             fn bitand_assign(&mut self, rhs: bool) {
 374                 *self &= Self::splat(rhs);
 375             }
 376         }
 377
 378         impl<const LANES: usize> core::ops::BitOrAssign for $name<LANES>
 379         where
 380             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 381         {
 382             #[inline]
 383             fn bitor_assign(&mut self, rhs: Self) {
 384                 self.0 = self.0 | rhs.0;
 385             }
 386         }
 387
 388         impl<const LANES: usize> core::ops::BitOrAssign<bool> for $name<LANES>
 389         where
 390             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 391         {
 392             #[inline]
 393             fn bitor_assign(&mut self, rhs: bool) {
 394                 *self |= Self::splat(rhs);
 395             }
 396         }
 397
 398         impl<const LANES: usize> core::ops::BitXorAssign for $name<LANES>
 399         where
 400             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 401         {
 402             #[inline]
 403             fn bitxor_assign(&mut self, rhs: Self) {
 404                 self.0 = self.0 ^ rhs.0;
 405             }
 406         }
 407
 408         impl<const LANES: usize> core::ops::BitXorAssign<bool> for $name<LANES>
 409         where
 410             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 411         {
 412             #[inline]
 413             fn bitxor_assign(&mut self, rhs: bool) {
 414                 *self ^= Self::splat(rhs);
 415             }
 416         }
 417     };
 418 }
 419
 420 define_opaque_mask! {
 421     /// Mask for vectors with `LANES` 8-bit elements.
 422     ///
 423     /// The layout of this type is unspecified.
 424     struct Mask8<const LANES: usize>(mask_impl::Mask8<LANES>);
 425     @bits SimdI8
 426 }
 427
 428 define_opaque_mask! {
 429     /// Mask for vectors with `LANES` 16-bit elements.
 430     ///
 431     /// The layout of this type is unspecified.
 432     struct Mask16<const LANES: usize>(mask_impl::Mask16<LANES>);
 433     @bits SimdI16
 434 }
 435
 436 define_opaque_mask! {
 437     /// Mask for vectors with `LANES` 32-bit elements.
 438     ///
 439     /// The layout of this type is unspecified.
 440     struct Mask32<const LANES: usize>(mask_impl::Mask32<LANES>);
 441     @bits SimdI32
 442 }
 443
 444 define_opaque_mask! {
 445     /// Mask for vectors with `LANES` 64-bit elements.
 446     ///
 447     /// The layout of this type is unspecified.
 448     struct Mask64<const LANES: usize>(mask_impl::Mask64<LANES>);
 449     @bits SimdI64
 450 }
 451
 452 define_opaque_mask! {
 453     /// Mask for vectors with `LANES` pointer-width elements.
 454     ///
 455     /// The layout of this type is unspecified.
 456     struct MaskSize<const LANES: usize>(mask_impl::MaskSize<LANES>);
 457     @bits SimdIsize
 458 }
 459
 460 /// Vector of eight 8-bit masks
 461 pub type mask8x8 = Mask8<8>;
 462
 463 /// Vector of 16 8-bit masks
 464 pub type mask8x16 = Mask8<16>;
 465
 466 /// Vector of 32 8-bit masks
 467 pub type mask8x32 = Mask8<32>;
 468
 469 /// Vector of 16 8-bit masks
 470 pub type mask8x64 = Mask8<64>;
 471
 472 /// Vector of four 16-bit masks
 473 pub type mask16x4 = Mask16<4>;
 474
 475 /// Vector of eight 16-bit masks
 476 pub type mask16x8 = Mask16<8>;
 477
 478 /// Vector of 16 16-bit masks
 479 pub type mask16x16 = Mask16<16>;
 480
 481 /// Vector of 32 16-bit masks
 482 pub type mask16x32 = Mask32<32>;
 483
 484 /// Vector of two 32-bit masks
 485 pub type mask32x2 = Mask32<2>;
 486
 487 /// Vector of four 32-bit masks
 488 pub type mask32x4 = Mask32<4>;
 489
 490 /// Vector of eight 32-bit masks
 491 pub type mask32x8 = Mask32<8>;
 492
 493 /// Vector of 16 32-bit masks
 494 pub type mask32x16 = Mask32<16>;
 495
 496 /// Vector of two 64-bit masks
 497 pub type mask64x2 = Mask64<2>;
 498
 499 /// Vector of four 64-bit masks
 500 pub type mask64x4 = Mask64<4>;
 501
 502 /// Vector of eight 64-bit masks
 503 pub type mask64x8 = Mask64<8>;
 504
 505 /// Vector of two pointer-width masks
 506 pub type masksizex2 = MaskSize<2>;
 507
 508 /// Vector of four pointer-width masks
 509 pub type masksizex4 = MaskSize<4>;
 510
 511 /// Vector of eight pointer-width masks
 512 pub type masksizex8 = MaskSize<8>;
 513
 514 macro_rules! impl_from {
 515     { $from:ident ($from_inner:ident) => $($to:ident ($to_inner:ident)),* } => {
 516         $(
 517         impl<const LANES: usize> From<$from<LANES>> for $to<LANES>
 518         where
 519             crate::LaneCount<LANES>: crate::SupportedLaneCount,
 520         {
 521             fn from(value: $from<LANES>) -> Self {
 522                 Self(value.0.into())
 523             }
 524         }
 525         )*
 526     }
 527 }
 528 impl_from! { Mask8 (SimdI8) => Mask16 (SimdI16), Mask32 (SimdI32), Mask64 (SimdI64), MaskSize (SimdIsize) }
 529 impl_from! { Mask16 (SimdI16) => Mask32 (SimdI32), Mask64 (SimdI64), MaskSize (SimdIsize), Mask8 (SimdI8) }
 530 impl_from! { Mask32 (SimdI32) => Mask64 (SimdI64), MaskSize (SimdIsize), Mask8 (SimdI8), Mask16 (SimdI16) }
 531 impl_from! { Mask64 (SimdI64) => MaskSize (SimdIsize), Mask8 (SimdI8), Mask16 (SimdI16), Mask32 (SimdI32) }
 532 impl_from! { MaskSize (SimdIsize) => Mask8 (SimdI8), Mask16 (SimdI16), Mask32 (SimdI32), Mask64 (SimdI64) }