From 3af016cc057393dce6ac1a3d1527fcaca188db88 Mon Sep 17 00:00:00 2001 From: jsign Date: Fri, 4 Sep 2026 11:00:37 -0300 Subject: [PATCH 1/3] Reduce prover memory traffic GKR computes the next round message from folded values in a small local buffer. PCS combines overlapping claims in paired lane windows, then computes the first sumcheck message from the complete weights. On eligible x86 targets, large extension NTTs reuse the base-field NTT through a coefficient view that requires no copy. The public API, protocol, and transcript order remain unchanged. --- crates/lean_vm/src/gkr.rs | 112 +++++++++++++++++++++++-- crates/pcs/src/ring_switch.rs | 32 +++++--- crates/pcs/src/stack_open.rs | 94 +++++++++++++++++++++ crates/pcs/src/stack_open/basis.rs | 109 ++++++++++++++++++++++++ crates/pcs/src/whir.rs | 128 +++++++++++++++++++++++++++-- crates/pcs/src/whir_ntt_ext.rs | 26 +++++- 6 files changed, 479 insertions(+), 22 deletions(-) create mode 100644 crates/pcs/src/stack_open/basis.rs diff --git a/crates/lean_vm/src/gkr.rs b/crates/lean_vm/src/gkr.rs index ee7b841b4..963308ec6 100644 --- a/crates/lean_vm/src/gkr.rs +++ b/crates/lean_vm/src/gkr.rs @@ -10,7 +10,6 @@ use crate::PAR_THRESHOLD; use crate::transcript::{Challenger, ProverState, Receiver, Transmitter, VerifierState}; use primitives::field::{F192, F192Unreduced, mul_unreduced4, mul2, mul4}; use primitives::multilinear::{eq_table, interp, poly_eval, shrink_eq_low}; -#[cfg(target_arch = "x86_64")] use primitives::stream::Stream; use zk_alloc::ArenaVec; @@ -226,6 +225,68 @@ impl QuaternaryLayerState { self.logical_rows /= 2; } + fn fold_and_message(&mut self, challenge: F192, equality: &[F192]) -> [F192; 4] { + let stored_rows = self.values.len() / 4; + let rows = stored_rows.div_ceil(2); + self.next.truncate(4 * rows); + let values = &self.values; + let dst = parallel::SendPtr(self.next.as_mut_ptr()); + const PAIRS: usize = 16; + let pairs = rows.div_ceil(2); + let task = |index: usize| { + let first = index * PAIRS; + let end = (first + PAIRS).min(pairs); + let mut stage = [F192::ZERO; 8 * PAIRS]; + let end_row = (2 * end).min(rows); + for row in 2 * first..end_row { + let lo = 8 * row; + let left = &values[lo..lo + 4]; + let right = values.get(lo + 4..lo + 8).unwrap_or(&[F192::ONE; 4]); + let product = mul4(std::array::from_fn(|i| left[i] + right[i]), [challenge; 4]); + let offset = 4 * (row - 2 * first); + for i in 0..4 { + stage[offset + i] = left[i] + product[i]; + } + } + let mut message = [F192Unreduced::ZERO; 4]; + for pair in first..end { + let lo = 8 * (pair - first); + let left = &stage[lo..lo + 4]; + let right = if 2 * pair + 1 < rows { + &stage[lo + 4..lo + 8] + } else { + &[F192::ONE; 4] + }; + let lines = std::array::from_fn(|i| [left[i], left[i] + right[i]]); + let terms = quartic_summand(lines, equality[pair]); + for i in 0..4 { + message[i] ^= terms[i]; + } + } + // The next round reads the destination; this round reads only the local stage. + let stream = Stream::new(); + let len = 4 * (end_row - 2 * first); + // SAFETY: tasks own disjoint initialized prefixes of the output, covering every row. + unsafe { stream.copy(dst.slice(8 * first, len), &stage[..len]) }; + message + }; + let xor = |mut a: [F192Unreduced; 4], b: [F192Unreduced; 4]| { + for i in 0..4 { + a[i] ^= b[i]; + } + a + }; + let tasks = pairs.div_ceil(PAIRS); + let message = if rows >= PAR_THRESHOLD { + parallel::map_reduce(tasks, || [F192Unreduced::ZERO; 4], task, xor) + } else { + (0..tasks).map(task).fold([F192Unreduced::ZERO; 4], xor) + }; + std::mem::swap(&mut self.values, &mut self.next); + self.logical_rows /= 2; + message.map(F192Unreduced::reduce) + } + fn children(&self) -> [F192; 4] { debug_assert_eq!(self.values.len(), 4); debug_assert_eq!(self.logical_rows, 1); @@ -314,8 +375,12 @@ pub fn prove_product_triple(leaves: [ArenaVec; 3], ps: &mut ProverState, s Vec::new() }; let mut round_point = Vec::with_capacity(round_count); - for _ in 0..round_count { - let messages = [0, 1, 2].map(|tree| trees[tree].round_message(&equality)); + let mut messages = if round_count > 0 { + trees.each_ref().map(|tree| tree.round_message(&equality)) + } else { + [[F192::ZERO; 4]; 3] + }; + for round in 0..round_count { let mut coeffs = [0, 1, 2, 3].map(|coefficient| { messages[0][coefficient] + lambda * (messages[1][coefficient] + lambda * messages[2][coefficient]) }); @@ -326,10 +391,14 @@ pub fn prove_product_triple(leaves: [ArenaVec; 3], ps: &mut ProverState, s ps.add_scalars(&coeffs); let challenge = ps.sample(); round_point.push(challenge); - for tree in &mut trees { - tree.fold(challenge); - } shrink_eq_low(&mut equality); + if round + 1 < round_count { + messages = trees.each_mut().map(|tree| tree.fold_and_message(challenge, &equality)); + } else { + for tree in &mut trees { + tree.fold(challenge); + } + } } for tree in &trees { @@ -475,6 +544,37 @@ mod tests { } } + #[test] + fn fused_fold_matches_separate_fold_and_message() { + for width in [4usize, 16, 1 << 14] { + for len in [1, 4, 5, 7, 8, 9, 31, 32, 33, 4 * width - 5, 4 * width - 1, 4 * width] { + if len > 4 * width { + continue; + } + let values: ArenaVec = (0..len) + .map(|i| F192::new((17 * i + 1) as u64, (i * i + 3) as u64, (5 * i + 7) as u64)) + .collect(); + let mut reference = QuaternaryLayerState::new(ArenaVec::from_slice(&values), width); + let mut fused = QuaternaryLayerState::new(values, width); + let point: Vec = (0..width.ilog2() - 1) + .map(|i| F192::new(31 + u64::from(i), 7, 11)) + .collect(); + let mut equality = eq_table(&point); + while reference.logical_rows > 2 { + let challenge = F192::new(reference.logical_rows as u64, 13, 19); + reference.fold(challenge); + shrink_eq_low(&mut equality); + let message = fused.fold_and_message(challenge, &equality); + assert_eq!(message, reference.round_message(&equality), "width={width}, len={len}"); + assert_eq!(&*fused.values, &*reference.values, "width={width}, len={len}"); + } + reference.fold(F192::Y); + fused.fold(F192::Y); + assert_eq!(fused.children(), reference.children()); + } + } + } + #[test] fn radix_four_roundtrip_at_even_and_odd_depths() { for mu in 0..=10 { diff --git a/crates/pcs/src/ring_switch.rs b/crates/pcs/src/ring_switch.rs index ff4e65ada..e43a73f7c 100644 --- a/crates/pcs/src/ring_switch.rs +++ b/crates/pcs/src/ring_switch.rs @@ -461,19 +461,31 @@ pub(crate) fn combine_deferred_into(outputs: &[DeferredRingSwitchOutput], out: & ); parallel::chunks_mut(out, block_len, |hi, out_block| { - for (claim_idx, claim) in outputs.iter().enumerate() { - let e_hi = claim.eq_hi[hi]; - if claim_idx == 0 { - for (slot, &e_lo) in out_block.iter_mut().zip(&claim.eq_lo) { - *slot = fold_one_slot_ext(e_lo * e_hi, &claim.table); - } - } else { - for (slot, &e_lo) in out_block.iter_mut().zip(&claim.eq_lo) { - *slot += fold_one_slot_ext(e_lo * e_hi, &claim.table); + combine_deferred_chunk(outputs, hi * block_len, out_block); + }); +} + +pub(crate) fn combine_deferred_chunk(outputs: &[DeferredRingSwitchOutput], start: usize, out: &mut [F192]) { + for (claim_idx, claim) in outputs.iter().enumerate() { + let block_len = claim.eq_lo.len(); + assert!(start + out.len() <= block_len * claim.eq_hi.len()); + let mut done = 0; + while done < out.len() { + let index = start + done; + let lo = index % block_len; + let len = (block_len - lo).min(out.len() - done); + let e_hi = claim.eq_hi[index / block_len]; + for (slot, &e_lo) in out[done..done + len].iter_mut().zip(&claim.eq_lo[lo..lo + len]) { + let value = fold_one_slot_ext(e_lo * e_hi, &claim.table); + if claim_idx == 0 { + *slot = value; + } else { + *slot += value; } } + done += len; } - }); + } } /// Split point for the factored eq build: low half sized ~n/2 (min 4, the diff --git a/crates/pcs/src/stack_open.rs b/crates/pcs/src/stack_open.rs index 36c5cba72..90d2498d3 100644 --- a/crates/pcs/src/stack_open.rs +++ b/crates/pcs/src/stack_open.rs @@ -66,6 +66,8 @@ use super::ring_switch; use super::whir::{ProverConfig, VerifierConfig}; use super::whir::{ProverData, recursive_prover_with_basis, recursive_verifier_with_basis_succinct}; +mod basis; + // --------------------------------------------------------------------------- // Claim types // --------------------------------------------------------------------------- @@ -413,6 +415,27 @@ pub fn open_batch_mixed_whir_stacked( let mut target = rs_outputs .iter() .fold(F192::ZERO, |acc, out| acc + out.batched_sumcheck_claim); + if stack.len() >= 1 << 16 { + target += point_claims + .iter() + .zip(lambdas_pd) + .fold(F192::ZERO, |sum, (claim, &lambda)| sum + lambda * claim.value()); + let lane_block = 1usize << (log_n - config.initial_k); + let (b_stack, message) = basis::build(stack, lane_block, point_claims, lambdas_pd, ring, &rs_outputs); + mark("basis + initial message", &mut t); + super::whir::recursive_prover_with_prepared_basis( + config, + log_n, + stack, + b_stack, + target, + &prover_data.codeword, + &prover_data.merkle_tree, + Some(message), + ps, + ); + return; + } // Parallel first-touch wins for the tower stack: its many scattered point // claims otherwise fault pages one claim at a time. A scatter that lands on // slots an earlier one already touched has to accumulate, so those slots @@ -563,6 +586,77 @@ mod tests { const DOMAIN: &[u8] = b"stack-open-test"; + #[test] + fn fused_basis_matches_dense_weights() { + let mut rng = Rng::new(0xBA515); + for (lane_vars, lanes) in [(6usize, 1usize), (6, 3), (10, 15), (10, 37)] { + let lane_block = 1 << lane_vars; + let stack: Vec = (0..lanes * lane_block).map(|_| F64(rng.next_u64())).collect(); + let qflock_vars = lane_vars + usize::from(lanes > 1); + let qflock_len = 1 << qflock_vars; + let offset = if stack.len() >= 2 * qflock_len { qflock_len } else { 0 }; + let ring = RingSwitchOpen { + offset, + qflock_vars, + claims: (0..2) + .map(|_| RingSwitchClaim { + suffix_point: rng.ext_vec(qflock_vars), + s_hat_v: None, + }) + .collect(), + }; + let coordinates = rng.ext_vec(192); + let rs_outputs: Vec<_> = ring + .claims + .iter() + .map(|claim| { + let state = + ring_switch::prove_prepare(&stack[offset..offset + qflock_len], &claim.suffix_point, None); + ring_switch::prove_finish_deferred(state, &coordinates, rng.ext()) + }) + .collect(); + let mut claims: Vec<_> = [ + (offset, qflock_vars), + ((lanes - 1) * lane_block, lane_vars), + (8, 3), + (0, 0), + ] + .into_iter() + .map(|(offset, vars)| StackClaim::Point { + offset, + low_point: rng.ext_vec(vars), + value: rng.ext(), + }) + .collect(); + for stride_log in [0, 1, 3, qflock_vars - 1, qflock_vars] { + claims.push(StackClaim::Strided { + offset, + slot: (1 << stride_log) - 1, + stride_log, + point: rng.ext_vec(qflock_vars - stride_log), + value: rng.ext(), + }); + } + let lambdas = rng.ext_vec(claims.len()); + let mut expected = vec![F192::ZERO; stack.len()]; + ring_switch::combine_deferred_into(&rs_outputs, &mut expected[offset..offset + qflock_len]); + let mut target = F192::ZERO; + fold_stacked_point_claims( + &mut expected, + &mut target, + &claims, + &lambdas, + &vec![false; claims.len()], + ); + let (actual, message) = basis::build(&stack, lane_block, &claims, &lambdas, &ring, &rs_outputs); + assert_eq!(&*actual, expected, "lane_vars={lane_vars}, lanes={lanes}"); + let (_, expected_message) = super::super::whir::build_initial_basis(&stack, lane_block, |start, dst| { + dst.copy_from_slice(&expected[start..start + dst.len()]); + }); + assert_eq!(message, expected_message); + } + } + struct Instance { vc: VerifierConfig, log_n: usize, diff --git a/crates/pcs/src/stack_open/basis.rs b/crates/pcs/src/stack_open/basis.rs new file mode 100644 index 000000000..853a5fe2b --- /dev/null +++ b/crates/pcs/src/stack_open/basis.rs @@ -0,0 +1,109 @@ +use std::mem::MaybeUninit; + +use primitives::field::{F64, F192}; +use zk_alloc::ArenaVec; + +use super::{RingSwitchOpen, StackClaim}; +use crate::ring_switch::{DeferredRingSwitchOutput, combine_deferred_chunk}; +use crate::whir::{INITIAL_BASIS_CHUNK, SumcheckMessage, build_eq_table_ext_seeded, build_initial_basis}; + +struct PointWeight<'a> { + offset: usize, + end: usize, + slot: usize, + stride: usize, + low: &'a [F192], + high: ArenaVec, +} + +impl<'a> PointWeight<'a> { + fn new(claim: &'a StackClaim, lambda: F192, chunk_log: usize) -> Self { + let (offset, slot, stride_log, point) = match claim { + StackClaim::Point { offset, low_point, .. } => (*offset, 0, 0, low_point.as_slice()), + StackClaim::Strided { + offset, + slot, + stride_log, + point, + .. + } => (*offset, *slot, *stride_log, point.as_slice()), + }; + let len = 1usize << (stride_log + point.len()); + let stride = 1usize << stride_log; + assert!(offset.is_multiple_of(len), "claim must be aligned to its support"); + assert!(slot < stride, "claim slot must fit the stride"); + let low_vars = point.len().min(chunk_log.saturating_sub(stride_log)); + let (low, high_point) = point.split_at(low_vars); + let mut high = zk_alloc::alloc_uninit(1 << high_point.len()); + build_eq_table_ext_seeded(high_point, lambda, &mut high); + // SAFETY: the seeded equality build initializes the whole table. + let high = unsafe { zk_alloc::assume_init(high) }; + Self { + offset, + end: offset + len, + slot, + stride, + low, + high, + } + } + + fn add(&self, start: usize, dst: &mut [F192], scratch: &mut [MaybeUninit]) { + let base = self.offset + self.slot; + let lo = start.max(base); + let hi = (start + dst.len()).min(self.end); + if lo >= hi { + return; + } + let first = (lo - base).div_ceil(self.stride); + let end = (hi - base).div_ceil(self.stride); + if first == end { + return; + } + let len = 1usize << self.low.len(); + assert!(first.is_multiple_of(len) && end - first == len); + build_eq_table_ext_seeded(self.low, self.high[first / len], &mut scratch[..len]); + // SAFETY: the build above initializes this prefix before the scatter reads it. + let eq = unsafe { std::slice::from_raw_parts(scratch.as_ptr().cast::(), len) }; + let dst_offset = base + first * self.stride - start; + for (i, &value) in eq.iter().enumerate() { + dst[dst_offset + i * self.stride] += value; + } + } +} + +pub(super) fn build( + stack: &[F64], + lane_block: usize, + claims: &[StackClaim], + lambdas: &[F192], + ring: &RingSwitchOpen, + rs_outputs: &[DeferredRingSwitchOutput], +) -> (ArenaVec, SumcheckMessage) { + assert_eq!(claims.len(), lambdas.len()); + let chunk_log = lane_block.min(INITIAL_BASIS_CHUNK).ilog2() as usize; + let weights: Vec<_> = claims + .iter() + .zip(lambdas) + .map(|(claim, &lambda)| PointWeight::new(claim, lambda, chunk_log)) + .collect(); + let mut by_lane = vec![Vec::new(); stack.len() / lane_block]; + for (index, weight) in weights.iter().enumerate() { + for lane in &mut by_lane[weight.offset / lane_block..weight.end.div_ceil(lane_block)] { + lane.push(index); + } + } + let ring_end = ring.offset + (1 << ring.qflock_vars); + build_initial_basis(stack, lane_block, |start, dst| { + dst.fill(F192::ZERO); + let lo = start.max(ring.offset); + let hi = (start + dst.len()).min(ring_end); + if lo < hi { + combine_deferred_chunk(rs_outputs, lo - ring.offset, &mut dst[lo - start..hi - start]); + } + let mut scratch = [MaybeUninit::uninit(); INITIAL_BASIS_CHUNK]; + for &index in &by_lane[start / lane_block] { + weights[index].add(start, dst, &mut scratch); + } + }) +} diff --git a/crates/pcs/src/whir.rs b/crates/pcs/src/whir.rs index 66e4ac7e3..05ad83a05 100644 --- a/crates/pcs/src/whir.rs +++ b/crates/pcs/src/whir.rs @@ -496,8 +496,8 @@ pub(crate) fn ligero_commit_ext( // lifts the witness into E and all later rounds are pure E. /// (u_0, u_2) per round in E. -#[derive(Clone, Copy, Debug)] -struct SumcheckMessage { +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct SumcheckMessage { u_0: F192, u_2: F192, } @@ -887,6 +887,52 @@ fn round_msg_blocks(f: &[T], b: &[F192], block: usize) -> Sumch } } +pub(crate) const INITIAL_BASIS_CHUNK: usize = 256; + +pub(crate) fn build_initial_basis( + f: &[F64], + block: usize, + fill: impl Fn(usize, &mut [F192]) + Sync, +) -> (ArenaVec, SumcheckMessage) { + assert!(block.is_power_of_two() && f.len().is_multiple_of(block)); + let n_blocks = f.len() / block; + let per = block.div_ceil(INITIAL_BASIS_CHUNK); + // SAFETY: each task fills and publishes its disjoint lane windows before returning. + let mut basis = unsafe { ArenaVec::::uninitialized(f.len()) }; + let dst = parallel::SendPtr(basis.as_mut_ptr()); + let task = |t: usize| { + let (pair, chunk) = (t / per, t % per); + let offset = chunk * INITIAL_BASIS_CHUNK; + let len = INITIAL_BASIS_CHUNK.min(block - offset); + let lo = 2 * pair * block + offset; + let mut b0 = [F192::ZERO; INITIAL_BASIS_CHUNK]; + let mut b1 = [F192::ZERO; INITIAL_BASIS_CHUNK]; + fill(lo, &mut b0[..len]); + let stream = Stream::new(); + let message = if 2 * pair + 1 < n_blocks { + let hi = lo + block; + fill(hi, &mut b1[..len]); + let message = msg_terms_pair(&f[lo..lo + len], &f[hi..hi + len], &b0[..len], &b1[..len]); + // SAFETY: this task owns the high lane window, disjoint from every other task. + unsafe { stream.copy(dst.slice(hi, len), &b1[..len]) }; + message + } else { + msg_terms_lone(&f[lo..lo + len], &b0[..len]) + }; + // SAFETY: this task owns the low lane window, disjoint from every other task. + unsafe { stream.copy(dst.slice(lo, len), &b0[..len]) }; + message + }; + let (u_0, u_2) = accumulate_msg(n_blocks.div_ceil(2) * per, f.len() / 2, F192BaseUnreduced::ZERO, task); + ( + basis, + SumcheckMessage { + u_0: u_0.reduce(), + u_2: u_2.reduce(), + }, + ) +} + /// Fused lane fold + next-round message. Mirror of [`fold_and_msg_lsb`] for the /// block pairing: a task owns one output *pair* (so four input blocks), because /// that is the smallest unit the next round's message is local to. @@ -1024,10 +1070,16 @@ impl<'a> SumcheckProver<'a> { /// `block` is the lane block size `2^(log_n - initial_k)`: the first /// `initial_k` rounds are the lane fold, so round 0's message already pairs /// whole blocks rather than adjacent words. - fn new(f: &'a [F64], b1: ArenaVec, h1: F192, block: usize) -> (Self, SumcheckMessage) { + fn new( + f: &'a [F64], + b1: ArenaVec, + h1: F192, + block: usize, + initial_message: Option, + ) -> (Self, SumcheckMessage) { let _span = tracing::info_span!("Sumcheck round", round = 0, log_size = f.len().ilog2()).entered(); assert_eq!(f.len(), b1.len()); - let msg = round_msg_blocks(f, &b1, block); + let msg = initial_message.unwrap_or_else(|| round_msg_blocks(f, &b1, block)); let inst = Self { f: Witness::Base(f), combined_basis: b1, @@ -1250,6 +1302,30 @@ pub fn recursive_prover_with_basis( l0_codeword: &[F64], l0_tree: &[Hash], ps: &mut impl Transmitter, +) { + recursive_prover_with_prepared_basis( + config, + log_n, + witness, + b_initial, + target, + l0_codeword, + l0_tree, + None, + ps, + ); +} + +pub(crate) fn recursive_prover_with_prepared_basis( + config: &ProverConfig, + log_n: usize, + witness: &[F64], + b_initial: ArenaVec, + target: F192, + l0_codeword: &[F64], + l0_tree: &[Hash], + initial_message: Option, + ps: &mut impl Transmitter, ) { let r = config.level_steps; let initial_k = config.initial_k; @@ -1322,7 +1398,7 @@ pub fn recursive_prover_with_basis( let _t = std::time::Instant::now(); let sumcheck_span = tracing::info_span!("Sumcheck"); let (mut sc_prover, start_msg) = - sumcheck_span.in_scope(|| SumcheckProver::new(witness, b_initial, target, lane_block)); + sumcheck_span.in_scope(|| SumcheckProver::new(witness, b_initial, target, lane_block, initial_message)); send_msg(ps, start_msg, target); let mut r_lane_fold = Vec::with_capacity(initial_k); @@ -2769,4 +2845,46 @@ mod tests { forward_transform_interleaved_ext_parallel_from_layer(&ntt, &mut b, lanes, 1); assert_eq!(a, b); } + + #[test] + fn ext_ntt_coefficient_view_matches_scalar() { + let mut rng = Rng::new(0xE192); + for (log_d, lanes) in [(3usize, 1usize), (8, 4), (12, 16), (14, 16)] { + let ntt = AdditiveNttF64::standard(log_d); + let original = rng.ext_vec((1 << log_d) * lanes); + for start_layer in [0, 1, 2, 3, 4, log_d / 2, log_d] { + if start_layer > log_d { + continue; + } + let mut expected = original.clone(); + forward_transform_interleaved_ext_scalar_from_layer(&ntt, &mut expected, lanes, start_layer); + let mut actual = original.clone(); + crate::whir_ntt_ext::forward_transform_interleaved_ext_via_base(&ntt, &mut actual, lanes, start_layer); + assert_eq!( + actual, expected, + "log_d={log_d}, lanes={lanes}, start_layer={start_layer}" + ); + let mut dispatched = original.clone(); + forward_transform_interleaved_ext_from_layer(&ntt, &mut dispatched, lanes, start_layer); + assert_eq!(dispatched, expected); + } + } + } + + #[test] + fn initial_basis_message_matches_materialized_weights() { + let mut rng = Rng::new(0xBA515); + for block in [1, 16, INITIAL_BASIS_CHUNK, 4 * INITIAL_BASIS_CHUNK] { + for lanes in [1, 2, 3, 37] { + let f: Vec = (0..block * lanes).map(|_| F64(rng.next_u64())).collect(); + let expected = rng.ext_vec(f.len()); + let expected_message = round_msg_blocks(&f, &expected, block); + let (actual, message) = build_initial_basis(&f, block, |start, out| { + out.copy_from_slice(&expected[start..start + out.len()]); + }); + assert_eq!(&*actual, expected, "block={block}, lanes={lanes}"); + assert_eq!(message, expected_message, "block={block}, lanes={lanes}"); + } + } + } } diff --git a/crates/pcs/src/whir_ntt_ext.rs b/crates/pcs/src/whir_ntt_ext.rs index de6730859..8280d0530 100644 --- a/crates/pcs/src/whir_ntt_ext.rs +++ b/crates/pcs/src/whir_ntt_ext.rs @@ -35,7 +35,31 @@ pub(crate) fn forward_transform_interleaved_ext_from_layer( assert!(log_d <= ntt.log_domain_size()); assert!(start_layer <= log_d); - forward_transform_interleaved_ext_parallel_from_layer(ntt, data, num_ntts, start_layer); + // Keep the dedicated NEON and AVX-512 kernels; other x86 targets share the base NTT's larger groups. + if log_d >= 12 + && cfg!(all( + target_arch = "x86_64", + not(all(target_feature = "vpclmulqdq", target_feature = "avx512f")) + )) + { + forward_transform_interleaved_ext_via_base(ntt, data, num_ntts, start_layer); + } else { + forward_transform_interleaved_ext_parallel_from_layer(ntt, data, num_ntts, start_layer); + } +} + +pub(crate) fn forward_transform_interleaved_ext_via_base( + ntt: &AdditiveNttF64, + data: &mut [F192], + num_ntts: usize, + start_layer: usize, +) { + const _: () = assert!(size_of::() == 3 * size_of::()); + const _: () = assert!(align_of::() == align_of::()); + // SAFETY: F192 is repr(C) over three u64 coefficients, and F64 is transparent over u64. + // Base-field twiddles act independently on each coefficient, preserving the row layout. + let coefficients = unsafe { std::slice::from_raw_parts_mut(data.as_mut_ptr().cast::(), 3 * data.len()) }; + ntt.forward_transform_interleaved_parallel_from_layer(coefficients, 3 * num_ntts, start_layer); } /// Scalar reference for the E-valued interleaved forward NTT (test oracle and From 675ad46ff5582040b8aa2d9075d9c9d1b54165e5 Mon Sep 17 00:00:00 2001 From: Tom Wambsgans Date: Fri, 4 Sep 2026 20:36:14 +0200 Subject: [PATCH 2/3] pcs: keep only the fused basis path, delete the dense one The size gate on the fused lifted-weight build never selected the dense path in practice: an aggregate opens the stack at 2^26 words and the recursion at 2^24, against a 2^16 threshold, and the pcs tests were the only caller below it. Routing those through the fused build passes, so the second path was dead weight. Deleting it takes claim_write_plan, fold_stacked_point_claims and add_eq_table_ext_seeded with it, and reduces combine_deferred_into to its combine_deferred_chunk core. fused_basis_matches_dense_weights keeps its differential coverage against a naive eq tensor written in the test, which shares no code with the build under test. Net 211 lines lighter than the parent, and 220 SPHINCS measures 2.944 s against the parent's 2.942 s on a Ryzen 7 PRO 8700GE. --- crates/pcs/src/ring_switch.rs | 24 +--- crates/pcs/src/stack_open.rs | 222 ++++++---------------------------- crates/pcs/src/whir.rs | 50 +------- 3 files changed, 42 insertions(+), 254 deletions(-) diff --git a/crates/pcs/src/ring_switch.rs b/crates/pcs/src/ring_switch.rs index e43a73f7c..75ec4a852 100644 --- a/crates/pcs/src/ring_switch.rs +++ b/crates/pcs/src/ring_switch.rs @@ -449,22 +449,6 @@ pub(crate) fn prove_finish_deferred( /// No per-claim dense vector is allocated or read back, and the first claim /// **writes** rather than accumulates, so the caller need not pre-zero `out`. /// -/// Use one pass per claim because interleaving tables increases lookup pressure. -pub(crate) fn combine_deferred_into(outputs: &[DeferredRingSwitchOutput], out: &mut [F192]) { - assert!(!outputs.is_empty()); - let block_len = outputs[0].eq_lo.len(); - assert!(block_len.is_power_of_two()); - assert!( - outputs - .iter() - .all(|o| { o.eq_lo.len() == block_len && o.eq_lo.len() * o.eq_hi.len() == out.len() }) - ); - - parallel::chunks_mut(out, block_len, |hi, out_block| { - combine_deferred_chunk(outputs, hi * block_len, out_block); - }); -} - pub(crate) fn combine_deferred_chunk(outputs: &[DeferredRingSwitchOutput], start: usize, out: &mut [F192]) { for (claim_idx, claim) in outputs.iter().enumerate() { let block_len = claim.eq_lo.len(); @@ -681,7 +665,7 @@ mod tests { .iter() .fold(F192::ZERO, |acc, out| acc + out.batched_sumcheck_claim); let mut deferred_basis = vec![F192::ZERO; expected_basis.len()]; - combine_deferred_into(&deferred, &mut deferred_basis); + combine_deferred_chunk(&deferred, 0, &mut deferred_basis); assert_eq!(deferred_target, expected_target); assert_eq!(deferred_basis, expected_basis); @@ -771,7 +755,7 @@ mod tests { assert_eq!(apply_composed_map(value, &challenges), expanded); } - /// The byte-table fold of a dense tensor: the kernel `combine_deferred_into` + /// The byte-table fold of a dense tensor: the kernel `combine_deferred_chunk` /// runs per slot, without its factored-eq slot generation. fn fold_dense(tensor: &[F192], coordinate_weights: &[F192]) -> Vec { let tables = build_fold_byte_table_ext(coordinate_weights); @@ -882,7 +866,7 @@ mod tests { ); } - /// The byte-table fold behind `combine_deferred_into` must match the naive + /// The byte-table fold behind `combine_deferred_chunk` must match the naive /// bit-scan on arbitrary (not necessarily eq-structured) input. #[test] fn rs_eq_ind_fast_matches_naive() { @@ -959,7 +943,7 @@ mod tests { let out = prove_finish_deferred(state, &coordinate_weights, F192::ONE); let sumcheck_claim = out.batched_sumcheck_claim; let mut rs_eq_ind = vec![F192::ZERO; packed.len()]; - combine_deferred_into(&[out], &mut rs_eq_ind); + combine_deferred_chunk(&[out], 0, &mut rs_eq_ind); assert_eq!(inner_product_base_ext(&packed, &rs_eq_ind), sumcheck_claim); recursive_prover_with_basis( &pc, diff --git a/crates/pcs/src/stack_open.rs b/crates/pcs/src/stack_open.rs index 90d2498d3..08fd4535b 100644 --- a/crates/pcs/src/stack_open.rs +++ b/crates/pcs/src/stack_open.rs @@ -64,7 +64,7 @@ use primitives::multilinear::eq_eval; use super::pack::PACKING_WIDTH; use super::ring_switch; use super::whir::{ProverConfig, VerifierConfig}; -use super::whir::{ProverData, recursive_prover_with_basis, recursive_verifier_with_basis_succinct}; +use super::whir::{ProverData, recursive_verifier_with_basis_succinct}; mod basis; @@ -168,128 +168,6 @@ fn claim_range(claim: &StackClaim) -> (usize, usize) { } } -/// Per-claim "this claim may WRITE its range instead of accumulating into it", -/// plus the ranges those writes cover. A `Point` claim qualifies when nothing -/// written earlier (the q_flock block, which `combine_deferred_into` fills, or -/// an earlier claim) lands anywhere in its range: claims are folded in list -/// order, so its slice is still untouched when its turn comes, and in char 2 -/// writing where a zero would have been is bit-identical. -/// -/// The returned ranges are pairwise disjoint, so the caller only has to zero -/// the gaps between them. -fn claim_write_plan(claims: &[StackClaim], qflock: (usize, usize)) -> (Vec, Vec<(usize, usize)>) { - let mut touched = vec![qflock]; - let mut written = vec![qflock]; - let mut write_first = Vec::with_capacity(claims.len()); - for claim in claims { - let range = claim_range(claim); - let free = touched.iter().all(|&(s, e)| range.1 <= s || e <= range.0); - let first = free && matches!(claim, StackClaim::Point { .. }); - if first { - written.push(range); - } - touched.push(range); - write_first.push(first); - } - (write_first, written) -} - -/// Fold the lambda-weighted point claims into the stack weight `b_stack` and -/// running `target` (pure: the caller has already observed the claim values -/// and sampled `lambdas` in transcript order). A `Point` builds eq over ONLY -/// its aligned slice, a `Strided` scatters the eq of its high coords at the -/// slot's stride. Every claim but the first one on a range scatters with `+=`, -/// so overlapping slices accumulate correctly; the OUTER loop therefore stays -/// serial (several bus claims can land on one column region), and parallelism -/// lives inside each claim: the lambda-seeded eq build (parallel above its level -/// floor) and the strided scatter. Small slices stay fully serial (with many -/// tiny point claims, pool dispatch would cost more than the fold itself). The -/// lambda seeding and the serial/parallel splits are exact-field/order-preserving, -/// so `b_stack`'s bytes (and hence the proof) are unchanged relative to the -/// build-then-multiply form. -/// -/// `write_first` comes from [`claim_write_plan`]: where it is set, the claim's -/// slice is uninitialized and the eq table is written straight into it by -/// [`super::whir::build_eq_table_ext_seeded`]; elsewhere -/// [`super::whir::add_eq_table_ext_seeded`] accumulates, expanding its last -/// coordinate straight into `b_stack` so the table's largest level is never -/// staged in scratch. -fn fold_stacked_point_claims( - b_stack: &mut [F192], - target: &mut F192, - claims: &[StackClaim], - lambdas: &[F192], - write_first: &[bool], -) { - // One reusable eq scratch: half the largest accumulating Point claim (the - // seeded add expands the last coordinate straight into `b_stack`, and a - // write-first Point needs no scratch at all), or the whole eq table of the - // largest Strided claim. A fresh multi-MB allocation per claim would pay the - // first-touch page faults anew. - let scratch_len = claims - .iter() - .zip(write_first) - .map(|(c, &first)| match c { - StackClaim::Point { .. } if first => 0, - StackClaim::Point { low_point, .. } => 1usize << low_point.len().saturating_sub(1), - StackClaim::Strided { point, .. } => 1usize << point.len(), - }) - .max() - .unwrap_or(0); - let mut scratch = zk_alloc::alloc_uninit(scratch_len); - for ((claim, g), &first) in claims.iter().zip(lambdas.iter()).zip(write_first) { - let g = *g; - match claim { - StackClaim::Point { - offset, - low_point, - value, - } => { - let len = 1usize << low_point.len(); - assert!( - offset % len == 0, - "StackClaim::Point: offset must be 2^|low_point|-aligned" - ); - let dst = &mut b_stack[*offset..*offset + len]; - if first { - super::whir::build_eq_table_ext_seeded(low_point, g, dst); - } else { - super::whir::add_eq_table_ext_seeded(low_point, g, &mut scratch, dst); - } - *target += g * *value; - } - StackClaim::Strided { - offset, - slot, - stride_log, - point, - value, - } => { - // Sparse: eq over the instance `point` (2^|point| entries), - // scattered at stride 2^stride_log from the slot's position. - // Identical b_stack contribution to the dense Point with - // low_point = slot_bits ++ point, at ~2^stride_log x less work. - let stride = 1usize << stride_log; - let block = 1usize << (stride_log + point.len()); - assert!(*slot < stride, "StackClaim::Strided: slot must fit the stride"); - assert!( - offset % block == 0, - "StackClaim::Strided: offset must be 2^(stride_log + |point|)-aligned" - ); - let base = *offset + *slot; - let len = 1usize << point.len(); - super::whir::build_eq_table_ext_seeded(point, g, &mut scratch[..len]); - // SAFETY: the build above initialized exactly this prefix. - let eq = unsafe { std::slice::from_raw_parts(scratch.as_ptr().cast::(), len) }; - for (j, &ej) in eq.iter().enumerate() { - b_stack[base + j * stride] += ej; - } - *target += g * *value; - } - } - } -} - /// The claim's weight `eq(full claim point, x)` at an arbitrary point `x` of /// the full stack cube. A `Point`'s full point is `[low_point, sel_bits]`, a /// `Strided`'s is `[slot_bits, point, sel_bits]`; neither is materialized. @@ -412,67 +290,24 @@ pub fn open_batch_mixed_whir_stacked( // 3. Combined target and lifted stack weight b_stack: the lambda-weighted // rs_eq_ind sum scattered at the q_flock slice, plus the point-claim // eq tensors scattered at their offsets. - let mut target = rs_outputs + let target = rs_outputs .iter() - .fold(F192::ZERO, |acc, out| acc + out.batched_sumcheck_claim); - if stack.len() >= 1 << 16 { - target += point_claims + .fold(F192::ZERO, |acc, out| acc + out.batched_sumcheck_claim) + + point_claims .iter() .zip(lambdas_pd) .fold(F192::ZERO, |sum, (claim, &lambda)| sum + lambda * claim.value()); - let lane_block = 1usize << (log_n - config.initial_k); - let (b_stack, message) = basis::build(stack, lane_block, point_claims, lambdas_pd, ring, &rs_outputs); - mark("basis + initial message", &mut t); - super::whir::recursive_prover_with_prepared_basis( - config, - log_n, - stack, - b_stack, - target, - &prover_data.codeword, - &prover_data.merkle_tree, - Some(message), - ps, - ); - return; - } - // Parallel first-touch wins for the tower stack: its many scattered point - // claims otherwise fault pages one claim at a time. A scatter that lands on - // slots an earlier one already touched has to accumulate, so those slots - // start at zero; a range whose first writer covers all of it does not, and - // zeroing it would be stores thrown away. `combine_deferred_into` writes the - // whole q_flock block, and `claim_write_plan` finds the point claims that - // likewise write a whole untouched range. - // - // SAFETY: every slot is written before it is read: the fill covers every gap - // between the written ranges, `combine_deferred_into` writes the q_flock - // block, and each `write_first` claim writes its whole slice before any - // later claim can accumulate into it. - let (write_first, mut written) = claim_write_plan(point_claims, (ring.offset, ring.offset + qflock_len)); - let mut b_stack = unsafe { zk_alloc::ArenaVec::::uninitialized(stack.len()) }; - { - const ZERO_CHUNK: usize = 1 << 16; - written.sort_unstable(); - let mut cursor = 0usize; - let zero = |part: &mut [F192]| parallel::chunks_mut(part, ZERO_CHUNK, |_, c| c.fill(F192::ZERO)); - for (start, end) in written { - if start > cursor { - zero(&mut b_stack[cursor..start]); - } - cursor = cursor.max(end); - } - zero(&mut b_stack[cursor..]); - mark("b_stack zero fill", &mut t); - let block = &mut b_stack[ring.offset..ring.offset + qflock_len]; - ring_switch::combine_deferred_into(&rs_outputs, block); - mark("rs_eq_ind scatter", &mut t); - } - fold_stacked_point_claims(&mut b_stack, &mut target, point_claims, lambdas_pd, &write_first); - mark("point-claim folds", &mut t); + + // The lifted weight is built and consumed in one pass: each lane window is + // filled from the ring-switch outputs and the point claims, then feeds + // round 0's message while it is still hot, so nothing re-reads the buffer. + let lane_block = 1usize << (log_n - config.initial_k); + let (b_stack, message) = basis::build(stack, lane_block, point_claims, lambdas_pd, ring, &rs_outputs); + mark("basis + initial message", &mut t); // 4. One WHIR over the full stack against the combined claim (the // stack is borrowed by the prover; no copy). - recursive_prover_with_basis( + super::whir::recursive_prover_with_prepared_basis( config, log_n, stack, @@ -480,8 +315,9 @@ pub fn open_batch_mixed_whir_stacked( target, &prover_data.codeword, &prover_data.merkle_tree, + Some(message), ps, - ) + ); } // --------------------------------------------------------------------------- @@ -638,16 +474,28 @@ mod tests { }); } let lambdas = rng.ext_vec(claims.len()); + // Oracle: the dense weight written out naively, one eq entry at a + // time, so it shares no code with the fused build under test. let mut expected = vec![F192::ZERO; stack.len()]; - ring_switch::combine_deferred_into(&rs_outputs, &mut expected[offset..offset + qflock_len]); - let mut target = F192::ZERO; - fold_stacked_point_claims( - &mut expected, - &mut target, - &claims, - &lambdas, - &vec![false; claims.len()], - ); + ring_switch::combine_deferred_chunk(&rs_outputs, 0, &mut expected[offset..offset + qflock_len]); + for (claim, &lambda) in claims.iter().zip(&lambdas) { + let (base, stride_log, point) = match claim { + StackClaim::Point { offset, low_point, .. } => (*offset, 0, low_point.as_slice()), + StackClaim::Strided { + offset, + slot, + stride_log, + point, + .. + } => (*offset + *slot, *stride_log, point.as_slice()), + }; + for j in 0..1usize << point.len() { + let w = point.iter().enumerate().fold(lambda, |w, (i, &p_i)| { + w * if (j >> i) & 1 == 1 { p_i } else { F192::ONE + p_i } + }); + expected[base + (j << stride_log)] += w; + } + } let (actual, message) = basis::build(&stack, lane_block, &claims, &lambdas, &ring, &rs_outputs); assert_eq!(&*actual, expected, "lane_vars={lane_vars}, lanes={lanes}"); let (_, expected_message) = super::super::whir::build_initial_basis(&stack, lane_block, |start, dst| { diff --git a/crates/pcs/src/whir.rs b/crates/pcs/src/whir.rs index 05ad83a05..d4dd60665 100644 --- a/crates/pcs/src/whir.rs +++ b/crates/pcs/src/whir.rs @@ -138,54 +138,10 @@ impl EqTableSlot for std::mem::MaybeUninit { } } -/// Add `seed * eq(point, .)` into `dst` (length `2^point.len()`), with `scratch` -/// holding the table for all but the last coordinate (length `2^(point.len()-1)`). -/// -/// The last doubling level is half the whole table, and it writes straight into -/// `dst`: materializing it in scratch and adding it afterwards would move that -/// half three times (write it, read it back, read-modify-write `dst`) where this -/// moves it once. Same field operations in the same order, so `dst` ends -/// bit-identical to the build-then-add form. -pub(crate) fn add_eq_table_ext_seeded( - point: &[F192], - seed: F192, - scratch: &mut [std::mem::MaybeUninit], - dst: &mut [F192], -) { - let n = point.len(); - assert_eq!(dst.len(), 1usize << n, "dst must have length 2^point.len()"); - let Some((&r, head)) = point.split_last() else { - dst[0] += seed; - return; - }; - let half = 1usize << head.len(); - build_eq_table_ext_seeded(head, seed, &mut scratch[..half]); - // SAFETY: the build above initialized exactly this prefix. - let eq = unsafe { std::slice::from_raw_parts(scratch.as_ptr().cast::(), half) }; - let (lo, hi) = dst.split_at_mut(half); - // Same floor as the seeded build: below it, dispatch costs more than the work. - const PAR_THRESHOLD: usize = 1 << 12; - let expand = |lo: &mut [F192], hi: &mut [F192], eq: &[F192]| { - for ((l, h), &v) in lo.iter_mut().zip(hi.iter_mut()).zip(eq) { - let high = v * r; - *h += high; - *l += v + high; - } - }; - if half < PAR_THRESHOLD { - expand(lo, hi, eq); - } else { - let chunk = parallel::recommended_chunk_size(half); - parallel::chunks_mut2(lo, hi, chunk, |ci, lo_c, hi_c| { - expand(lo_c, hi_c, &eq[ci * chunk..ci * chunk + lo_c.len()]); - }); - } -} - /// In-place seeded core of [`build_eq_table_ext_parallel`]: fills -/// `out[..2^point.len()]` with `seed * eq(point, .)`. Write-only, so it also -/// serves the first claim landing on a range of `stack_open`'s `b_stack`, which -/// then needs no prior zeroing. +/// `out[..2^point.len()]` with `seed * eq(point, .)`. Write-only, which is what +/// lets `stack_open`'s fused basis build each claim's table straight into a +/// reused scratch buffer. /// /// Seeding folds a batching scalar into the table for free: every entry is /// `seed` times a product of point factors, and field multiplication is From 8626945ca77e661d26a4027cfd541d3586566e93 Mon Sep 17 00:00:00 2001 From: Tom Wambsgans Date: Fri, 4 Sep 2026 20:51:20 +0200 Subject: [PATCH 3/3] pcs: drop the dead write-first branch in combine_deferred_chunk Every caller now arrives with a zeroed destination: `basis::build` fills its window before folding into it, and the three test sites allocate zeroed vectors. The first-claim write that let a caller skip pre-zeroing had no remaining beneficiary, and unlike the hoisted version it came from, the branch sat inside the innermost loop. Also records why `QuaternaryLayerState::fold` cannot fold into `fold_and_message`: the last `shrink_eq_low` empties the equality table, so the final round has nothing to weight a message by. Replacing it panics with a zero-length table in the lean_compiler suite. --- crates/lean_vm/src/gkr.rs | 3 +++ crates/pcs/src/ring_switch.rs | 17 ++++++----------- 2 files changed, 9 insertions(+), 11 deletions(-) diff --git a/crates/lean_vm/src/gkr.rs b/crates/lean_vm/src/gkr.rs index 963308ec6..8191d86e2 100644 --- a/crates/lean_vm/src/gkr.rs +++ b/crates/lean_vm/src/gkr.rs @@ -395,6 +395,9 @@ pub fn prove_product_triple(leaves: [ArenaVec; 3], ps: &mut ProverState, s if round + 1 < round_count { messages = trees.each_mut().map(|tree| tree.fold_and_message(challenge, &equality)); } else { + // The last shrink exhausts `equality`, so the final round has no + // table to weight a message by and needs the fold alone. Both + // kernels stay for that reason; `fold` is not dead. for tree in &mut trees { tree.fold(challenge); } diff --git a/crates/pcs/src/ring_switch.rs b/crates/pcs/src/ring_switch.rs index 75ec4a852..12d88d53f 100644 --- a/crates/pcs/src/ring_switch.rs +++ b/crates/pcs/src/ring_switch.rs @@ -445,12 +445,12 @@ pub(crate) fn prove_finish_deferred( } } -/// Fold several deferred claims directly into their final combined dense basis. -/// No per-claim dense vector is allocated or read back, and the first claim -/// **writes** rather than accumulates, so the caller need not pre-zero `out`. -/// +/// Fold several deferred claims into `out[start..]` of their combined dense +/// basis, accumulating, so `out` arrives zeroed. No per-claim dense vector is +/// allocated or read back. `start` is an offset into the basis, which lets a +/// caller cover it one cache-resident window at a time. pub(crate) fn combine_deferred_chunk(outputs: &[DeferredRingSwitchOutput], start: usize, out: &mut [F192]) { - for (claim_idx, claim) in outputs.iter().enumerate() { + for claim in outputs { let block_len = claim.eq_lo.len(); assert!(start + out.len() <= block_len * claim.eq_hi.len()); let mut done = 0; @@ -460,12 +460,7 @@ pub(crate) fn combine_deferred_chunk(outputs: &[DeferredRingSwitchOutput], start let len = (block_len - lo).min(out.len() - done); let e_hi = claim.eq_hi[index / block_len]; for (slot, &e_lo) in out[done..done + len].iter_mut().zip(&claim.eq_lo[lo..lo + len]) { - let value = fold_one_slot_ext(e_lo * e_hi, &claim.table); - if claim_idx == 0 { - *slot = value; - } else { - *slot += value; - } + *slot += fold_one_slot_ext(e_lo * e_hi, &claim.table); } done += len; }