From d513647645122b81b92b75e5aefdd0d93d6f4986 Mon Sep 17 00:00:00 2001
From: Caleb Zulawski <caleb.zulawski@gmail.com>
Date: Thu, 6 Jun 2024 21:09:37 -0400
Subject: [PATCH] Add aarch64 workarounds

---
 crates/core_simd/src/ops.rs            | 29 ++++++++++++++++++++++++--
 crates/core_simd/src/simd/num/float.rs | 28 +++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/crates/core_simd/src/ops.rs b/crates/core_simd/src/ops.rs
index d8e10eeaa1a..5e026f1ddae 100644
--- a/crates/core_simd/src/ops.rs
+++ b/crates/core_simd/src/ops.rs
@@ -96,8 +96,33 @@ macro_rules! int_divrem_guard {
                 // Nice base case to make it easy to const-fold away the other branch.
                 $rhs
             };
-            // Safety: $lhs and rhs are vectors
-            unsafe { core::intrinsics::simd::$simd_call($lhs, rhs) }
+
+            // aarch64 fails for arbitrary `v % 0` for non-powers-of-two
+            const { assert!(Self::LEN <= 64) };
+            if cfg!(not(target_arch = "aarch64")) || Self::LEN <= 2 {
+                // Safety: $lhs and rhs are vectors
+                unsafe { core::intrinsics::simd::$simd_call($lhs, rhs) }
+            } else if Self::LEN <= 4 {
+                // Safety: $lhs and rhs are vectors
+                let x: Simd::<_, 4> = unsafe { core::intrinsics::simd::$simd_call($lhs.resize::<4>(Default::default()), rhs.resize::<4>(Default::default())) };
+                x.resize(Default::default())
+            } else if Self::LEN <= 8 {
+                // Safety: $lhs and rhs are vectors
+                let x: Simd::<_, 8> = unsafe { core::intrinsics::simd::$simd_call($lhs.resize::<8>(Default::default()), rhs.resize::<8>(Default::default())) };
+                x.resize(Default::default())
+            } else if Self::LEN <= 16 {
+                // Safety: $lhs and rhs are vectors
+                let x: Simd::<_, 16> = unsafe { core::intrinsics::simd::$simd_call($lhs.resize::<16>(Default::default()), rhs.resize::<16>(Default::default())) };
+                x.resize(Default::default())
+            } else if Self::LEN <= 32 {
+                // Safety: $lhs and rhs are vectors
+                let x: Simd::<_, 32> = unsafe { core::intrinsics::simd::$simd_call($lhs.resize::<32>(Default::default()), rhs.resize::<32>(Default::default())) };
+                x.resize(Default::default())
+            } else {
+                // Safety: $lhs and rhs are vectors
+                let x: Simd::<_, 64> = unsafe { core::intrinsics::simd::$simd_call($lhs.resize::<64>(Default::default()), rhs.resize::<64>(Default::default())) };
+                x.resize(Default::default())
+            }
         }
     };
 }
diff --git a/crates/core_simd/src/simd/num/float.rs b/crates/core_simd/src/simd/num/float.rs
index 59e43851ea8..a5d6b429f40 100644
--- a/crates/core_simd/src/simd/num/float.rs
+++ b/crates/core_simd/src/simd/num/float.rs
@@ -255,6 +255,7 @@ macro_rules! impl_trait {
             type Bits = Simd<$bits_ty, N>;
             type Cast<T: SimdElement> = Simd<T, N>;
 
+            #[cfg(not(target_arch = "aarch64"))]
             #[inline]
             fn cast<T: SimdCast>(self) -> Self::Cast<T>
             {
@@ -262,6 +263,33 @@ macro_rules! impl_trait {
                 unsafe { core::intrinsics::simd::simd_as(self) }
             }
 
+            // https://github.com/llvm/llvm-project/issues/94694
+            #[cfg(target_arch = "aarch64")]
+            #[inline]
+            fn cast<T: SimdCast>(self) -> Self::Cast<T>
+            {
+                const { assert!(N <= 64) };
+                if N <= 2 {
+                    // Safety: supported types are guaranteed by SimdCast
+                    unsafe { core::intrinsics::simd::simd_as(self) }
+                } else if N <= 4 {
+                    let x = self.resize::<4>(Default::default()).cast();
+                    x.resize::<N>(x[0])
+                } else if N <= 8 {
+                    let x = self.resize::<8>(Default::default()).cast();
+                    x.resize::<N>(x[0])
+                } else if N <= 16 {
+                    let x = self.resize::<16>(Default::default()).cast();
+                    x.resize::<N>(x[0])
+                } else if N <= 32 {
+                    let x = self.resize::<32>(Default::default()).cast();
+                    x.resize::<N>(x[0])
+                } else {
+                    let x = self.resize::<64>(Default::default()).cast();
+                    x.resize::<N>(x[0])
+                }
+            }
+
             #[inline]
             #[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
             unsafe fn to_int_unchecked<I: SimdCast>(self) -> Self::Cast<I>