Guarantee Mask has the same layout as Simd. Implement select as a trait that also supports bitmasks.

Author: calebzulawski

crates/core_simd/src/masks.rs

@@ -2,17 +2,10 @@
 //! Types representing
 #![allow(non_camel_case_types)]
 
-#[cfg_attr(
-    not(all(target_arch = "x86_64", target_feature = "avx512f")),
-    path = "masks/full_masks.rs"
-)]
-#[cfg_attr(
-    all(target_arch = "x86_64", target_feature = "avx512f"),
-    path = "masks/bitmask.rs"
-)]
+#[path = "masks/full_masks.rs"]
 mod mask_impl;
 
-use crate::simd::{LaneCount, Simd, SimdCast, SimdElement, SupportedLaneCount};
+use crate::simd::{LaneCount, Select, Simd, SimdCast, SimdElement, SupportedLaneCount};
 use core::cmp::Ordering;
 use core::{fmt, mem};
 
@@ -105,9 +98,8 @@ impl_element! { isize, usize }
 ///
 /// Masks represent boolean inclusion/exclusion on a per-element basis.
 ///
-/// The layout of this type is unspecified, and may change between platforms
-/// and/or Rust versions, and code should not assume that it is equivalent to
-/// `[T; N]`.
+/// The layout of this type is equivalent to `Simd<T, N>`, but elements
+/// are guaranteed to be either 0 or -1.
 #[repr(transparent)]
 pub struct Mask<T, const N: usize>(mask_impl::Mask<T, N>)
 where

crates/core_simd/src/mod.rs

