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