wai-quantum 0.6.0

//! The vector unit's kernels: the gate's fixed-point arithmetic in its
//! 128-byte vectors, thirty-two lanes of 32 bits. Each lane's 64-bit value is a
//! pair of words, high and low; every step on it is 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 core::arch::hexagon::v128::*; // allow-name: the toolchain's module of the vector unit's intrinsics

/// One 128-byte vector.
type V = HvxVector; // allow-name: the toolchain's vector type

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

/// Each lane's 64-bit value as its high and low words.
#[derive(Clone, Copy)]
struct Wide {
    hi: V,
    lo: V,
}

// Carries and borrows are bit logic on the words, not predicates: the carry out
// of `s = a + b` is bit 31 of `(a & b) | ((a | b) & !s)`, and the borrow out of
// `d = a − b` is bit 31 of `(!a & b) | ((!a | b) & d)`.

/// `a + b` in each lane: the low words' carry goes into the high word.
#[inline(always)]
unsafe fn add(a: Wide, b: Wide) -> Wide {
    unsafe {
        let lo = Q6_Vw_vadd_VwVw(a.lo, b.lo);
        let carry = Q6_Vuw_vlsr_VuwR(Q6_V_vor_VV(Q6_V_vand_VV(a.lo, b.lo), Q6_V_vand_VV(Q6_V_vor_VV(a.lo, b.lo), Q6_V_vnot_V(lo))), 31);
        Wide { hi: Q6_Vw_vadd_VwVw(Q6_Vw_vadd_VwVw(a.hi, b.hi), carry), lo }
    }
}

/// `a − b` in each lane: the low words' borrow comes out of the high word.
#[inline(always)]
unsafe fn sub(a: Wide, b: Wide) -> Wide {
    unsafe {
        let lo = Q6_Vw_vsub_VwVw(a.lo, b.lo);
        let na = Q6_V_vnot_V(a.lo);
        let borrow = Q6_Vuw_vlsr_VuwR(Q6_V_vor_VV(Q6_V_vand_VV(na, b.lo), Q6_V_vand_VV(Q6_V_vor_VV(na, b.lo), lo)), 31);
        Wide { hi: Q6_Vw_vsub_VwVw(Q6_Vw_vsub_VwVw(a.hi, b.hi), borrow), lo }
    }
}

/// `a·b` in each lane, exactly: `a` times `b`'s low half (unsigned), plus `a`
/// times `b`'s high half (signed) sixteen bits up.
#[inline(always)]
unsafe fn product(a: V, b: V) -> Wide {
    unsafe {
        let p = Q6_W_vmpyoacc_WVwVh(Q6_W_vmpye_VwVuh(a, b), a, b);
        Wide { hi: Q6_V_hi_W(p), lo: Q6_V_lo_W(p) }
    }
}

/// `(e·a + RND) >> SH` in each lane, the shift arithmetic over all 64 bits.
#[inline(always)]
unsafe fn fx<const RND: i32, const SH: i32>(e: V, a: V) -> Wide {
    unsafe {
        let r = add(product(a, e), Wide { hi: Q6_V_vzero(), lo: Q6_V_vsplat_R(RND) });
        Wide { hi: Q6_Vw_vasr_VwR(r.hi, SH), lo: Q6_V_vor_VV(Q6_Vuw_vlsr_VuwR(r.lo, SH), Q6_Vw_vasl_VwR(r.hi, 32 - SH)) }
    }
}

/// The vector of `s` from lane `k`.
#[inline(always)]
unsafe fn load(s: &[i32], k: usize) -> V {
    debug_assert!(k + LANES <= s.len());
    // SAFETY: the caller's slices hold whole vectors; the load is unaligned.
    unsafe { core::ptr::read_unaligned(s.as_ptr().add(k).cast::<V>()) }
}

/// Writes `v` to `d` from lane `k`.
#[inline(always)]
unsafe fn store(d: &mut [i32], k: usize, v: V) {
    debug_assert!(k + LANES <= d.len());
    // SAFETY: as for `load`.
    unsafe { core::ptr::write_unaligned(d.as_mut_ptr().add(k).cast::<V>(), v) }
}

/// 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); row `r` (0r, 0i, 1r, 1i) goes to
/// `hi[r]` and `lo[r]`.
///
/// # Safety
///
/// The processor has the vector unit in 128-byte mode; every slice has the same
/// length, a whole number of vectors.
pub unsafe fn rows<const RND: i32, const SH: i32, const LAST_ADD: bool>(m: &[i32; 8], cols: [&[i32]; 4], hi: [&mut [i32]; 4], lo: [&mut [i32]; 4]) {
    let n = cols[0].len();
    debug_assert!(n.is_multiple_of(LANES) && cols.iter().all(|c| c.len() == n));
    unsafe {
        let mut e = [Q6_V_vzero(); 8];
        for (v, &x) in e.iter_mut().zip(m) {
            *v = Q6_V_vsplat_R(x);
        }
        let [hi0r, hi0i, hi1r, hi1i] = hi;
        let [lo0r, lo0i, lo1r, lo1i] = lo;
        let f = fx::<RND, SH>;
        let mut k = 0;
        while k < n {
            let (a0r, a0i, a1r, a1i) = (load(cols[0], k), load(cols[1], k), load(cols[2], k), load(cols[3], k));
            let t = add(sub(f(e[0], a0r), f(e[1], a0i)), f(e[2], a1r));
            let o0r = if LAST_ADD { add(t, f(e[3], a1i)) } else { sub(t, f(e[3], a1i)) };
            let o0i = add(add(add(f(e[0], a0i), f(e[1], a0r)), f(e[2], a1i)), f(e[3], a1r));
            let o1r = sub(add(sub(f(e[4], a0r), f(e[5], a0i)), f(e[6], a1r)), f(e[7], a1i));
            let o1i = add(add(add(f(e[4], a0i), f(e[5], a0r)), f(e[6], a1i)), f(e[7], a1r));
            for (h, l, o) in [(&mut *hi0r, &mut *lo0r, o0r), (&mut *hi0i, &mut *lo0i, o0i), (&mut *hi1r, &mut *lo1r, o1r), (&mut *hi1i, &mut *lo1i, o1i)] {
                store(h, k, o.hi);
                store(l, k, o.lo);
            }
            k += LANES;
        }
    }
}

/// Each pair's `a·b` and `(a·b + RND) >> SH`, as high and low words.
///
/// # Safety
///
/// As for [`rows`].
pub unsafe fn probe<const RND: i32, const SH: i32>(a: &[i32], b: &[i32], prod: [&mut [i32]; 2], fixed: [&mut [i32]; 2]) {
    let n = a.len();
    debug_assert!(n.is_multiple_of(LANES) && b.len() == n);
    let [ph, pl] = prod;
    let [fh, fl] = fixed;
    unsafe {
        let mut k = 0;
        while k < n {
            let (va, vb) = (load(a, k), load(b, k));
            let p = product(va, vb);
            let f = fx::<RND, SH>(vb, va);
            store(ph, k, p.hi);
            store(pl, k, p.lo);
            store(fh, k, f.hi);
            store(fl, k, f.lo);
            k += LANES;
        }
    }
}