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