hermes_simd/dispatch/

complex.rs

//! Generic interleaved complex kernels.
//!
//! The public surface accepts primitive lane slices in `[re, im, ...]` order so
//! domain crates can use their own complex storage types without Hermes taking a
//! dependency on a concrete complex-number crate.
//!
//! Arithmetic stays in vector registers for every `(T, Arch)` pair via the
//! adjacent-pair shuffle and alternating-FMA primitives on
//! [`SimdKernel`]:
//!
//! - `a * b`       = `fmaddsub(dup_even(a), b, mul(dup_odd(a), swap_adjacent(b)))`
//! - `a * conj(b)` = `fmsubadd(dup_odd(a), swap_adjacent(b), mul(dup_even(a), b))`
//!
//! Runtime architecture selection is generated by `#[runtime_dispatch]`, the
//! same mechanism used by the dense kernels — no per-type provider impls and no
//! `OnceLock` feature caching.

use hermes_simd_core::{arch::SimdArch, kernel::SimdKernel, scalar::Scalar, view::SimdError};
use hermes_simd_macros::runtime_dispatch;

const MAX_STACK_LANES: usize = 128;

#[inline]
fn mul_pair<T, const CONJ_B: bool>(ar: T, ai: T, br: T, bi: T) -> (T, T)
where
    T: Scalar,
{
    if CONJ_B {
        // a * conj(b): (ar+i·ai)(br-i·bi) = (ar·br + ai·bi) + i·(ai·br - ar·bi)
        (ar * br + ai * bi, ai * br - ar * bi)
    } else {
        // a * b: (ar+i·ai)(br+i·bi) = (ar·br - ai·bi) + i·(ar·bi + ai·br)
        (ar * br - ai * bi, ar * bi + ai * br)
    }
}

/// Computes one register of interleaved complex products: `a[k] * b[k]`
/// (or `a[k] * conj(b[k])` when `CONJ_B`) over adjacent lane pairs.
///
/// # Safety
/// Processor must support the target features required by `A`.
#[inline(always)]
unsafe fn complex_mul_vector<T, A, const CONJ_B: bool>(av: A::Vector, bv: A::Vector) -> A::Vector
where
    T: Scalar,
    A: SimdArch + SimdKernel<T>,
{
    let b_sw = A::swap_adjacent(bv);
    if CONJ_B {
        // even: ai*bi + ar*br = re ; odd: ai*br - ar*bi = im
        A::fmsubadd(A::dup_odd(av), b_sw, A::mul(A::dup_even(av), bv))
    } else {
        // even: ar*br - ai*bi = re ; odd: ar*bi + ai*br = im
        A::fmaddsub(A::dup_even(av), bv, A::mul(A::dup_odd(av), b_sw))
    }
}

