core/portable-simd/crates/core_simd/src/swizzle.rs
1use crate::simd::{LaneCount, Mask, MaskElement, Simd, SimdElement, SupportedLaneCount};
2
3/// Constructs a new SIMD vector by copying elements from selected elements in other vectors.
4///
5/// When swizzling one vector, elements are selected like [`Swizzle::swizzle`].
6///
7/// When swizzling two vectors, elements are selected like [`Swizzle::concat_swizzle`].
8///
9/// # Examples
10///
11/// With a single SIMD vector, the const array specifies element indices in that vector:
12/// ```
13/// # #![feature(portable_simd)]
14/// # use core::simd::{u32x2, u32x4, simd_swizzle};
15/// let v = u32x4::from_array([10, 11, 12, 13]);
16///
17/// // Keeping the same size
18/// let r: u32x4 = simd_swizzle!(v, [3, 0, 1, 2]);
19/// assert_eq!(r.to_array(), [13, 10, 11, 12]);
20///
21/// // Changing the number of elements
22/// let r: u32x2 = simd_swizzle!(v, [3, 1]);
23/// assert_eq!(r.to_array(), [13, 11]);
24/// ```
25///
26/// With two input SIMD vectors, the const array specifies element indices in the concatenation of
27/// those vectors:
28/// ```
29/// # #![feature(portable_simd)]
30/// # #[cfg(feature = "as_crate")] use core_simd::simd;
31/// # #[cfg(not(feature = "as_crate"))] use core::simd;
32/// # use simd::{u32x2, u32x4, simd_swizzle};
33/// let a = u32x4::from_array([0, 1, 2, 3]);
34/// let b = u32x4::from_array([4, 5, 6, 7]);
35///
36/// // Keeping the same size
37/// let r: u32x4 = simd_swizzle!(a, b, [0, 1, 6, 7]);
38/// assert_eq!(r.to_array(), [0, 1, 6, 7]);
39///
40/// // Changing the number of elements
41/// let r: u32x2 = simd_swizzle!(a, b, [0, 4]);
42/// assert_eq!(r.to_array(), [0, 4]);
43/// ```
44#[allow(unused_macros)]
45pub macro simd_swizzle {
46 (
47 $vector:expr, $index:expr $(,)?
48 ) => {
49 {
50 use $crate::simd::Swizzle;
51 struct Impl;
52 impl Swizzle<{$index.len()}> for Impl {
53 const INDEX: [usize; {$index.len()}] = $index;
54 }
55 Impl::swizzle($vector)
56 }
57 },
58 (
59 $first:expr, $second:expr, $index:expr $(,)?
60 ) => {
61 {
62 use $crate::simd::Swizzle;
63 struct Impl;
64 impl Swizzle<{$index.len()}> for Impl {
65 const INDEX: [usize; {$index.len()}] = $index;
66 }
67 Impl::concat_swizzle($first, $second)
68 }
69 }
70}
71
72/// Creates a vector from the elements of another vector.
73pub trait Swizzle<const N: usize> {
74 /// Map from the elements of the input vector to the output vector.
75 const INDEX: [usize; N];
76
77 /// Creates a new vector from the elements of `vector`.
78 ///
79 /// Lane `i` of the output is `vector[Self::INDEX[i]]`.
80 #[inline]
81 #[must_use = "method returns a new vector and does not mutate the original inputs"]
82 fn swizzle<T, const M: usize>(vector: Simd<T, M>) -> Simd<T, N>
83 where
84 T: SimdElement,
85 LaneCount<N>: SupportedLaneCount,
86 LaneCount<M>: SupportedLaneCount,
87 {
88 // Safety: `vector` is a vector, and the index is a const vector of u32.
89 unsafe {
90 core::intrinsics::simd::simd_shuffle(
91 vector,
92 vector,
93 const {
94 let mut output = [0; N];
95 let mut i = 0;
96 while i < N {
97 let index = Self::INDEX[i];
98 assert!(index as u32 as usize == index);
99 assert!(
100 index < M,
101 "source element index exceeds input vector length"
102 );
103 output[i] = index as u32;
104 i += 1;
105 }
106
107 // The index list needs to be returned as a vector.
108 #[repr(simd)]
109 struct SimdShuffleIdx<const LEN: usize>([u32; LEN]);
110 SimdShuffleIdx(output)
111 },
112 )
113 }
114 }
115
116 /// Creates a new vector from the elements of `first` and `second`.
117 ///
118 /// Lane `i` of the output is `concat[Self::INDEX[i]]`, where `concat` is the concatenation of
119 /// `first` and `second`.
120 #[inline]
121 #[must_use = "method returns a new vector and does not mutate the original inputs"]
122 fn concat_swizzle<T, const M: usize>(first: Simd<T, M>, second: Simd<T, M>) -> Simd<T, N>
123 where
124 T: SimdElement,
125 LaneCount<N>: SupportedLaneCount,
126 LaneCount<M>: SupportedLaneCount,
127 {
128 // Safety: `first` and `second` are vectors, and the index is a const vector of u32.
129 unsafe {
130 core::intrinsics::simd::simd_shuffle(
131 first,
132 second,
133 const {
134 let mut output = [0; N];
135 let mut i = 0;
136 while i < N {
137 let index = Self::INDEX[i];
138 assert!(index as u32 as usize == index);
139 assert!(
140 index < 2 * M,
141 "source element index exceeds input vector length"
142 );
143 output[i] = index as u32;
144 i += 1;
145 }
146
147 // The index list needs to be returned as a vector.
148 #[repr(simd)]
149 struct SimdShuffleIdx<const LEN: usize>([u32; LEN]);
150 SimdShuffleIdx(output)
151 },
152 )
153 }
154 }
155
156 /// Creates a new mask from the elements of `mask`.
157 ///
158 /// Element `i` of the output is `mask[Self::INDEX[i]]`.
159 #[inline]
160 #[must_use = "method returns a new mask and does not mutate the original inputs"]
161 fn swizzle_mask<T, const M: usize>(mask: Mask<T, M>) -> Mask<T, N>
162 where
163 T: MaskElement,
164 LaneCount<N>: SupportedLaneCount,
165 LaneCount<M>: SupportedLaneCount,
166 {
167 // SAFETY: all elements of this mask come from another mask
168 unsafe { Mask::from_int_unchecked(Self::swizzle(mask.to_int())) }
169 }
170
171 /// Creates a new mask from the elements of `first` and `second`.
172 ///
173 /// Element `i` of the output is `concat[Self::INDEX[i]]`, where `concat` is the concatenation of
174 /// `first` and `second`.
175 #[inline]
176 #[must_use = "method returns a new mask and does not mutate the original inputs"]
177 fn concat_swizzle_mask<T, const M: usize>(first: Mask<T, M>, second: Mask<T, M>) -> Mask<T, N>
178 where
179 T: MaskElement,
180 LaneCount<N>: SupportedLaneCount,
181 LaneCount<M>: SupportedLaneCount,
182 {
183 // SAFETY: all elements of this mask come from another mask
184 unsafe { Mask::from_int_unchecked(Self::concat_swizzle(first.to_int(), second.to_int())) }
185 }
186}
187
188impl<T, const N: usize> Simd<T, N>
189where
190 T: SimdElement,
191 LaneCount<N>: SupportedLaneCount,
192{
193 /// Reverse the order of the elements in the vector.
194 #[inline]
195 #[must_use = "method returns a new vector and does not mutate the original inputs"]
196 pub fn reverse(self) -> Self {
197 struct Reverse;
198
199 impl<const N: usize> Swizzle<N> for Reverse {
200 const INDEX: [usize; N] = const {
201 let mut index = [0; N];
202 let mut i = 0;
203 while i < N {
204 index[i] = N - i - 1;
205 i += 1;
206 }
207 index
208 };
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 `self.len() - OFFSET` elements move to the front. After calling `rotate_elements_left`,
216 /// the element previously at index `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_elements_left<const OFFSET: usize>(self) -> Self {
220 struct Rotate<const OFFSET: usize>;
221
222 impl<const OFFSET: usize, const N: usize> Swizzle<N> for Rotate<OFFSET> {
223 const INDEX: [usize; N] = const {
224 let offset = OFFSET % N;
225 let mut index = [0; N];
226 let mut i = 0;
227 while i < N {
228 index[i] = (i + offset) % N;
229 i += 1;
230 }
231 index
232 };
233 }
234
235 Rotate::<OFFSET>::swizzle(self)
236 }
237
238 /// Rotates the vector such that the first `self.len() - OFFSET` elements of the vector move to
239 /// the end while the last `OFFSET` elements move to the front. After calling `rotate_elements_right`,
240 /// the element previously at index `self.len() - OFFSET` will become the first element in the slice.
241 #[inline]
242 #[must_use = "method returns a new vector and does not mutate the original inputs"]
243 pub fn rotate_elements_right<const OFFSET: usize>(self) -> Self {
244 struct Rotate<const OFFSET: usize>;
245
246 impl<const OFFSET: usize, const N: usize> Swizzle<N> for Rotate<OFFSET> {
247 const INDEX: [usize; N] = const {
248 let offset = N - OFFSET % N;
249 let mut index = [0; N];
250 let mut i = 0;
251 while i < N {
252 index[i] = (i + offset) % N;
253 i += 1;
254 }
255 index
256 };
257 }
258
259 Rotate::<OFFSET>::swizzle(self)
260 }
261
262 /// Shifts the vector elements to the left by `OFFSET`, filling in with
263 /// `padding` from the right.
264 #[inline]
265 #[must_use = "method returns a new vector and does not mutate the original inputs"]
266 pub fn shift_elements_left<const OFFSET: usize>(self, padding: T) -> Self {
267 struct Shift<const OFFSET: usize>;
268
269 impl<const OFFSET: usize, const N: usize> Swizzle<N> for Shift<OFFSET> {
270 const INDEX: [usize; N] = const {
271 let mut index = [N; N];
272 let mut i = 0;
273 while i + OFFSET < N {
274 index[i] = i + OFFSET;
275 i += 1;
276 }
277 index
278 };
279 }
280
281 Shift::<OFFSET>::concat_swizzle(self, Simd::splat(padding))
282 }
283
284 /// Shifts the vector elements to the right by `OFFSET`, filling in with
285 /// `padding` from the left.
286 #[inline]
287 #[must_use = "method returns a new vector and does not mutate the original inputs"]
288 pub fn shift_elements_right<const OFFSET: usize>(self, padding: T) -> Self {
289 struct Shift<const OFFSET: usize>;
290
291 impl<const OFFSET: usize, const N: usize> Swizzle<N> for Shift<OFFSET> {
292 const INDEX: [usize; N] = const {
293 let mut index = [N; N];
294 let mut i = OFFSET;
295 while i < N {
296 index[i] = i - OFFSET;
297 i += 1;
298 }
299 index
300 };
301 }
302
303 Shift::<OFFSET>::concat_swizzle(self, Simd::splat(padding))
304 }
305
306 /// Interleave two vectors.
307 ///
308 /// The resulting vectors contain elements taken alternatively from `self` and `other`, first
309 /// filling the first result, and then the second.
310 ///
311 /// The reverse of this operation is [`Simd::deinterleave`].
312 ///
313 /// ```
314 /// # #![feature(portable_simd)]
315 /// # use core::simd::Simd;
316 /// let a = Simd::from_array([0, 1, 2, 3]);
317 /// let b = Simd::from_array([4, 5, 6, 7]);
318 /// let (x, y) = a.interleave(b);
319 /// assert_eq!(x.to_array(), [0, 4, 1, 5]);
320 /// assert_eq!(y.to_array(), [2, 6, 3, 7]);
321 /// ```
322 #[inline]
323 #[must_use = "method returns a new vector and does not mutate the original inputs"]
324 pub fn interleave(self, other: Self) -> (Self, Self) {
325 const fn interleave<const N: usize>(high: bool) -> [usize; N] {
326 let mut idx = [0; N];
327 let mut i = 0;
328 while i < N {
329 let dst_index = if high { i + N } else { i };
330 let src_index = dst_index / 2 + (dst_index % 2) * N;
331 idx[i] = src_index;
332 i += 1;
333 }
334 idx
335 }
336
337 struct Lo;
338 struct Hi;
339
340 impl<const N: usize> Swizzle<N> for Lo {
341 const INDEX: [usize; N] = interleave::<N>(false);
342 }
343
344 impl<const N: usize> Swizzle<N> for Hi {
345 const INDEX: [usize; N] = interleave::<N>(true);
346 }
347
348 (
349 Lo::concat_swizzle(self, other),
350 Hi::concat_swizzle(self, other),
351 )
352 }
353
354 /// Deinterleave two vectors.
355 ///
356 /// The first result takes every other element of `self` and then `other`, starting with
357 /// the first element.
358 ///
359 /// The second result takes every other element of `self` and then `other`, starting with
360 /// the second element.
361 ///
362 /// The reverse of this operation is [`Simd::interleave`].
363 ///
364 /// ```
365 /// # #![feature(portable_simd)]
366 /// # #[cfg(feature = "as_crate")] use core_simd::simd;
367 /// # #[cfg(not(feature = "as_crate"))] use core::simd;
368 /// # use simd::Simd;
369 /// let a = Simd::from_array([0, 4, 1, 5]);
370 /// let b = Simd::from_array([2, 6, 3, 7]);
371 /// let (x, y) = a.deinterleave(b);
372 /// assert_eq!(x.to_array(), [0, 1, 2, 3]);
373 /// assert_eq!(y.to_array(), [4, 5, 6, 7]);
374 /// ```
375 #[inline]
376 #[must_use = "method returns a new vector and does not mutate the original inputs"]
377 pub fn deinterleave(self, other: Self) -> (Self, Self) {
378 const fn deinterleave<const N: usize>(second: bool) -> [usize; N] {
379 let mut idx = [0; N];
380 let mut i = 0;
381 while i < N {
382 idx[i] = i * 2 + second as usize;
383 i += 1;
384 }
385 idx
386 }
387
388 struct Even;
389 struct Odd;
390
391 impl<const N: usize> Swizzle<N> for Even {
392 const INDEX: [usize; N] = deinterleave::<N>(false);
393 }
394
395 impl<const N: usize> Swizzle<N> for Odd {
396 const INDEX: [usize; N] = deinterleave::<N>(true);
397 }
398
399 (
400 Even::concat_swizzle(self, other),
401 Odd::concat_swizzle(self, other),
402 )
403 }
404
405 /// Resize a vector.
406 ///
407 /// If `M` > `N`, extends the length of a vector, setting the new elements to `value`.
408 /// If `M` < `N`, truncates the vector to the first `M` elements.
409 ///
410 /// ```
411 /// # #![feature(portable_simd)]
412 /// # #[cfg(feature = "as_crate")] use core_simd::simd;
413 /// # #[cfg(not(feature = "as_crate"))] use core::simd;
414 /// # use simd::u32x4;
415 /// let x = u32x4::from_array([0, 1, 2, 3]);
416 /// assert_eq!(x.resize::<8>(9).to_array(), [0, 1, 2, 3, 9, 9, 9, 9]);
417 /// assert_eq!(x.resize::<2>(9).to_array(), [0, 1]);
418 /// ```
419 #[inline]
420 #[must_use = "method returns a new vector and does not mutate the original inputs"]
421 pub fn resize<const M: usize>(self, value: T) -> Simd<T, M>
422 where
423 LaneCount<M>: SupportedLaneCount,
424 {
425 struct Resize<const N: usize>;
426 impl<const N: usize, const M: usize> Swizzle<M> for Resize<N> {
427 const INDEX: [usize; M] = const {
428 let mut index = [0; M];
429 let mut i = 0;
430 while i < M {
431 index[i] = if i < N { i } else { N };
432 i += 1;
433 }
434 index
435 };
436 }
437 Resize::<N>::concat_swizzle(self, Simd::splat(value))
438 }
439
440 /// Extract a vector from another vector.
441 ///
442 /// ```
443 /// # #![feature(portable_simd)]
444 /// # #[cfg(feature = "as_crate")] use core_simd::simd;
445 /// # #[cfg(not(feature = "as_crate"))] use core::simd;
446 /// # use simd::u32x4;
447 /// let x = u32x4::from_array([0, 1, 2, 3]);
448 /// assert_eq!(x.extract::<1, 2>().to_array(), [1, 2]);
449 /// ```
450 #[inline]
451 #[must_use = "method returns a new vector and does not mutate the original inputs"]
452 pub fn extract<const START: usize, const LEN: usize>(self) -> Simd<T, LEN>
453 where
454 LaneCount<LEN>: SupportedLaneCount,
455 {
456 struct Extract<const N: usize, const START: usize>;
457 impl<const N: usize, const START: usize, const LEN: usize> Swizzle<LEN> for Extract<N, START> {
458 const INDEX: [usize; LEN] = const {
459 assert!(START + LEN <= N, "index out of bounds");
460 let mut index = [0; LEN];
461 let mut i = 0;
462 while i < LEN {
463 index[i] = START + i;
464 i += 1;
465 }
466 index
467 };
468 }
469 Extract::<N, START>::swizzle(self)
470 }
471}
472
473impl<T, const N: usize> Mask<T, N>
474where
475 T: MaskElement,
476 LaneCount<N>: SupportedLaneCount,
477{
478 /// Reverse the order of the elements in the mask.
479 #[inline]
480 #[must_use = "method returns a new vector and does not mutate the original inputs"]
481 pub fn reverse(self) -> Self {
482 // Safety: swizzles are safe for masks
483 unsafe { Self::from_int_unchecked(self.to_int().reverse()) }
484 }
485
486 /// Rotates the mask such that the first `OFFSET` elements of the slice move to the end
487 /// while the last `self.len() - OFFSET` elements move to the front. After calling `rotate_elements_left`,
488 /// the element previously at index `OFFSET` will become the first element in the slice.
489 #[inline]
490 #[must_use = "method returns a new vector and does not mutate the original inputs"]
491 pub fn rotate_elements_left<const OFFSET: usize>(self) -> Self {
492 // Safety: swizzles are safe for masks
493 unsafe { Self::from_int_unchecked(self.to_int().rotate_elements_left::<OFFSET>()) }
494 }
495
496 /// Rotates the mask such that the first `self.len() - OFFSET` elements of the mask move to
497 /// the end while the last `OFFSET` elements move to the front. After calling `rotate_elements_right`,
498 /// the element previously at index `self.len() - OFFSET` will become the first element in the slice.
499 #[inline]
500 #[must_use = "method returns a new vector and does not mutate the original inputs"]
501 pub fn rotate_elements_right<const OFFSET: usize>(self) -> Self {
502 // Safety: swizzles are safe for masks
503 unsafe { Self::from_int_unchecked(self.to_int().rotate_elements_right::<OFFSET>()) }
504 }
505
506 /// Shifts the mask elements to the left by `OFFSET`, filling in with
507 /// `padding` from the right.
508 #[inline]
509 #[must_use = "method returns a new mask and does not mutate the original inputs"]
510 pub fn shift_elements_left<const OFFSET: usize>(self, padding: bool) -> Self {
511 // Safety: swizzles are safe for masks
512 unsafe {
513 Self::from_int_unchecked(self.to_int().shift_elements_left::<OFFSET>(if padding {
514 T::TRUE
515 } else {
516 T::FALSE
517 }))
518 }
519 }
520
521 /// Shifts the mask elements to the right by `OFFSET`, filling in with
522 /// `padding` from the left.
523 #[inline]
524 #[must_use = "method returns a new mask and does not mutate the original inputs"]
525 pub fn shift_elements_right<const OFFSET: usize>(self, padding: bool) -> Self {
526 // Safety: swizzles are safe for masks
527 unsafe {
528 Self::from_int_unchecked(self.to_int().shift_elements_right::<OFFSET>(if padding {
529 T::TRUE
530 } else {
531 T::FALSE
532 }))
533 }
534 }
535
536 /// Interleave two masks.
537 ///
538 /// The resulting masks contain elements taken alternatively from `self` and `other`, first
539 /// filling the first result, and then the second.
540 ///
541 /// The reverse of this operation is [`Mask::deinterleave`].
542 ///
543 /// ```
544 /// # #![feature(portable_simd)]
545 /// # #[cfg(feature = "as_crate")] use core_simd::simd;
546 /// # #[cfg(not(feature = "as_crate"))] use core::simd;
547 /// # use simd::mask32x4;
548 /// let a = mask32x4::from_array([false, true, false, true]);
549 /// let b = mask32x4::from_array([false, false, true, true]);
550 /// let (x, y) = a.interleave(b);
551 /// assert_eq!(x.to_array(), [false, false, true, false]);
552 /// assert_eq!(y.to_array(), [false, true, true, true]);
553 /// ```
554 #[inline]
555 #[must_use = "method returns a new vector and does not mutate the original inputs"]
556 pub fn interleave(self, other: Self) -> (Self, Self) {
557 let (lo, hi) = self.to_int().interleave(other.to_int());
558 // Safety: swizzles are safe for masks
559 unsafe { (Self::from_int_unchecked(lo), Self::from_int_unchecked(hi)) }
560 }
561
562 /// Deinterleave two masks.
563 ///
564 /// The first result takes every other element of `self` and then `other`, starting with
565 /// the first element.
566 ///
567 /// The second result takes every other element of `self` and then `other`, starting with
568 /// the second element.
569 ///
570 /// The reverse of this operation is [`Mask::interleave`].
571 ///
572 /// ```
573 /// # #![feature(portable_simd)]
574 /// # #[cfg(feature = "as_crate")] use core_simd::simd;
575 /// # #[cfg(not(feature = "as_crate"))] use core::simd;
576 /// # use simd::mask32x4;
577 /// let a = mask32x4::from_array([false, true, false, true]);
578 /// let b = mask32x4::from_array([false, false, true, true]);
579 /// let (x, y) = a.deinterleave(b);
580 /// assert_eq!(x.to_array(), [false, false, false, true]);
581 /// assert_eq!(y.to_array(), [true, true, false, true]);
582 /// ```
583 #[inline]
584 #[must_use = "method returns a new vector and does not mutate the original inputs"]
585 pub fn deinterleave(self, other: Self) -> (Self, Self) {
586 let (even, odd) = self.to_int().deinterleave(other.to_int());
587 // Safety: swizzles are safe for masks
588 unsafe {
589 (
590 Self::from_int_unchecked(even),
591 Self::from_int_unchecked(odd),
592 )
593 }
594 }
595
596 /// Resize a mask.
597 ///
598 /// If `M` > `N`, extends the length of a mask, setting the new elements to `value`.
599 /// If `M` < `N`, truncates the mask to the first `M` elements.
600 ///
601 /// ```
602 /// # #![feature(portable_simd)]
603 /// # #[cfg(feature = "as_crate")] use core_simd::simd;
604 /// # #[cfg(not(feature = "as_crate"))] use core::simd;
605 /// # use simd::mask32x4;
606 /// let x = mask32x4::from_array([false, true, true, false]);
607 /// assert_eq!(x.resize::<8>(true).to_array(), [false, true, true, false, true, true, true, true]);
608 /// assert_eq!(x.resize::<2>(true).to_array(), [false, true]);
609 /// ```
610 #[inline]
611 #[must_use = "method returns a new vector and does not mutate the original inputs"]
612 pub fn resize<const M: usize>(self, value: bool) -> Mask<T, M>
613 where
614 LaneCount<M>: SupportedLaneCount,
615 {
616 // Safety: swizzles are safe for masks
617 unsafe {
618 Mask::<T, M>::from_int_unchecked(self.to_int().resize::<M>(if value {
619 T::TRUE
620 } else {
621 T::FALSE
622 }))
623 }
624 }
625
626 /// Extract a vector from another vector.
627 ///
628 /// ```
629 /// # #![feature(portable_simd)]
630 /// # #[cfg(feature = "as_crate")] use core_simd::simd;
631 /// # #[cfg(not(feature = "as_crate"))] use core::simd;
632 /// # use simd::mask32x4;
633 /// let x = mask32x4::from_array([false, true, true, false]);
634 /// assert_eq!(x.extract::<1, 2>().to_array(), [true, true]);
635 /// ```
636 #[inline]
637 #[must_use = "method returns a new vector and does not mutate the original inputs"]
638 pub fn extract<const START: usize, const LEN: usize>(self) -> Mask<T, LEN>
639 where
640 LaneCount<LEN>: SupportedLaneCount,
641 {
642 // Safety: swizzles are safe for masks
643 unsafe { Mask::<T, LEN>::from_int_unchecked(self.to_int().extract::<START, LEN>()) }
644 }
645}