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