Determine SIMD at runtime

ext/crates/fp/build.rs

@@ -29,8 +29,6 @@ fn main() -> io::Result<()> {
     writer.add_value("NUM_PRIMES", "usize", num_primes);
     writer.add_raw("/// The `NUM_PRIMES`th prime number.");
     writer.add_value("MAX_PRIME", "usize", max_prime);
-    // `NOT_A_PRIME` is never used if odd-primes is disabled.
-    writer.add_raw("#[allow(dead_code)]");
     writer.add_raw(
         "/// A sentinel value. `PRIME_TO_INDEX_MAP[i] == NOT_A_PRIME` if and only if `i` is not",
     );

ext/crates/fp/src/constants.rs

@@ -1,3 +1,5 @@
+#![allow(dead_code)]
+
 use build_const::build_const;
 
 build_const!("constants");

ext/crates/fp/src/simd/generic.rs

@@ -1,15 +1,7 @@
 use crate::limb::Limb;
 
-pub(crate) type SimdLimb = Limb;
-
-pub(crate) unsafe fn load(limb: *const Limb) -> SimdLimb {
-    *limb
-}
-
-pub(crate) unsafe fn store(limb: *mut Limb, val: SimdLimb) {
-    *limb = val;
-}
-
-pub(crate) unsafe fn xor(left: SimdLimb, right: SimdLimb) -> SimdLimb {
-    left ^ right
+pub(super) fn add_simd(target: &mut [Limb], source: &[Limb], min_limb: usize) {
+    for (target_limb, source_limb) in target.iter_mut().zip(source.iter()).skip(min_limb) {
+        *target_limb ^= source_limb
+    }
 }

ext/crates/fp/src/simd/mod.rs

@@ -1,34 +1,16 @@
-cfg_if::cfg_if! {
-    if #[cfg(target_arch = "x86_64")] {
-        mod x86_64;
-        use x86_64::*;
-    } else {
-        mod generic;
-        use generic::*;
-    }
-}
+use crate::limb::Limb;
 
-use super::limb::Limb;
+mod generic;
 
-const LIMBS_PER_SIMD: usize = std::mem::size_of::<SimdLimb>() / crate::constants::BYTES_PER_LIMB;
+#[cfg(target_arch = "x86_64")]
+mod x86_64;
 
 pub(crate) fn add_simd(target: &mut [Limb], source: &[Limb], min_limb: usize) {
-    let max_limb = target.len();
-    let target = target.as_mut_ptr();
-    let source = source.as_ptr();
-    let chunks = (max_limb - min_limb) / LIMBS_PER_SIMD;
-    for i in 0..chunks {
-        unsafe {
-            let mut target_chunk = load(target.add(LIMBS_PER_SIMD * i + min_limb));
-            let source_chunk = load(source.add(LIMBS_PER_SIMD * i + min_limb));
-            target_chunk = xor(target_chunk, source_chunk);
-            store(target.add(LIMBS_PER_SIMD * i + min_limb), target_chunk);
-        }
-    }
-    for i in (min_limb + LIMBS_PER_SIMD * chunks)..max_limb {
-        unsafe {
-            // pointer arithmetic
-            *target.add(i) = *target.add(i) ^ *source.add(i);
+    cfg_if::cfg_if! {
+        if #[cfg(target_arch = "x86_64")] {
+            x86_64::add_simd(target, source, min_limb)
+        } else {
+            generic::add_simd(target, source, min_limb)
         }
     }
 }

ext/crates/fp/src/simd/x86_64/avx.rs

@@ -0,0 +1,22 @@
+use std::arch::x86_64;
+
+use crate::limb::Limb;
+
+type SimdLimb = x86_64::__m256;
+
+#[target_feature(enable = "avx")]
+fn load(limb: *const Limb) -> SimdLimb {
+    unsafe { x86_64::_mm256_loadu_ps(limb as *const f32) }
+}
+
+#[target_feature(enable = "avx")]
+fn store(limb: *mut Limb, val: SimdLimb) {
+    unsafe { x86_64::_mm256_storeu_ps(limb as *mut f32, val) }
+}
+
+#[target_feature(enable = "avx")]
+fn xor(left: SimdLimb, right: SimdLimb) -> SimdLimb {
+    x86_64::_mm256_xor_ps(left, right)
+}
+
+super::add_simd_arch!("avx");

ext/crates/fp/src/simd/x86_64/avx2.rs

@@ -0,0 +1,22 @@
+use std::arch::x86_64;
+
+use crate::limb::Limb;
+
+type SimdLimb = x86_64::__m256i;
+
+#[target_feature(enable = "avx2")]
+fn load(limb: *const Limb) -> SimdLimb {
+    unsafe { x86_64::_mm256_loadu_si256(limb as *const SimdLimb) }
+}
+
+#[target_feature(enable = "avx2")]
+fn store(limb: *mut Limb, val: SimdLimb) {
+    unsafe { x86_64::_mm256_storeu_si256(limb as *mut SimdLimb, val) }
+}
+
+#[target_feature(enable = "avx2")]
+fn xor(left: SimdLimb, right: SimdLimb) -> SimdLimb {
+    x86_64::_mm256_xor_si256(left, right)
+}
+
+super::add_simd_arch!("avx2");

ext/crates/fp/src/simd/x86_64/avx512.rs

@@ -0,0 +1,22 @@
+use std::arch::x86_64;
+
+use crate::limb::Limb;
+
+type SimdLimb = x86_64::__m512i;
+
+#[target_feature(enable = "avx512f")]
+fn load(limb: *const Limb) -> SimdLimb {
+    unsafe { x86_64::_mm512_loadu_si512(limb as *const SimdLimb) }
+}
+
+#[target_feature(enable = "avx512f")]
+fn store(limb: *mut Limb, val: SimdLimb) {
+    unsafe { x86_64::_mm512_storeu_si512(limb as *mut SimdLimb, val) }
+}
+
+#[target_feature(enable = "avx512f")]
+fn xor(left: SimdLimb, right: SimdLimb) -> SimdLimb {
+    x86_64::_mm512_xor_si512(left, right)
+}
+
+super::add_simd_arch!("avx512f");

ext/crates/fp/src/simd/x86_64/mod.rs

@@ -0,0 +1,51 @@
+mod avx;
+mod avx2;
+mod avx512;
+mod sse2;
+
+use crate::limb::Limb;
+
+macro_rules! add_simd_arch {
+    ($arch:tt) => {
+        const LIMBS_PER_SIMD: usize =
+            std::mem::size_of::<SimdLimb>() / crate::constants::BYTES_PER_LIMB;
+
+        #[target_feature(enable = $arch)]
+        pub(super) fn add_simd(target: &mut [Limb], source: &[Limb], min_limb: usize) {
+            let max_limb = target.len();
+            let target = target.as_mut_ptr();
+            let source = source.as_ptr();
+            let chunks = (max_limb - min_limb) / LIMBS_PER_SIMD;
+            for i in 0..chunks {
+                unsafe {
+                    let mut target_chunk = load(target.add(LIMBS_PER_SIMD * i + min_limb));
+                    let source_chunk = load(source.add(LIMBS_PER_SIMD * i + min_limb));
+                    target_chunk = xor(target_chunk, source_chunk);
+                    store(target.add(LIMBS_PER_SIMD * i + min_limb), target_chunk);
+                }
+            }
+            for i in (min_limb + LIMBS_PER_SIMD * chunks)..max_limb {
+                unsafe {
+                    // pointer arithmetic
+                    *target.add(i) = *target.add(i) ^ *source.add(i);
+                }
+            }
+        }
+    };
+}
+
+use add_simd_arch;
+
+pub(super) fn add_simd(target: &mut [Limb], source: &[Limb], min_limb: usize) {
+    if is_x86_feature_detected!("avx512f") {
+        unsafe { avx512::add_simd(target, source, min_limb) }
+    } else if is_x86_feature_detected!("avx2") {
+        unsafe { avx2::add_simd(target, source, min_limb) }
+    } else if is_x86_feature_detected!("avx") {
+        unsafe { avx::add_simd(target, source, min_limb) }
+    } else if is_x86_feature_detected!("sse2") {
+        unsafe { sse2::add_simd(target, source, min_limb) }
+    } else {
+        super::generic::add_simd(target, source, min_limb)
+    }
+}

ext/crates/fp/src/simd/x86_64/sse2.rs

@@ -0,0 +1,22 @@
+use std::arch::x86_64;
+
+use crate::limb::Limb;
+
+type SimdLimb = x86_64::__m128i;
+
+#[target_feature(enable = "sse2")]
+fn load(limb: *const Limb) -> SimdLimb {
+    unsafe { x86_64::_mm_loadu_si128(limb as *const SimdLimb) }
+}
+
+#[target_feature(enable = "sse2")]
+fn store(limb: *mut Limb, val: SimdLimb) {
+    unsafe { x86_64::_mm_storeu_si128(limb as *mut SimdLimb, val) }
+}
+
+#[target_feature(enable = "sse2")]
+fn xor(left: SimdLimb, right: SimdLimb) -> SimdLimb {
+    x86_64::_mm_xor_si128(left, right)
+}
+
+super::add_simd_arch!("sse2");