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

   1 use crate::simd::intrinsics;
   2 use crate::simd::{LaneCount, Simd, SimdElement, SupportedLaneCount};
   3
   4 /// Constructs a new vector by selecting values from the lanes of the source vector or vectors to use.
   5 ///
   6 /// When swizzling one vector, the indices of the result vector are indicated by a `const` array
   7 /// of `usize`, like [`Swizzle`].
   8 /// When swizzling two vectors, the indices are indicated by a `const` array of [`Which`], like
   9 /// [`Swizzle2`].
  10 ///
  11 /// # Examples
  12 /// ## One source vector
  13 /// ```
  14 /// # #![feature(portable_simd)]
  15 /// # #[cfg(feature = "std")] use core_simd::{Simd, simd_swizzle};
  16 /// # #[cfg(not(feature = "std"))] use core::simd::{Simd, simd_swizzle};
  17 /// let v = Simd::<f32, 4>::from_array([0., 1., 2., 3.]);
  18 ///
  19 /// // Keeping the same size
  20 /// let r = simd_swizzle!(v, [3, 0, 1, 2]);
  21 /// assert_eq!(r.to_array(), [3., 0., 1., 2.]);
  22 ///
  23 /// // Changing the number of lanes
  24 /// let r = simd_swizzle!(v, [3, 1]);
  25 /// assert_eq!(r.to_array(), [3., 1.]);
  26 /// ```
  27 ///
  28 /// ## Two source vectors
  29 /// ```
  30 /// # #![feature(portable_simd)]
  31 /// # #[cfg(feature = "std")] use core_simd::{Simd, simd_swizzle, Which};
  32 /// # #[cfg(not(feature = "std"))] use core::simd::{Simd, simd_swizzle, Which};
  33 /// use Which::*;
  34 /// let a = Simd::<f32, 4>::from_array([0., 1., 2., 3.]);
  35 /// let b = Simd::<f32, 4>::from_array([4., 5., 6., 7.]);
  36 ///
  37 /// // Keeping the same size
  38 /// let r = simd_swizzle!(a, b, [First(0), First(1), Second(2), Second(3)]);
  39 /// assert_eq!(r.to_array(), [0., 1., 6., 7.]);
  40 ///
  41 /// // Changing the number of lanes
  42 /// let r = simd_swizzle!(a, b, [First(0), Second(0)]);
  43 /// assert_eq!(r.to_array(), [0., 4.]);
  44 /// ```
  45 #[allow(unused_macros)]
  46 pub macro simd_swizzle {
  47     (
  48         $vector:expr, $index:expr $(,)?
  49     ) => {
  50         {
  51             use $crate::simd::Swizzle;
  52             struct Impl;
  53             impl<const LANES: usize> Swizzle<LANES, {$index.len()}> for Impl {
  54                 const INDEX: [usize; {$index.len()}] = $index;
  55             }
  56             Impl::swizzle($vector)
  57         }
  58     },
  59     (
  60         $first:expr, $second:expr, $index:expr $(,)?
  61     ) => {
  62         {
  63             use $crate::simd::{Which, Swizzle2};
  64             struct Impl;
  65             impl<const LANES: usize> Swizzle2<LANES, {$index.len()}> for Impl {
  66                 const INDEX: [Which; {$index.len()}] = $index;
  67             }
  68             Impl::swizzle2($first, $second)
  69         }
  70     }
  71 }
  72
  73 /// An index into one of two vectors.
  74 #[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
  75 pub enum Which {
  76     /// Indexes the first vector.
  77     First(usize),
  78     /// Indexes the second vector.
  79     Second(usize),
  80 }
  81
  82 /// Create a vector from the elements of another vector.
  83 pub trait Swizzle<const INPUT_LANES: usize, const OUTPUT_LANES: usize> {
  84     /// Map from the lanes of the input vector to the output vector.
  85     const INDEX: [usize; OUTPUT_LANES];
  86
  87     /// Create a new vector from the lanes of `vector`.
  88     ///
  89     /// Lane `i` of the output is `vector[Self::INDEX[i]]`.
  90     #[inline]
  91     #[must_use = "method returns a new vector and does not mutate the original inputs"]
  92     fn swizzle<T>(vector: Simd<T, INPUT_LANES>) -> Simd<T, OUTPUT_LANES>
  93     where
  94         T: SimdElement,
  95         LaneCount<INPUT_LANES>: SupportedLaneCount,
  96         LaneCount<OUTPUT_LANES>: SupportedLaneCount,
  97     {
  98         unsafe { intrinsics::simd_shuffle(vector, vector, Self::INDEX_IMPL) }
  99     }
 100 }
 101
 102 /// Create a vector from the elements of two other vectors.
 103 pub trait Swizzle2<const INPUT_LANES: usize, const OUTPUT_LANES: usize> {
 104     /// Map from the lanes of the input vectors to the output vector
 105     const INDEX: [Which; OUTPUT_LANES];
 106
 107     /// Create a new vector from the lanes of `first` and `second`.
 108     ///
 109     /// Lane `i` is `first[j]` when `Self::INDEX[i]` is `First(j)`, or `second[j]` when it is
 110     /// `Second(j)`.
 111     #[inline]
 112     #[must_use = "method returns a new vector and does not mutate the original inputs"]
 113     fn swizzle2<T>(
 114         first: Simd<T, INPUT_LANES>,
 115         second: Simd<T, INPUT_LANES>,
 116     ) -> Simd<T, OUTPUT_LANES>
 117     where
 118         T: SimdElement,
 119         LaneCount<INPUT_LANES>: SupportedLaneCount,
 120         LaneCount<OUTPUT_LANES>: SupportedLaneCount,
 121     {
 122         unsafe { intrinsics::simd_shuffle(first, second, Self::INDEX_IMPL) }
 123     }
 124 }
 125
 126 /// The `simd_shuffle` intrinsic expects `u32`, so do error checking and conversion here.
 127 /// This trait hides `INDEX_IMPL` from the public API.
 128 trait SwizzleImpl<const INPUT_LANES: usize, const OUTPUT_LANES: usize> {
 129     const INDEX_IMPL: [u32; OUTPUT_LANES];
 130 }
 131
 132 impl<T, const INPUT_LANES: usize, const OUTPUT_LANES: usize> SwizzleImpl<INPUT_LANES, OUTPUT_LANES>
 133     for T
 134 where
 135     T: Swizzle<INPUT_LANES, OUTPUT_LANES> + ?Sized,
 136 {
 137     const INDEX_IMPL: [u32; OUTPUT_LANES] = {
 138         let mut output = [0; OUTPUT_LANES];
 139         let mut i = 0;
 140         while i < OUTPUT_LANES {
 141             let index = Self::INDEX[i];
 142             assert!(index as u32 as usize == index);
 143             assert!(index < INPUT_LANES, "source lane exceeds input lane count",);
 144             output[i] = index as u32;
 145             i += 1;
 146         }
 147         output
 148     };
 149 }
 150
 151 /// The `simd_shuffle` intrinsic expects `u32`, so do error checking and conversion here.
 152 /// This trait hides `INDEX_IMPL` from the public API.
 153 trait Swizzle2Impl<const INPUT_LANES: usize, const OUTPUT_LANES: usize> {
 154     const INDEX_IMPL: [u32; OUTPUT_LANES];
 155 }
 156
 157 impl<T, const INPUT_LANES: usize, const OUTPUT_LANES: usize> Swizzle2Impl<INPUT_LANES, OUTPUT_LANES>
 158     for T
 159 where
 160     T: Swizzle2<INPUT_LANES, OUTPUT_LANES> + ?Sized,
 161 {
 162     const INDEX_IMPL: [u32; OUTPUT_LANES] = {
 163         let mut output = [0; OUTPUT_LANES];
 164         let mut i = 0;
 165         while i < OUTPUT_LANES {
 166             let (offset, index) = match Self::INDEX[i] {
 167                 Which::First(index) => (false, index),
 168                 Which::Second(index) => (true, index),
 169             };
 170             assert!(index < INPUT_LANES, "source lane exceeds input lane count",);
 171
 172             // lanes are indexed by the first vector, then second vector
 173             let index = if offset { index + INPUT_LANES } else { index };
 174             assert!(index as u32 as usize == index);
 175             output[i] = index as u32;
 176             i += 1;
 177         }
 178         output
 179     };
 180 }
 181
 182 impl<T, const LANES: usize> Simd<T, LANES>
 183 where
 184     T: SimdElement,
 185     LaneCount<LANES>: SupportedLaneCount,
 186 {
 187     /// Reverse the order of the lanes in the vector.
 188     #[inline]
 189     #[must_use = "method returns a new vector and does not mutate the original inputs"]
 190     pub fn reverse(self) -> Self {
 191         const fn reverse_index<const LANES: usize>() -> [usize; LANES] {
 192             let mut index = [0; LANES];
 193             let mut i = 0;
 194             while i < LANES {
 195                 index[i] = LANES - i - 1;
 196                 i += 1;
 197             }
 198             index
 199         }
 200
 201         struct Reverse;
 202
 203         impl<const LANES: usize> Swizzle<LANES, LANES> for Reverse {
 204             const INDEX: [usize; LANES] = reverse_index::<LANES>();
 205         }
 206
 207         Reverse::swizzle(self)
 208     }
 209
 210     /// Rotates the vector such that the first `OFFSET` elements of the slice move to the end
 211     /// while the last `LANES - OFFSET` elements move to the front. After calling `rotate_lanes_left`,
 212     /// the element previously in lane `OFFSET` will become the first element in the slice.
 213     #[inline]
 214     #[must_use = "method returns a new vector and does not mutate the original inputs"]
 215     pub fn rotate_lanes_left<const OFFSET: usize>(self) -> Self {
 216         const fn rotate_index<const OFFSET: usize, const LANES: usize>() -> [usize; LANES] {
 217             let offset = OFFSET % LANES;
 218             let mut index = [0; LANES];
 219             let mut i = 0;
 220             while i < LANES {
 221                 index[i] = (i + offset) % LANES;
 222                 i += 1;
 223             }
 224             index
 225         }
 226
 227         struct Rotate<const OFFSET: usize>;
 228
 229         impl<const OFFSET: usize, const LANES: usize> Swizzle<LANES, LANES> for Rotate<OFFSET> {
 230             const INDEX: [usize; LANES] = rotate_index::<OFFSET, LANES>();
 231         }
 232
 233         Rotate::<OFFSET>::swizzle(self)
 234     }
 235
 236     /// Rotates the vector such that the first `LANES - OFFSET` elements of the vector move to
 237     /// the end while the last `OFFSET` elements move to the front. After calling `rotate_lanes_right`,
 238     /// the element previously at index `LANES - OFFSET` will become the first element in the slice.
 239     #[inline]
 240     #[must_use = "method returns a new vector and does not mutate the original inputs"]
 241     pub fn rotate_lanes_right<const OFFSET: usize>(self) -> Self {
 242         const fn rotate_index<const OFFSET: usize, const LANES: usize>() -> [usize; LANES] {
 243             let offset = LANES - OFFSET % LANES;
 244             let mut index = [0; LANES];
 245             let mut i = 0;
 246             while i < LANES {
 247                 index[i] = (i + offset) % LANES;
 248                 i += 1;
 249             }
 250             index
 251         }
 252
 253         struct Rotate<const OFFSET: usize>;
 254
 255         impl<const OFFSET: usize, const LANES: usize> Swizzle<LANES, LANES> for Rotate<OFFSET> {
 256             const INDEX: [usize; LANES] = rotate_index::<OFFSET, LANES>();
 257         }
 258
 259         Rotate::<OFFSET>::swizzle(self)
 260     }
 261
 262     /// Interleave two vectors.
 263     ///
 264     /// Produces two vectors with lanes taken alternately from `self` and `other`.
 265     ///
 266     /// The first result contains the first `LANES / 2` lanes from `self` and `other`,
 267     /// alternating, starting with the first lane of `self`.
 268     ///
 269     /// The second result contains the last `LANES / 2` lanes from `self` and `other`,
 270     /// alternating, starting with the lane `LANES / 2` from the start of `self`.
 271     ///
 272     /// ```
 273     /// #![feature(portable_simd)]
 274     /// # #[cfg(feature = "std")] use core_simd::Simd;
 275     /// # #[cfg(not(feature = "std"))] use core::simd::Simd;
 276     /// let a = Simd::from_array([0, 1, 2, 3]);
 277     /// let b = Simd::from_array([4, 5, 6, 7]);
 278     /// let (x, y) = a.interleave(b);
 279     /// assert_eq!(x.to_array(), [0, 4, 1, 5]);
 280     /// assert_eq!(y.to_array(), [2, 6, 3, 7]);
 281     /// ```
 282     #[inline]
 283     #[must_use = "method returns a new vector and does not mutate the original inputs"]
 284     pub fn interleave(self, other: Self) -> (Self, Self) {
 285         const fn lo<const LANES: usize>() -> [Which; LANES] {
 286             let mut idx = [Which::First(0); LANES];
 287             let mut i = 0;
 288             while i < LANES {
 289                 let offset = i / 2;
 290                 idx[i] = if i % 2 == 0 {
 291                     Which::First(offset)
 292                 } else {
 293                     Which::Second(offset)
 294                 };
 295                 i += 1;
 296             }
 297             idx
 298         }
 299         const fn hi<const LANES: usize>() -> [Which; LANES] {
 300             let mut idx = [Which::First(0); LANES];
 301             let mut i = 0;
 302             while i < LANES {
 303                 let offset = (LANES + i) / 2;
 304                 idx[i] = if i % 2 == 0 {
 305                     Which::First(offset)
 306                 } else {
 307                     Which::Second(offset)
 308                 };
 309                 i += 1;
 310             }
 311             idx
 312         }
 313
 314         struct Lo;
 315         struct Hi;
 316
 317         impl<const LANES: usize> Swizzle2<LANES, LANES> for Lo {
 318             const INDEX: [Which; LANES] = lo::<LANES>();
 319         }
 320
 321         impl<const LANES: usize> Swizzle2<LANES, LANES> for Hi {
 322             const INDEX: [Which; LANES] = hi::<LANES>();
 323         }
 324
 325         (Lo::swizzle2(self, other), Hi::swizzle2(self, other))
 326     }
 327
 328     /// Deinterleave two vectors.
 329     ///
 330     /// The first result takes every other lane of `self` and then `other`, starting with
 331     /// the first lane.
 332     ///
 333     /// The second result takes every other lane of `self` and then `other`, starting with
 334     /// the second lane.
 335     ///
 336     /// ```
 337     /// #![feature(portable_simd)]
 338     /// # #[cfg(feature = "std")] use core_simd::Simd;
 339     /// # #[cfg(not(feature = "std"))] use core::simd::Simd;
 340     /// let a = Simd::from_array([0, 4, 1, 5]);
 341     /// let b = Simd::from_array([2, 6, 3, 7]);
 342     /// let (x, y) = a.deinterleave(b);
 343     /// assert_eq!(x.to_array(), [0, 1, 2, 3]);
 344     /// assert_eq!(y.to_array(), [4, 5, 6, 7]);
 345     /// ```
 346     #[inline]
 347     #[must_use = "method returns a new vector and does not mutate the original inputs"]
 348     pub fn deinterleave(self, other: Self) -> (Self, Self) {
 349         const fn even<const LANES: usize>() -> [Which; LANES] {
 350             let mut idx = [Which::First(0); LANES];
 351             let mut i = 0;
 352             while i < LANES / 2 {
 353                 idx[i] = Which::First(2 * i);
 354                 idx[i + LANES / 2] = Which::Second(2 * i);
 355                 i += 1;
 356             }
 357             idx
 358         }
 359         const fn odd<const LANES: usize>() -> [Which; LANES] {
 360             let mut idx = [Which::First(0); LANES];
 361             let mut i = 0;
 362             while i < LANES / 2 {
 363                 idx[i] = Which::First(2 * i + 1);
 364                 idx[i + LANES / 2] = Which::Second(2 * i + 1);
 365                 i += 1;
 366             }
 367             idx
 368         }
 369
 370         struct Even;
 371         struct Odd;
 372
 373         impl<const LANES: usize> Swizzle2<LANES, LANES> for Even {
 374             const INDEX: [Which; LANES] = even::<LANES>();
 375         }
 376
 377         impl<const LANES: usize> Swizzle2<LANES, LANES> for Odd {
 378             const INDEX: [Which; LANES] = odd::<LANES>();
 379         }
 380
 381         (Even::swizzle2(self, other), Odd::swizzle2(self, other))
 382     }
 383 }