lattice_embed/simd/

normalize.rs

//! SIMD in-place L2 normalization kernels.
//!
//! See docs/simd.md for the two-pass algorithm and backend differences.

#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::*;

#[cfg(target_arch = "aarch64")]
use std::arch::aarch64::*;

#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
use std::arch::wasm32::*;

use super::simd_config;

#[cfg(target_arch = "x86_64")]
use super::dot_product::{horizontal_sum_avx2, horizontal_sum_avx512};

#[cfg(target_arch = "aarch64")]
use super::dot_product::horizontal_sum_neon;

#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
use super::dot_product::horizontal_sum_simd128;

/// **Unstable**: SIMD dispatch layer; use `lattice_embed::utils::normalize` for the stable wrapper.
#[inline]
pub fn normalize(vector: &mut [f32]) {
    let config = simd_config();

    #[cfg(target_arch = "x86_64")]
    {
        if config.avx512f_enabled {
            // SAFETY: Runtime feature detection verified AVX-512F. The mutable
            // slice is valid for the call lifetime; the callee uses unaligned
            // loads/stores and chunk/remainder bounds that stay inside the slice.
            return unsafe { normalize_avx512_unrolled(vector) };
        }
        if config.avx2_enabled && config.fma_enabled {
            // SAFETY: Runtime feature detection verified AVX2+FMA. The mutable
            // slice is valid for the call lifetime; the callee uses unaligned
            // loads/stores and chunk/remainder bounds that stay inside the slice.
            return unsafe { normalize_avx2_unrolled(vector) };
        }
    }

    #[cfg(target_arch = "aarch64")]
    {
        if config.neon_enabled {
            // SAFETY: NEON is available on aarch64. The mutable slice is valid
            // for the call lifetime; the callee uses unaligned loads/stores and
            // bounded chunk/remainder loops that stay inside the slice.
            return unsafe { normalize_neon_unrolled(vector) };
        }
    }

    #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
    {
        if config.simd128_enabled() {
            // SAFETY: compiled with wasm32 simd128 (compile-time gate, see
            // `SimdConfig::simd128_enabled`). The mutable slice is valid for
            // the call lifetime; the callee uses alignment-free loads/stores
            // and bounded chunk/remainder loops that stay inside the slice.
            return unsafe { normalize_simd128_unrolled(vector) };
        }
    }

    normalize_scalar(vector)
}

/// Scalar normalization.
pub(crate) fn normalize_scalar(vector: &mut [f32]) {
    let norm: f32 = vector.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm > 0.0 {
        let inv_norm = 1.0 / norm;
        vector.iter_mut().for_each(|x| *x *= inv_norm);
    }
}

