Skip to main content

rten_simd/
ops.rs

1//! Traits for operations on SIMD vectors.
2//!
3//! The entry point is the [`Isa`] trait, an implementation of which is passed
4//! to SIMD operations when evaluated. This has methods for each of the
5//! supported element types which returns the implementation of operations on
6//! SIMD vectors with that element type.
7//!
8//! The [`BitOps`] trait provides operations available on all element types
9//! which only require treating a value as a sequence of bits (load, store,
10//! splat, bitwise ops, select etc.). The [`NumOps`] sub-trait adds arithmetic
11//! operations (add, subtract, multiply, comparison) which require interpreting
12//! the bits as numbers. The sub-traits [`FloatOps`] and [`SignedIntOps`]
13//! provide additional operations on float and signed integer element types.
14//! Additionally there are traits for individual operations such as [`Extend`]
15//! or [`NarrowSaturate`] which are available on a subset of element types.
16
17use std::mem::MaybeUninit;
18
19use crate::elem::Elem;
20use crate::f16;
21use crate::simd::{Mask, Simd};
22
23/// Entry point for performing SIMD operations using a particular Instruction
24/// Set Architecture (ISA).
25///
26/// Implementations of this trait are types which can only be instantiated
27/// if the instruction set is available. They are usually zero-sized and thus
28/// free to copy.
29///
30/// # Safety
31///
32/// Implementations must ensure they can only be constructed if the
33/// instruction set is supported on the current system.
34pub unsafe trait Isa: Copy {
35    /// SIMD vector with an unspecified element type. This is used for
36    /// bitwise casting between different vector types.
37    type Bits: Simd;
38
39    /// Mask vector for 32-bit lanes.
40    type M32: Mask;
41
42    /// Mask vector for 16-bit lanes.
43    type M16: Mask;
44
45    /// Mask vector for 8-bit lanes.
46    type M8: Mask;
47
48    /// SIMD vector with `f32` elements.
49    type F32: Simd<Elem = f32, Isa = Self, Mask = Self::M32>;
50
51    /// SIMD vector with `i32` elements.
52    type I32: Simd<Elem = i32, Isa = Self, Mask = Self::M32>;
53
54    /// SIMD vector with `i16` elements.
55    type I16: Simd<Elem = i16, Isa = Self, Mask = Self::M16>;
56
57    /// SIMD vector with `i8` elements.
58    type I8: Simd<Elem = i8, Isa = Self, Mask = Self::M8>;
59
60    /// SIMD vector with `u8` elements.
61    type U8: Simd<Elem = u8, Isa = Self, Mask = Self::M8>;
62
63    /// SIMD vector with `u16` elements.
64    type U16: Simd<Elem = u16, Isa = Self, Mask = Self::M16>;
65
66    /// SIMD vector with `u32` elements.
67    type U32: Simd<Elem = u32, Isa = Self, Mask = Self::M32>;
68
69    /// SIMD vector with `f16` elements.
70    type F16: Simd<Elem = f16, Isa = Self, Mask = Self::M16>;
71
72    /// Operations on SIMD vectors with `f32` elements.
73    fn f32(
74        self,
75    ) -> impl FloatOps<f32, Simd = Self::F32, Int = Self::I32>
76    + NarrowSaturate<f32, f16, Output = Self::F16>;
77
78    /// Operations on SIMD vectors with `f16` elements.
79    ///
80    /// Only bit-level operations ([`BitOps`]) plus conversion to `f32` (via
81    /// [`Extend`]) are supported. Conversion from `f32` is available via the
82    /// [`NarrowSaturate`] implementation returned by [`f32`](Isa::f32).
83    fn f16(self) -> impl Extend<f16, Output = Self::F32, Simd = Self::F16>;
84
85    /// Operations on SIMD vectors with `i32` elements.
86    fn i32(
87        self,
88    ) -> impl SignedIntOps<i32, Simd = Self::I32>
89    + NarrowSaturate<i32, i16, Output = Self::I16>
90    + Concat<i32>
91    + ToFloat<i32, Output = Self::F32>;
92
93    /// Operations on SIMD vectors with `i16` elements.
94    fn i16(
95        self,
96    ) -> impl SignedIntOps<i16, Simd = Self::I16>
97    + NarrowSaturate<i16, u8, Output = Self::U8>
98    + Extend<i16, Output = Self::I32>
99    + Interleave<i16>;
100
101    /// Operations on SIMD vectors with `i8` elements.
102    fn i8(
103        self,
104    ) -> impl SignedIntOps<i8, Simd = Self::I8> + Extend<i8, Output = Self::I16> + Interleave<i8>;
105
106    /// Operations on SIMD vectors with `u8` elements.
107    fn u8(
108        self,
109    ) -> impl IntOps<u8, Simd = Self::U8> + Extend<u8, Output = Self::U16> + Interleave<u8>;
110
111    /// Operations on SIMD vectors with `u16` elements.
112    fn u16(self) -> impl IntOps<u16, Simd = Self::U16>;
113
114    /// Operations on mask vectors for 32-bit lanes.
115    fn m32(self) -> impl MaskOps<Self::M32>;
116
117    /// Operations on mask vectors for 16-bit lanes.
118    fn m16(self) -> impl MaskOps<Self::M16>;
119
120    /// Operations on mask vectors for 8-bit lanes.
121    fn m8(self) -> impl MaskOps<Self::M8>;
122}
123
124/// Get the [`NumOps`] implementation from an [`Isa`] for a given element type.
125///
126/// This trait is useful for writing SIMD operations which are generic over the
127/// element type. It is implemented for all of the element types supported in
128/// SIMD vectors.
129///
130/// # Example
131///
132/// This example shows how to use [`GetNumOps`] to write a vectorized `Sum`
133/// operation.
134///
135/// ```
136/// use rten_simd::{Isa, SimdIterable, SimdOp};
137/// use rten_simd::ops::{BitOps, GetNumOps, NumOps};
138///
139/// struct Sum<'a, T>(&'a [T]);
140///
141/// impl<T: std::ops::Add<Output=T> + GetNumOps> SimdOp for Sum<'_, T> {
142///   type Output = T;
143///   
144///   #[inline(always)]
145///   fn eval<I: Isa>(self, isa: I) -> Self::Output {
146///     let ops = T::num_ops(isa);
147///
148///     // Build `ops.len()` partial sums in parallel. If the slice length is
149///     // not a multiple of `ops.len()` it will be padded with zeros.
150///     let mut sum = ops.zero();
151///     for chunk in self.0.simd_iter_pad(ops) {
152///         sum = ops.add(sum, chunk);
153///     }
154///
155///     // Horizontally reduce the SIMD vector containing partial sums to a
156///     // single value.
157///     ops.sum(sum)
158///   }
159/// }
160///
161/// let vals: Vec<_> = (1..20i32).collect();
162/// let sum = Sum(&vals).dispatch();
163/// assert_eq!(sum, vals.iter().sum());
164/// ```
165pub trait GetNumOps
166where
167    Self: GetSimd + 'static,
168{
169    /// Return the [`NumOps`] implementation from a SIMD [`Isa`] that provides
170    /// operations on vectors containing elements of type `Self`.
171    fn num_ops<I: Isa>(isa: I) -> impl NumOps<Self, Simd = Self::Simd<I>>;
172}
173
174macro_rules! impl_get_ops {
175    ($trait:ty, $method:ident, $ops:ident, $type:ident) => {
176        impl $trait for $type {
177            fn $method<I: Isa>(isa: I) -> impl $ops<Self, Simd = <Self as GetSimd>::Simd<I>> {
178                isa.$type()
179            }
180        }
181    };
182}
183impl_get_ops!(GetNumOps, num_ops, NumOps, f32);
184impl_get_ops!(GetNumOps, num_ops, NumOps, i16);
185impl_get_ops!(GetNumOps, num_ops, NumOps, i32);
186impl_get_ops!(GetNumOps, num_ops, NumOps, i8);
187impl_get_ops!(GetNumOps, num_ops, NumOps, u16);
188impl_get_ops!(GetNumOps, num_ops, NumOps, u8);
189
190/// Get the [`BitOps`] implementation from an [`Isa`] for a given element type.
191///
192/// This is the bit-level counterpart of [`GetNumOps`]. It is implemented for
193/// all element types which support bit-level SIMD operations, which is a
194/// superset of the types that support arithmetic operations.
195pub trait GetBitOps
196where
197    Self: GetSimd + 'static,
198{
199    /// Return the [`BitOps`] implementation from a SIMD [`Isa`] that provides
200    /// operations on vectors containing elements of type `Self`.
201    fn bit_ops<I: Isa>(isa: I) -> impl BitOps<Self, Simd = Self::Simd<I>>;
202}
203impl_get_ops!(GetBitOps, bit_ops, BitOps, f16);
204impl_get_ops!(GetBitOps, bit_ops, BitOps, f32);
205impl_get_ops!(GetBitOps, bit_ops, BitOps, i16);
206impl_get_ops!(GetBitOps, bit_ops, BitOps, i32);
207impl_get_ops!(GetBitOps, bit_ops, BitOps, i8);
208impl_get_ops!(GetBitOps, bit_ops, BitOps, u16);
209impl_get_ops!(GetBitOps, bit_ops, BitOps, u8);
210
211/// Get the [`Simd`] implementation from an [`Isa`] for a given element type.
212///
213/// For example the type `<f32 as GetSimd>::Simd<I>` yields `I::F32` where
214/// `I` is an `Isa`. This trait is used for example by
215/// [`SimdUnaryOp`](crate::SimdUnaryOp) to determine the type of SIMD vector
216/// that corresponds to the element type.
217pub trait GetSimd: Elem {
218    type Simd<I: Isa>: Simd<Elem = Self, Isa = I>;
219}
220
221macro_rules! impl_getsimd {
222    ($ty:ty, $simd:ident) => {
223        impl GetSimd for $ty {
224            type Simd<I: Isa> = I::$simd;
225        }
226    };
227}
228impl_getsimd!(f16, F16);
229impl_getsimd!(f32, F32);
230impl_getsimd!(i16, I16);
231impl_getsimd!(i32, I32);
232impl_getsimd!(i8, I8);
233impl_getsimd!(u16, U16);
234impl_getsimd!(u8, U8);
235
236/// Get the [`FloatOps`] implementation from an [`Isa`] for a given element type.
237///
238/// This is a specialization of [`GetNumOps`] for float element types.
239pub trait GetFloatOps
240where
241    Self: GetSimd,
242{
243    fn float_ops<I: Isa>(isa: I) -> impl FloatOps<Self, Simd = Self::Simd<I>>;
244}
245impl_get_ops!(GetFloatOps, float_ops, FloatOps, f32);
246
247/// Get the [`IntOps`] implementation from an [`Isa`] for a given element type.
248///
249/// This is a specialization of [`GetNumOps`] for signed integer element types.
250pub trait GetIntOps
251where
252    Self: GetSimd,
253{
254    fn int_ops<I: Isa>(isa: I) -> impl IntOps<Self, Simd = Self::Simd<I>>;
255}
256impl_get_ops!(GetIntOps, int_ops, IntOps, i16);
257impl_get_ops!(GetIntOps, int_ops, IntOps, i32);
258impl_get_ops!(GetIntOps, int_ops, IntOps, i8);
259impl_get_ops!(GetIntOps, int_ops, IntOps, u8);
260impl_get_ops!(GetIntOps, int_ops, IntOps, u16);
261
262/// Get the [`SignedIntOps`] implementation from an [`Isa`] for a given element type.
263///
264/// This is a specialization of [`GetNumOps`] for signed integer element types.
265pub trait GetSignedIntOps
266where
267    Self: GetSimd,
268{
269    fn signed_int_ops<I: Isa>(isa: I) -> impl SignedIntOps<Self, Simd = Self::Simd<I>>;
270}
271impl_get_ops!(GetSignedIntOps, signed_int_ops, SignedIntOps, i32);
272impl_get_ops!(GetSignedIntOps, signed_int_ops, SignedIntOps, i16);
273impl_get_ops!(GetSignedIntOps, signed_int_ops, SignedIntOps, i8);
274
275/// SIMD operations on a [`Mask`] vector.
276///
277/// # Safety
278///
279/// Implementations must ensure they can only be constructed if the
280/// instruction set is supported on the current system.
281pub unsafe trait MaskOps<M: Mask>: Copy {
282    /// Compute `x & y`.
283    fn and(self, x: M, y: M) -> M;
284
285    /// Return true if any lanes are true.
286    fn any(self, x: M) -> bool;
287
288    /// Return true if all lanes are false.
289    fn all_false(self, x: M) -> bool {
290        !self.any(x)
291    }
292
293    /// Return true if all lanes are true.
294    fn all(self, x: M) -> bool;
295}
296
297/// Bit-level operations available on all SIMD vector types.
298///
299/// This trait provides operations which treat a SIMD vector as a sequence of
300/// bits, without interpreting those bits as numbers of a particular type:
301///
302/// - Load from and store into memory
303/// - Creating a new vector filled with zeros or a specific value
304/// - Combining elements from two vectors according to a mask
305/// - Bitwise operations (and, or, xor, not)
306///
307/// Arithmetic operations which require interpreting the bits as numbers are
308/// provided by the [`NumOps`] sub-trait. Splitting these operations allows
309/// supporting data types for which only bit-level operations are available. For
310/// example many x86 and Arm CPUs support loading, storing and shuffling `f16`
311/// vectors, but not arithmetic on them.
312///
313/// # Safety
314///
315/// Implementations must ensure they can only be constructed if the
316/// instruction set is supported on the current system.
317#[allow(clippy::len_without_is_empty)]
318pub unsafe trait BitOps<T: Elem>: Copy {
319    /// SIMD vector containing lanes of type `T`.
320    type Simd: Simd<Elem = T>;
321
322    /// Convert `x` to an untyped vector of the same width.
323    #[allow(clippy::wrong_self_convention)]
324    fn from_bits(self, x: <<Self::Simd as Simd>::Isa as Isa>::Bits) -> Self::Simd {
325        Self::Simd::from_bits(x)
326    }
327
328    /// Return the number of elements in the vector.
329    fn len(self) -> usize;
330
331    /// Create a new vector with all lanes set to zero.
332    fn zero(self) -> Self::Simd {
333        self.splat(T::default())
334    }
335
336    /// Broadcast the element from one lane of a vector to all lanes of a new
337    /// vector.
338    fn broadcast_lane<const LANE: i32>(self, x: Self::Simd) -> Self::Simd {
339        let val = x.to_array()[LANE as usize];
340        self.splat(val)
341    }
342
343    /// Return the bitwise AND of `x` and `y`.
344    fn and(self, x: Self::Simd, y: Self::Simd) -> Self::Simd;
345
346    /// Return the bitwise NOT of `x`.
347    fn not(self, x: Self::Simd) -> Self::Simd;
348
349    /// Return the bitwise OR of `x` and `y`.
350    fn or(self, x: Self::Simd, y: Self::Simd) -> Self::Simd;
351
352    /// Return the bitwise XOR of `x` and `y`.
353    fn xor(self, x: Self::Simd, y: Self::Simd) -> Self::Simd;
354
355    /// Create a new vector with all lanes set to `x`.
356    fn splat(self, x: T) -> Self::Simd;
357
358    /// Reduce the elements in `x` to a single value using `f`, then
359    /// return a new vector with the accumulated value broadcast to each lane.
360    #[inline]
361    fn fold_splat<F: Fn(T, T) -> T>(self, x: Self::Simd, accum: T, f: F) -> Self::Simd {
362        let reduced = x.to_array().into_iter().fold(accum, f);
363        self.splat(reduced)
364    }
365
366    /// Return a mask with the first `n` lanes set to true.
367    fn first_n_mask(self, n: usize) -> <Self::Simd as Simd>::Mask;
368
369    /// Load the first `self.len()` elements from a slice into a vector.
370    ///
371    /// Panics if `xs.len() < self.len()`.
372    #[inline]
373    #[track_caller]
374    fn load(self, xs: &[T]) -> Self::Simd {
375        assert!(
376            xs.len() >= self.len(),
377            "slice length {} too short for SIMD vector width {}",
378            xs.len(),
379            self.len()
380        );
381        unsafe { self.load_ptr(xs.as_ptr()) }
382    }
383
384    /// Load `N` vectors from consecutive sub-slices of `xs`.
385    ///
386    /// Panics if `xs.len() < self.len() * N`.
387    #[inline]
388    #[track_caller]
389    fn load_many<const N: usize>(self, xs: &[T]) -> [Self::Simd; N] {
390        let v_len = self.len();
391        assert!(
392            xs.len() >= v_len * N,
393            "slice length {} too short for {} * SIMD vector width {}",
394            xs.len(),
395            N,
396            v_len
397        );
398        // Safety: `xs.add(i * v_len)` points to at least `v_len` elements.
399        std::array::from_fn(|i| unsafe { self.load_ptr(xs.as_ptr().add(i * v_len)) })
400    }
401
402    /// Load elements from `xs` into a vector.
403    ///
404    /// If the vector length exceeds `xs.len()`, the tail is padded with zeros.
405    ///
406    /// Returns the padded vector and a mask of the lanes which were set.
407    #[inline]
408    fn load_pad(self, xs: &[T]) -> (Self::Simd, <Self::Simd as Simd>::Mask) {
409        let n = xs.len().min(self.len());
410        let mask = self.first_n_mask(n);
411
412        // Safety: `xs.add(i)` is valid for all positions where mask is set
413        let vec = unsafe { self.load_ptr_mask(xs.as_ptr(), mask) };
414
415        (vec, mask)
416    }
417
418    /// Load vector of elements from `ptr`.
419    ///
420    /// `ptr` is not required to have any particular alignment.
421    ///
422    /// # Safety
423    ///
424    /// `ptr` must point to `self.len()` initialized elements of type `T`.
425    unsafe fn load_ptr(self, ptr: *const T) -> Self::Simd;
426
427    /// Load vector elements from `ptr` using a mask.
428    ///
429    /// `ptr` is not required to have any particular alignment.
430    ///
431    /// # Safety
432    ///
433    /// For each mask position `i` which is true, `ptr.add(i)` must point to
434    /// an initialized element of type `T`.
435    unsafe fn load_ptr_mask(self, ptr: *const T, mask: <Self::Simd as Simd>::Mask) -> Self::Simd;
436
437    /// Select elements from `x` or `y` according to a mask.
438    ///
439    /// Elements are selected from `x` where the corresponding mask element
440    /// is one or `y` if zero.
441    fn select(self, x: Self::Simd, y: Self::Simd, mask: <Self::Simd as Simd>::Mask) -> Self::Simd;
442
443    /// Store the values in this vector to a memory location.
444    ///
445    /// # Safety
446    ///
447    /// `ptr` must point to `self.len()` elements.
448    unsafe fn store_ptr(self, x: Self::Simd, ptr: *mut T);
449
450    /// Store `x` into the first `self.len()` elements of `xs`.
451    #[inline]
452    fn store(self, x: Self::Simd, xs: &mut [T]) {
453        assert!(xs.len() >= self.len());
454        unsafe { self.store_ptr(x, xs.as_mut_ptr()) }
455    }
456
457    /// Store `x` into the first `self.len()` elements of `xs`.
458    ///
459    /// This is a variant of [`store`](BitOps::store) which takes an
460    /// uninitialized slice as input and returns the initialized portion of the
461    /// slice.
462    #[inline]
463    fn store_uninit(self, x: Self::Simd, xs: &mut [MaybeUninit<T>]) -> &mut [T] {
464        let len = self.len();
465        let xs_ptr = xs.as_mut_ptr() as *mut T;
466        assert!(xs.len() >= len);
467        unsafe {
468            self.store_ptr(x, xs_ptr);
469
470            // Safety: `store_ptr` initialized `len` elements of `xs`.
471            std::slice::from_raw_parts_mut(xs_ptr, len)
472        }
473    }
474
475    /// Store `N` vectors into consecutive sub-slices of `xs`, returning the
476    /// initialized portion.
477    ///
478    /// This can be faster than [`store_uninit`](Self::store_uninit) when
479    /// storing several vectors as it only performs a single bounds check.
480    ///
481    /// Panics if `xs.len() < self.len() * N`.
482    #[inline(always)]
483    fn store_many_uninit<const N: usize>(
484        self,
485        vecs: [Self::Simd; N],
486        xs: &mut [MaybeUninit<T>],
487    ) -> &mut [T] {
488        let v_len = self.len();
489        let total = v_len * N;
490
491        // Bounds-check the whole batch up front. The panic message is kept
492        // simple with no arguments. Adding arguments was observed to cause
493        // stack writes, slowing down calls to this function in hot loops.
494        assert!(
495            xs.len() >= total,
496            "slice length too short for SIMD vector array"
497        );
498        let dest_ptr = xs.as_mut_ptr() as *mut T;
499
500        for (i, vec) in vecs.into_iter().enumerate() {
501            // Safety: `xs` holds at least `total = N * v_len` elements, so the
502            // store of `v_len` elements at offset `i * v_len` (i in `0..N`)
503            // stays in bounds.
504            unsafe { self.store_ptr(vec, dest_ptr.add(i * v_len)) };
505        }
506        // Safety: the loop initialized `total` elements of `xs`.
507        unsafe { std::slice::from_raw_parts_mut(dest_ptr, total) }
508    }
509
510    /// Store the values in this vector to a memory location, where the
511    /// corresponding mask element is set.
512    ///
513    /// # Safety
514    ///
515    /// For each position `i` in the mask which is true, `ptr.add(i)` must point
516    /// to a valid element of type `Self::Elem`.
517    unsafe fn store_ptr_mask(self, x: Self::Simd, ptr: *mut T, mask: <Self::Simd as Simd>::Mask);
518
519    /// Pre-fetch the cache line containing `ptr` for reading.
520    fn prefetch(self, ptr: *const T) {
521        // Default implementation does nothing
522        let _ = ptr;
523    }
524
525    /// Pre-fetch the cache line containing `ptr` for writing.
526    fn prefetch_write(self, ptr: *mut T) {
527        // Default implementation does nothing
528        let _ = ptr;
529    }
530}
531
532/// Arithmetic operations available on all numeric SIMD vector types.
533///
534/// This trait extends [`BitOps`] with operations which interpret the bits of a
535/// SIMD vector as numbers of type `T`:
536///
537/// - Add, subtract and multiply
538/// - Comparison (equality, less than, greater than etc.)
539///
540/// # Safety
541///
542/// Implementations must ensure they can only be constructed if the
543/// instruction set is supported on the current system.
544pub unsafe trait NumOps<T: Elem>: BitOps<T> {
545    /// Compute `x + y`.
546    fn add(self, x: Self::Simd, y: Self::Simd) -> Self::Simd;
547
548    /// Compute `x - y`.
549    fn sub(self, x: Self::Simd, y: Self::Simd) -> Self::Simd;
550
551    /// Compute `x * y`.
552    fn mul(self, x: Self::Simd, y: Self::Simd) -> Self::Simd;
553
554    /// Create a new vector with all lanes set to one.
555    fn one(self) -> Self::Simd {
556        self.splat(T::one())
557    }
558
559    /// Compute `a * b + c`.
560    ///
561    /// This will use fused multiply-add instructions if available. For float
562    /// element types, this may use one or two roundings.
563    fn mul_add(self, a: Self::Simd, b: Self::Simd, c: Self::Simd) -> Self::Simd {
564        self.add(self.mul(a, b), c)
565    }
566
567    /// Evaluate a polynomial using Horner's method.
568    ///
569    /// Computes `x * coeffs[0] + x^2 * coeffs[1] ... x^n * coeffs[N]`
570    #[inline]
571    fn poly_eval(self, x: Self::Simd, coeffs: &[Self::Simd]) -> Self::Simd {
572        let mut y = coeffs[coeffs.len() - 1];
573        for i in (0..coeffs.len() - 1).rev() {
574            y = self.mul_add(y, x, coeffs[i]);
575        }
576        self.mul(y, x)
577    }
578
579    /// Return a mask indicating whether elements in `x` are less than `y`.
580    #[inline]
581    fn lt(self, x: Self::Simd, y: Self::Simd) -> <Self::Simd as Simd>::Mask {
582        self.gt(y, x)
583    }
584
585    /// Return a mask indicating whether elements in `x` are less or equal to `y`.
586    #[inline]
587    fn le(self, x: Self::Simd, y: Self::Simd) -> <Self::Simd as Simd>::Mask {
588        self.ge(y, x)
589    }
590
591    /// Return a mask indicating whether elements in `x` are equal to `y`.
592    fn eq(self, x: Self::Simd, y: Self::Simd) -> <Self::Simd as Simd>::Mask;
593
594    /// Return a mask indicating whether elements in `x` are greater or equal to `y`.
595    fn ge(self, x: Self::Simd, y: Self::Simd) -> <Self::Simd as Simd>::Mask;
596
597    /// Return a mask indicating whether elements in `x` are greater than `y`.
598    fn gt(self, x: Self::Simd, y: Self::Simd) -> <Self::Simd as Simd>::Mask;
599
600    /// Return the minimum of `x` and `y` for each lane.
601    fn min(self, x: Self::Simd, y: Self::Simd) -> Self::Simd {
602        self.select(x, y, self.le(x, y))
603    }
604
605    /// Return the maximum of `x` and `y` for each lane.
606    fn max(self, x: Self::Simd, y: Self::Simd) -> Self::Simd {
607        self.select(x, y, self.ge(x, y))
608    }
609
610    /// Clamp values in `x` to minimum and maximum values from corresponding
611    /// lanes in `min` and `max`.
612    fn clamp(self, x: Self::Simd, min: Self::Simd, max: Self::Simd) -> Self::Simd {
613        self.min(self.max(x, min), max)
614    }
615
616    /// Horizontally sum the elements in a vector.
617    ///
618    /// If the sum overflows, it will wrap. This choice was made to enable
619    /// consistency between native intrinsics for horizontal addition and the
620    /// generic implementation.
621    fn sum(self, x: Self::Simd) -> T {
622        let mut sum = T::default();
623        for elem in x.to_array() {
624            sum = sum.wrapping_add(elem);
625        }
626        sum
627    }
628}
629
630/// Operations available on SIMD vectors with float elements.
631pub trait FloatOps<T: Elem>: NumOps<T> {
632    /// Integer SIMD vector of the same bit-width as this vector.
633    type Int: Simd;
634
635    /// Compute x / y
636    fn div(self, x: Self::Simd, y: Self::Simd) -> Self::Simd;
637
638    /// Compute 1. / x
639    fn reciprocal(self, x: Self::Simd) -> Self::Simd {
640        self.div(self.one(), x)
641    }
642
643    /// Compute `-x`
644    fn neg(self, x: Self::Simd) -> Self::Simd {
645        self.sub(self.zero(), x)
646    }
647
648    /// Compute the absolute value of `x`
649    fn abs(self, x: Self::Simd) -> Self::Simd {
650        self.select(self.neg(x), x, self.lt(x, self.zero()))
651    }
652
653    /// Round `x` to the nearest integer value, with ties to even.
654    ///
655    /// This is like [`f32::round_ties_even`].
656    fn round_ties_even(self, x: Self::Simd) -> Self::Simd;
657
658    /// Compute `c - a * b`.
659    fn mul_sub_from(self, a: Self::Simd, b: Self::Simd, c: Self::Simd) -> Self::Simd {
660        self.sub(c, self.mul(a, b))
661    }
662
663    /// Convert each lane to an integer of the same width, rounding towards zero.
664    fn to_int_trunc(self, x: Self::Simd) -> Self::Int;
665
666    /// Convert each lane to an integer of the same width, rounding to nearest
667    /// with ties to even.
668    fn to_int_round(self, x: Self::Simd) -> Self::Int;
669}
670
671/// Operations on SIMD vectors with integer elements.
672pub trait IntOps<T: Elem>: NumOps<T> {
673    /// Shift each lane in `x` left by `SHIFT` bits.
674    fn shift_left<const SHIFT: i32>(self, x: Self::Simd) -> Self::Simd;
675
676    /// Shift each lane in `x` right by `SHIFT` bits.
677    ///
678    /// For signed integer types this is an arithmetic shift, so shifting a
679    /// negative number right will preserve the sign, like the `>>` operator.
680    fn shift_right<const SHIFT: i32>(self, x: Self::Simd) -> Self::Simd;
681}
682
683/// Operations on SIMD vectors with signed integer elements.
684pub trait SignedIntOps<T: Elem>: IntOps<T> {
685    /// Compute the absolute value of `x`
686    fn abs(self, x: Self::Simd) -> Self::Simd {
687        self.select(self.neg(x), x, self.lt(x, self.zero()))
688    }
689
690    /// Return `-x`.
691    fn neg(self, x: Self::Simd) -> Self::Simd {
692        self.sub(self.zero(), x)
693    }
694}
695
696/// Widen lanes to a type with twice the width.
697///
698/// For integer types, the extended type has the same signed-ness. For `f16`,
699/// the extended type is `f32`.
700pub trait Extend<T: Elem>: BitOps<T> {
701    /// SIMD vector type with elements that have twice the bit-width of
702    /// those in `Self::SIMD`.
703    type Output;
704
705    /// Extend each lane in the low half of the input to a type with twice the
706    /// width.
707    fn extend_low(self, x: Self::Simd) -> Self::Output;
708
709    /// Extend each lane in the high half of the input to a type with twice the
710    /// width.
711    fn extend_high(self, x: Self::Simd) -> Self::Output;
712}
713
714/// Interleave elements from the low or high halves of two vectors to form a
715/// new vector.
716pub trait Interleave<T: Elem>: NumOps<T> {
717    /// Interleave elements from the low halves of two vectors.
718    fn interleave_low(self, a: Self::Simd, b: Self::Simd) -> Self::Simd;
719
720    /// Interleave elements from the high halves of two vectors.
721    fn interleave_high(self, a: Self::Simd, b: Self::Simd) -> Self::Simd;
722}
723
724/// Concatenate elements from the low or high halves of two vectors to form a
725/// new vector.
726pub trait Concat<T: Elem>: NumOps<T> {
727    /// Concatenate elements from the low halves of two vectors.
728    fn concat_low(self, a: Self::Simd, b: Self::Simd) -> Self::Simd;
729
730    /// Concatenate elements from the high halves of two vectors.
731    fn concat_high(self, a: Self::Simd, b: Self::Simd) -> Self::Simd;
732}
733
734/// Convert each lane to a float with the same bit width.
735pub trait ToFloat<T: Elem>: NumOps<T> {
736    type Output;
737
738    fn to_float(self, x: Self::Simd) -> Self::Output;
739}
740
741/// Narrow lanes to one with half the bit-width, using truncation.
742///
743/// For integer types, the narrowed type has the same signed-ness.
744#[cfg(target_arch = "x86_64")]
745pub(crate) trait Narrow<S: Simd> {
746    type Output;
747
748    /// Truncate each lane in a pair of vectors to one with half the bit-width.
749    ///
750    /// Returns a vector containing the concatenation of the narrowed lanes
751    /// from `low` followed by the narrowed lanes from `high`.
752    fn narrow_truncate(self, low: S, high: S) -> Self::Output;
753}
754
755/// Narrow lanes to one with half the bit-width, using saturation.
756///
757/// Conceptually, this converts each element from `S1::Elem` to `S2::Elem` using
758/// `x.clamp(S2::Elem::MIN as S1::Elem, S2::Elem::MAX as S1::Elem) as S2::Elem`.
759pub trait NarrowSaturate<T: Elem, U: Elem>: NumOps<T> {
760    type Output: Simd<Elem = U>;
761
762    /// Narrow each lane in a pair of vectors to one with half the bit-width.
763    ///
764    /// Returns a vector containing the concatenation of the narrowed lanes
765    /// from `low` followed by the narrowed lanes from `high`.
766    fn narrow_saturate(self, low: Self::Simd, high: Self::Simd) -> Self::Output;
767}
768
769#[cfg(test)]
770mod tests {
771    use crate::elem::WrappingAdd;
772    use crate::f16;
773    use crate::ops::{
774        BitOps, Concat, Extend, FloatOps, IntOps, Interleave, MaskOps, NarrowSaturate, NumOps,
775        SignedIntOps, ToFloat,
776    };
777    use crate::{Isa, Mask, Simd, SimdOp, assert_simd_eq, assert_simd_ne, test_simd_op};
778
779    // Generate tests for operations available on all numeric types.
780    macro_rules! test_num_ops {
781        ($modname:ident, $elem:ident, $mask_elem:ident) => {
782            mod $modname {
783                use super::{
784                    BitOps, Isa, MaskOps, NumOps, Simd, SimdOp, WrappingAdd, assert_simd_eq,
785                    assert_simd_ne, test_simd_op,
786                };
787
788                #[test]
789                fn test_load_store() {
790                    test_simd_op!(isa, {
791                        let ops = isa.$elem();
792
793                        let src: Vec<_> = (0..ops.len() * 4).map(|x| x as $elem).collect();
794                        let mut dst = vec![0 as $elem; src.len()];
795
796                        for (src_chunk, dst_chunk) in
797                            src.chunks(ops.len()).zip(dst.chunks_mut(ops.len()))
798                        {
799                            let x = ops.load(src_chunk);
800                            ops.store(x, dst_chunk);
801                        }
802
803                        assert_eq!(dst, src);
804                    })
805                }
806
807                #[test]
808                fn test_store_uninit() {
809                    test_simd_op!(isa, {
810                        let ops = isa.$elem();
811
812                        let src: Vec<_> = (0..ops.len() + 3).map(|x| x as $elem).collect();
813                        let mut dest = Vec::with_capacity(src.len());
814
815                        let x = ops.load(&src);
816
817                        let init = ops.store_uninit(x, dest.spare_capacity_mut());
818                        assert_eq!(init, &src[0..ops.len()]);
819                    })
820                }
821
822                #[test]
823                fn test_load_many() {
824                    test_simd_op!(isa, {
825                        let ops = isa.$elem();
826
827                        let src: Vec<_> = (0..ops.len() * 2).map(|x| x as $elem).collect();
828
829                        let xs = ops.load_many::<2>(&src);
830                        assert_simd_eq!(xs[0], ops.load(&src));
831                        assert_simd_eq!(xs[1], ops.load(&src[ops.len()..]));
832                    })
833                }
834
835                #[test]
836                fn test_store_many_uninit() {
837                    test_simd_op!(isa, {
838                        let ops = isa.$elem();
839
840                        let src: Vec<_> = (0..ops.len() * 2).map(|x| x as $elem).collect();
841                        let xs = ops.load_many::<2>(&src);
842
843                        let mut dest = Vec::with_capacity(src.len());
844                        let init = ops.store_many_uninit(xs, dest.spare_capacity_mut());
845                        assert_eq!(init, &src[..]);
846                    })
847                }
848
849                #[test]
850                fn test_load_pad() {
851                    test_simd_op!(isa, {
852                        let ops = isa.$elem();
853
854                        // Array which is shorter than vector length for all ISAs.
855                        let src = [0, 1, 2].map(|x| x as $elem);
856
857                        let (vec, _mask) = ops.load_pad(&src);
858                        let vec_array = vec.to_array();
859                        let vec_slice = vec_array.as_ref();
860
861                        assert_eq!(&vec_slice[..src.len()], &src);
862                        for i in ops.len()..vec_slice.len() {
863                            assert_eq!(vec_array[i], 0 as $elem);
864                        }
865                    })
866                }
867
868                #[test]
869                fn test_bin_ops() {
870                    test_simd_op!(isa, {
871                        let ops = isa.$elem();
872
873                        let a = 2 as $elem;
874                        let b = 3 as $elem;
875
876                        let x = ops.splat(a);
877                        let y = ops.splat(b);
878
879                        // Add
880                        let expected = ops.splat(a + b);
881                        let actual = ops.add(x, y);
882                        assert_simd_eq!(actual, expected);
883
884                        // Sub
885                        let expected = ops.splat(b - a);
886                        let actual = ops.sub(y, x);
887                        assert_simd_eq!(actual, expected);
888
889                        // Mul
890                        let expected = ops.splat(a * b);
891                        let actual = ops.mul(x, y);
892                        assert_simd_eq!(actual, expected);
893                    })
894                }
895
896                #[test]
897                fn test_cmp_ops() {
898                    test_simd_op!(isa, {
899                        let ops = isa.$elem();
900                        let mo = isa.$mask_elem();
901
902                        let x = ops.splat(1 as $elem);
903                        let y = ops.splat(2 as $elem);
904
905                        assert!(mo.all(ops.eq(x, x)));
906                        assert!(mo.all_false(ops.eq(x, y)));
907                        assert!(mo.all(ops.le(x, x)));
908                        assert!(mo.all(ops.le(x, y)));
909                        assert!(mo.all_false(ops.le(y, x)));
910                        assert!(mo.all(ops.ge(x, x)));
911                        assert!(mo.all_false(ops.ge(x, y)));
912                        assert!(mo.all_false(ops.gt(x, y)));
913                        assert!(mo.all(ops.gt(y, x)));
914                    })
915                }
916
917                #[test]
918                fn test_mul_add() {
919                    test_simd_op!(isa, {
920                        let ops = isa.$elem();
921
922                        let a = ops.splat(2 as $elem);
923                        let b = ops.splat(3 as $elem);
924                        let c = ops.splat(4 as $elem);
925
926                        let actual = ops.mul_add(a, b, c);
927                        let expected = ops.splat(((2. * 3.) + 4.) as $elem);
928
929                        assert_simd_eq!(actual, expected);
930                    })
931                }
932
933                #[test]
934                fn test_min_max() {
935                    test_simd_op!(isa, {
936                        let ops = isa.$elem();
937
938                        let x = ops.splat(3 as $elem);
939
940                        // Min
941                        let y_min = ops.min(x, ops.splat(2 as $elem));
942                        let y_min_2 = ops.min(ops.splat(2 as $elem), x);
943                        assert_simd_eq!(y_min, y_min_2);
944                        assert_simd_eq!(y_min, ops.splat(2 as $elem));
945
946                        // Max
947                        let y_max = ops.max(x, ops.splat(4 as $elem));
948                        let y_max_2 = ops.max(ops.splat(4 as $elem), x);
949                        assert_simd_eq!(y_max, y_max_2);
950                        assert_simd_eq!(y_max, ops.splat(4 as $elem));
951
952                        // Clamp
953                        let y_clamped = ops.clamp(x, ops.splat(0 as $elem), ops.splat(4 as $elem));
954                        assert_simd_eq!(y_clamped, ops.splat(3 as $elem));
955                    })
956                }
957
958                #[test]
959                fn test_and() {
960                    test_simd_op!(isa, {
961                        let ops = isa.$elem();
962                        let zeros = ops.zero();
963                        let ones = ops.not(zeros);
964
965                        // Cast to bits here because all-ones is a NaN if elements
966                        // are floats, and NaNs are not equal to themselves.
967                        assert_simd_eq!(ops.and(zeros, zeros).to_bits(), zeros.to_bits());
968                        assert_simd_eq!(ops.and(zeros, ones).to_bits(), zeros.to_bits());
969                        assert_simd_eq!(ops.and(ones, zeros).to_bits(), zeros.to_bits());
970                        assert_simd_eq!(ops.and(ones, ones).to_bits(), ones.to_bits());
971                    })
972                }
973
974                #[test]
975                fn test_not() {
976                    test_simd_op!(isa, {
977                        let ops = isa.$elem();
978                        let zeros = ops.zero();
979                        let ones = ops.not(zeros);
980                        assert_simd_ne!(zeros, ones);
981
982                        let zeros_2 = ops.not(ones);
983                        assert_simd_eq!(zeros_2, zeros);
984                    })
985                }
986
987                #[test]
988                fn test_or() {
989                    test_simd_op!(isa, {
990                        let ops = isa.$elem();
991                        let zeros = ops.zero();
992                        let ones = ops.not(zeros);
993
994                        // Cast to bits here because all-ones is a NaN if elements
995                        // are floats, and NaNs are not equal to themselves.
996                        assert_simd_eq!(ops.or(zeros, zeros).to_bits(), zeros.to_bits());
997                        assert_simd_eq!(ops.or(zeros, ones).to_bits(), ones.to_bits());
998                        assert_simd_eq!(ops.or(ones, zeros).to_bits(), ones.to_bits());
999                        assert_simd_eq!(ops.or(ones, ones).to_bits(), ones.to_bits());
1000                    })
1001                }
1002
1003                #[test]
1004                fn test_xor() {
1005                    test_simd_op!(isa, {
1006                        let ops = isa.$elem();
1007
1008                        let zeros = ops.zero();
1009                        let ones = ops.not(zeros);
1010
1011                        // Cast to bits here because all-ones is a NaN if the
1012                        // element type is a float, and NaNs are not equal to
1013                        // themselves.
1014                        assert_simd_eq!(ops.xor(zeros, zeros).to_bits(), zeros.to_bits());
1015                        assert_simd_eq!(ops.xor(ones, ones).to_bits(), zeros.to_bits());
1016                        assert_simd_eq!(ops.xor(zeros, ones).to_bits(), ones.to_bits());
1017                        assert_simd_eq!(ops.xor(ones, zeros).to_bits(), ones.to_bits());
1018                    })
1019                }
1020
1021                #[test]
1022                fn test_sum() {
1023                    test_simd_op!(isa, {
1024                        let ops = isa.$elem();
1025
1026                        let vec: Vec<_> = (0..ops.len()).map(|x| x as $elem).collect();
1027                        let expected = vec
1028                            .iter()
1029                            .fold(0 as $elem, |sum, x| WrappingAdd::wrapping_add(sum, *x));
1030
1031                        let x = ops.load(&vec);
1032                        let y = ops.sum(x);
1033
1034                        assert_eq!(y, expected);
1035                    })
1036                }
1037
1038                #[test]
1039                fn test_poly_eval() {
1040                    test_simd_op!(isa, {
1041                        let ops = isa.$elem();
1042
1043                        let coeffs = [2, 3, 4].map(|x| x as $elem);
1044                        let x = 2 as $elem;
1045                        let y = ops.poly_eval(ops.splat(x), &coeffs.map(|c| ops.splat(c)));
1046
1047                        let expected =
1048                            (x * coeffs[0]) + (x * x * coeffs[1]) + (x * x * x * coeffs[2]);
1049                        assert_simd_eq!(y, ops.splat(expected));
1050                    })
1051                }
1052
1053                #[test]
1054                fn test_broadcast_lane() {
1055                    test_simd_op!(isa, {
1056                        let ops = isa.$elem();
1057
1058                        let vec: Vec<_> = (0..ops.len()).map(|x| x as $elem).collect();
1059                        let xs = ops.load(&vec);
1060
1061                        let ys = ops.broadcast_lane::<0>(xs);
1062                        assert_simd_eq!(ops.splat(0 as $elem), ys);
1063
1064                        let ys = ops.broadcast_lane::<1>(xs);
1065                        assert_simd_eq!(ops.splat(1 as $elem), ys);
1066
1067                        let ys = ops.broadcast_lane::<2>(xs);
1068                        assert_simd_eq!(ops.splat(2 as $elem), ys);
1069
1070                        let ys = ops.broadcast_lane::<3>(xs);
1071                        assert_simd_eq!(ops.splat(3 as $elem), ys);
1072                    });
1073                }
1074            }
1075        };
1076    }
1077
1078    test_num_ops!(num_ops_f32, f32, m32);
1079    test_num_ops!(num_ops_i32, i32, m32);
1080    test_num_ops!(num_ops_i16, i16, m16);
1081    test_num_ops!(num_ops_i8, i8, m8);
1082    test_num_ops!(num_ops_u8, u8, m8);
1083    test_num_ops!(num_ops_u16, u16, m16);
1084
1085    // Test that x8 multiply truncates result as expected.
1086    #[test]
1087    fn test_i8_mul_truncate() {
1088        test_simd_op!(isa, {
1089            let ops = isa.i8();
1090
1091            let x = 17i8;
1092            let y = 19i8;
1093
1094            let x_vec = ops.splat(x);
1095            let y_vec = ops.splat(y);
1096            let expected = ops.splat(x.wrapping_mul(y));
1097            let actual = ops.mul(x_vec, y_vec);
1098
1099            assert_simd_eq!(actual, expected);
1100        })
1101    }
1102
1103    #[test]
1104    fn test_u8_mul_truncate() {
1105        test_simd_op!(isa, {
1106            let ops = isa.u8();
1107
1108            let x = 17u8;
1109            let y = 19u8;
1110
1111            let x_vec = ops.splat(x);
1112            let y_vec = ops.splat(y);
1113            let expected = ops.splat(x.wrapping_mul(y));
1114            let actual = ops.mul(x_vec, y_vec);
1115
1116            assert_simd_eq!(actual, expected);
1117        })
1118    }
1119
1120    // Generate tests for operations available on all float types.
1121    macro_rules! test_float_ops {
1122        ($modname:ident, $elem:ident, $int_elem:ident) => {
1123            mod $modname {
1124                use super::{BitOps, FloatOps, Isa, Simd, SimdOp, assert_simd_eq, test_simd_op};
1125
1126                #[test]
1127                fn test_div() {
1128                    test_simd_op!(isa, {
1129                        let ops = isa.$elem();
1130
1131                        let x = ops.splat(1.);
1132                        let y = ops.splat(2.);
1133                        let expected = ops.splat(0.5);
1134                        let actual = ops.div(x, y);
1135                        assert_simd_eq!(actual, expected);
1136                    })
1137                }
1138
1139                #[test]
1140                fn test_reciprocal() {
1141                    test_simd_op!(isa, {
1142                        let ops = isa.$elem();
1143
1144                        let vals = [-5., -2., 2., 5.];
1145                        for v in vals {
1146                            let x = ops.splat(v);
1147                            let y = ops.reciprocal(x);
1148                            let expected = ops.splat(1. / v);
1149                            assert_simd_eq!(y, expected);
1150                        }
1151                    })
1152                }
1153
1154                #[test]
1155                fn test_abs() {
1156                    test_simd_op!(isa, {
1157                        let ops = isa.$elem();
1158
1159                        let vals = [-1., 0., 1.];
1160                        for v in vals {
1161                            let x = ops.splat(v);
1162                            let y = ops.abs(x);
1163                            let expected = ops.splat(v.abs());
1164                            assert_simd_eq!(y, expected);
1165                        }
1166                    })
1167                }
1168
1169                #[test]
1170                fn test_neg() {
1171                    test_simd_op!(isa, {
1172                        let ops = isa.$elem();
1173
1174                        let x = ops.splat(3 as $elem);
1175
1176                        let expected = ops.splat(-3 as $elem);
1177                        let actual = ops.neg(x);
1178                        assert_simd_eq!(actual, expected);
1179                    })
1180                }
1181
1182                #[test]
1183                fn test_mul_sub_from() {
1184                    test_simd_op!(isa, {
1185                        let ops = isa.$elem();
1186
1187                        let a = ops.splat(2 as $elem);
1188                        let b = ops.splat(3 as $elem);
1189                        let c = ops.splat(4 as $elem);
1190
1191                        let actual = ops.mul_sub_from(a, b, c);
1192                        let expected = ops.splat((-(2. * 3.) + 4.) as $elem);
1193
1194                        assert_simd_eq!(actual, expected);
1195                    })
1196                }
1197
1198                #[test]
1199                fn test_round_ties_even() {
1200                    test_simd_op!(isa, {
1201                        let ops = isa.$elem();
1202
1203                        let x = ops.splat(3.5 as $elem);
1204
1205                        let expected = ops.splat(4 as $elem);
1206                        let actual = ops.round_ties_even(x);
1207                        assert_simd_eq!(actual, expected);
1208                    })
1209                }
1210
1211                #[test]
1212                fn test_to_int_trunc() {
1213                    test_simd_op!(isa, {
1214                        let ops = isa.$elem();
1215
1216                        let x = ops.splat(12.345);
1217                        let y = ops.to_int_trunc(x);
1218                        let expected = isa.$int_elem().splat(12);
1219                        assert_simd_eq!(y, expected);
1220                    })
1221                }
1222            }
1223        };
1224    }
1225
1226    test_float_ops!(float_ops_f32, f32, i32);
1227
1228    // Generate tests for operations available on unsigned integer types.
1229    macro_rules! test_unsigned_int_ops {
1230        ($modname:ident, $elem:ident) => {
1231            mod $modname {
1232                use super::{BitOps, IntOps, Isa, Simd, SimdOp, assert_simd_eq, test_simd_op};
1233
1234                #[test]
1235                fn test_shift_left() {
1236                    test_simd_op!(isa, {
1237                        let ops = isa.$elem();
1238
1239                        let x = ops.splat(42);
1240                        let y = ops.shift_left::<1>(x);
1241                        let expected = ops.splat(42 << 1);
1242                        assert_simd_eq!(y, expected);
1243                    })
1244                }
1245
1246                #[test]
1247                fn test_shift_right() {
1248                    test_simd_op!(isa, {
1249                        let ops = isa.$elem();
1250
1251                        let x = ops.splat(42);
1252                        let y = ops.shift_right::<1>(x);
1253                        let expected = ops.splat(42 >> 1);
1254                        assert_simd_eq!(y, expected);
1255                    })
1256                }
1257            }
1258        };
1259    }
1260
1261    test_unsigned_int_ops!(uint_ops_u16, u16);
1262
1263    // Generate tests for operations available on signed integer types.
1264    macro_rules! test_signed_int_ops {
1265        ($modname:ident, $elem:ident) => {
1266            mod $modname {
1267                use super::{
1268                    BitOps, IntOps, Isa, NumOps, SignedIntOps, Simd, SimdOp, assert_simd_eq,
1269                    test_simd_op,
1270                };
1271
1272                #[test]
1273                fn test_abs() {
1274                    test_simd_op!(isa, {
1275                        let ops = isa.$elem();
1276
1277                        let vals = [-1, 0, 1];
1278                        for v in vals {
1279                            let x = ops.splat(v);
1280                            let y = ops.abs(x);
1281                            let expected = ops.splat(v.abs());
1282                            assert_simd_eq!(y, expected);
1283                        }
1284                    })
1285                }
1286
1287                // Add / Sub / Mul with a negative argument.
1288                #[test]
1289                fn test_bin_ops_neg() {
1290                    test_simd_op!(isa, {
1291                        let ops = isa.$elem();
1292
1293                        let a = -2 as $elem;
1294                        let b = 3 as $elem;
1295
1296                        let x = ops.splat(a);
1297                        let y = ops.splat(b);
1298
1299                        // Add
1300                        let expected = ops.splat(a + b);
1301                        let actual = ops.add(x, y);
1302                        assert_simd_eq!(actual, expected);
1303
1304                        // Sub
1305                        let expected = ops.splat(b - a);
1306                        let actual = ops.sub(y, x);
1307                        assert_simd_eq!(actual, expected);
1308
1309                        // Mul
1310                        let expected = ops.splat(a * b);
1311                        let actual = ops.mul(x, y);
1312                        assert_simd_eq!(actual, expected);
1313                    })
1314                }
1315
1316                #[test]
1317                fn test_shift_left() {
1318                    test_simd_op!(isa, {
1319                        let ops = isa.$elem();
1320
1321                        let x = ops.splat(42);
1322                        let y = ops.shift_left::<1>(x);
1323                        let expected = ops.splat(42 << 1);
1324                        assert_simd_eq!(y, expected);
1325                    })
1326                }
1327
1328                #[test]
1329                fn test_shift_right() {
1330                    test_simd_op!(isa, {
1331                        let ops = isa.$elem();
1332
1333                        let x = ops.splat(42);
1334                        let y = ops.shift_right::<1>(x);
1335                        let expected = ops.splat(42 >> 1);
1336                        assert_simd_eq!(y, expected);
1337
1338                        // `shift_right` is an arithmetic right shift, so it
1339                        // preserves the sign.
1340                        let x = ops.splat(-128);
1341                        let y = ops.shift_right::<1>(x);
1342                        let expected = ops.splat(-64);
1343                        assert_simd_eq!(y, expected);
1344                    })
1345                }
1346
1347                #[test]
1348                fn test_neg() {
1349                    test_simd_op!(isa, {
1350                        let ops = isa.$elem();
1351
1352                        let x = ops.splat(3 as $elem);
1353
1354                        let expected = ops.splat(-3 as $elem);
1355                        let actual = ops.neg(x);
1356                        assert_simd_eq!(actual, expected);
1357                    })
1358                }
1359            }
1360        };
1361    }
1362
1363    test_signed_int_ops!(int_ops_i32, i32);
1364    test_signed_int_ops!(int_ops_i16, i16);
1365    test_signed_int_ops!(int_ops_i8, i8);
1366
1367    // For small positive values, signed comparison ops will work on unsigned
1368    // values. Make sure we really are using unsigned comparison.
1369    #[test]
1370    fn test_cmp_gt_ge_u16() {
1371        test_simd_op!(isa, {
1372            let ops = isa.u16();
1373            let m16 = isa.m16();
1374
1375            let x = ops.splat(i16::MAX as u16);
1376            let y = ops.splat(i16::MAX as u16 + 1);
1377
1378            assert!(m16.all(ops.gt(y, x)));
1379            assert!(m16.all(ops.ge(y, x)));
1380        });
1381    }
1382
1383    #[test]
1384    fn test_cmp_gt_ge_u8() {
1385        test_simd_op!(isa, {
1386            let ops = isa.u8();
1387            let m8 = isa.m8();
1388
1389            let x = ops.splat(i8::MAX as u8);
1390            let y = ops.splat(i8::MAX as u8 + 1);
1391
1392            assert!(m8.all(ops.gt(y, x)));
1393            assert!(m8.all(ops.ge(y, x)));
1394        });
1395    }
1396
1397    macro_rules! test_mask_ops {
1398        ($elem_type:ident, $mask_type:ident) => {
1399            test_simd_op!(isa, {
1400                let ops = isa.$elem_type();
1401                let mask_ops = isa.$mask_type();
1402
1403                // First-n mask
1404                let ones = ops.first_n_mask(ops.len());
1405                let zeros = ops.first_n_mask(0);
1406                let first = ops.first_n_mask(1);
1407
1408                // Bitwise and
1409                assert_simd_eq!(mask_ops.and(ones, ones), ones);
1410                assert_simd_eq!(mask_ops.and(first, ones), first);
1411                assert_simd_eq!(mask_ops.and(first, zeros), zeros);
1412
1413                // Any
1414                assert!(mask_ops.any(ones));
1415                assert!(mask_ops.any(first));
1416                assert!(!mask_ops.any(zeros));
1417
1418                // All
1419                assert!(mask_ops.all(ones));
1420                assert!(!mask_ops.all(zeros));
1421                assert!(!mask_ops.all(first));
1422
1423                // All false
1424                assert!(mask_ops.all_false(zeros));
1425                assert!(!mask_ops.all_false(ones));
1426                assert!(!mask_ops.all_false(first));
1427            });
1428        };
1429    }
1430
1431    #[test]
1432    fn test_mask_ops_m32() {
1433        test_mask_ops!(i32, m32);
1434    }
1435
1436    #[test]
1437    fn test_mask_ops_m16() {
1438        test_mask_ops!(i16, m16);
1439    }
1440
1441    #[test]
1442    fn test_mask_ops_m8() {
1443        test_mask_ops!(i8, m8);
1444    }
1445
1446    macro_rules! test_narrow_saturate {
1447        ($test_name:ident, $src:ident, $dest:ident) => {
1448            #[test]
1449            fn $test_name() {
1450                test_simd_op!(isa, {
1451                    let ops = isa.$src();
1452
1453                    let src: Vec<$src> = (0..ops.len() * 2).map(|x| x as $src).collect();
1454                    let expected: Vec<$dest> = src
1455                        .iter()
1456                        .map(|&x| x.clamp($dest::MIN as $src, $dest::MAX as $src) as $dest)
1457                        .collect();
1458
1459                    let x_low = ops.load(&src[..ops.len()]);
1460                    let x_high = ops.load(&src[ops.len()..]);
1461                    let y = ops.narrow_saturate(x_low, x_high);
1462
1463                    assert_eq!(y.to_array().as_ref(), expected);
1464                });
1465            }
1466        };
1467    }
1468
1469    test_narrow_saturate!(test_narrow_i32_i16, i32, i16);
1470    test_narrow_saturate!(test_narrow_u16_u8, i16, u8);
1471
1472    macro_rules! test_extend {
1473        ($test_name:ident, $src:ident, $dest:ident) => {
1474            #[test]
1475            fn $test_name() {
1476                test_simd_op!(isa, {
1477                    let ops = isa.$src();
1478                    let dst_ops = isa.$dest();
1479
1480                    let src: Vec<$src> = (0..ops.len()).map(|x| x as $src).collect();
1481                    let expected: Vec<$dest> = src.iter().map(|&x| x as $dest).collect();
1482
1483                    let x = ops.load(&src);
1484                    let y_low = ops.extend_low(x);
1485                    let y_high = ops.extend_high(x);
1486                    assert_eq!(y_low.to_array().as_ref(), &expected[..dst_ops.len()]);
1487                    assert_eq!(y_high.to_array().as_ref(), &expected[dst_ops.len()..]);
1488                });
1489            }
1490        };
1491    }
1492    test_extend!(test_extend_i8_i16, i8, i16);
1493    test_extend!(test_extend_i16_i32, i16, i32);
1494    test_extend!(test_extend_u8_u16, u8, u16);
1495
1496    macro_rules! test_interleave {
1497        ($test_name:ident, $elem:ident) => {
1498            #[test]
1499            fn $test_name() {
1500                test_simd_op!(isa, {
1501                    let ops = isa.$elem();
1502
1503                    let even: Vec<_> = (0..ops.len()).map(|x| x as $elem * 2).collect();
1504                    let even = ops.load(&even);
1505
1506                    let odd: Vec<_> = (0..ops.len()).map(|x| 1 + (x as $elem * 2)).collect();
1507                    let odd = ops.load(&odd);
1508
1509                    let expected_low: Vec<_> = (0..ops.len()).map(|x| x as $elem).collect();
1510                    let expected_high: Vec<_> = (0..ops.len())
1511                        .map(|x| ops.len() as $elem + x as $elem)
1512                        .collect();
1513
1514                    let y_low = ops.interleave_low(even, odd);
1515                    let y_high = ops.interleave_high(even, odd);
1516                    assert_eq!(y_low.to_array().as_ref(), expected_low);
1517                    assert_eq!(y_high.to_array().as_ref(), expected_high);
1518                });
1519            }
1520        };
1521    }
1522    test_interleave!(test_interleave_i16, i16);
1523    test_interleave!(test_interleave_i8, i8);
1524    test_interleave!(test_interleave_u8, u8);
1525
1526    #[test]
1527    fn test_concat_i32() {
1528        test_simd_op!(isa, {
1529            let ops = isa.i32();
1530            let src: Vec<_> = (0..ops.len() * 2).map(|x| x as i32).collect();
1531            let src_a = &src[..ops.len()];
1532            let src_b = &src[ops.len()..];
1533
1534            let half_len = ops.len() / 2;
1535            let expected_low: Vec<_> = (0..ops.len())
1536                .map(|i| {
1537                    if i < half_len {
1538                        src_a[i]
1539                    } else {
1540                        src_b[i - half_len]
1541                    }
1542                })
1543                .collect();
1544            let expected_hi: Vec<_> = (0..ops.len())
1545                .map(|i| {
1546                    if i < half_len {
1547                        src_a[half_len + i]
1548                    } else {
1549                        src_b[i]
1550                    }
1551                })
1552                .collect();
1553
1554            let a = ops.load(&src_a);
1555            let b = ops.load(&src_b);
1556            let ab_lo = ops.concat_low(a, b);
1557            let ab_hi = ops.concat_high(a, b);
1558
1559            assert_eq!(ab_lo.to_array().as_ref(), expected_low);
1560            assert_eq!(ab_hi.to_array().as_ref(), expected_hi);
1561        });
1562    }
1563
1564    #[test]
1565    fn test_reinterpret_cast() {
1566        test_simd_op!(isa, {
1567            let x = 1.456f32;
1568            let x_i32 = x.to_bits() as i32;
1569
1570            let x_vec = isa.f32().splat(x);
1571            let y_vec: I::I32 = x_vec.reinterpret_cast();
1572
1573            let expected = isa.i32().splat(x_i32);
1574            assert_simd_eq!(y_vec, expected);
1575        })
1576    }
1577
1578    #[test]
1579    fn test_to_float_i32() {
1580        test_simd_op!(isa, {
1581            let x = isa.i32().splat(42);
1582            let y = isa.i32().to_float(x);
1583
1584            let expected = isa.f32().splat(42.0);
1585            assert_simd_eq!(y, expected);
1586        });
1587    }
1588
1589    #[test]
1590    fn test_f16_bit_ops() {
1591        test_simd_op!(isa, {
1592            let ops = isa.f16();
1593
1594            // Load / store round trip across several vectors.
1595            let src: Vec<f16> = (0..ops.len() * 2)
1596                .map(|i| f16::from_f32(i as f32))
1597                .collect();
1598            let mut dst = vec![f16::default(); src.len()];
1599            for (s, d) in src.chunks(ops.len()).zip(dst.chunks_mut(ops.len())) {
1600                let v = ops.load(s);
1601                ops.store(v, d);
1602            }
1603            assert_eq!(dst, src);
1604
1605            // Splat
1606            let val = f16::from_f32(2.5);
1607            let splatted = ops.splat(val).to_array();
1608            assert!(splatted.as_ref().iter().all(|&x| x == val));
1609
1610            // Bitwise ops
1611            let zeros = ops.zero();
1612            let ones = ops.not(zeros);
1613            assert_eq!(
1614                ops.and(ones, zeros).to_array().as_ref(),
1615                zeros.to_array().as_ref()
1616            );
1617            assert_eq!(
1618                ops.or(zeros, ones).to_array().as_ref(),
1619                ones.to_array().as_ref()
1620            );
1621            assert_eq!(
1622                ops.xor(ones, ones).to_array().as_ref(),
1623                zeros.to_array().as_ref()
1624            );
1625        });
1626    }
1627
1628    #[test]
1629    fn test_f16_f32_conversion() {
1630        use crate::f16;
1631
1632        test_simd_op!(isa, {
1633            let f32_ops = isa.f32();
1634            let f16_ops = isa.f16();
1635
1636            let n = f32_ops.len() * 2;
1637            let src: Vec<f32> = (0..n).map(|i| (i as f32) * 0.25 - 3.0).collect();
1638
1639            // f32 -> f16
1640            let lo = f32_ops.load(&src[..f32_ops.len()]);
1641            let hi = f32_ops.load(&src[f32_ops.len()..]);
1642            let half = f32_ops.narrow_saturate(lo, hi);
1643
1644            let expected: Vec<f16> = src.iter().map(|&x| f16::from_f32(x)).collect();
1645            assert_eq!(half.to_array().as_ref(), expected.as_slice());
1646
1647            // f16 -> f32
1648            let back_lo = f16_ops.extend_low(half);
1649            let back_hi = f16_ops.extend_high(half);
1650            let expected_lo: Vec<f32> = expected[..f32_ops.len()]
1651                .iter()
1652                .map(|h| h.to_f32())
1653                .collect();
1654            let expected_hi: Vec<f32> = expected[f32_ops.len()..]
1655                .iter()
1656                .map(|h| h.to_f32())
1657                .collect();
1658            assert_eq!(back_lo.to_array().as_ref(), expected_lo.as_slice());
1659            assert_eq!(back_hi.to_array().as_ref(), expected_hi.as_slice());
1660        });
1661    }
1662}