wai-quantum 0.6.0

//! The x86-64 vector unit's kernels (AVX-512, its foundation instructions):
//! the gate's fixed-point arithmetic in 512-bit vectors, eight 64-bit lanes,
//! every step exact.
//!
//! This file's text is the implementation id's input. `RND` is the rounding
//! constant, `SH` the shift and `LAST_ADD` whether row 0's last real term is
//! added (a mutant) instead of subtracted.
use std::arch::x86_64::*;

/// Lanes per vector.
pub const LANES: usize = 8;

/// Eight operands of `s` from `k`, sign-extended to 64 bits.
#[inline]
#[target_feature(enable = "avx512f")]
unsafe fn load(s: &[i32], k: usize) -> __m512i {
    debug_assert!(k + LANES <= s.len());
    // SAFETY: eight values from `k` are in bounds; the load is unaligned.
    _mm512_cvtepi32_epi64(unsafe { _mm256_loadu_si256(s.as_ptr().add(k).cast()) })
}

/// Writes the eight lanes of `v` to `d` from `k`.
#[inline]
#[target_feature(enable = "avx512f")]
unsafe fn store(d: &mut [i64], k: usize, v: __m512i) {
    debug_assert!(k + LANES <= d.len());
    // SAFETY: as for `load`.
    unsafe { _mm512_storeu_si512(d.as_mut_ptr().add(k).cast(), v) }
}

/// `(e·a + RND) >> SH` in each lane. The multiply takes each lane's low 32
/// bits as signed and forms the exact 64-bit product; the unit's 64-bit
/// arithmetic shift is the shift.
#[inline]
#[target_feature(enable = "avx512f")]
fn fx<const RND: i64, const SH: u32>(e: __m512i, a: __m512i) -> __m512i {
    _mm512_srai_epi64::<SH>(_mm512_add_epi64(_mm512_mul_epi32(a, e), _mm512_set1_epi64(RND)))
}

/// A gate's four rows for every pair: `m` the gate's entries, `cols` the pairs'
/// columns (re a0, im a0, re a1, im a1), `out` the rows (0r, 0i, 1r, 1i).
///
/// # Safety
///
/// The processor has AVX-512's foundation instructions; every slice has the
/// same length, a multiple of [`LANES`].
#[target_feature(enable = "avx512f")]
pub unsafe fn rows<const RND: i64, const SH: u32, const LAST_ADD: bool>(m: &[i32; 8], cols: [&[i32]; 4], out: [&mut [i64]; 4]) {
    let n = cols[0].len();
    debug_assert!(n.is_multiple_of(LANES) && cols.iter().all(|c| c.len() == n) && out.iter().all(|o| o.len() == n));
    let mut e = [_mm512_setzero_si512(); 8];
    for (v, &x) in e.iter_mut().zip(m) {
        *v = _mm512_set1_epi64(i64::from(x));
    }
    let [o0r, o0i, o1r, o1i] = out;
    macro_rules! f {
        ($e:expr, $a:expr) => {
            fx::<RND, SH>($e, $a)
        };
    }
    let mut k = 0;
    while k < n {
        // SAFETY: the caller's slices hold whole vectors.
        let (a0r, a0i, a1r, a1i) = unsafe { (load(cols[0], k), load(cols[1], k), load(cols[2], k), load(cols[3], k)) };
        let t = _mm512_add_epi64(_mm512_sub_epi64(f!(e[0], a0r), f!(e[1], a0i)), f!(e[2], a1r));
        let r0r = if LAST_ADD { _mm512_add_epi64(t, f!(e[3], a1i)) } else { _mm512_sub_epi64(t, f!(e[3], a1i)) };
        let r0i = _mm512_add_epi64(_mm512_add_epi64(_mm512_add_epi64(f!(e[0], a0i), f!(e[1], a0r)), f!(e[2], a1i)), f!(e[3], a1r));
        let r1r = _mm512_sub_epi64(_mm512_add_epi64(_mm512_sub_epi64(f!(e[4], a0r), f!(e[5], a0i)), f!(e[6], a1r)), f!(e[7], a1i));
        let r1i = _mm512_add_epi64(_mm512_add_epi64(_mm512_add_epi64(f!(e[4], a0i), f!(e[5], a0r)), f!(e[6], a1i)), f!(e[7], a1r));
        // SAFETY: as for the loads.
        unsafe {
            store(o0r, k, r0r);
            store(o0i, k, r0i);
            store(o1r, k, r1r);
            store(o1i, k, r1i);
        }
        k += LANES;
    }
}

/// Each pair's `a·b` and `(a·b + RND) >> SH`.
///
/// # Safety
///
/// As for [`rows`].
#[target_feature(enable = "avx512f")]
pub unsafe fn probe<const RND: i64, const SH: u32>(a: &[i32], b: &[i32], prod: &mut [i64], fixed: &mut [i64]) {
    let n = a.len();
    debug_assert!(n.is_multiple_of(LANES) && b.len() == n && prod.len() == n && fixed.len() == n);
    let mut k = 0;
    while k < n {
        // SAFETY: the caller's slices hold whole vectors.
        unsafe {
            let (va, vb) = (load(a, k), load(b, k));
            store(prod, k, _mm512_mul_epi32(va, vb));
            store(fixed, k, fx::<RND, SH>(vb, va));
        }
        k += LANES;
    }
}