lightweight-simd 0.1.1

//! Assembly probes for the vectorization-verification framework (step 3).
//
// Most or all contents generated by AI (model glm-5.3)
//
// This file is NOT part of the cargo build/test surface: `scripts/check_asm.py`
// compiles it directly with `rustc --emit asm` against the release rlib of the
// crate (built with --all-features). Each probe is a `#[no_mangle]`
// free function so the check script can locate its body by symbol label; every
// input and the result are fenced with `std::hint::black_box` to defeat
// constant folding and dead-code elimination. Expectation table lives in the
// script. Naming: `probe_<op>_<type>`.

use std::hint::black_box;

use lightweight_simd::{c64x8, f16x8, f32x16, f64x8, f64x16, i32x16, i64x8};

/* #region REQUIRED: lane-wise f64 arithmetic */

#[no_mangle]
pub fn probe_add_f64x8(a: f64x8, b: f64x8) -> f64x8 {
    black_box(black_box(a) + black_box(b))
}

#[no_mangle]
pub fn probe_mul_f64x8(a: f64x8, b: f64x8) -> f64x8 {
    black_box(black_box(a) * black_box(b))
}

#[no_mangle]
pub fn probe_div_f64x8(a: f64x8, b: f64x8) -> f64x8 {
    black_box(black_box(a) / black_box(b))
}

#[no_mangle]
pub fn probe_neg_f64x8(a: f64x8) -> f64x8 {
    black_box(-black_box(a))
}

#[no_mangle]
pub fn probe_min_f64x8(a: f64x8, b: f64x8) -> f64x8 {
    let (a, b) = (black_box(a), black_box(b));
    black_box(a.min(b))
}

#[no_mangle]
pub fn probe_max_f64x8(a: f64x8, b: f64x8) -> f64x8 {
    let (a, b) = (black_box(a), black_box(b));
    black_box(a.max(b))
}

#[no_mangle]
pub fn probe_copysign_f64x8(a: f64x8, b: f64x8) -> f64x8 {
    let (a, b) = (black_box(a), black_box(b));
    black_box(a.copysign(b))
}

#[no_mangle]
pub fn probe_sqrt_f64x8(a: f64x8) -> f64x8 {
    black_box(black_box(a).sqrt())
}

#[no_mangle]
pub fn probe_abs_f64x8(a: f64x8) -> f64x8 {
    black_box(black_box(a).abs())
}

#[no_mangle]
pub fn probe_floor_f64x8(a: f64x8) -> f64x8 {
    black_box(black_box(a).floor())
}

/* #endregion */

/* #region REQUIRED: fused multiply-add */

#[no_mangle]
pub fn probe_mul_add_f64x8(a: f64x8, b: f64x8, c: f64x8) -> f64x8 {
    let (a, b, c) = (black_box(a), black_box(b), black_box(c));
    black_box(a.mul_add(b, c))
}

#[no_mangle]
pub fn probe_fma_from_f64x8(mut a: f64x8, b: f64x8, c: f64x8) -> f64x8 {
    let (b, c) = (black_box(b), black_box(c));
    a.fma_from(b, c);
    black_box(a)
}

#[no_mangle]
pub fn probe_mul_add_f64x16(a: f64x16, b: f64x16, c: f64x16) -> f64x16 {
    let (a, b, c) = (black_box(a), black_box(b), black_box(c));
    black_box(a.mul_add(b, c))
}

/* #endregion */

/* #region REQUIRED: f32 and integer lanes */

#[no_mangle]
pub fn probe_add_f32x16(a: f32x16, b: f32x16) -> f32x16 {
    black_box(black_box(a) + black_box(b))
}

#[no_mangle]
pub fn probe_mul_add_f32x16(a: f32x16, b: f32x16, c: f32x16) -> f32x16 {
    let (a, b, c) = (black_box(a), black_box(b), black_box(c));
    black_box(a.mul_add(b, c))
}

#[no_mangle]
pub fn probe_add_i64x8(a: i64x8, b: i64x8) -> i64x8 {
    black_box(black_box(a) + black_box(b))
}

#[no_mangle]
pub fn probe_bitwise_i32x16(a: i32x16, b: i32x16, c: i32x16) -> i32x16 {
    let (a, b, c) = (black_box(a), black_box(b), black_box(c));
    black_box((a & b) | (b ^ c))
}

#[no_mangle]
pub fn probe_shift_i64x8(a: i64x8, s: usize) -> i64x8 {
    let (a, s) = (black_box(a), black_box(s));
    black_box((a << s) | (a >> s))
}

/* #endregion */

/* #region REQUIRED: comparisons, masks, memory */

#[no_mangle]
pub fn probe_cmp_select_f64x8(a: f64x8, b: f64x8, c: f64x8) -> f64x8 {
    let (a, b, c) = (black_box(a), black_box(b), black_box(c));
    let m = a.simd_gt(b);
    black_box(a.mask_select(m, c))
}

#[no_mangle]
pub fn probe_eq_any_all_f64x8(a: f64x8, b: f64x8) -> (f64x8, bool, bool) {
    let (a, b) = (black_box(a), black_box(b));
    let m = a.simd_eq(b);
    let n = a.simd_ne(b);
    // the mask_select forces the masks to be consumed as vectors: masks that
    // only feed reductions or bool-array conversions get scalarized by LLVM
    let v = a.mask_select(m, b);
    black_box((v, m.any_true(), n.all_true()))
}

#[no_mangle]
pub fn probe_slice_io_unaligned(src: &[f64], dst: &mut [f64], k: f64) {
    let (src, dst, k) = (black_box(src), black_box(dst), black_box(k));
    let v = f64x8::from_slice(src);
    let v = v * f64x8::splat(k);
    v.store_slice(dst);
}

#[no_mangle]
pub fn probe_slice_io_aligned(a: &f64x8, b: &f64x8, out: &mut f64x8) {
    let (a, b) = (black_box(a), black_box(b));
    *black_box(out) = *a + *b;
}

#[no_mangle]
pub fn probe_store_slice_partial(v: &f64x8, dst: &mut [f64]) {
    let (v, dst) = (black_box(v), black_box(dst));
    v.store_slice_partial(dst);
}

/* #endregion */

/* #region INFORMATIONAL: reported, never fails the run */

#[no_mangle]
pub fn probe_reduce_sum_f64x8(a: f64x8) -> f64 {
    black_box(black_box(a).reduce_sum())
}

/// Naive fused expression: NOT fused under default rustc flags; the script
/// also compiles an extra row with `-C llvm-args=-fp-contract=fast` where it
/// must fuse, documenting the opt-in for downstream users.
#[no_mangle]
pub fn probe_expr_fma_f64x8(a: f64x8, b: f64x8, c: f64x8) -> f64x8 {
    black_box(black_box(a) * black_box(b) + black_box(c))
}

/// Soft-float via `half`; known-scalar on x86 (verified even with an
/// AVX512-FP16 target). Expectations are per-arch in the check script.
#[no_mangle]
pub fn probe_f16_add_f16x8(a: f16x8, b: f16x8) -> f16x8 {
    black_box(black_box(a) + black_box(b))
}

/// Generic `num-complex` multiply; vectorizes on x86-64
/// (vmulpd/vshufpd/vaddsubpd). Expectations are per-arch in the check
/// script.
#[no_mangle]
pub fn probe_complex_mul_c64x8(a: c64x8, b: c64x8) -> c64x8 {
    black_box(black_box(a) * black_box(b))
}

/* #endregion */