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

   1 mod float;
   2 mod int;
   3 mod uint;
   4
   5 pub use float::*;
   6 pub use int::*;
   7 pub use uint::*;
   8
   9 // Vectors of pointers are not for public use at the current time.
  10 pub(crate) mod ptr;
  11
  12 use crate::simd::intrinsics;
  13 use crate::simd::{LaneCount, Mask, MaskElement, SupportedLaneCount};
  14
  15 /// A SIMD vector of `LANES` elements of type `T`.
  16 #[repr(simd)]
  17 pub struct Simd<T, const LANES: usize>([T; LANES])
  18 where
  19     T: SimdElement,
  20     LaneCount<LANES>: SupportedLaneCount;
  21
  22 impl<T, const LANES: usize> Simd<T, LANES>
  23 where
  24     LaneCount<LANES>: SupportedLaneCount,
  25     T: SimdElement,
  26 {
  27     /// Number of lanes in this vector.
  28     pub const LANES: usize = LANES;
  29
  30     /// Get the number of lanes in this vector.
  31     pub const fn lanes(&self) -> usize {
  32         LANES
  33     }
  34
  35     /// Construct a SIMD vector by setting all lanes to the given value.
  36     pub const fn splat(value: T) -> Self {
  37         Self([value; LANES])
  38     }
  39
  40     /// Returns an array reference containing the entire SIMD vector.
  41     pub const fn as_array(&self) -> &[T; LANES] {
  42         &self.0
  43     }
  44
  45     /// Returns a mutable array reference containing the entire SIMD vector.
  46     pub fn as_mut_array(&mut self) -> &mut [T; LANES] {
  47         &mut self.0
  48     }
  49
  50     /// Converts an array to a SIMD vector.
  51     pub const fn from_array(array: [T; LANES]) -> Self {
  52         Self(array)
  53     }
  54
  55     /// Converts a SIMD vector to an array.
  56     pub const fn to_array(self) -> [T; LANES] {
  57         self.0
  58     }
  59
  60     /// Converts a slice to a SIMD vector containing `slice[..LANES]`
  61     /// # Panics
  62     /// `from_slice` will panic if the slice's `len` is less than the vector's `Simd::LANES`.
  63     #[must_use]
  64     pub const fn from_slice(slice: &[T]) -> Self {
  65         assert!(
  66             slice.len() >= LANES,
  67             "slice length must be at least the number of lanes"
  68         );
  69         let mut array = [slice[0]; LANES];
  70         let mut i = 0;
  71         while i < LANES {
  72             array[i] = slice[i];
  73             i += 1;
  74         }
  75         Self(array)
  76     }
  77
  78     /// Performs lanewise conversion of a SIMD vector's elements to another SIMD-valid type.
  79     /// This follows the semantics of Rust's `as` conversion for casting
  80     /// integers to unsigned integers (interpreting as the other type, so `-1` to `MAX`),
  81     /// and from floats to integers (truncating, or saturating at the limits) for each lane,
  82     /// or vice versa.
  83     ///
  84     /// # Examples
  85     /// ```
  86     /// # #![feature(portable_simd)]
  87     /// # #[cfg(feature = "std")] use core_simd::Simd;
  88     /// # #[cfg(not(feature = "std"))] use core::simd::Simd;
  89     /// let floats: Simd<f32, 4> = Simd::from_array([1.9, -4.5, f32::INFINITY, f32::NAN]);
  90     /// let ints = floats.cast::<i32>();
  91     /// assert_eq!(ints, Simd::from_array([1, -4, i32::MAX, 0]));
  92     ///
  93     /// // Formally equivalent, but `Simd::cast` can optimize better.
  94     /// assert_eq!(ints, Simd::from_array(floats.to_array().map(|x| x as i32)));
  95     ///
  96     /// // The float conversion does not round-trip.
  97     /// let floats_again = ints.cast();
  98     /// assert_ne!(floats, floats_again);
  99     /// assert_eq!(floats_again, Simd::from_array([1.0, -4.0, 2147483647.0, 0.0]));
 100     /// ```
 101     #[must_use]
 102     #[inline]
 103     #[cfg(not(bootstrap))]
 104     pub fn cast<U: SimdElement>(self) -> Simd<U, LANES> {
 105         unsafe { intrinsics::simd_as(self) }
 106     }
 107
 108     /// Reads from potentially discontiguous indices in `slice` to construct a SIMD vector.
 109     /// If an index is out-of-bounds, the lane is instead selected from the `or` vector.
 110     ///
 111     /// # Examples
 112     /// ```
 113     /// # #![feature(portable_simd)]
 114     /// # #[cfg(feature = "std")] use core_simd::Simd;
 115     /// # #[cfg(not(feature = "std"))] use core::simd::Simd;
 116     /// let vec: Vec<i32> = vec![10, 11, 12, 13, 14, 15, 16, 17, 18];
 117     /// let idxs = Simd::from_array([9, 3, 0, 5]);
 118     /// let alt = Simd::from_array([-5, -4, -3, -2]);
 119     ///
 120     /// let result = Simd::gather_or(&vec, idxs, alt); // Note the lane that is out-of-bounds.
 121     /// assert_eq!(result, Simd::from_array([-5, 13, 10, 15]));
 122     /// ```
 123     #[must_use]
 124     #[inline]
 125     pub fn gather_or(slice: &[T], idxs: Simd<usize, LANES>, or: Self) -> Self {
 126         Self::gather_select(slice, Mask::splat(true), idxs, or)
 127     }
 128
 129     /// Reads from potentially discontiguous indices in `slice` to construct a SIMD vector.
 130     /// If an index is out-of-bounds, the lane is set to the default value for the type.
 131     ///
 132     /// # Examples
 133     /// ```
 134     /// # #![feature(portable_simd)]
 135     /// # #[cfg(feature = "std")] use core_simd::Simd;
 136     /// # #[cfg(not(feature = "std"))] use core::simd::Simd;
 137     /// let vec: Vec<i32> = vec![10, 11, 12, 13, 14, 15, 16, 17, 18];
 138     /// let idxs = Simd::from_array([9, 3, 0, 5]);
 139     ///
 140     /// let result = Simd::gather_or_default(&vec, idxs); // Note the lane that is out-of-bounds.
 141     /// assert_eq!(result, Simd::from_array([0, 13, 10, 15]));
 142     /// ```
 143     #[must_use]
 144     #[inline]
 145     pub fn gather_or_default(slice: &[T], idxs: Simd<usize, LANES>) -> Self
 146     where
 147         T: Default,
 148     {
 149         Self::gather_or(slice, idxs, Self::splat(T::default()))
 150     }
 151
 152     /// Reads from potentially discontiguous indices in `slice` to construct a SIMD vector.
 153     /// The mask `enable`s all `true` lanes and disables all `false` lanes.
 154     /// If an index is disabled or is out-of-bounds, the lane is selected from the `or` vector.
 155     ///
 156     /// # Examples
 157     /// ```
 158     /// # #![feature(portable_simd)]
 159     /// # #[cfg(feature = "std")] use core_simd::{Simd, Mask};
 160     /// # #[cfg(not(feature = "std"))] use core::simd::{Simd, Mask};
 161     /// let vec: Vec<i32> = vec![10, 11, 12, 13, 14, 15, 16, 17, 18];
 162     /// let idxs = Simd::from_array([9, 3, 0, 5]);
 163     /// let alt = Simd::from_array([-5, -4, -3, -2]);
 164     /// let enable = Mask::from_array([true, true, true, false]); // Note the mask of the last lane.
 165     ///
 166     /// let result = Simd::gather_select(&vec, enable, idxs, alt); // Note the lane that is out-of-bounds.
 167     /// assert_eq!(result, Simd::from_array([-5, 13, 10, -2]));
 168     /// ```
 169     #[must_use]
 170     #[inline]
 171     pub fn gather_select(
 172         slice: &[T],
 173         enable: Mask<isize, LANES>,
 174         idxs: Simd<usize, LANES>,
 175         or: Self,
 176     ) -> Self {
 177         let enable: Mask<isize, LANES> = enable & idxs.lanes_lt(Simd::splat(slice.len()));
 178         // SAFETY: We have masked-off out-of-bounds lanes.
 179         unsafe { Self::gather_select_unchecked(slice, enable, idxs, or) }
 180     }
 181
 182     /// Reads from potentially discontiguous indices in `slice` to construct a SIMD vector.
 183     /// The mask `enable`s all `true` lanes and disables all `false` lanes.
 184     /// If an index is disabled, the lane is selected from the `or` vector.
 185     ///
 186     /// # Safety
 187     ///
 188     /// Calling this function with an `enable`d out-of-bounds index is *[undefined behavior]*
 189     /// even if the resulting value is not used.
 190     ///
 191     /// # Examples
 192     /// ```
 193     /// # #![feature(portable_simd)]
 194     /// # #[cfg(feature = "std")] use core_simd::{Simd, Mask};
 195     /// # #[cfg(not(feature = "std"))] use core::simd::{Simd, Mask};
 196     /// let vec: Vec<i32> = vec![10, 11, 12, 13, 14, 15, 16, 17, 18];
 197     /// let idxs = Simd::from_array([9, 3, 0, 5]);
 198     /// let alt = Simd::from_array([-5, -4, -3, -2]);
 199     /// let enable = Mask::from_array([true, true, true, false]); // Note the final mask lane.
 200     /// // If this mask was used to gather, it would be unsound. Let's fix that.
 201     /// let enable = enable & idxs.lanes_lt(Simd::splat(vec.len()));
 202     ///
 203     /// // We have masked the OOB lane, so it's safe to gather now.
 204     /// let result = unsafe { Simd::gather_select_unchecked(&vec, enable, idxs, alt) };
 205     /// assert_eq!(result, Simd::from_array([-5, 13, 10, -2]));
 206     /// ```
 207     /// [undefined behavior]: https://doc.rust-lang.org/reference/behavior-considered-undefined.html
 208     #[must_use]
 209     #[inline]
 210     pub unsafe fn gather_select_unchecked(
 211         slice: &[T],
 212         enable: Mask<isize, LANES>,
 213         idxs: Simd<usize, LANES>,
 214         or: Self,
 215     ) -> Self {
 216         let base_ptr = crate::simd::ptr::SimdConstPtr::splat(slice.as_ptr());
 217         // Ferris forgive me, I have done pointer arithmetic here.
 218         let ptrs = base_ptr.wrapping_add(idxs);
 219         // SAFETY: The ptrs have been bounds-masked to prevent memory-unsafe reads insha'allah
 220         unsafe { intrinsics::simd_gather(or, ptrs, enable.to_int()) }
 221     }
 222
 223     /// Writes the values in a SIMD vector to potentially discontiguous indices in `slice`.
 224     /// If two lanes in the scattered vector would write to the same index
 225     /// only the last lane is guaranteed to actually be written.
 226     ///
 227     /// # Examples
 228     /// ```
 229     /// # #![feature(portable_simd)]
 230     /// # #[cfg(feature = "std")] use core_simd::Simd;
 231     /// # #[cfg(not(feature = "std"))] use core::simd::Simd;
 232     /// let mut vec: Vec<i32> = vec![10, 11, 12, 13, 14, 15, 16, 17, 18];
 233     /// let idxs = Simd::from_array([9, 3, 0, 0]);
 234     /// let vals = Simd::from_array([-27, 82, -41, 124]);
 235     ///
 236     /// vals.scatter(&mut vec, idxs); // index 0 receives two writes.
 237     /// assert_eq!(vec, vec![124, 11, 12, 82, 14, 15, 16, 17, 18]);
 238     /// ```
 239     #[inline]
 240     pub fn scatter(self, slice: &mut [T], idxs: Simd<usize, LANES>) {
 241         self.scatter_select(slice, Mask::splat(true), idxs)
 242     }
 243
 244     /// Writes the values in a SIMD vector to multiple potentially discontiguous indices in `slice`.
 245     /// The mask `enable`s all `true` lanes and disables all `false` lanes.
 246     /// If an enabled index is out-of-bounds, the lane is not written.
 247     /// If two enabled lanes in the scattered vector would write to the same index,
 248     /// only the last lane is guaranteed to actually be written.
 249     ///
 250     /// # Examples
 251     /// ```
 252     /// # #![feature(portable_simd)]
 253     /// # #[cfg(feature = "std")] use core_simd::{Simd, Mask};
 254     /// # #[cfg(not(feature = "std"))] use core::simd::{Simd, Mask};
 255     /// let mut vec: Vec<i32> = vec![10, 11, 12, 13, 14, 15, 16, 17, 18];
 256     /// let idxs = Simd::from_array([9, 3, 0, 0]);
 257     /// let vals = Simd::from_array([-27, 82, -41, 124]);
 258     /// let enable = Mask::from_array([true, true, true, false]); // Note the mask of the last lane.
 259     ///
 260     /// vals.scatter_select(&mut vec, enable, idxs); // index 0's second write is masked, thus omitted.
 261     /// assert_eq!(vec, vec![-41, 11, 12, 82, 14, 15, 16, 17, 18]);
 262     /// ```
 263     #[inline]
 264     pub fn scatter_select(
 265         self,
 266         slice: &mut [T],
 267         enable: Mask<isize, LANES>,
 268         idxs: Simd<usize, LANES>,
 269     ) {
 270         let enable: Mask<isize, LANES> = enable & idxs.lanes_lt(Simd::splat(slice.len()));
 271         // SAFETY: We have masked-off out-of-bounds lanes.
 272         unsafe { self.scatter_select_unchecked(slice, enable, idxs) }
 273     }
 274
 275     /// Writes the values in a SIMD vector to multiple potentially discontiguous indices in `slice`.
 276     /// The mask `enable`s all `true` lanes and disables all `false` lanes.
 277     /// If two enabled lanes in the scattered vector would write to the same index,
 278     /// only the last lane is guaranteed to actually be written.
 279     ///
 280     /// # Safety
 281     ///
 282     /// Calling this function with an enabled out-of-bounds index is *[undefined behavior]*,
 283     /// and may lead to memory corruption.
 284     ///
 285     /// # Examples
 286     /// ```
 287     /// # #![feature(portable_simd)]
 288     /// # #[cfg(feature = "std")] use core_simd::{Simd, Mask};
 289     /// # #[cfg(not(feature = "std"))] use core::simd::{Simd, Mask};
 290     /// let mut vec: Vec<i32> = vec![10, 11, 12, 13, 14, 15, 16, 17, 18];
 291     /// let idxs = Simd::from_array([9, 3, 0, 0]);
 292     /// let vals = Simd::from_array([-27, 82, -41, 124]);
 293     /// let enable = Mask::from_array([true, true, true, false]); // Note the mask of the last lane.
 294     /// // If this mask was used to scatter, it would be unsound. Let's fix that.
 295     /// let enable = enable & idxs.lanes_lt(Simd::splat(vec.len()));
 296     ///
 297     /// // We have masked the OOB lane, so it's safe to scatter now.
 298     /// unsafe { vals.scatter_select_unchecked(&mut vec, enable, idxs); }
 299     /// // index 0's second write is masked, thus was omitted.
 300     /// assert_eq!(vec, vec![-41, 11, 12, 82, 14, 15, 16, 17, 18]);
 301     /// ```
 302     /// [undefined behavior]: https://doc.rust-lang.org/reference/behavior-considered-undefined.html
 303     #[inline]
 304     pub unsafe fn scatter_select_unchecked(
 305         self,
 306         slice: &mut [T],
 307         enable: Mask<isize, LANES>,
 308         idxs: Simd<usize, LANES>,
 309     ) {
 310         // SAFETY: This block works with *mut T derived from &mut 'a [T],
 311         // which means it is delicate in Rust's borrowing model, circa 2021:
 312         // &mut 'a [T] asserts uniqueness, so deriving &'a [T] invalidates live *mut Ts!
 313         // Even though this block is largely safe methods, it must be exactly this way
 314         // to prevent invalidating the raw ptrs while they're live.
 315         // Thus, entering this block requires all values to use being already ready:
 316         // 0. idxs we want to write to, which are used to construct the mask.
 317         // 1. enable, which depends on an initial &'a [T] and the idxs.
 318         // 2. actual values to scatter (self).
 319         // 3. &mut [T] which will become our base ptr.
 320         unsafe {
 321             // Now Entering ☢️ *mut T Zone
 322             let base_ptr = crate::simd::ptr::SimdMutPtr::splat(slice.as_mut_ptr());
 323             // Ferris forgive me, I have done pointer arithmetic here.
 324             let ptrs = base_ptr.wrapping_add(idxs);
 325             // The ptrs have been bounds-masked to prevent memory-unsafe writes insha'allah
 326             intrinsics::simd_scatter(self, ptrs, enable.to_int())
 327             // Cleared ☢️ *mut T Zone
 328         }
 329     }
 330 }
 331
 332 impl<T, const LANES: usize> Copy for Simd<T, LANES>
 333 where
 334     T: SimdElement,
 335     LaneCount<LANES>: SupportedLaneCount,
 336 {
 337 }
 338
 339 impl<T, const LANES: usize> Clone for Simd<T, LANES>
 340 where
 341     T: SimdElement,
 342     LaneCount<LANES>: SupportedLaneCount,
 343 {
 344     fn clone(&self) -> Self {
 345         *self
 346     }
 347 }
 348
 349 impl<T, const LANES: usize> Default for Simd<T, LANES>
 350 where
 351     LaneCount<LANES>: SupportedLaneCount,
 352     T: SimdElement + Default,
 353 {
 354     #[inline]
 355     fn default() -> Self {
 356         Self::splat(T::default())
 357     }
 358 }
 359
 360 impl<T, const LANES: usize> PartialEq for Simd<T, LANES>
 361 where
 362     LaneCount<LANES>: SupportedLaneCount,
 363     T: SimdElement + PartialEq,
 364 {
 365     #[inline]
 366     fn eq(&self, other: &Self) -> bool {
 367         // TODO use SIMD equality
 368         self.to_array() == other.to_array()
 369     }
 370 }
 371
 372 impl<T, const LANES: usize> PartialOrd for Simd<T, LANES>
 373 where
 374     LaneCount<LANES>: SupportedLaneCount,
 375     T: SimdElement + PartialOrd,
 376 {
 377     #[inline]
 378     fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
 379         // TODO use SIMD equality
 380         self.to_array().partial_cmp(other.as_ref())
 381     }
 382 }
 383
 384 impl<T, const LANES: usize> Eq for Simd<T, LANES>
 385 where
 386     LaneCount<LANES>: SupportedLaneCount,
 387     T: SimdElement + Eq,
 388 {
 389 }
 390
 391 impl<T, const LANES: usize> Ord for Simd<T, LANES>
 392 where
 393     LaneCount<LANES>: SupportedLaneCount,
 394     T: SimdElement + Ord,
 395 {
 396     #[inline]
 397     fn cmp(&self, other: &Self) -> core::cmp::Ordering {
 398         // TODO use SIMD equality
 399         self.to_array().cmp(other.as_ref())
 400     }
 401 }
 402
 403 impl<T, const LANES: usize> core::hash::Hash for Simd<T, LANES>
 404 where
 405     LaneCount<LANES>: SupportedLaneCount,
 406     T: SimdElement + core::hash::Hash,
 407 {
 408     #[inline]
 409     fn hash<H>(&self, state: &mut H)
 410     where
 411         H: core::hash::Hasher,
 412     {
 413         self.as_array().hash(state)
 414     }
 415 }
 416
 417 // array references
 418 impl<T, const LANES: usize> AsRef<[T; LANES]> for Simd<T, LANES>
 419 where
 420     LaneCount<LANES>: SupportedLaneCount,
 421     T: SimdElement,
 422 {
 423     #[inline]
 424     fn as_ref(&self) -> &[T; LANES] {
 425         &self.0
 426     }
 427 }
 428
 429 impl<T, const LANES: usize> AsMut<[T; LANES]> for Simd<T, LANES>
 430 where
 431     LaneCount<LANES>: SupportedLaneCount,
 432     T: SimdElement,
 433 {
 434     #[inline]
 435     fn as_mut(&mut self) -> &mut [T; LANES] {
 436         &mut self.0
 437     }
 438 }
 439
 440 // slice references
 441 impl<T, const LANES: usize> AsRef<[T]> for Simd<T, LANES>
 442 where
 443     LaneCount<LANES>: SupportedLaneCount,
 444     T: SimdElement,
 445 {
 446     #[inline]
 447     fn as_ref(&self) -> &[T] {
 448         &self.0
 449     }
 450 }
 451
 452 impl<T, const LANES: usize> AsMut<[T]> for Simd<T, LANES>
 453 where
 454     LaneCount<LANES>: SupportedLaneCount,
 455     T: SimdElement,
 456 {
 457     #[inline]
 458     fn as_mut(&mut self) -> &mut [T] {
 459         &mut self.0
 460     }
 461 }
 462
 463 // vector/array conversion
 464 impl<T, const LANES: usize> From<[T; LANES]> for Simd<T, LANES>
 465 where
 466     LaneCount<LANES>: SupportedLaneCount,
 467     T: SimdElement,
 468 {
 469     fn from(array: [T; LANES]) -> Self {
 470         Self(array)
 471     }
 472 }
 473
 474 impl<T, const LANES: usize> From<Simd<T, LANES>> for [T; LANES]
 475 where
 476     LaneCount<LANES>: SupportedLaneCount,
 477     T: SimdElement,
 478 {
 479     fn from(vector: Simd<T, LANES>) -> Self {
 480         vector.to_array()
 481     }
 482 }
 483
 484 mod sealed {
 485     pub trait Sealed {}
 486 }
 487 use sealed::Sealed;
 488
 489 /// Marker trait for types that may be used as SIMD vector elements.
 490 /// SAFETY: This trait, when implemented, asserts the compiler can monomorphize
 491 /// `#[repr(simd)]` structs with the marked type as an element.
 492 /// Strictly, it is valid to impl if the vector will not be miscompiled.
 493 /// Practically, it is user-unfriendly to impl it if the vector won't compile,
 494 /// even when no soundness guarantees are broken by allowing the user to try.
 495 pub unsafe trait SimdElement: Sealed + Copy {
 496     /// The mask element type corresponding to this element type.
 497     type Mask: MaskElement;
 498 }
 499
 500 impl Sealed for u8 {}
 501 unsafe impl SimdElement for u8 {
 502     type Mask = i8;
 503 }
 504
 505 impl Sealed for u16 {}
 506 unsafe impl SimdElement for u16 {
 507     type Mask = i16;
 508 }
 509
 510 impl Sealed for u32 {}
 511 unsafe impl SimdElement for u32 {
 512     type Mask = i32;
 513 }
 514
 515 impl Sealed for u64 {}
 516 unsafe impl SimdElement for u64 {
 517     type Mask = i64;
 518 }
 519
 520 impl Sealed for usize {}
 521 unsafe impl SimdElement for usize {
 522     type Mask = isize;
 523 }
 524
 525 impl Sealed for i8 {}
 526 unsafe impl SimdElement for i8 {
 527     type Mask = i8;
 528 }
 529
 530 impl Sealed for i16 {}
 531 unsafe impl SimdElement for i16 {
 532     type Mask = i16;
 533 }
 534
 535 impl Sealed for i32 {}
 536 unsafe impl SimdElement for i32 {
 537     type Mask = i32;
 538 }
 539
 540 impl Sealed for i64 {}
 541 unsafe impl SimdElement for i64 {
 542     type Mask = i64;
 543 }
 544
 545 impl Sealed for isize {}
 546 unsafe impl SimdElement for isize {
 547     type Mask = isize;
 548 }
 549
 550 impl Sealed for f32 {}
 551 unsafe impl SimdElement for f32 {
 552     type Mask = i32;
 553 }
 554
 555 impl Sealed for f64 {}
 556 unsafe impl SimdElement for f64 {
 557     type Mask = i64;
 558 }