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