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