/// Multiplies interleaved complex values in-place using architecture `A`.
///
/// `a` and `b` must have identical even lengths. The operation is value
/// preserving with respect to scalar complex multiplication over adjacent
/// primitive lane pairs:
///
/// `a[k] = a[k] * b[k]` when `CONJ_B == false`, and
/// `a[k] = a[k] * conj(b[k])` when `CONJ_B == true`.
///
/// # Examples
///
/// ```
/// use hermes_simd::{interleaved_complex_mul_assign, Scalar};
///
/// let mut lhs = [1.0_f64, 2.0, 3.0, -1.0];
/// let rhs = [4.0_f64, -2.0, 0.5, 5.0];
///
/// interleaved_complex_mul_assign::<f64, Scalar, false>(&mut lhs, &rhs).unwrap();
///
/// assert_eq!(lhs, [8.0, 6.0, 6.5, 14.5]);
/// ```
#[inline]
pub fn interleaved_complex_mul_assign<T, A, const CONJ_B: bool>(
    a: &mut [T],
    b: &[T],
) -> Result<(), SimdError>
where
    T: Scalar,
    A: SimdArch + SimdKernel<T>,
{
    if a.len() != b.len() || (a.len() & 1) != 0 {
        return Err(SimdError::LengthMismatch);
    }

    if A::REGISTER_WIDTH_BITS == 0 && a.len() >= 32_768 {
        let mut lane = 0usize;
        while lane + 8 <= a.len() {
            let (re0, im0) = mul_pair::<T, CONJ_B>(a[lane], a[lane + 1], b[lane], b[lane + 1]);
            a[lane] = re0;
            a[lane + 1] = im0;

            let (re1, im1) =
                mul_pair::<T, CONJ_B>(a[lane + 2], a[lane + 3], b[lane + 2], b[lane + 3]);
            a[lane + 2] = re1;
            a[lane + 3] = im1;

            let (re2, im2) =
                mul_pair::<T, CONJ_B>(a[lane + 4], a[lane + 5], b[lane + 4], b[lane + 5]);
            a[lane + 4] = re2;
            a[lane + 5] = im2;

            let (re3, im3) =
                mul_pair::<T, CONJ_B>(a[lane + 6], a[lane + 7], b[lane + 6], b[lane + 7]);
            a[lane + 6] = re3;
            a[lane + 7] = im3;

            lane += 8;
        }
        while lane < a.len() {
            let (re, im) = mul_pair::<T, CONJ_B>(a[lane], a[lane + 1], b[lane], b[lane + 1]);
            a[lane] = re;
            a[lane + 1] = im;
            lane += 2;
        }
        return Ok(());
    }

    let lanes = A::LANE_COUNT;
    let mut offset = 0usize;

    // The vectorized path requires an even lane count so every register holds
    // whole (re, im) pairs; all provided backends satisfy this.
    if lanes >= 2 && lanes & 1 == 0 {
        while offset + 2 * lanes <= a.len() {
            // SAFETY: offset + 2*lanes <= len was checked above; `a` and `b`
            // are valid for reads (and `a` for writes) of `2*lanes`
            // primitive values. `A`'s target features are guaranteed by the
            // dispatching caller.
            unsafe {
                let av0 = A::load_unaligned(a.as_ptr().add(offset));
                let bv0 = A::load_unaligned(b.as_ptr().add(offset));
                let res0 = complex_mul_vector::<T, A, CONJ_B>(av0, bv0);
                A::store_unaligned(a.as_mut_ptr().add(offset), res0);

                let next = offset + lanes;
                let av1 = A::load_unaligned(a.as_ptr().add(next));
                let bv1 = A::load_unaligned(b.as_ptr().add(next));
                let res1 = complex_mul_vector::<T, A, CONJ_B>(av1, bv1);
                A::store_unaligned(a.as_mut_ptr().add(next), res1);
            }
            offset += 2 * lanes;
        }

        while offset + lanes <= a.len() {
            // SAFETY: offset + lanes <= len was checked above; `a` and `b` are
            // valid for reads (and `a` for writes) of `lanes` primitive values.
            // `A`'s target features are guaranteed by the dispatching caller.
            unsafe {
                let av = A::load_unaligned(a.as_ptr().add(offset));
                let bv = A::load_unaligned(b.as_ptr().add(offset));
                let res = complex_mul_vector::<T, A, CONJ_B>(av, bv);
                A::store_unaligned(a.as_mut_ptr().add(offset), res);
            }
            offset += lanes;
        }
    }

    let mut lane = offset;
    while lane < a.len() {
        let (re, im) = mul_pair::<T, CONJ_B>(a[lane], a[lane + 1], b[lane], b[lane + 1]);
        a[lane] = re;
        a[lane + 1] = im;
        lane += 2;
    }

    Ok(())
}

