fix : SIMD

Author: minjaekimartist

src/simd/generic/avx512.rs

@@ -1,20 +1,104 @@
-//! AVX-512 SIMD backend — 16-wide f32 using __m512.
+//! AVX-512 SIMD backend — 16-wide f32 using `__m512`.
 //!
-//! Enabled with `--features avx512` on x86_64 targets.
-//! Requires nightly Rust for AVX-512 intrinsics.
-
-// AVX-512 support requires nightly and is a placeholder for future implementation.
-// The trait implementation follows the same pattern as AVX2 but with __m512 / 16 lanes.
-
-// When stabilized, this will use:
-// - _mm512_load_ps, _mm512_store_ps
-// - _mm512_set1_ps
-// - _mm512_add_ps, _mm512_sub_ps, _mm512_mul_ps, _mm512_div_ps
-// - _mm512_fmadd_ps
-// - _mm512_max_ps, _mm512_min_ps
-// - _mm512_abs_ps
-// - _mm512_cmp_ps_mask + _mm512_mask_blend_ps
-
-// For now, re-export scalar as a placeholder if AVX-512 intrinsics aren't available.
-#[cfg(not(all(target_arch = "x86_64", target_feature = "avx512f")))]
-pub type F32x16 = super::scalar::ScalarFloat;
+//! Requires the `avx512f` target feature at the call site (guarded by
+//! `#[target_feature]` + `is_x86_feature_detected!` in `crate::simd::x86_64`
+//! - an internal module - or by the `sim-avx512` opt-in feature assuming
+//! target hardware support in [`crate::sim`]).
+
+use core::arch::x86_64::*;
+
+use super::traits::SimdFloat;
+
+/// 16-wide f32 SIMD type using AVX-512F.
+#[derive(Copy, Clone)]
+#[repr(transparent)]
+pub struct F32x16(__m512);
+
+impl SimdFloat for F32x16 {
+    const WIDTH: usize = 16;
+
+    #[inline(always)]
+    unsafe fn load(ptr: *const f32) -> Self {
+        F32x16(unsafe { _mm512_loadu_ps(ptr) })
+    }
+
+    #[inline(always)]
+    unsafe fn store(self, ptr: *mut f32) {
+        unsafe { _mm512_storeu_ps(ptr, self.0) };
+    }
+
+    #[inline(always)]
+    fn splat(v: f32) -> Self {
+        F32x16(unsafe { _mm512_set1_ps(v) })
+    }
+
+    #[inline(always)]
+    fn add(self, rhs: Self) -> Self {
+        F32x16(unsafe { _mm512_add_ps(self.0, rhs.0) })
+    }
+
+    #[inline(always)]
+    fn sub(self, rhs: Self) -> Self {
+        F32x16(unsafe { _mm512_sub_ps(self.0, rhs.0) })
+    }
+
+    #[inline(always)]
+    fn mul(self, rhs: Self) -> Self {
+        F32x16(unsafe { _mm512_mul_ps(self.0, rhs.0) })
+    }
+
+    #[inline(always)]
+    fn div(self, rhs: Self) -> Self {
+        F32x16(unsafe { _mm512_div_ps(self.0, rhs.0) })
+    }
+
+    #[inline(always)]
+    fn fma(self, b: Self, c: Self) -> Self {
+        F32x16(unsafe { _mm512_fmadd_ps(self.0, b.0, c.0) })
+    }
+
+    #[inline(always)]
+    fn max(self, rhs: Self) -> Self {
+        F32x16(unsafe { _mm512_max_ps(self.0, rhs.0) })
+    }
+
+    #[inline(always)]
+    fn min(self, rhs: Self) -> Self {
+        F32x16(unsafe { _mm512_min_ps(self.0, rhs.0) })
+    }
+
+    #[inline(always)]
+    fn abs(self) -> Self {
+        let mask = unsafe { _mm512_castsi512_ps(_mm512_set1_epi32(0x7FFF_FFFF_u32 as i32)) };
+        F32x16(unsafe { _mm512_and_ps(self.0, mask) })
+    }
+
+    #[inline(always)]
+    fn neg(self) -> Self {
+        let zero = unsafe { _mm512_setzero_ps() };
+        F32x16(unsafe { _mm512_sub_ps(zero, self.0) })
+    }
+
+    #[inline(always)]
+    fn cmp_ge(self, rhs: Self) -> Self {
+        // AVX-512 compares produce a 16-bit mask register, not a vector -
+        // broadcast it back into a full lane-wise all-1s/all-0s vector so
+        // this matches every other backend's `SimdFloat::cmp_ge` contract.
+        let mask = unsafe { _mm512_cmp_ps_mask(self.0, rhs.0, _CMP_GE_OQ) };
+        let all_ones = unsafe { _mm512_castsi512_ps(_mm512_set1_epi32(-1)) };
+        F32x16(unsafe { _mm512_maskz_mov_ps(mask, all_ones) })
+    }
+
+    #[inline(always)]
+    fn blend(mask: Self, a: Self, b: Self) -> Self {
+        // Inverse of `cmp_ge`: collapse the lane-wise vector mask back into
+        // a mask register for `_mm512_mask_blend_ps`.
+        let k = unsafe { _mm512_movepi32_mask(_mm512_castps_si512(mask.0)) };
+        F32x16(unsafe { _mm512_mask_blend_ps(k, b.0, a.0) })
+    }
+
+    #[inline(always)]
+    fn first(self) -> f32 {
+        unsafe { _mm_cvtss_f32(_mm512_castps512_ps128(self.0)) }
+    }
+}