/// Normalizes in two passes with AVX-512F.
///
/// # Safety
/// Caller must provide AVX-512F; chunked unaligned access stays in bounds.
/// See [`docs/simd.md`](../../docs/simd.md#kernel-safety-boundary) for the shared kernel invariant.
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f")]
unsafe fn normalize_avx512_unrolled(vector: &mut [f32]) {
    const SIMD_WIDTH: usize = 16;
    const UNROLL: usize = 4;
    const CHUNK_SIZE: usize = SIMD_WIDTH * UNROLL;

    let n = vector.len();
    let chunks = n / CHUNK_SIZE;
    let main_processed = chunks * CHUNK_SIZE;
    let remaining = n - main_processed;
    let remaining_chunks = remaining / SIMD_WIDTH;

    // First pass: compute L2 norm with 4 accumulators
    let mut norm0 = _mm512_setzero_ps();
    let mut norm1 = _mm512_setzero_ps();
    let mut norm2 = _mm512_setzero_ps();
    let mut norm3 = _mm512_setzero_ps();

    for i in 0..chunks {
        let base = i * CHUNK_SIZE;

        let v0 = _mm512_loadu_ps(vector.as_ptr().add(base));
        norm0 = _mm512_fmadd_ps(v0, v0, norm0);

        let v1 = _mm512_loadu_ps(vector.as_ptr().add(base + SIMD_WIDTH));
        norm1 = _mm512_fmadd_ps(v1, v1, norm1);

        let v2 = _mm512_loadu_ps(vector.as_ptr().add(base + SIMD_WIDTH * 2));
        norm2 = _mm512_fmadd_ps(v2, v2, norm2);

        let v3 = _mm512_loadu_ps(vector.as_ptr().add(base + SIMD_WIDTH * 3));
        norm3 = _mm512_fmadd_ps(v3, v3, norm3);
    }

    let norm_vec = _mm512_add_ps(_mm512_add_ps(norm0, norm1), _mm512_add_ps(norm2, norm3));

    // Remainder for norm calculation with single-register AVX-512F loop
    let mut norm_remainder = _mm512_setzero_ps();
    for i in 0..remaining_chunks {
        let offset = main_processed + i * SIMD_WIDTH;
        let v = _mm512_loadu_ps(vector.as_ptr().add(offset));
        norm_remainder = _mm512_fmadd_ps(v, v, norm_remainder);
    }

    let mut norm_sq = horizontal_sum_avx512(norm_vec) + horizontal_sum_avx512(norm_remainder);

    // Scalar tail for norm (recomputed inline to avoid cross-pass variable dependency)
    for i in (main_processed + remaining_chunks * SIMD_WIDTH)..n {
        norm_sq += vector[i] * vector[i];
    }

    let norm = norm_sq.sqrt();
    // Match `normalize_scalar`, which only scales when `norm > 0.0`. Rejecting
    // NaN here too (a NaN element makes `norm` NaN) leaves the vector
    // byte-identical to the scalar path instead of scaling by a NaN inv_norm.
    // `is_nan() || <= 0.0` is the lint-clean equivalent of `!(norm > 0.0)`.
    if norm.is_nan() || norm <= 0.0 {
        return;
    }

    let inv_norm = 1.0 / norm;
    let inv_norm_vec = _mm512_set1_ps(inv_norm);

    // Second pass: scale by inverse norm with 4x unrolling
    for i in 0..chunks {
        let base = i * CHUNK_SIZE;

        let v0 = _mm512_loadu_ps(vector.as_ptr().add(base));
        _mm512_storeu_ps(
            vector.as_mut_ptr().add(base),
            _mm512_mul_ps(v0, inv_norm_vec),
        );

        let v1 = _mm512_loadu_ps(vector.as_ptr().add(base + SIMD_WIDTH));
        _mm512_storeu_ps(
            vector.as_mut_ptr().add(base + SIMD_WIDTH),
            _mm512_mul_ps(v1, inv_norm_vec),
        );

        let v2 = _mm512_loadu_ps(vector.as_ptr().add(base + SIMD_WIDTH * 2));
        _mm512_storeu_ps(
            vector.as_mut_ptr().add(base + SIMD_WIDTH * 2),
            _mm512_mul_ps(v2, inv_norm_vec),
        );

        let v3 = _mm512_loadu_ps(vector.as_ptr().add(base + SIMD_WIDTH * 3));
        _mm512_storeu_ps(
            vector.as_mut_ptr().add(base + SIMD_WIDTH * 3),
            _mm512_mul_ps(v3, inv_norm_vec),
        );
    }

    // Remainder for scaling with single-register AVX-512F loop
    for i in 0..remaining_chunks {
        let offset = main_processed + i * SIMD_WIDTH;
        let v = _mm512_loadu_ps(vector.as_ptr().add(offset));
        _mm512_storeu_ps(
            vector.as_mut_ptr().add(offset),
            _mm512_mul_ps(v, inv_norm_vec),
        );
    }

    // Final scalar remainder (recomputed inline to avoid cross-pass variable dependency)
    for i in (main_processed + remaining_chunks * SIMD_WIDTH)..n {
        vector[i] *= inv_norm;
    }
}