@@ -29,6 +29,7 @@ pub mod simd {
     pub use crate::core_simd::cast::*;
     pub use crate::core_simd::lane_count::{LaneCount, SupportedLaneCount};
     pub use crate::core_simd::masks::*;
+    pub use crate::core_simd::select::*;
     pub use crate::core_simd::swizzle::*;
     pub use crate::core_simd::to_bytes::ToBytes;
     pub use crate::core_simd::vector::*;

crates/core_simd/src/ops.rs

@@ -1,4 +1,4 @@
-use crate::simd::{LaneCount, Simd, SimdElement, SupportedLaneCount, cmp::SimdPartialEq};
+use crate::simd::{LaneCount, Select, Simd, SimdElement, SupportedLaneCount, cmp::SimdPartialEq};
 use core::ops::{Add, Mul};
 use core::ops::{BitAnd, BitOr, BitXor};
 use core::ops::{Div, Rem, Sub};

crates/core_simd/src/select.rs

@@ -1,54 +1,167 @@
 use crate::simd::{LaneCount, Mask, MaskElement, Simd, SimdElement, SupportedLaneCount};
 
-impl<T, const N: usize> Mask<T, N>
+/// Choose elements from two vectors using a mask.
+///
+/// For each element in the mask, choose the corresponding element from `true_values` if
+/// that element mask is true, and `false_values` if that element mask is false.
+///
+/// If the mask is `u64`, it's treated as a bitmask with the least significant bit
+/// corresponding to the first element.
+///
+/// # Examples
+///
+/// ## Selecting values from `Simd`
+/// ```
+/// # #![feature(portable_simd)]
+/// # #[cfg(feature = "as_crate")] use core_simd::simd;
+/// # #[cfg(not(feature = "as_crate"))] use core::simd;
+/// # use simd::{Simd, Mask, Select};
+/// let a = Simd::from_array([0, 1, 2, 3]);
+/// let b = Simd::from_array([4, 5, 6, 7]);
+/// let mask = Mask::<i32, 4>::from_array([true, false, false, true]);
+/// let c = mask.select(a, b);
+/// assert_eq!(c.to_array(), [0, 5, 6, 3]);
+/// ```
+///
+/// ## Selecting values from `Mask`
+/// ```
+/// # #![feature(portable_simd)]
+/// # #[cfg(feature = "as_crate")] use core_simd::simd;
+/// # #[cfg(not(feature = "as_crate"))] use core::simd;
+/// # use simd::{Mask, Select};
+/// let a = Mask::<i32, 4>::from_array([true, true, false, false]);
+/// let b = Mask::<i32, 4>::from_array([false, false, true, true]);
+/// let mask = Mask::<i32, 4>::from_array([true, false, false, true]);
+/// let c = mask.select(a, b);
+/// assert_eq!(c.to_array(), [true, false, true, false]);
+/// ```
+///
+/// ## Selecting with a bitmask
+/// ```
+/// # #![feature(portable_simd)]
+/// # #[cfg(feature = "as_crate")] use core_simd::simd;
+/// # #[cfg(not(feature = "as_crate"))] use core::simd;
+/// # use simd::{Mask, Select};
+/// let a = Mask::<i32, 4>::from_array([true, true, false, false]);
+/// let b = Mask::<i32, 4>::from_array([false, false, true, true]);
+/// let mask = 0b1001;
+/// let c = mask.select(a, b);
+/// assert_eq!(c.to_array(), [true, false, true, false]);
+/// ```
+pub trait Select<T> {
+    /// Choose elements
+    fn select(self, true_values: T, false_values: T) -> T;
+}
+
+impl<T, U, const N: usize> Select<Simd<T, N>> for Mask<U, N>
+where
+    T: SimdElement,
+    U: MaskElement,
+    LaneCount<N>: SupportedLaneCount,
+{
+    #[inline]
+    fn select(self, true_values: Simd<T, N>, false_values: Simd<T, N>) -> Simd<T, N> {
+        // Safety:
+        // simd_as between masks is always safe (they're vectors of ints).
+        // simd_select uses a mask that matches the width and number of elements
+        unsafe {
+            let mask: Simd<T::Mask, N> = core::intrinsics::simd::simd_as(self.to_simd());
+            core::intrinsics::simd::simd_select(mask, true_values, false_values)
+        }
+    }
+}
+
+impl<T, const N: usize> Select<Simd<T, N>> for u64
+where
+    T: SimdElement,
+    LaneCount<N>: SupportedLaneCount,
+{
+    #[inline]
+    fn select(self, true_values: Simd<T, N>, false_values: Simd<T, N>) -> Simd<T, N> {
+        const {
+            assert!(N <= 64, "number of elements can't be greater than 64");
+        }
+
+        // LLVM assumes bit order should match endianness
+        let bitmask = if cfg!(target_endian = "big") {
+            let rev = self.reverse_bits();
+            if N < 64 {
+                // Shift things back to the right
+                rev >> (64 - N)
+            } else {
+                rev
+            }
+        } else {
+            self
+        };
+
+        #[inline]
+        unsafe fn select_impl<T, U, const M: usize, const N: usize>(
+            bitmask: U,
+            true_values: Simd<T, N>,
+            false_values: Simd<T, N>,
+        ) -> Simd<T, N>
+        where
+            T: SimdElement,
+            LaneCount<M>: SupportedLaneCount,
+            LaneCount<N>: SupportedLaneCount,
+        {
+            let default = true_values[0];
+            let true_values = true_values.resize::<M>(default);
+            let false_values = false_values.resize::<M>(default);
+
+            // Safety: the caller guarantees that the size of U matches M
+            let selected = unsafe {
+                core::intrinsics::simd::simd_select_bitmask(bitmask, true_values, false_values)
+            };
+
+            selected.resize::<N>(default)
+        }
+
+        // TODO modify simd_bitmask_select to truncate input, making this unnecessary
+        if N <= 8 {
+            // Safety: bitmask matches length
+            unsafe { select_impl::<T, u8, 8, N>(bitmask as u8, true_values, false_values) }
+        } else if N <= 16 {
+            // Safety: bitmask matches length
+            unsafe { select_impl::<T, u16, 16, N>(bitmask as u16, true_values, false_values) }
+        } else if N <= 32 {
+            // Safety: bitmask matches length
+            unsafe { select_impl::<T, u32, 32, N>(bitmask as u32, true_values, false_values) }
+        } else {
+            // Safety: bitmask matches length
+            unsafe { select_impl::<T, u64, 64, N>(bitmask, true_values, false_values) }
+        }
+    }
+}
+
+impl<T, U, const N: usize> Select<Mask<T, N>> for Mask<U, N>
 where
     T: MaskElement,
+    U: MaskElement,
     LaneCount<N>: SupportedLaneCount,
 {
-    /// Choose elements from two vectors.
-    ///
-    /// For each element in the mask, choose the corresponding element from `true_values` if
-    /// that element mask is true, and `false_values` if that element mask is false.
-    ///
-    /// # Examples
-    /// ```
-    /// # #![feature(portable_simd)]
-    /// # use core::simd::{Simd, Mask};
-    /// let a = Simd::from_array([0, 1, 2, 3]);
-    /// let b = Simd::from_array([4, 5, 6, 7]);
-    /// let mask = Mask::from_array([true, false, false, true]);
-    /// let c = mask.select(a, b);
-    /// assert_eq!(c.to_array(), [0, 5, 6, 3]);
-    /// ```
     #[inline]
-    #[must_use = "method returns a new vector and does not mutate the original inputs"]
-    pub fn select<U>(self, true_values: Simd<U, N>, false_values: Simd<U, N>) -> Simd<U, N>
-    where
-        U: SimdElement<Mask = T>,
-    {
-        // Safety: The mask has been cast to a vector of integers,
-        // and the operands to select between are vectors of the same type and length.
-        unsafe { core::intrinsics::simd::simd_select(self.to_simd(), true_values, false_values) }
+    fn select(self, true_values: Mask<T, N>, false_values: Mask<T, N>) -> Mask<T, N> {
+        let selected: Simd<T, N> =
+            Select::select(self, true_values.to_simd(), false_values.to_simd());
+
+        // Safety: all values come from masks
+        unsafe { Mask::from_simd_unchecked(selected) }
     }
+}
 
-    /// Choose elements from two masks.
-    ///
-    /// For each element in the mask, choose the corresponding element from `true_values` if
-    /// that element mask is true, and `false_values` if that element mask is false.
-    ///
-    /// # Examples
-    /// ```
-    /// # #![feature(portable_simd)]
-    /// # use core::simd::Mask;
-    /// let a = Mask::<i32, 4>::from_array([true, true, false, false]);
-    /// let b = Mask::<i32, 4>::from_array([false, false, true, true]);
-    /// let mask = Mask::<i32, 4>::from_array([true, false, false, true]);
-    /// let c = mask.select_mask(a, b);
-    /// assert_eq!(c.to_array(), [true, false, true, false]);
-    /// ```
+impl<T, const N: usize> Select<Mask<T, N>> for u64
+where
+    T: MaskElement,
+    LaneCount<N>: SupportedLaneCount,
+{
     #[inline]
-    #[must_use = "method returns a new mask and does not mutate the original inputs"]
-    pub fn select_mask(self, true_values: Self, false_values: Self) -> Self {
-        self & true_values | !self & false_values
+    fn select(self, true_values: Mask<T, N>, false_values: Mask<T, N>) -> Mask<T, N> {
+        let selected: Simd<T, N> =
+            Select::select(self, true_values.to_simd(), false_values.to_simd());
+
+        // Safety: all values come from masks
+        unsafe { Mask::from_simd_unchecked(selected) }
     }
 }

crates/core_simd/src/simd/cmp/ord.rs

@@ -1,5 +1,5 @@
 use crate::simd::{
-    LaneCount, Mask, Simd, SupportedLaneCount,
+    LaneCount, Mask, Select, Simd, SupportedLaneCount,
     cmp::SimdPartialEq,
     ptr::{SimdConstPtr, SimdMutPtr},
 };
@@ -194,12 +194,12 @@ macro_rules! impl_mask {
         {
             #[inline]
             fn simd_max(self, other: Self) -> Self {
-                self.simd_gt(other).select_mask(other, self)
+                self.simd_gt(other).select(other, self)
             }
 
             #[inline]
             fn simd_min(self, other: Self) -> Self {
-                self.simd_lt(other).select_mask(other, self)
+                self.simd_lt(other).select(other, self)
             }
 
             #[inline]

crates/core_simd/src/simd/num/float.rs

@@ -1,6 +1,6 @@
 use super::sealed::Sealed;
 use crate::simd::{
-    LaneCount, Mask, Simd, SimdCast, SimdElement, SupportedLaneCount,
+    LaneCount, Mask, Select, Simd, SimdCast, SimdElement, SupportedLaneCount,
     cmp::{SimdPartialEq, SimdPartialOrd},
 };
 

crates/core_simd/src/simd/num/int.rs

@@ -1,6 +1,6 @@
 use super::sealed::Sealed;
 use crate::simd::{
-    LaneCount, Mask, Simd, SimdCast, SimdElement, SupportedLaneCount, cmp::SimdOrd,
+    LaneCount, Mask, Select, Simd, SimdCast, SimdElement, SupportedLaneCount, cmp::SimdOrd,
     cmp::SimdPartialOrd, num::SimdUint,
 };
 

crates/core_simd/src/swizzle_dyn.rs

@@ -1,4 +1,4 @@
-use crate::simd::{LaneCount, Simd, SupportedLaneCount};
+use crate::simd::{LaneCount, Select, Simd, SupportedLaneCount};
 use core::mem;
 
 impl<const N: usize> Simd<u8, N>