Fix endianness correction

Author: calebzulawski

crates/core_simd/src/masks.rs

@@ -296,23 +296,23 @@ where
         }
 
         // TODO modify simd_bitmask to zero-extend output, making this unnecessary
-        let bitmask = if N <= 8 {
+        let (bitmask, n_used) = if N <= 8 {
             // Safety: bitmask matches length
-            unsafe { to_bitmask_impl::<T, u8, 8, N>(self) as u64 }
+            unsafe { (to_bitmask_impl::<T, u8, 8, N>(self) as u64, 8) }
         } else if N <= 16 {
             // Safety: bitmask matches length
-            unsafe { to_bitmask_impl::<T, u16, 16, N>(self) as u64 }
+            unsafe { (to_bitmask_impl::<T, u16, 16, N>(self) as u64, 16) }
         } else if N <= 32 {
             // Safety: bitmask matches length
-            unsafe { to_bitmask_impl::<T, u32, 32, N>(self) as u64 }
+            unsafe { (to_bitmask_impl::<T, u32, 32, N>(self) as u64, 32) }
         } else {
             // Safety: bitmask matches length
-            unsafe { to_bitmask_impl::<T, u64, 64, N>(self) }
+            unsafe { (to_bitmask_impl::<T, u64, 64, N>(self), 64) }
         };
 
         // LLVM assumes bit order should match endianness
         if cfg!(target_endian = "big") {
-            bitmask.reverse_bits() >> (64 - N.min(64))
+            bitmask.reverse_bits() >> (n_used - N.min(64))
         } else {
             bitmask
         }

crates/core_simd/src/select.rs

@@ -82,22 +82,15 @@ 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: core::ops::Shr<usize, Output = U>,
+            const M: usize,
+            const N: usize,
+        >(
             bitmask: U,
+            bitmask_reversed: U,
             true_values: Simd<T, N>,
             false_values: Simd<T, N>,
         ) -> Simd<T, N>
@@ -110,6 +103,13 @@ 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 = if cfg!(target_endian = "big") {
+                bitmask_reversed >> (M - N)
+            } else {
+                bitmask
+            };
+
             // 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,17 +120,49 @@ 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,
+                    bitmask.reverse_bits(),
+                    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,
+                    bitmask.reverse_bits(),
+                    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,
+                    bitmask.reverse_bits(),
+                    true_values,
+                    false_values,
+                )
+            }
         } else {
+            let bitmask = self as u64;
             // Safety: bitmask matches length
-            unsafe { select_impl::<T, u64, 64, N>(bitmask, true_values, false_values) }
+            unsafe {
+                select_impl::<T, u64, 64, N>(
+                    bitmask,
+                    bitmask.reverse_bits(),
+                    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))
 }