/// Normalizes in two passes with AVX2 and FMA.
///
/// # Safety
/// Caller must provide AVX2 and FMA; chunked unaligned access stays in bounds.
/// See [`docs/simd.md`](../../docs/simd.md#kernel-safety-boundary) for the shared kernel invariant.
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2", enable = "fma")]
unsafe fn normalize_avx2_unrolled(vector: &mut [f32]) {
    const SIMD_WIDTH: usize = 8;
    const UNROLL: usize = 4;
    const CHUNK_SIZE: usize = SIMD_WIDTH * UNROLL;
    let n = vector.len();
    let chunks = n / CHUNK_SIZE;

    // First pass: compute L2 norm with 4 accumulators
    let mut norm0 = _mm256_setzero_ps();
    let mut norm1 = _mm256_setzero_ps();
    let mut norm2 = _mm256_setzero_ps();
    let mut norm3 = _mm256_setzero_ps();

    for i in 0..chunks {
        let base = i * CHUNK_SIZE;

        let v0 = _mm256_loadu_ps(vector.as_ptr().add(base));
        norm0 = _mm256_fmadd_ps(v0, v0, norm0);

        let v1 = _mm256_loadu_ps(vector.as_ptr().add(base + SIMD_WIDTH));
        norm1 = _mm256_fmadd_ps(v1, v1, norm1);

        let v2 = _mm256_loadu_ps(vector.as_ptr().add(base + SIMD_WIDTH * 2));
        norm2 = _mm256_fmadd_ps(v2, v2, norm2);

        let v3 = _mm256_loadu_ps(vector.as_ptr().add(base + SIMD_WIDTH * 3));
        norm3 = _mm256_fmadd_ps(v3, v3, norm3);
    }

    let norm_vec = _mm256_add_ps(_mm256_add_ps(norm0, norm1), _mm256_add_ps(norm2, norm3));
    let mut norm_sq = horizontal_sum_avx2(norm_vec);

    // Remainder for norm calculation
    for i in (chunks * CHUNK_SIZE)..n {
        norm_sq += vector[i] * vector[i];
    }

    let norm = norm_sq.sqrt();
    // Match `normalize_scalar`: reject 0.0 and NaN alike so a NaN-containing
    // vector is left unchanged rather than scaled by NaN. `is_nan() || <= 0.0`
    // is the lint-clean equivalent of `!(norm > 0.0)`.
    if norm.is_nan() || norm <= 0.0 {
        return;
    }

    let inv_norm = 1.0 / norm;
    let inv_norm_vec = _mm256_set1_ps(inv_norm);

    // Second pass: divide by norm with 4x unrolling
    for i in 0..chunks {
        let base = i * CHUNK_SIZE;

        let v0 = _mm256_loadu_ps(vector.as_ptr().add(base));
        _mm256_storeu_ps(
            vector.as_mut_ptr().add(base),
            _mm256_mul_ps(v0, inv_norm_vec),
        );

        let v1 = _mm256_loadu_ps(vector.as_ptr().add(base + SIMD_WIDTH));
        _mm256_storeu_ps(
            vector.as_mut_ptr().add(base + SIMD_WIDTH),
            _mm256_mul_ps(v1, inv_norm_vec),
        );

        let v2 = _mm256_loadu_ps(vector.as_ptr().add(base + SIMD_WIDTH * 2));
        _mm256_storeu_ps(
            vector.as_mut_ptr().add(base + SIMD_WIDTH * 2),
            _mm256_mul_ps(v2, inv_norm_vec),
        );

        let v3 = _mm256_loadu_ps(vector.as_ptr().add(base + SIMD_WIDTH * 3));
        _mm256_storeu_ps(
            vector.as_mut_ptr().add(base + SIMD_WIDTH * 3),
            _mm256_mul_ps(v3, inv_norm_vec),
        );
    }

    // Remainder for scaling
    for i in (chunks * CHUNK_SIZE)..n {
        vector[i] *= inv_norm;
    }
}

