diff --git a/AGENTS.md b/AGENTS.md index 470429067..2985987ab 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -218,7 +218,7 @@ Resident memory (`peak-memory`) therefore mixes what the code needs with what je A buffer every slot of which the caller writes before reading skips the zero-fill: it is `Box::new_uninit_slice` (or a `Vec`'s `spare_capacity_mut`), written through its `MaybeUninit` slots and made a vector by `assume_init` (or `set_len`) only once every slot is written. A kernel typed over `&mut [T]` that writes before it reads takes those slots through `primitives::write_only` (`T: Copy`). -The opening's weight is never stored: the first WHIR pass fills it one chunk at a time, and the first fold builds the lane-folded weight in closed form (a point claim's lanes contract to one seed of its eq table, a ring-switched claim's to one GF(2)-linear map of its in-lane eq table), so neither keeps a full-size weight. +The opening's weight is never stored: the first WHIR pass fills it one chunk at a time, and the first fold builds the lane-folded weight in closed form (a point claim's lanes contract to one seed of its eq table, a ring-switched claim's to one GF(2)-linear map of its in-lane eq table), so neither keeps a full-size weight. Weight generation uses bounded aligned windows, capped by the lane block, to amortize map visits. The ARM EOR3 byte-table backend uses longer windows; scalar and x86 byte-sliced backends keep smaller scratch. The first pass initializes its per-worker lane-window scratch directly on the heap, including the zeros needed for padded lanes, rather than constructing a large stack temporary. ## The thread pool (`parallel`) diff --git a/crates/pcs/src/whir/sumcheck/first_pass.rs b/crates/pcs/src/whir/sumcheck/first_pass.rs index 2c2d59fa1..8988eff25 100644 --- a/crates/pcs/src/whir/sumcheck/first_pass.rs +++ b/crates/pcs/src/whir/sumcheck/first_pass.rs @@ -794,13 +794,10 @@ fn grid_pass_with(f: &[F64], block: usize, w: &dyn InitialWeight }; let n_tasks = lanes.len().div_ceil(group) * per; - let new_scratch = || { - Box::new(Scratch { - raw: [[F192::ZERO; INITIAL_BASIS_CHUNK]; GROUP], - fg: [[0; ROW]; GRID], - bg: [WeightRow::default(); GRID], - }) - }; + // Initialize in place: the lane windows must not become a large stack temporary. + // SAFETY: Scratch contains only arrays of integer words, including F192 and WeightRow, + // and their zero bit patterns are the zero values required for padded lanes. + let new_scratch = || unsafe { Box::::new_zeroed().assume_init() }; let new_acc = || Box::new([ProductRow::default(); GRID]); let acc = if lanes.len() * block < FIRST_PASS_PAR_THRESHOLD { let (mut scratch, mut acc) = (new_scratch(), new_acc()); @@ -852,10 +849,16 @@ impl LaneWeight { /// - vectors 0 and 1: coefficients 0 and 1 of the even words, 128-bit lane `m` for word `2m`; /// - vectors 2 and 3: the same for the odd words, lane `m` for word `2m + 1`; /// - vectors 4 and 5: coefficient 2, of the even words, then of the odd words. -#[derive(Clone, Copy, Default)] +#[derive(Clone, Copy)] #[repr(C, align(64))] pub(super) struct WeightFold([[[u64; ROW]; 6]; INITIAL_BASIS_CHUNK / ROW]); +impl Default for WeightFold { + fn default() -> Self { + Self([[[0; ROW]; 6]; INITIAL_BASIS_CHUNK / ROW]) + } +} + /// `xs` by rows of `ROW`, the last one zero-padded. #[inline(always)] fn padded_rows(xs: &[T]) -> impl Iterator + '_ { diff --git a/crates/pcs/src/whir/sumcheck/mod.rs b/crates/pcs/src/whir/sumcheck/mod.rs index a6a8856cc..55e0a7f3f 100644 --- a/crates/pcs/src/whir/sumcheck/mod.rs +++ b/crates/pcs/src/whir/sumcheck/mod.rs @@ -42,8 +42,13 @@ const ROUND_CHUNK: usize = 2048; /// Words of the initial weight one fill call writes: one chunk, aligned to its size. /// -/// A lane block smaller than this is filled whole. -pub(crate) const INITIAL_BASIS_CHUNK: usize = 256; +/// A lane block smaller than this is filled whole. The ARM EOR3 byte-table maps use +/// longer runs before switching maps; other backends keep their compact scratch. +pub(crate) const INITIAL_BASIS_CHUNK: usize = if cfg!(all(target_arch = "aarch64", target_feature = "sha3")) { + 1024 +} else { + 256 +}; // Sumcheck over `E`, the witness in `K` until the first fold and in `E` after it. //