diff --git a/crates/lean_vm/src/gkr.rs b/crates/lean_vm/src/gkr.rs index ee7b841b4..8191d86e2 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,17 @@ 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 { + // 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); + } + } } for tree in &trees { @@ -475,6 +547,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..12d88d53f 100644 --- a/crates/pcs/src/ring_switch.rs +++ b/crates/pcs/src/ring_switch.rs @@ -445,35 +445,26 @@ 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`. -/// -/// 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| { - 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); - } +/// 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 in outputs { + 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]) { + *slot += fold_one_slot_ext(e_lo * e_hi, &claim.table); } + done += len; } - }); + } } /// Split point for the factored eq build: low half sized ~n/2 (min 4, the @@ -669,7 +660,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); @@ -759,7 +750,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); @@ -870,7 +861,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() { @@ -947,7 +938,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 36c5cba72..08fd4535b 100644 --- a/crates/pcs/src/stack_open.rs +++ b/crates/pcs/src/stack_open.rs @@ -64,7 +64,9 @@ 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; // --------------------------------------------------------------------------- // Claim types @@ -166,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. @@ -410,46 +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); - // 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); + .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()); + + // 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, @@ -457,8 +315,9 @@ pub fn open_batch_mixed_whir_stacked( target, &prover_data.codeword, &prover_data.merkle_tree, + Some(message), ps, - ) + ); } // --------------------------------------------------------------------------- @@ -563,6 +422,89 @@ 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()); + // 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_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| { + 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..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 @@ -496,8 +452,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 +843,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 +1026,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 +1258,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 +1354,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 +2801,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