/// Normalizes in two passes with NEON reciprocal-square-root refinement.
///
/// # Safety
/// Caller must run on aarch64; chunked unaligned access stays in bounds.
/// See [`docs/simd.md`](../../docs/simd.md#kernel-safety-boundary) for convergence and fallback semantics.
#[cfg(target_arch = "aarch64")]
#[inline]
unsafe fn normalize_neon_unrolled(vector: &mut [f32]) {
    const SIMD_WIDTH: usize = 4;
    const UNROLL: usize = 4;
    const CHUNK_SIZE: usize = SIMD_WIDTH * UNROLL;
    let n = vector.len();
    let chunks = n / CHUNK_SIZE;

    // First pass: compute L2 norm with 4 accumulators
    let mut norm0 = vdupq_n_f32(0.0);
    let mut norm1 = vdupq_n_f32(0.0);
    let mut norm2 = vdupq_n_f32(0.0);
    let mut norm3 = vdupq_n_f32(0.0);

    for i in 0..chunks {
        let base = i * CHUNK_SIZE;

        let v0 = vld1q_f32(vector.as_ptr().add(base));
        norm0 = vfmaq_f32(norm0, v0, v0);

        let v1 = vld1q_f32(vector.as_ptr().add(base + SIMD_WIDTH));
        norm1 = vfmaq_f32(norm1, v1, v1);

        let v2 = vld1q_f32(vector.as_ptr().add(base + SIMD_WIDTH * 2));
        norm2 = vfmaq_f32(norm2, v2, v2);

        let v3 = vld1q_f32(vector.as_ptr().add(base + SIMD_WIDTH * 3));
        norm3 = vfmaq_f32(norm3, v3, v3);
    }

    let norm_vec = vaddq_f32(vaddq_f32(norm0, norm1), vaddq_f32(norm2, norm3));
    let mut norm_sq = horizontal_sum_neon(norm_vec);

    for val in vector.iter().skip(chunks * CHUNK_SIZE) {
        norm_sq += val * val;
    }

    // Match `normalize_scalar`: reject 0.0 and NaN alike. A subnormal-but-positive
    // norm_sq still passes here and is handled by the finite-fallback below; only
    // zero/NaN short-circuit to leave the vector unchanged, keeping NEON
    // byte-consistent with the scalar path. `is_nan() || <= 0.0` is the lint-clean
    // equivalent of `!(norm_sq > 0.0)`.
    if norm_sq.is_nan() || norm_sq <= 0.0 {
        return;
    }

    // vrsqrteq_f32 gives ~8-bit estimate; two Newton–Raphson steps reach full f32
    // precision (~23 bits), eliminating any residual above the 1e-5 accuracy gate.
    // vrsqrtsq_f32(a, b) = (3 - a*b) / 2  →  y' = y * vrsqrtsq_f32(x, y*y)
    let norm_sq_v = vdupq_n_f32(norm_sq);
    let y0 = vrsqrteq_f32(norm_sq_v);
    let y1 = vmulq_f32(y0, vrsqrtsq_f32(norm_sq_v, vmulq_f32(y0, y0)));
    let y2 = vmulq_f32(y1, vrsqrtsq_f32(norm_sq_v, vmulq_f32(y1, y1)));
    // SAFETY: y2 has 4 identical lanes (norm_sq_v is a broadcast), so lane 0 is the
    // scalar inv_norm used for both the NEON broadcast and the 1-3 element tail.
    let mut inv_norm = vgetq_lane_f32(y2, 0);
    // vrsqrte/Newton overflow to inf for a subnormal-but-nonzero norm_sq (‖v‖ ≲ 7e-20),
    // where the AVX2/scalar lanes stay finite via 1.0/sqrt; fall back to keep NEON
    // byte-consistent with them rather than scaling the vector to inf/NaN.
    if !inv_norm.is_finite() {
        inv_norm = 1.0 / norm_sq.sqrt();
    }
    let inv_norm_vec = vdupq_n_f32(inv_norm);

    // Second pass: scale by inverse norm with 4x unrolling
    for i in 0..chunks {
        let base = i * CHUNK_SIZE;

        let v0 = vld1q_f32(vector.as_ptr().add(base));
        vst1q_f32(vector.as_mut_ptr().add(base), vmulq_f32(v0, inv_norm_vec));

        let v1 = vld1q_f32(vector.as_ptr().add(base + SIMD_WIDTH));
        vst1q_f32(
            vector.as_mut_ptr().add(base + SIMD_WIDTH),
            vmulq_f32(v1, inv_norm_vec),
        );

        let v2 = vld1q_f32(vector.as_ptr().add(base + SIMD_WIDTH * 2));
        vst1q_f32(
            vector.as_mut_ptr().add(base + SIMD_WIDTH * 2),
            vmulq_f32(v2, inv_norm_vec),
        );

        let v3 = vld1q_f32(vector.as_ptr().add(base + SIMD_WIDTH * 3));
        vst1q_f32(
            vector.as_mut_ptr().add(base + SIMD_WIDTH * 3),
            vmulq_f32(v3, inv_norm_vec),
        );
    }

    // Remainder for scaling
    for val in vector.iter_mut().skip(chunks * CHUNK_SIZE) {
        *val *= inv_norm;
    }
}

