1 | use 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)] |
45 | pub 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 | /// Create a vector from the elements of another vector. |
73 | pub 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 | /// Create 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 array 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 | output |
107 | }, |
108 | ) |
109 | } |
110 | } |
111 | |
112 | /// Create a new vector from the elements of `first` and `second`. |
113 | /// |
114 | /// Lane `i` of the output is `concat[Self::INDEX[i]]`, where `concat` is the concatenation of |
115 | /// `first` and `second`. |
116 | #[inline ] |
117 | #[must_use = "method returns a new vector and does not mutate the original inputs" ] |
118 | fn concat_swizzle<T, const M: usize>(first: Simd<T, M>, second: Simd<T, M>) -> Simd<T, N> |
119 | where |
120 | T: SimdElement, |
121 | LaneCount<N>: SupportedLaneCount, |
122 | LaneCount<M>: SupportedLaneCount, |
123 | { |
124 | // Safety: `first` and `second` are vectors, and the index is a const array of u32. |
125 | unsafe { |
126 | core::intrinsics::simd::simd_shuffle( |
127 | first, |
128 | second, |
129 | const { |
130 | let mut output = [0; N]; |
131 | let mut i = 0; |
132 | while i < N { |
133 | let index = Self::INDEX[i]; |
134 | assert!(index as u32 as usize == index); |
135 | assert!( |
136 | index < 2 * M, |
137 | "source element index exceeds input vector length" |
138 | ); |
139 | output[i] = index as u32; |
140 | i += 1; |
141 | } |
142 | output |
143 | }, |
144 | ) |
145 | } |
146 | } |
147 | |
148 | /// Create a new mask from the elements of `mask`. |
149 | /// |
150 | /// Element `i` of the output is `concat[Self::INDEX[i]]`, where `concat` is the concatenation of |
151 | /// `first` and `second`. |
152 | #[inline ] |
153 | #[must_use = "method returns a new mask and does not mutate the original inputs" ] |
154 | fn swizzle_mask<T, const M: usize>(mask: Mask<T, M>) -> Mask<T, N> |
155 | where |
156 | T: MaskElement, |
157 | LaneCount<N>: SupportedLaneCount, |
158 | LaneCount<M>: SupportedLaneCount, |
159 | { |
160 | // SAFETY: all elements of this mask come from another mask |
161 | unsafe { Mask::from_int_unchecked(Self::swizzle(mask.to_int())) } |
162 | } |
163 | |
164 | /// Create a new mask from the elements of `first` and `second`. |
165 | /// |
166 | /// Element `i` of the output is `concat[Self::INDEX[i]]`, where `concat` is the concatenation of |
167 | /// `first` and `second`. |
168 | #[inline ] |
169 | #[must_use = "method returns a new mask and does not mutate the original inputs" ] |
170 | fn concat_swizzle_mask<T, const M: usize>(first: Mask<T, M>, second: Mask<T, M>) -> Mask<T, N> |
171 | where |
172 | T: MaskElement, |
173 | LaneCount<N>: SupportedLaneCount, |
174 | LaneCount<M>: SupportedLaneCount, |
175 | { |
176 | // SAFETY: all elements of this mask come from another mask |
177 | unsafe { Mask::from_int_unchecked(Self::concat_swizzle(first.to_int(), second.to_int())) } |
178 | } |
179 | } |
180 | |
181 | impl<T, const N: usize> Simd<T, N> |
182 | where |
183 | T: SimdElement, |
184 | LaneCount<N>: SupportedLaneCount, |
185 | { |
186 | /// Reverse the order of the elements in the vector. |
187 | #[inline ] |
188 | #[must_use = "method returns a new vector and does not mutate the original inputs" ] |
189 | pub fn reverse(self) -> Self { |
190 | struct Reverse; |
191 | |
192 | impl<const N: usize> Swizzle<N> for Reverse { |
193 | const INDEX: [usize; N] = const { |
194 | let mut index = [0; N]; |
195 | let mut i = 0; |
196 | while i < N { |
197 | index[i] = N - i - 1; |
198 | i += 1; |
199 | } |
200 | index |
201 | }; |
202 | } |
203 | |
204 | Reverse::swizzle(self) |
205 | } |
206 | |
207 | /// Rotates the vector such that the first `OFFSET` elements of the slice move to the end |
208 | /// while the last `self.len() - OFFSET` elements move to the front. After calling `rotate_elements_left`, |
209 | /// the element previously at index `OFFSET` will become the first element in the slice. |
210 | #[inline ] |
211 | #[must_use = "method returns a new vector and does not mutate the original inputs" ] |
212 | pub fn rotate_elements_left<const OFFSET: usize>(self) -> Self { |
213 | struct Rotate<const OFFSET: usize>; |
214 | |
215 | impl<const OFFSET: usize, const N: usize> Swizzle<N> for Rotate<OFFSET> { |
216 | const INDEX: [usize; N] = const { |
217 | let offset = OFFSET % N; |
218 | let mut index = [0; N]; |
219 | let mut i = 0; |
220 | while i < N { |
221 | index[i] = (i + offset) % N; |
222 | i += 1; |
223 | } |
224 | index |
225 | }; |
226 | } |
227 | |
228 | Rotate::<OFFSET>::swizzle(self) |
229 | } |
230 | |
231 | /// Rotates the vector such that the first `self.len() - OFFSET` elements of the vector move to |
232 | /// the end while the last `OFFSET` elements move to the front. After calling `rotate_elements_right`, |
233 | /// the element previously at index `self.len() - OFFSET` will become the first element in the slice. |
234 | #[inline ] |
235 | #[must_use = "method returns a new vector and does not mutate the original inputs" ] |
236 | pub fn rotate_elements_right<const OFFSET: usize>(self) -> Self { |
237 | struct Rotate<const OFFSET: usize>; |
238 | |
239 | impl<const OFFSET: usize, const N: usize> Swizzle<N> for Rotate<OFFSET> { |
240 | const INDEX: [usize; N] = const { |
241 | let offset = N - OFFSET % N; |
242 | let mut index = [0; N]; |
243 | let mut i = 0; |
244 | while i < N { |
245 | index[i] = (i + offset) % N; |
246 | i += 1; |
247 | } |
248 | index |
249 | }; |
250 | } |
251 | |
252 | Rotate::<OFFSET>::swizzle(self) |
253 | } |
254 | |
255 | /// Interleave two vectors. |
256 | /// |
257 | /// The resulting vectors contain elements taken alternatively from `self` and `other`, first |
258 | /// filling the first result, and then the second. |
259 | /// |
260 | /// The reverse of this operation is [`Simd::deinterleave`]. |
261 | /// |
262 | /// ``` |
263 | /// # #![feature (portable_simd)] |
264 | /// # use core::simd::Simd; |
265 | /// let a = Simd::from_array([0, 1, 2, 3]); |
266 | /// let b = Simd::from_array([4, 5, 6, 7]); |
267 | /// let (x, y) = a.interleave(b); |
268 | /// assert_eq!(x.to_array(), [0, 4, 1, 5]); |
269 | /// assert_eq!(y.to_array(), [2, 6, 3, 7]); |
270 | /// ``` |
271 | #[inline ] |
272 | #[must_use = "method returns a new vector and does not mutate the original inputs" ] |
273 | pub fn interleave(self, other: Self) -> (Self, Self) { |
274 | const fn interleave<const N: usize>(high: bool) -> [usize; N] { |
275 | let mut idx = [0; N]; |
276 | let mut i = 0; |
277 | while i < N { |
278 | let dst_index = if high { i + N } else { i }; |
279 | let src_index = dst_index / 2 + (dst_index % 2) * N; |
280 | idx[i] = src_index; |
281 | i += 1; |
282 | } |
283 | idx |
284 | } |
285 | |
286 | struct Lo; |
287 | struct Hi; |
288 | |
289 | impl<const N: usize> Swizzle<N> for Lo { |
290 | const INDEX: [usize; N] = interleave::<N>(false); |
291 | } |
292 | |
293 | impl<const N: usize> Swizzle<N> for Hi { |
294 | const INDEX: [usize; N] = interleave::<N>(true); |
295 | } |
296 | |
297 | ( |
298 | Lo::concat_swizzle(self, other), |
299 | Hi::concat_swizzle(self, other), |
300 | ) |
301 | } |
302 | |
303 | /// Deinterleave two vectors. |
304 | /// |
305 | /// The first result takes every other element of `self` and then `other`, starting with |
306 | /// the first element. |
307 | /// |
308 | /// The second result takes every other element of `self` and then `other`, starting with |
309 | /// the second element. |
310 | /// |
311 | /// The reverse of this operation is [`Simd::interleave`]. |
312 | /// |
313 | /// ``` |
314 | /// # #![feature (portable_simd)] |
315 | /// # use core::simd::Simd; |
316 | /// let a = Simd::from_array([0, 4, 1, 5]); |
317 | /// let b = Simd::from_array([2, 6, 3, 7]); |
318 | /// let (x, y) = a.deinterleave(b); |
319 | /// assert_eq!(x.to_array(), [0, 1, 2, 3]); |
320 | /// assert_eq!(y.to_array(), [4, 5, 6, 7]); |
321 | /// ``` |
322 | #[inline ] |
323 | #[must_use = "method returns a new vector and does not mutate the original inputs" ] |
324 | pub fn deinterleave(self, other: Self) -> (Self, Self) { |
325 | const fn deinterleave<const N: usize>(second: bool) -> [usize; N] { |
326 | let mut idx = [0; N]; |
327 | let mut i = 0; |
328 | while i < N { |
329 | idx[i] = i * 2 + second as usize; |
330 | i += 1; |
331 | } |
332 | idx |
333 | } |
334 | |
335 | struct Even; |
336 | struct Odd; |
337 | |
338 | impl<const N: usize> Swizzle<N> for Even { |
339 | const INDEX: [usize; N] = deinterleave::<N>(false); |
340 | } |
341 | |
342 | impl<const N: usize> Swizzle<N> for Odd { |
343 | const INDEX: [usize; N] = deinterleave::<N>(true); |
344 | } |
345 | |
346 | ( |
347 | Even::concat_swizzle(self, other), |
348 | Odd::concat_swizzle(self, other), |
349 | ) |
350 | } |
351 | |
352 | /// Resize a vector. |
353 | /// |
354 | /// If `M` > `N`, extends the length of a vector, setting the new elements to `value`. |
355 | /// If `M` < `N`, truncates the vector to the first `M` elements. |
356 | /// |
357 | /// ``` |
358 | /// # #![feature (portable_simd)] |
359 | /// # #[cfg (feature = "as_crate" )] use core_simd::simd; |
360 | /// # #[cfg (not(feature = "as_crate" ))] use core::simd; |
361 | /// # use simd::u32x4; |
362 | /// let x = u32x4::from_array([0, 1, 2, 3]); |
363 | /// assert_eq!(x.resize::<8>(9).to_array(), [0, 1, 2, 3, 9, 9, 9, 9]); |
364 | /// assert_eq!(x.resize::<2>(9).to_array(), [0, 1]); |
365 | /// ``` |
366 | #[inline ] |
367 | #[must_use = "method returns a new vector and does not mutate the original inputs" ] |
368 | pub fn resize<const M: usize>(self, value: T) -> Simd<T, M> |
369 | where |
370 | LaneCount<M>: SupportedLaneCount, |
371 | { |
372 | struct Resize<const N: usize>; |
373 | impl<const N: usize, const M: usize> Swizzle<M> for Resize<N> { |
374 | const INDEX: [usize; M] = const { |
375 | let mut index = [0; M]; |
376 | let mut i = 0; |
377 | while i < M { |
378 | index[i] = if i < N { i } else { N }; |
379 | i += 1; |
380 | } |
381 | index |
382 | }; |
383 | } |
384 | Resize::<N>::concat_swizzle(self, Simd::splat(value)) |
385 | } |
386 | } |
387 | |