/// Computes an interleaved complex dot product using architecture `A`.
///
/// Inputs must have identical even lengths in `[re0, im0, re1, im1, ...]`
/// primitive lane order. The returned tuple is `(re, im)` for
/// `sum(a[k] * b[k])`; when `CONJ_B` is true, the operation is
/// `sum(a[k] * conj(b[k]))`.
///
/// # Examples
///
/// ```
/// use hermes_simd::{interleaved_complex_dot, Scalar};
///
/// let lhs = [1.0_f64, 2.0, 3.0, 4.0];
/// let rhs = [5.0_f64, 6.0, 7.0, 8.0];
///
/// let product_sum = interleaved_complex_dot::<f64, Scalar, false>(&lhs, &rhs).unwrap();
///
/// assert_eq!(product_sum, (-18.0, 68.0));
/// ```
#[inline]
pub fn interleaved_complex_dot<T, A, const CONJ_B: bool>(
    a: &[T],
    b: &[T],
) -> Result<(T, T), SimdError>
where
    T: Scalar,
    A: SimdArch + SimdKernel<T>,
{
    if a.len() != b.len() || (a.len() & 1) != 0 {
        return Err(SimdError::LengthMismatch);
    }

    let lanes = A::LANE_COUNT;
    let mut offset = 0usize;
    let mut re = T::ZERO;
    let mut im = T::ZERO;

    if lanes >= 2 && lanes & 1 == 0 && offset + lanes <= a.len() {
        assert!(
            lanes <= MAX_STACK_LANES,
            "SIMD lane count exceeds stack buffer"
        );
        // Two independent accumulators break the loop-carried add dependency
        // so consecutive FMAs can overlap in the pipeline.
        // SAFETY: feature availability is guaranteed by the dispatching caller.
        let (mut acc0, mut acc1) = unsafe { (A::zero(), A::zero()) };
        while offset + 2 * lanes <= a.len() {
            // SAFETY: offset + 2*lanes <= len was checked above; `a` and `b`
            // are valid for reads of `2*lanes` primitive values.
            unsafe {
                let av0 = A::load_unaligned(a.as_ptr().add(offset));
                let bv0 = A::load_unaligned(b.as_ptr().add(offset));
                acc0 = A::add(acc0, complex_mul_vector::<T, A, CONJ_B>(av0, bv0));
                let av1 = A::load_unaligned(a.as_ptr().add(offset + lanes));
                let bv1 = A::load_unaligned(b.as_ptr().add(offset + lanes));
                acc1 = A::add(acc1, complex_mul_vector::<T, A, CONJ_B>(av1, bv1));
            }
            offset += 2 * lanes;
        }
        while offset + lanes <= a.len() {
            // SAFETY: offset + lanes <= len was checked above; `a` and `b` are
            // valid for reads of `lanes` primitive values.
            unsafe {
                let av = A::load_unaligned(a.as_ptr().add(offset));
                let bv = A::load_unaligned(b.as_ptr().add(offset));
                acc0 = A::add(acc0, complex_mul_vector::<T, A, CONJ_B>(av, bv));
            }
            offset += lanes;
        }
        // SAFETY: same feature contract as above.
        let acc = unsafe { A::add(acc0, acc1) };

        // Single spill: fold even lanes into re, odd lanes into im.
        let mut buf = [T::ZERO; MAX_STACK_LANES];
        // SAFETY: buf holds at least `lanes` elements (asserted above).
        unsafe { A::store_unaligned(buf.as_mut_ptr(), acc) };
        let mut lane = 0usize;
        while lane < lanes {
            re = re + buf[lane];
            im = im + buf[lane + 1];
            lane += 2;
        }
    }

    let mut lane = offset;
    while lane < a.len() {
        let (prod_re, prod_im) = mul_pair::<T, CONJ_B>(a[lane], a[lane + 1], b[lane], b[lane + 1]);
        re = re + prod_re;
        im = im + prod_im;
        lane += 2;
    }

    Ok((re, im))
}

#[runtime_dispatch(avx512f, avx2, neon, scalar)]
pub(super) fn dispatch_interleaved_complex_mul_assign_impl<T, const CONJ_B: bool, A>(
    a: &mut [T],
    b: &[T],
) -> Result<(), SimdError>
where
    T: Scalar,
    A: SimdArch + SimdKernel<T>,
{
    interleaved_complex_mul_assign::<T, A, CONJ_B>(a, b)
}

#[runtime_dispatch(avx512f, avx2, neon, scalar)]
pub(super) fn dispatch_interleaved_complex_dot_impl<T, const CONJ_B: bool, A>(
    a: &[T],
    b: &[T],
) -> Result<(T, T), SimdError>
where
    T: Scalar,
    A: SimdArch + SimdKernel<T>,
{
    interleaved_complex_dot::<T, A, CONJ_B>(a, b)
}