/// Normalizes in two passes with wasm32 SIMD128.
///
/// # Safety
/// This function requires the compile-time `simd128` target feature; bounds are chunked.
/// See [`docs/simd.md`](../../docs/simd.md#kernel-safety-boundary) for wasm alignment semantics.
#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
#[inline]
unsafe fn normalize_simd128_unrolled(vector: &mut [f32]) {
    const SIMD_WIDTH: usize = 4;
    const UNROLL: usize = 4;
    const CHUNK_SIZE: usize = SIMD_WIDTH * UNROLL;
    let n = vector.len();
    let chunks = n / CHUNK_SIZE;

    // First pass: compute L2 norm with 4 accumulators
    let mut norm0 = f32x4_splat(0.0);
    let mut norm1 = f32x4_splat(0.0);
    let mut norm2 = f32x4_splat(0.0);
    let mut norm3 = f32x4_splat(0.0);

    for i in 0..chunks {
        let base = i * CHUNK_SIZE;

        let v0 = v128_load(vector.as_ptr().add(base) as *const v128);
        norm0 = f32x4_add(norm0, f32x4_mul(v0, v0));

        let v1 = v128_load(vector.as_ptr().add(base + SIMD_WIDTH) as *const v128);
        norm1 = f32x4_add(norm1, f32x4_mul(v1, v1));

        let v2 = v128_load(vector.as_ptr().add(base + SIMD_WIDTH * 2) as *const v128);
        norm2 = f32x4_add(norm2, f32x4_mul(v2, v2));

        let v3 = v128_load(vector.as_ptr().add(base + SIMD_WIDTH * 3) as *const v128);
        norm3 = f32x4_add(norm3, f32x4_mul(v3, v3));
    }

    let norm_vec = f32x4_add(f32x4_add(norm0, norm1), f32x4_add(norm2, norm3));
    let mut norm_sq = horizontal_sum_simd128(norm_vec);

    // Remainder for norm calculation
    for i in (chunks * CHUNK_SIZE)..n {
        norm_sq += vector[i] * vector[i];
    }

    let norm = norm_sq.sqrt();
    // Match `normalize_scalar`: reject 0.0 and NaN alike so a NaN-containing
    // vector is left unchanged rather than scaled by NaN. `is_nan() || <= 0.0`
    // is the lint-clean equivalent of `!(norm > 0.0)`.
    if norm.is_nan() || norm <= 0.0 {
        return;
    }

    let inv_norm = 1.0 / norm;
    let inv_norm_vec = f32x4_splat(inv_norm);

    // Second pass: divide by norm with 4x unrolling
    for i in 0..chunks {
        let base = i * CHUNK_SIZE;

        let v0 = v128_load(vector.as_ptr().add(base) as *const v128);
        v128_store(
            vector.as_mut_ptr().add(base) as *mut v128,
            f32x4_mul(v0, inv_norm_vec),
        );

        let v1 = v128_load(vector.as_ptr().add(base + SIMD_WIDTH) as *const v128);
        v128_store(
            vector.as_mut_ptr().add(base + SIMD_WIDTH) as *mut v128,
            f32x4_mul(v1, inv_norm_vec),
        );

        let v2 = v128_load(vector.as_ptr().add(base + SIMD_WIDTH * 2) as *const v128);
        v128_store(
            vector.as_mut_ptr().add(base + SIMD_WIDTH * 2) as *mut v128,
            f32x4_mul(v2, inv_norm_vec),
        );

        let v3 = v128_load(vector.as_ptr().add(base + SIMD_WIDTH * 3) as *const v128);
        v128_store(
            vector.as_mut_ptr().add(base + SIMD_WIDTH * 3) as *mut v128,
            f32x4_mul(v3, inv_norm_vec),
        );
    }

    // Remainder for scaling
    for i in (chunks * CHUNK_SIZE)..n {
        vector[i] *= inv_norm;
    }
}