src/simd/generic/dispatch.rs

@@ -1,29 +1,41 @@
 // Runtime SIMD dispatch — selects the best available backend.
 //
-// Currently returns the compile-time selected backend.
-// Platform-specific backends use `is_x86_feature_detected!` etc.
-// at initialization time, then store the result in a static.
+// Considers every backend compiled in: unconditionally under the `simd`
+// feature (this crate's own runtime-detected `x86_64`/`aarch64` layer,
+// which must work correctly on whatever CPU the binary actually ends up
+// running on), or under the matching `sim-avx2`/`sim-avx512`/`sim-neon`
+// opt-in feature (`crate::sim`'s compile-time-selected layer, which
+// assumes the build targets hardware known to support it - `sim`'s own
+// callers don't invoke this function to make that choice, they just
+// monomorphize against one backend directly, but `simd`'s callers do).
 
-/// Query which SIMD backend is available at runtime.
+/// Query which SIMD backend is available at runtime, among those compiled in.
 pub fn detect_backend() -> SimdBackend {
     #[cfg(target_arch = "x86_64")]
     {
-        #[cfg(feature = "sim-avx512")]
+        #[cfg(any(feature = "simd", feature = "sim-avx512"))]
         if is_x86_feature_detected!("avx512f") {
             return SimdBackend::Avx512;
         }
 
-        #[cfg(feature = "sim-avx2")]
+        #[cfg(any(feature = "simd", feature = "sim-avx2"))]
         if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
             return SimdBackend::Avx2;
         }
+
+        // SSE2 is guaranteed present on every x86_64 target, so no runtime
+        // check is needed - it's the floor beneath AVX2/AVX-512 for the
+        // `simd` feature's portable, runtime-detected layer specifically
+        // (`sim` has no equivalent opt-in feature for it).
+        #[cfg(feature = "simd")]
+        return SimdBackend::Sse2;
     }
 
     // NEON is always available on aarch64
-    #[cfg(all(target_arch = "aarch64", feature = "sim-neon"))]
+    #[cfg(all(target_arch = "aarch64", any(feature = "simd", feature = "sim-neon")))]
     return SimdBackend::Neon;
 
-    #[cfg(not(all(target_arch = "aarch64", feature = "sim-neon")))]
+    #[allow(unreachable_code)]
     SimdBackend::Scalar
 }
 
