Fix endianness correction

Author: calebzulawski

crates/core_simd/src/masks.rs

@@ -6,6 +6,36 @@ use crate::simd::{LaneCount, Select, Simd, SimdCast, SimdElement, SupportedLaneC
 use core::cmp::Ordering;
 use core::{fmt, mem};
 
+pub(crate) trait FixEndianness {
+    fn fix_endianness(self, elements: usize) -> Self;
+}
+
+macro_rules! impl_fix_endianness {
+    { $($int:ty),* } => {
+        $(
+        impl FixEndianness for $int {
+            #[inline(always)]
+            fn fix_endianness(self, elements: usize) -> Self {
+                if cfg!(target_endian = "big") {
+                    let rev = <$int>::reverse_bits(self);
+                    let bitsize = size_of::<$int>() * 8;
+                    if elements < bitsize {
+                        // Shift things back to the right
+                        rev >> (bitsize - elements)
+                    } else {
+                        rev
+                    }
+                } else {
+                    self
+                }
+            }
+        }
+        )*
+    }
+}
+
+impl_fix_endianness! { u8, u16, u32, u64 }
+
 mod sealed {
     use super::*;
 
@@ -283,7 +313,9 @@ where
     #[must_use = "method returns a new integer and does not mutate the original value"]
     pub fn to_bitmask(self) -> u64 {
         #[inline]
-        unsafe fn to_bitmask_impl<T, U, const M: usize, const N: usize>(mask: Mask<T, N>) -> U
+        unsafe fn to_bitmask_impl<T, U: FixEndianness, const M: usize, const N: usize>(
+            mask: Mask<T, N>,
+        ) -> U
         where
             T: MaskElement,
             LaneCount<M>: SupportedLaneCount,
@@ -292,11 +324,14 @@ where
             let resized = mask.resize::<M>(false);
 
             // Safety: `resized` is an integer vector with length M, which must match T
-            unsafe { core::intrinsics::simd::simd_bitmask(resized.0) }
+            let bitmask: U = unsafe { core::intrinsics::simd::simd_bitmask(resized.0) };
+
+            // LLVM assumes bit order should match endianness
+            bitmask.fix_endianness(N)
         }
 
         // TODO modify simd_bitmask to zero-extend output, making this unnecessary
-        let bitmask = if N <= 8 {
+        if N <= 8 {
             // Safety: bitmask matches length
             unsafe { to_bitmask_impl::<T, u8, 8, N>(self) as u64 }
         } else if N <= 16 {
@@ -308,13 +343,6 @@ where
         } else {
             // Safety: bitmask matches length
             unsafe { to_bitmask_impl::<T, u64, 64, N>(self) }
-        };
-
-        // LLVM assumes bit order should match endianness
-        if cfg!(target_endian = "big") {
-            bitmask.reverse_bits() >> (64 - N.min(64))
-        } else {
-            bitmask
         }
     }
 

crates/core_simd/src/select.rs

@@ -1,4 +1,6 @@
-use crate::simd::{LaneCount, Mask, MaskElement, Simd, SimdElement, SupportedLaneCount};
+use crate::simd::{
+    FixEndianness, LaneCount, Mask, MaskElement, Simd, SimdElement, SupportedLaneCount,
+};
 
 /// Choose elements from two vectors using a mask.
 ///
@@ -82,21 +84,8 @@ where
             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>(
+        unsafe fn select_impl<T, U: FixEndianness, const M: usize, const N: usize>(
             bitmask: U,
             true_values: Simd<T, N>,
             false_values: Simd<T, N>,
@@ -110,6 +99,9 @@ where
             let true_values = true_values.resize::<M>(default);
             let false_values = false_values.resize::<M>(default);
 
+            // LLVM assumes bit order should match endianness
+            let bitmask = bitmask.fix_endianness(N);
+
             // 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)
@@ -120,15 +112,19 @@ where
 
         // TODO modify simd_bitmask_select to truncate input, making this unnecessary
         if N <= 8 {
+            let bitmask = self as u8;
             // Safety: bitmask matches length
-            unsafe { select_impl::<T, u8, 8, N>(bitmask as u8, true_values, false_values) }
+            unsafe { select_impl::<T, u8, 8, N>(bitmask, true_values, false_values) }
         } else if N <= 16 {
+            let bitmask = self as u16;
             // Safety: bitmask matches length
-            unsafe { select_impl::<T, u16, 16, N>(bitmask as u16, true_values, false_values) }
+            unsafe { select_impl::<T, u16, 16, N>(bitmask, true_values, false_values) }
         } else if N <= 32 {
+            let bitmask = self as u32;
             // Safety: bitmask matches length
-            unsafe { select_impl::<T, u32, 32, N>(bitmask as u32, true_values, false_values) }
+            unsafe { select_impl::<T, u32, 32, N>(bitmask, true_values, false_values) }
         } else {
+            let bitmask = self;
             // Safety: bitmask matches length
             unsafe { select_impl::<T, u64, 64, N>(bitmask, true_values, false_values) }
         }

crates/core_simd/src/swizzle_dyn.rs

@@ -139,7 +139,7 @@ unsafe fn armv7_neon_swizzle_u8x16(bytes: Simd<u8, 16>, idxs: Simd<u8, 16>) -> S
 #[inline]
 #[allow(clippy::let_and_return)]
 unsafe fn avx2_pshufb(bytes: Simd<u8, 32>, idxs: Simd<u8, 32>) -> Simd<u8, 32> {
-    use crate::simd::{cmp::SimdPartialOrd, Select};
+    use crate::simd::{Select, cmp::SimdPartialOrd};
     #[cfg(target_arch = "x86")]
     use core::arch::x86;
     #[cfg(target_arch = "x86_64")]
@@ -200,7 +200,7 @@ fn zeroing_idxs<const N: usize>(idxs: Simd<u8, N>) -> Simd<u8, N>
 where
     LaneCount<N>: SupportedLaneCount,
 {
-    use crate::simd::{cmp::SimdPartialOrd, Select};
+    use crate::simd::{Select, cmp::SimdPartialOrd};
     idxs.simd_lt(Simd::splat(N as u8))
         .select(idxs, Simd::splat(u8::MAX))
 }