diff --git a/crates/pcs/src/whir/sumcheck/first_pass.rs b/crates/pcs/src/whir/sumcheck/first_pass.rs index cfca183bf..37afa5dcc 100644 --- a/crates/pcs/src/whir/sumcheck/first_pass.rs +++ b/crates/pcs/src/whir/sumcheck/first_pass.rs @@ -9,6 +9,12 @@ //! drawn, and one more pass folds all `R` lane bits at once (the small-value //! precomputation of Bagad, Dao, Domb and Thaler, : //! the witness is in `K`, so every product is a mixed one). +//! +//! On PMULL targets, the grid's final extension digit is consumed directly by +//! mixed-product accumulation. Its zero, one and infinity sums stay in registers +//! across the offsets in a row, so neither the infinity weights nor intermediate +//! partial sums need a separate memory pass. Wider x86 kernels retain their +//! lane-parallel layout. use super::{ Basis, BasisFill, FIRST_PASS_PAR_THRESHOLD, INITIAL_BASIS_CHUNK, KEEP_WEIGHT_MAX_THREADS, PRECOMPUTED_ROUNDS, @@ -64,6 +70,43 @@ fn extend_grid(grid: &mut [T; GRID], add: impl Fn(&T, & } } +#[cfg(all(target_arch = "aarch64", target_feature = "aes"))] +#[inline(always)] +fn extend_grid_digits(grid: &mut [T; GRID], digits: usize, add: impl Fn(&T, &T) -> T) { + let mut stride = 1; + for i in 0..digits { + for &high in &LANE_IN_GRID[..1 << (R - 1 - i)] { + let base = 3 * stride * high; + for at in base..base + stride { + grid[at + 2 * stride] = add(&grid[at], &grid[at + stride]); + } + } + stride *= 3; + } +} + +/// Finish each row's extension at the point where its products are consumed. +#[cfg(all(target_arch = "aarch64", target_feature = "aes"))] +#[inline(always)] +fn accumulate_grid( + fg: &mut [[u64; ROW]; GRID], + bg: &mut [WeightRow; GRID], + acc: &mut [ProductRow; GRID], +) { + if R > 0 { + extend_grid_digits::<_, R>(fg, R - 1, |a, b| std::array::from_fn(|i| a[i] ^ b[i])); + extend_grid_digits::<_, R>(bg, R - 1, WeightRow::add); + // SAFETY: PMULL is enabled at compile time; the helper stays within the grid. + unsafe { lanes::finish_grid::(fg, bg, acc) }; + return; + } + extend_grid::<_, R>(fg, |a, b| std::array::from_fn(|i| a[i] ^ b[i])); + extend_grid::<_, R>(bg, WeightRow::add); + for (a, (k, w)) in acc.iter_mut().zip(fg.iter().zip(bg.iter())).take(3usize.pow(R as u32)) { + a.mul_acc(w, k); + } +} + /// A row of `ROW` weights, laid out so that both qwords of every 128-bit lane meet a word: /// lane `j` of `lo` is `[c0, c1]` of weight `2j`, of `hi` the same of weight `2j+1`, and /// `c2` holds every weight's last coefficient in order. @@ -300,6 +343,8 @@ mod avx512 { ))] #[cfg_attr(all(target_feature = "vpclmulqdq", target_feature = "avx512f"), allow(dead_code))] mod lanes { + #[cfg(target_arch = "aarch64")] + use super::GRID; use super::{F64, F192, LaneWeight, ProductRow, ROW, WeightRow}; #[cfg(target_arch = "aarch64")] use core::arch::aarch64::*; @@ -513,6 +558,72 @@ mod lanes { } } + /// Consume the final extension digit and keep each point's sums in registers across the row. + /// + /// Only the first pair of each coefficient stores this path's sum; the remaining pairs stay zero. + /// Reducing or combining a ProductRow therefore uses the same representation as the other paths. + #[cfg(all(target_arch = "aarch64", target_feature = "aes"))] + #[inline(always)] + pub(super) unsafe fn finish_grid( + fg: &[[u64; ROW]; GRID], + bg: &[WeightRow; GRID], + acc: &mut [ProductRow; GRID], + ) { + #[inline(always)] + fn update(s: (Neon, Neon, Neon), k: Neon, lo: Neon, hi: Neon, c2: Neon) -> (Neon, Neon, Neon) { + ( + s.0.xor3(lo.mul::<0x00>(k), hi.mul::<0x10>(k)), + s.1.xor3(lo.mul::<0x01>(k), hi.mul::<0x11>(k)), + s.2.xor3(c2.mul::<0x00>(k), c2.mul::<0x11>(k)), + ) + } + #[inline(always)] + unsafe fn load(acc: &ProductRow) -> (Neon, Neon, Neon) { + // SAFETY: every coefficient has a complete pair; only the first pair holds this path's sum. + unsafe { + ( + Neon::load(acc.0[0].as_ptr()), + Neon::load(acc.0[1].as_ptr()), + Neon::load(acc.0[2].as_ptr()), + ) + } + } + #[inline(always)] + unsafe fn store(s: (Neon, Neon, Neon), acc: &mut ProductRow) { + // SAFETY: as for load; remaining pairs stay zero. + unsafe { + s.0.store(acc.0[0].as_mut_ptr()); + s.1.store(acc.0[1].as_mut_ptr()); + s.2.store(acc.0[2].as_mut_ptr()); + } + } + let stride = 3usize.pow((R - 1) as u32); + for at in 0..stride { + // SAFETY: R is in 1..=PRECOMPUTED_ROUNDS; the three points and each pair lie within their rows. + unsafe { + let mut s0 = load(&acc[at]); + let mut s1 = load(&acc[at + stride]); + let mut s2 = load(&acc[at + 2 * stride]); + for c in (0..ROW).step_by(2) { + let k0 = Neon::load(fg[at].as_ptr().add(c)); + let k1 = Neon::load(fg[at + stride].as_ptr().add(c)); + let lo0 = Neon::load(bg[at].lo.as_ptr().add(c)); + let lo1 = Neon::load(bg[at + stride].lo.as_ptr().add(c)); + let hi0 = Neon::load(bg[at].hi.as_ptr().add(c)); + let hi1 = Neon::load(bg[at + stride].hi.as_ptr().add(c)); + let c20 = Neon::load(bg[at].c2.as_ptr().add(c)); + let c21 = Neon::load(bg[at + stride].c2.as_ptr().add(c)); + s0 = update(s0, k0, lo0, hi0, c20); + s1 = update(s1, k1, lo1, hi1, c21); + s2 = update(s2, k0.xor(k1), lo0.xor(lo1), hi0.xor(hi1), c20.xor(c21)); + } + store(s0, &mut acc[at]); + store(s1, &mut acc[at + stride]); + store(s2, &mut acc[at + 2 * stride]); + } + } + } + /// The products of the AVX-512 `fold_acc`, a register at a time, each coefficient packed from scalars. /// /// # Safety @@ -760,10 +871,15 @@ fn grid_pass_with( } bg[LANE_IN_GRID[l]] = WeightRow::pack(&bs[l][x..x + width]); } - extend_grid::<_, R>(fg, |a, b| std::array::from_fn(|i| a[i] ^ b[i])); - extend_grid::<_, R>(bg, WeightRow::add); - for (a, (k, w)) in acc.iter_mut().zip(fg.iter().zip(bg.iter())).take(points) { - a.mul_acc(w, k); + #[cfg(all(target_arch = "aarch64", target_feature = "aes"))] + accumulate_grid::(fg, bg, acc); + #[cfg(not(all(target_arch = "aarch64", target_feature = "aes")))] + { + extend_grid::<_, R>(fg, |a, b| std::array::from_fn(|i| a[i] ^ b[i])); + extend_grid::<_, R>(bg, WeightRow::add); + for (a, (k, w)) in acc.iter_mut().zip(fg.iter().zip(bg.iter())).take(points) { + a.mul_acc(w, k); + } } } }; @@ -935,6 +1051,58 @@ mod tests { use crate::whir::INITIAL_FOLDING_FACTOR; use primitives::test_util::Rng; + #[test] + fn grid_extension_products_match_scalar_definition() { + fn scalar_mul(mut a: u64, mut b: u64) -> u64 { + let mut p = 0; + for _ in 0..64 { + if b & 1 != 0 { + p ^= a; + } + let carry = a >> 63; + a <<= 1; + if carry != 0 { + a ^= 0x1b; + } + b >>= 1; + } + p + } + let mut rng = Rng::new(0x671D); + for rounds in 1..=PRECOMPUTED_ROUNDS { + for block in [1, 2, 4, 8, 16] { + for lanes in [1, (1 << rounds) - 1, 1 << rounds, (1 << rounds) + 3] { + let f: Vec = (0..block * lanes).map(|_| F64(rng.next_u64())).collect(); + let weight = rng.ext_vec(f.len()); + let got = initial_rounds(&f, block, rounds, &Basis::Dense(weight.clone())); + for (point, &sum) in got.grid.iter().enumerate() { + let mut want = F192::ZERO; + for start in (0..lanes).step_by(1 << rounds) { + for x in 0..block { + let (mut k, mut e) = (0, F192::ZERO); + for lane in 0..(1 << rounds).min(lanes - start) { + let included = (0..rounds).all(|bit| { + let digit = point / 3usize.pow(bit as u32) % 3; + digit == 2 || digit == (lane >> bit) & 1 + }); + if included { + k ^= f[(start + lane) * block + x].0; + e += weight[(start + lane) * block + x]; + } + } + want += F192::new(scalar_mul(e.c0, k), scalar_mul(e.c1, k), scalar_mul(e.c2, k)); + } + } + assert_eq!( + sum, want, + "rounds={rounds}, block={block}, lanes={lanes}, point={point}" + ); + } + } + } + } + } + #[test] fn the_kept_weight_is_the_filled_one() { // Invariant: writing the weight out leaves the sums alone and keeps every word of every lane, partial groups