@@ -32,6 +44,8 @@ pub fn detect_backend() -> SimdBackend {
 pub enum SimdBackend {
     /// No hardware SIMD; one lane at a time.
     Scalar,
+    /// x86_64 baseline SSE2 (4 `f32` lanes).
+    Sse2,
     /// x86_64 AVX2 (8 `f32` lanes).
     Avx2,
     /// x86_64 AVX-512 (16 `f32` lanes).
@@ -45,9 +59,21 @@ impl SimdBackend {
     pub fn width(self) -> usize {
         match self {
             SimdBackend::Scalar => 1,
+            SimdBackend::Sse2 => 4,
             SimdBackend::Avx2 => 8,
             SimdBackend::Avx512 => 16,
             SimdBackend::Neon => 4,
         }
     }
+
+    /// Short identifying name (`"scalar"`, `"sse2"`, `"avx2"`, `"avx512"`, `"neon"`).
+    pub fn as_str(self) -> &'static str {
+        match self {
+            SimdBackend::Scalar => "scalar",
+            SimdBackend::Sse2 => "sse2",
+            SimdBackend::Avx2 => "avx2",
+            SimdBackend::Avx512 => "avx512",
+            SimdBackend::Neon => "neon",
+        }
+    }
 }

src/simd/generic/mod.rs

@@ -1,25 +1,30 @@
 //! SIMD abstraction layer and fast math functions.
 //!
 //! All numerical computation in this module is written generically over the
-//! [`SimdFloat`] trait. The default [`ScalarFloat`] backend (WIDTH=1) works
-//! everywhere; platform backends are activated via feature flags:
+//! [`crate::simd::generic::SimdFloat`] trait. The default
+//! [`crate::simd::generic::ScalarFloat`] backend (WIDTH=1) works everywhere;
+//! platform backends are activated by [`crate::sim`]'s compile-time-selected
+//! opt-in features, or unconditionally by [`crate::simd`]'s own
+//! runtime-detected `simd` feature (see
+//! [`crate::simd::generic::dispatch::detect_backend`]):
 //!
-//! | Feature       | Type      | WIDTH | Platform  |
-//! |---------------|-----------|-------|-----------|
-//! | (none)        | `ScalarFloat` | 1 | all       |
-//! | `sim-avx2`    | `F32x8`   | 8     | x86_64    |
-//! | `sim-avx512`  | `F32x16`  | 16    | x86_64    |
-//! | `sim-neon`    | `F32x4`   | 4     | aarch64   |
+//! | Feature       | Type       | WIDTH | Platform  |
+//! |---------------|------------|-------|-----------|
+//! | (none)        | `ScalarFloat` | 1  | all       |
+//! | `simd`        | `F32x4Sse` | 4     | x86_64 (baseline floor)|
+//! | `sim-avx2`    | `F32x8`    | 8     | x86_64    |
+//! | `sim-avx512`  | `F32x16`   | 16    | x86_64    |
+//! | `sim-neon`    | `F32x4`    | 4     | aarch64   |
 //!
 //! # Fast Math
 //!
 //! The following functions are provided with < 1e-5 relative error:
 //!
-//! - [`simd_exp`] -- exponential via range reduction + degree-5 polynomial
-//! - [`simd_log`] -- natural log via atanh series expansion
-//! - [`simd_log1pexp`] -- numerically stable softplus
-//! - [`simd_sigmoid`] -- logistic sigmoid without overflow
-//! - [`simd_recip`], [`simd_rsqrt`] -- Newton-refined reciprocal/inverse sqrt
+//! - [`crate::simd::generic::simd_exp`] -- exponential via range reduction + degree-5 polynomial
+//! - [`crate::simd::generic::simd_log`] -- natural log via atanh series expansion
+//! - [`crate::simd::generic::simd_log1pexp`] -- numerically stable softplus
+//! - [`crate::simd::generic::simd_sigmoid`] -- logistic sigmoid without overflow
+//! - [`crate::simd::generic::simd_recip`], [`crate::simd::generic::simd_rsqrt`] -- Newton-refined reciprocal/inverse sqrt
 
 // `exp`'s Cody-Waite range-reduction constants are intentionally given to
 // more decimal digits than `f32` can represent exactly (for readability
@@ -36,7 +41,7 @@ pub mod reciprocal;
 pub mod scalar;
 /// SIMD logistic sigmoid (`simd_sigmoid`).
 pub mod sigmoid;
-/// The [`SimdFloat`] abstraction trait itself.
+/// The [`crate::simd::generic::SimdFloat`] abstraction trait itself.
 pub mod traits;
 
 /// AVX2 `F32x8` backend, `x86_64` only. Compiled whenever either
@@ -46,29 +51,35 @@ pub mod traits;
 #[cfg(all(target_arch = "x86_64", any(feature = "simd", feature = "sim-avx2")))]
 pub mod avx2;
 
-/// AVX-512 `F32x16` backend (`sim-avx512` feature, `x86_64` only).
-#[cfg(all(target_arch = "x86_64", feature = "sim-avx512"))]
+/// AVX-512 `F32x16` backend, `x86_64` only. Compiled whenever either
+/// [`crate::sim`]'s compile-time-selected `sim-avx512` backend or
+/// [`crate::simd`]'s runtime-dispatched `simd` feature needs it, mirroring
+/// [`crate::simd::generic::avx2`]'s dual use.
+#[cfg(all(target_arch = "x86_64", any(feature = "simd", feature = "sim-avx512")))]
 pub mod avx512;
 
 /// ARM NEON `F32x4` backend, `aarch64` only. Compiled whenever either
 /// [`crate::sim`]'s compile-time-selected `sim-neon` backend or
 /// [`crate::simd`]'s runtime-dispatched `simd` feature needs it, mirroring
-/// [`avx2`]'s dual use.
+/// [`crate::simd::generic::avx2`]'s dual use.
 #[cfg(all(target_arch = "aarch64", any(feature = "simd", feature = "sim-neon")))]
 pub mod neon;
 
-/// Baseline SSE2 `F32x4Sse` backend, `x86_64` only - see [`sse2`]'s module
-/// docs for why it needs no opt-in feature of its own.
+/// Baseline SSE2 `F32x4Sse` backend, `x86_64` only - see
+/// [`crate::simd::generic::sse2`]'s module docs for why it needs no opt-in
+/// feature of its own.
 #[cfg(target_arch = "x86_64")]
 pub mod sse2;
 
 /// Width-agnostic `dot`/`mul_elementwise`/`mix_scalar` shared by every
-/// [`SimdFloat`] backend - see [`ops`]'s module docs.
+/// [`crate::simd::generic::SimdFloat`] backend - see
+/// [`crate::simd::generic::ops`]'s module docs.
 pub mod ops;
 
-/// Runtime backend selection ([`dispatch::detect_backend`]).
+/// Runtime backend selection ([`crate::simd::generic::dispatch::detect_backend`]).
 pub mod dispatch;
 
+pub use dispatch::SimdBackend;
 pub use scalar::ScalarFloat;
 pub use traits::SimdFloat;
 
@@ -78,23 +89,36 @@ pub use reciprocal::{simd_recip, simd_rsqrt};
 pub use sigmoid::simd_sigmoid;
 
 /// Select the best available SIMD backend at compile time.
-/// Returns a string identifying the active backend.
+///
+/// This is [`crate::sim`]'s own compile-time choice (which `sim-*` feature
+/// was enabled), distinct from
+/// [`crate::simd::generic::dispatch::detect_backend`]'s runtime CPU check -
+/// see that function's docs for why the two differ. Returns a string
+/// identifying the active backend, from the same vocabulary as
+/// [`SimdBackend::as_str`].
 pub fn active_backend() -> &'static str {
     #[cfg(all(target_arch = "x86_64", feature = "sim-avx512"))]
     {
-        return "avx512";
+        return SimdBackend::Avx512.as_str();
     }
 
-    #[cfg(all(target_arch = "x86_64", feature = "sim-avx2"))]
+    // `not(sim-avx512)` keeps this mutually exclusive with the branch above
+    // so `--all-features` (which enables every `sim-*` feature at once)
+    // doesn't produce two unconditional `return`s in a row.
+    #[cfg(all(
+        target_arch = "x86_64",
+        feature = "sim-avx2",
+        not(feature = "sim-avx512")
+    ))]
     {
-        return "avx2";
+        return SimdBackend::Avx2.as_str();
     }
 
     #[cfg(all(target_arch = "aarch64", feature = "sim-neon"))]
     {
-        return "neon";
+        return SimdBackend::Neon.as_str();
     }
 
     #[allow(unreachable_code)]
-    "scalar"
+    SimdBackend::Scalar.as_str()
 }

src/simd/generic/sse2.rs

@@ -2,7 +2,8 @@
 //!
 //! SSE2 is guaranteed present on every `x86_64` target (it's part of the
 //! baseline ABI), so this needs no runtime feature detection and no opt-in
-//! Cargo feature - it's the always-available floor beneath [`super::avx2`].
+//! Cargo feature - it's the always-available floor beneath
+//! [`crate::simd::generic::avx2`].
 
 use core::arch::x86_64::*;