From 50696ccca317afc6ad5f98f8bb50be22f2e25379 Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Sun, 4 Oct 2026 22:44:51 +0100 Subject: [PATCH 01/18] pcs: cheaper AVX-512 butterfly, register radix-8/4, walked twiddles - mul_lanes_avx512: the reduction's spill is one vpshufb lookup of the top nibble and g(x) takes doublings, so fewer ops share the shift/shuffle pipes with the unpacks (VPCLMULQDQ zmm issues every other cycle on Zen 4). - radix-8 and radix-4 groups keep their rows in registers: each row is loaded and stored once per group, not once per layer. - run_layers puts the leftover one or two layers first (largest blocks) and walks twiddles block to block with one XOR (TwiddleWalk) instead of a subset sum per block; twiddles_radix8 is gone. Byte-identical: fibonacci 2M rate 1 digest 87a8a891167e3b9a, leanxmss 400 744c5cfc65b1157d. cargo test --release -p pcs passes. 1T (LEANVM_NUM_THREADS=1, taskset -c 2), base / this, two runs each: leanxmss 400: Commit 602, 601 -> 580, 572 ms; PCS open Commit L1-L4 380, 381 -> 349, 349 ms; Prove 5500, 5490 -> 5470, 5390 ms. fibonacci 2M: Commit 1160, 1150 -> 1110 ms (other run disturbed, load 11); L1-L4 798, 796 -> 746 ms. Kernel probe, 48 lanes in L1: 0.343 -> 0.308 ns per lane-butterfly. 32T leanxmss 400: Commit 67.6, 66.9 -> 65.6, 71.7 ms (second run had every stage slower, machine noise); proving unchanged within noise. --- crates/pcs/src/ntt/additive_ntt_f64.rs | 371 ++++++++++++++++++------- 1 file changed, 263 insertions(+), 108 deletions(-) diff --git a/crates/pcs/src/ntt/additive_ntt_f64.rs b/crates/pcs/src/ntt/additive_ntt_f64.rs index 311434030..31492a5f1 100644 --- a/crates/pcs/src/ntt/additive_ntt_f64.rs +++ b/crates/pcs/src/ntt/additive_ntt_f64.rs @@ -63,6 +63,40 @@ fn span_get(basis: &[F64], idx: usize) -> F64 { acc } +/// The twiddles of consecutive blocks of one layer, one addition a block. +/// +/// A twiddle is `Σ_j bit_j(block) · row[offset + j]`, so from block `b` to `b + 1` it changes by the +/// entries under the bits that flip: `b`'s trailing ones and the zero above them, a prefix of the row. +struct TwiddleWalk { + /// The current block's twiddle. + t: F64, + /// `steps[k] = Σ_{j <= k} row[offset + j]`, the change past `k` trailing ones. + steps: [F64; 64], +} + +impl TwiddleWalk { + /// The walk over a layer of `2^layer` blocks, from block `first`. + fn new(row: &[F64], offset: usize, layer: usize, first: usize) -> Self { + let mut walk = Self { + t: span_get(&row[offset..offset + layer], first), + steps: [F64::ZERO; 64], + }; + let mut acc = F64::ZERO; + for (step, &entry) in walk.steps.iter_mut().zip(&row[offset..offset + layer]) { + acc += entry; + *step = acc; + } + walk + } + + /// Move from block `block` to `block + 1`. + #[inline] + fn step(&mut self, block: usize) { + // Past the layer's last block the index has `layer` trailing ones, which reads a zero step. + self.t += self.steps[block.trailing_ones() as usize]; + } +} + /// Receives a finished block of codeword rows, as `(first_row, rows)`. pub type RowSink<'a> = dyn Fn(usize, &[F64]) + Sync + 'a; @@ -104,31 +138,6 @@ impl AdditiveNttF64 { span_get(&v[1..], block) } - /// The seven twiddles a radix-8 group needs, breadth-first: layer `layer`, - /// then `layer + 1` (one per half), then `layer + 2` (one per quarter). - /// - /// `span_get` is F_2-linear in the block index, so the six deeper twiddles are - /// the block's own contribution plus a fixed correction per sub-block index: - /// one scan of the three basis rows replaces seven. - pub(crate) fn twiddles_radix8(&self, layer: usize, block: usize) -> [F64; 7] { - let l = self.log_domain_size(); - let (v0, v1, v2) = ( - &self.evals[l - layer - 1], - &self.evals[l - layer - 2], - &self.evals[l - layer - 3], - ); - let (mut t0, mut a, mut c) = (F64::ZERO, F64::ZERO, F64::ZERO); - for j in 0..layer { - if (block >> j) & 1 == 1 { - t0 += v0[1 + j]; - a += v1[2 + j]; - c += v2[3 + j]; - } - } - let (d, e0, e1) = (v1[1], v2[1], v2[2]); - [t0, a, a + d, c, c + e0, c + e1, c + e0 + e1] - } - /// RS-encode a message already stored in the codeword's first replica. /// /// # Overview @@ -456,7 +465,12 @@ impl AdditiveNttF64 { /// /// - The domain has `2^d` rows and splits into `2^o` equal sub-blocks. /// - The buffer is one of them; its index fixes the global block index, and so the twiddle, of each block. - /// - Three layers fuse into one radix-8 sweep where blocks are wide enough, then two, then one. + /// - Three layers fuse into one radix-8 sweep, the one or two layers left over into a first sweep of their own. + /// + /// # Why the leftover goes first + /// + /// - A sweep costs a twiddle computation per block, and the first layers have the fewest, largest blocks. + /// - Last, a lone layer would pair two-row blocks, each with its own twiddle for one butterfly row. #[allow(clippy::too_many_arguments)] fn run_layers( &self, @@ -474,55 +488,69 @@ impl AdditiveNttF64 { let num_blocks_in_buf = 1usize << (layer - outer_log); let block_size = 1usize << (log_d - layer); let block_elems = block_size * num_ntts; - // A block's index in the whole domain, which picks its twiddle. - let global = |block_in_buf: usize| sub_idx * num_blocks_in_buf + block_in_buf; - - if layer + 2 < end_layer && block_size >= 8 { - // Three layers to go and blocks of at least 8 rows: one radix-8 sweep. - let eighth = block_size >> 3; - for block_in_buf in 0..num_blocks_in_buf { - let t = self.twiddles_radix8(layer, global(block_in_buf)); - let start = block_in_buf * block_elems; - butterfly_interleaved_fused_3layer(&mut buf[start..start + block_elems], &t, eighth, num_ntts); - } - layer += 3; - } else if layer + 1 < end_layer && block_size >= 4 { - // Two layers to go: one radix-4 sweep. - let quarter = block_size >> 2; - for block_in_buf in 0..num_blocks_in_buf { - let global_block = global(block_in_buf); - let t_outer = self.twiddle(layer, global_block); - let t_inner_a = self.twiddle(layer + 1, 2 * global_block); - let t_inner_b = self.twiddle(layer + 1, 2 * global_block + 1); - let start = block_in_buf * block_elems; - butterfly_interleaved_fused_2layer( - &mut buf[start..start + block_elems], - t_outer, - t_inner_a, - t_inner_b, - quarter, - num_ntts, + // The global index of the buffer's first block, which picks its twiddle; the others follow it. + let first = sub_idx * num_blocks_in_buf; + // Blocks span at least the layers left, since the end layer is at most `log_d`. + let blocks = buf.chunks_exact_mut(block_elems).zip(first..); + + match (end_layer - layer) % 3 { + 0 => { + // A radix-8 sweep, its seven twiddles breadth-first. + let (mut t0, mut a, mut c) = ( + TwiddleWalk::new(self.row(layer), 1, layer, first), + TwiddleWalk::new(self.row(layer + 1), 2, layer, first), + TwiddleWalk::new(self.row(layer + 2), 3, layer, first), ); + let (d, e0, e1) = (self.row(layer + 1)[1], self.row(layer + 2)[1], self.row(layer + 2)[2]); + for (block, global) in blocks { + let (a_t, c_t) = (a.t, c.t); + let t = [t0.t, a_t, a_t + d, c_t, c_t + e0, c_t + e1, c_t + e0 + e1]; + butterfly_interleaved_fused_3layer(block, &t, block_size >> 3, num_ntts); + t0.step(global); + a.step(global); + c.step(global); + } + layer += 3; } - layer += 2; - } else { - // One layer: a plain butterfly sweep. - let block_size_half = block_size >> 1; - for block_in_buf in 0..num_blocks_in_buf { - let twiddle = self.twiddle(layer, global(block_in_buf)); - let start = block_in_buf * block_elems; - butterfly_interleaved_block( - &mut buf[start..start + block_elems], - twiddle, - block_size_half, - num_ntts, + 2 => { + // A radix-4 sweep: one twiddle for the block, then one per half. + let (mut outer, mut inner) = ( + TwiddleWalk::new(self.row(layer), 1, layer, first), + TwiddleWalk::new(self.row(layer + 1), 2, layer, first), ); + let d = self.row(layer + 1)[1]; + for (block, global) in blocks { + butterfly_interleaved_fused_2layer( + block, + outer.t, + inner.t, + inner.t + d, + block_size >> 2, + num_ntts, + ); + outer.step(global); + inner.step(global); + } + layer += 2; + } + _ => { + // One layer: a plain butterfly sweep. + let mut walk = TwiddleWalk::new(self.row(layer), 1, layer, first); + for (block, global) in blocks { + butterfly_interleaved_block(block, walk.t, block_size >> 1, num_ntts); + walk.step(global); + } + layer += 1; } - layer += 1; } } } + /// The table row a layer's twiddles are subset sums of. + fn row(&self, layer: usize) -> &[F64] { + &self.evals[self.log_domain_size() - layer - 1] + } + /// Recover novel-basis coefficients from evaluations with a scalar inverse NTT. pub fn inverse_transform(&self, data: &mut [F64]) { let log_d = log2_strict_usize(data.len()); @@ -598,6 +626,23 @@ fn fused_rows( /// - The eight rows stay in L1 across all twelve butterflies. /// - The seven twiddles are breadth-first: one for layer L, two for L+1, four for L+2. fn butterfly_interleaved_fused_3layer(block: &mut [F64], t: &[F64; 7], eighth: usize, num_ntts: usize) { + #[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] + { + // SAFETY: the arm is compiled only with AVX-512F; registers only. + let tw = t.map(|x| unsafe { _mm512_set1_epi64(x.0 as i64) }); + let vectors = num_ntts / 8; + let done = 8 * vectors; + fused_rows::<8>(block, eighth, num_ntts, |rows| { + // SAFETY: the target features are enabled at compile time, and each row has `num_ntts` words. + unsafe { radix8_avx512(std::array::from_fn(|i| rows[i].as_mut_ptr()), &tw, vectors) }; + // The lane no vector covers, if any. + if done < num_ntts { + let mut tail = rows.each_mut().map(|row| &mut row[done..]); + radix8_butterflies(&mut tail, t); + } + }); + } + #[cfg(not(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f")))] fused_rows::<8>(block, eighth, num_ntts, |rows| radix8_butterflies(rows, t)); } @@ -798,6 +843,26 @@ fn butterfly_interleaved_fused_2layer( quarter: usize, num_ntts: usize, ) { + #[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] + { + // SAFETY: the arm is compiled only with AVX-512F; registers only. + let tw = [t_outer, t_inner_a, t_inner_b].map(|x| unsafe { _mm512_set1_epi64(x.0 as i64) }); + let vectors = num_ntts / 8; + let done = 8 * vectors; + fused_rows::<4>(block, quarter, num_ntts, |rows| { + // SAFETY: the target features are enabled at compile time, and each row has `num_ntts` words. + unsafe { radix4_avx512(std::array::from_fn(|i| rows[i].as_mut_ptr()), &tw, vectors) }; + // The lane no vector covers, if any. + if done < num_ntts { + let [row_a, row_b, row_c, row_d] = rows.each_mut().map(|row| &mut row[done..]); + butterfly_lanes(row_a, row_c, t_outer); + butterfly_lanes(row_b, row_d, t_outer); + butterfly_lanes(row_a, row_b, t_inner_a); + butterfly_lanes(row_c, row_d, t_inner_b); + } + }); + } + #[cfg(not(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f")))] fused_rows::<4>(block, quarter, num_ntts, |rows| { let [row_a, row_b, row_c, row_d] = rows; // Layer L: rows 2 apart, one twiddle for the block. @@ -832,14 +897,15 @@ fn butterfly_lanes(top: &mut [F64], bot: &mut [F64], twiddle: F64) { #[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] { let vectors = top.len() / 8; - // SAFETY: the target features are enabled at compile time and each - // iteration reads and writes exactly eight elements from both rows. + let done = 8 * vectors; + // SAFETY: the target features are enabled at compile time, and both rows have the lanes covered. unsafe { - for i in 0..vectors { - butterfly_lanes_avx512(top.as_mut_ptr().add(8 * i), bot.as_mut_ptr().add(8 * i), twiddle.0); - } + let tw = _mm512_set1_epi64(twiddle.0 as i64); + columns_avx512([top.as_mut_ptr(), bot.as_mut_ptr()], vectors, |r| { + butterfly_regs_avx512(r, 0, 1, tw); + }); } - for lane in 8 * vectors..top.len() { + for lane in done..top.len() { let v = bot[lane]; let new_u = top[lane] + v * twiddle; top[lane] = new_u; @@ -905,7 +971,7 @@ fn butterfly_lanes(top: &mut [F64], bot: &mut [F64], twiddle: F64) { } } -/// [`butterfly_lanes_avx512`] at half the width, for a machine with VPCLMULQDQ +/// [`mul_lanes_avx512`]'s butterfly at half the width, for a machine with VPCLMULQDQ /// but no AVX-512. Without it the base encode's innermost loop is scalar there, /// which costs it about half again as much. /// @@ -1022,7 +1088,7 @@ unsafe fn butterfly_lanes_neon_8(top: *mut F64, bot: *mut F64, twiddle: u64) { } } -/// Eight F64 butterflies with a shared twiddle, one per 64-bit lane of an AVX-512 register. +/// The products `v * t` of eight F64 lanes by one broadcast twiddle, reduced. /// /// # Algorithm /// @@ -1034,40 +1100,38 @@ unsafe fn butterfly_lanes_neon_8(top: *mut F64, bot: *mut F64, twiddle: u64) { /// /// reduce: p = lo + hi * x^64, x^64 = x^4 + x^3 + x + 1 /// p mod f = lo ^ g(hi ^ spill), g(y) = y ^ y<<1 ^ y<<3 ^ y<<4 -/// spill = hi>>63 ^ hi>>61 ^ hi>>60, the bits g pushes past x^63 +/// spill = n ^ n>>1 ^ n>>3 with n = hi>>60, the bits g pushes past x^63 /// ``` /// /// - The unpacks stay inside 128-bit lanes, so no shuffle crosses lanes. -/// - The reduction takes shifts and three-way XORs, `vpternlogq 0x96`, instead of two more carry-less multiplies. -/// - Zen 5 has one carry-less multiplier per core, so those two would dominate the kernel. -/// -/// # Safety -/// -/// - Requires VPCLMULQDQ and AVX-512F. -/// - Each pointer must address eight readable and writable words. +/// - The reduction takes no further carry-less multiply: the multiplier issues one every other cycle. +/// - Shifts and shuffles share two pipes with the unpacks, adds have four, so the left shifts of `g` are doublings. +/// - The spill is one table lookup of the top nibble where AVX-512BW has `vpshufb`. #[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] -#[inline] -#[target_feature(enable = "vpclmulqdq", enable = "avx512f", enable = "avx2")] -unsafe fn butterfly_lanes_avx512(top: *mut F64, bot: *mut F64, twiddle: u64) { - // SAFETY: - // - The caller supplies two valid eight-word rows. - // - This function's target features cover every intrinsic below. +#[inline(always)] +fn mul_lanes_avx512(v: __m512i, tw: __m512i) -> __m512i { + const XOR3: i32 = 0x96; + // SAFETY: the arm is compiled only with VPCLMULQDQ and AVX-512F (and BW where it uses it); registers only. unsafe { - // Load both rows and broadcast the twiddle to every lane. - let u = _mm512_loadu_si512(top.cast()); - let v = _mm512_loadu_si512(bot.cast()); - let tw = _mm512_set1_epi64(twiddle as i64); - - // Products v * t: even lanes, then odd lanes, one 128-bit product per 128-bit lane. let even = _mm512_clmulepi64_epi128::<0x00>(v, tw); let odd = _mm512_clmulepi64_epi128::<0x11>(v, tw); - // Back to lane order: qword i of lo / hi is the low / high half of lane i's product. let lo = _mm512_unpacklo_epi64(even, odd); let hi = _mm512_unpackhi_epi64(even, odd); - - // Reduce modulo x^64 + x^4 + x^3 + x + 1 with shifts and three-way XORs. - const XOR3: i32 = 0x96; - // The bits of hi * (x^4 + x^3 + x + 1) that land past x^63. + #[cfg(target_feature = "avx512bw")] + let spill = { + const SPILL: [u8; 16] = { + let mut table = [0u8; 16]; + let mut n = 0; + while n < 16 { + table[n] = (n ^ (n >> 1) ^ (n >> 3)) as u8; + n += 1; + } + table + }; + let table = _mm512_broadcast_i32x4(_mm_loadu_si128(SPILL.as_ptr().cast())); + _mm512_shuffle_epi8(table, _mm512_srli_epi64::<60>(hi)) + }; + #[cfg(not(target_feature = "avx512bw"))] let spill = _mm512_ternarylogic_epi64::( _mm512_srli_epi64::<63>(hi), _mm512_srli_epi64::<61>(hi), @@ -1075,14 +1139,105 @@ unsafe fn butterfly_lanes_avx512(top: *mut F64, bot: *mut F64, twiddle: u64) { ); // Both hi and spill are multiplied by the same constant, so fold them first. let x = _mm512_xor_si512(hi, spill); - // g(x) = x ^ x<<1 ^ x<<3 ^ x<<4, split across two three-way XORs. - let fx = _mm512_ternarylogic_epi64::(x, _mm512_slli_epi64::<1>(x), _mm512_slli_epi64::<3>(x)); - // Butterfly top: u' = u + v * t, with the product's lo and g(x) folded in one step. - let new_u = _mm512_ternarylogic_epi64::(u, lo, _mm512_xor_si512(fx, _mm512_slli_epi64::<4>(x))); - // Butterfly bottom: v' = v + u'. - let new_v = _mm512_xor_si512(v, new_u); - _mm512_storeu_si512(top.cast(), new_u); - _mm512_storeu_si512(bot.cast(), new_v); + let x2 = _mm512_add_epi64(x, x); + let x8 = { + let x4 = _mm512_add_epi64(x2, x2); + _mm512_add_epi64(x4, x4) + }; + let x16 = _mm512_add_epi64(x8, x8); + _mm512_ternarylogic_epi64::(_mm512_ternarylogic_epi64::(lo, x, x2), x8, x16) + } +} + +/// One butterfly of eight lanes held in registers: `u' = u + v * t`, `v' = v + u'`. +#[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] +#[inline(always)] +fn butterfly_regs_avx512(rows: &mut [__m512i], top: usize, bot: usize, tw: __m512i) { + // SAFETY: the arm is compiled only with AVX-512F; registers only. + unsafe { + let new_u = _mm512_xor_si512(rows[top], mul_lanes_avx512(rows[bot], tw)); + rows[bot] = _mm512_xor_si512(rows[bot], new_u); + rows[top] = new_u; + } +} + +/// Run `body` on columns of eight lanes of `N` rows held in registers, each row loaded and stored once. +/// +/// The lanes past the last whole vector are the caller's. +/// +/// # Safety +/// +/// - Requires AVX-512F. +/// - Each row pointer must address `8 * vectors` readable and writable words. +#[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] +#[inline(always)] +unsafe fn columns_avx512(rows: [*mut F64; N], vectors: usize, body: impl Fn(&mut [__m512i; N])) { + // SAFETY: every access is inside the words the caller guarantees. + unsafe { + for c in 0..vectors { + let mut r: [__m512i; N] = std::array::from_fn(|i| _mm512_loadu_si512(rows[i].add(8 * c).cast())); + body(&mut r); + for (i, row) in r.iter().enumerate() { + _mm512_storeu_si512(rows[i].add(8 * c).cast(), *row); + } + } + } +} + +/// A radix-8 row group with the eight rows in registers. +/// +/// Each row is loaded and stored once for its twelve butterflies, not three times. +/// +/// # Safety +/// +/// - Requires VPCLMULQDQ and AVX-512F. +/// - Each row pointer must address `8 * vectors` readable and writable words. +#[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] +#[inline] +#[target_feature(enable = "vpclmulqdq", enable = "avx512f", enable = "avx2")] +unsafe fn radix8_avx512(rows: [*mut F64; 8], t: &[__m512i; 7], vectors: usize) { + // SAFETY: forwarded from the caller. + unsafe { + columns_avx512(rows, vectors, |r| { + // Layer L: rows 4 apart, one twiddle for the whole block. + butterfly_regs_avx512(r, 0, 4, t[0]); + butterfly_regs_avx512(r, 1, 5, t[0]); + butterfly_regs_avx512(r, 2, 6, t[0]); + butterfly_regs_avx512(r, 3, 7, t[0]); + // Layer L+1: rows 2 apart, one twiddle per half. + butterfly_regs_avx512(r, 0, 2, t[1]); + butterfly_regs_avx512(r, 1, 3, t[1]); + butterfly_regs_avx512(r, 4, 6, t[2]); + butterfly_regs_avx512(r, 5, 7, t[2]); + // Layer L+2: adjacent rows, one twiddle per quarter. + butterfly_regs_avx512(r, 0, 1, t[3]); + butterfly_regs_avx512(r, 2, 3, t[4]); + butterfly_regs_avx512(r, 4, 5, t[5]); + butterfly_regs_avx512(r, 6, 7, t[6]); + }); + } +} + +/// A radix-4 row group with the four rows in registers. +/// +/// # Safety +/// +/// - Requires VPCLMULQDQ and AVX-512F. +/// - Each row pointer must address `8 * vectors` readable and writable words. +#[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] +#[inline] +#[target_feature(enable = "vpclmulqdq", enable = "avx512f", enable = "avx2")] +unsafe fn radix4_avx512(rows: [*mut F64; 4], t: &[__m512i; 3], vectors: usize) { + // SAFETY: forwarded from the caller. + unsafe { + columns_avx512(rows, vectors, |r| { + // Layer L: rows 2 apart, one twiddle for the block. + butterfly_regs_avx512(r, 0, 2, t[0]); + butterfly_regs_avx512(r, 1, 3, t[0]); + // Layer L+1: adjacent rows, one twiddle per half. + butterfly_regs_avx512(r, 0, 1, t[1]); + butterfly_regs_avx512(r, 2, 3, t[2]); + }); } } From c8e282f78b23c8ce443b7f3438465026689412c7 Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Sun, 4 Oct 2026 22:48:20 +0100 Subject: [PATCH 02/18] pcs: one worker's gathered groups take adjacent residues; commit reads the lane-major stack in place - With one worker (gated on the planned worker count; 32 threads keep the old plan) a gathered group takes up to 64 adjacent residues: they share every twiddle, so the group is one transform of R*n lanes, its copies in and out are R*n-word runs (whole lines, so streaming stores stop splitting lines) and the 41- or 35-lane scalar tail disappears. LONE_GATHER_WORDS bounds the scratch (32 MiB). - encode_lane_major_with: commit no longer transposes the stack into the codeword. When the first gathered pass takes >= 8 residues, it reads each lane's 8-word lines straight from the lane-major message and turns 8x8 blocks around in registers (gather_lanes, AVX-512). Otherwise (32T, small domains, other targets) it transposes first as before. - transform_for takes the worker count, so tests pin the one-worker plan. Byte-identical (same digests). cargo test --release -p pcs passes; tests cover the lone plan for in-place, no-message and lane-major encodes. Measured together with the next commit (see there). --- crates/pcs/src/ntt/additive_ntt_f64.rs | 385 ++++++++++++++++++++----- crates/pcs/src/whir/commit.rs | 7 +- 2 files changed, 311 insertions(+), 81 deletions(-) diff --git a/crates/pcs/src/ntt/additive_ntt_f64.rs b/crates/pcs/src/ntt/additive_ntt_f64.rs index 31492a5f1..4da83aa8d 100644 --- a/crates/pcs/src/ntt/additive_ntt_f64.rs +++ b/crates/pcs/src/ntt/additive_ntt_f64.rs @@ -100,6 +100,15 @@ impl TwiddleWalk { /// Receives a finished block of codeword rows, as `(first_row, rows)`. pub type RowSink<'a> = dyn Fn(usize, &[F64]) + Sync + 'a; +/// Where an encode's message is read from, while the codeword does not hold it yet. +#[derive(Clone, Copy)] +enum Message { + /// Row-major, `num_ntts` words a row; it may be the codeword's own first replica. + Rows(SendPtr), + /// Lane-major, in its own buffer: block `b` of the rows is codeword lane `num_ntts - 1 - b`. + Lanes(SendPtr), +} + /// Additive NTT over F_{2^64} with the standard polynomial-basis subspace /// `{1, x, x², …}`: the F_2-subspace is `{0, 1, …, 2^ℓ−1}` under the natural /// integer encoding, exactly as in the extension-field version (whose domain already @@ -153,14 +162,12 @@ impl AdditiveNttF64 { /// - The first pass therefore reads its rows from the message itself. /// - No pass fills the replicas only to read them back. /// - /// # Why the transpose stays separate + /// # Lane-major callers /// - /// - A lane-major caller transposes its message into place first. - /// - The transpose wants one long contiguous run per lane. - /// - The first pass wants hundreds of scattered rows at once. + /// - A lane-major message goes to [`Self::encode_lane_major_with`] instead, which reads it in place. pub fn encode_interleaved_in_place(&self, data: &mut [F64], num_ntts: usize, log_inv_rate: usize) { // The message is the buffer's own first replica. - let msg = SendPtr(data.as_mut_ptr()); + let msg = Message::Rows(SendPtr(data.as_mut_ptr())); self.transform(data, num_ntts, log_inv_rate, Some(msg), None); } @@ -177,7 +184,7 @@ impl AdditiveNttF64 { log_inv_rate: usize, on_rows: &RowSink<'_>, ) { - let msg = SendPtr(data.as_mut_ptr()); + let msg = Message::Rows(SendPtr(data.as_mut_ptr())); self.transform(data, num_ntts, log_inv_rate, Some(msg), Some(on_rows)); } @@ -204,7 +211,47 @@ impl AdditiveNttF64 { "the codeword is 2^log_inv_rate messages" ); // Read-only from here on: the pointer only feeds the first pass's reads. - let msg = SendPtr(msg.as_ptr().cast_mut()); + let msg = Message::Rows(SendPtr(msg.as_ptr().cast_mut())); + self.transform(data, num_ntts, log_inv_rate, Some(msg), Some(on_rows)); + } + + /// RS-encode a lane-major message, handing `on_rows` every finished block of rows. + /// + /// # Layout + /// + /// ```text + /// message, lane-major: [ block 0 | block 1 | ... | block n-1 ] each 2^(d - r) words + /// codeword lane t encodes message block n - 1 - t + /// ``` + /// + /// - The lane order is [`transpose_lane_major`]'s: the result equals that transpose, then the in-place encode. + /// - Every codeword word is written before it is read, so the codeword may start uninitialized. + /// - The blocks are as for the in-place encode. + /// + /// # Why the message is read in place + /// + /// - The first pass reads its rows from the message, so a row-major copy of it would be written only to be read once. + /// - Where the plan's first pass takes eight neighbouring rows at once, each lane gives it whole cache lines. + /// - Otherwise the message is transposed into the first replica, then encoded in place. + /// + /// # Panics + /// + /// Panics unless the codeword is exactly `2^r` messages long. + pub fn encode_lane_major_with( + &self, + data: &mut [F64], + msg: &[F64], + num_ntts: usize, + log_inv_rate: usize, + on_rows: &RowSink<'_>, + ) { + assert_eq!( + msg.len() << log_inv_rate, + data.len(), + "the codeword is 2^log_inv_rate messages" + ); + // Read-only from here on: the pointer only feeds the first pass's reads, or the transpose. + let msg = Message::Lanes(SendPtr(msg.as_ptr().cast_mut())); self.transform(data, num_ntts, log_inv_rate, Some(msg), Some(on_rows)); } @@ -243,13 +290,33 @@ impl AdditiveNttF64 { /// /// - Each deep sub-block is final, and hands its rows to `on_rows` while they are in L2. /// - Without a split deep pass, the rows go over in parallel blocks at the end. + /// + /// # One worker + /// + /// - A lone worker has its core's whole L3, so a gathered group takes up to `2^MAX_LOG_RESIDUES` residues. + /// - Their rows are adjacent, so each copy in or out is one long run, and a row group has that many times + /// the lanes. + /// - Several workers keep one residue a group: their scratch then stays in their share of L2. fn transform( &self, data: &mut [F64], num_ntts: usize, start: usize, - msg: Option>, + msg: Option, + on_rows: Option<&RowSink<'_>>, + ) { + self.transform_for(data, num_ntts, start, msg, on_rows, parallel::num_threads()); + } + + /// [`Self::transform`], planned for `workers` workers. + fn transform_for( + &self, + data: &mut [F64], + num_ntts: usize, + start: usize, + msg: Option, on_rows: Option<&RowSink<'_>>, + workers: usize, ) { // The buffer is 2^log_d rows of `num_ntts` words. assert!(num_ntts > 0); @@ -271,8 +338,7 @@ impl AdditiveNttF64 { // From 2^12 rows on, cut at least 2^LOG_SUBS_PER_WORKER deep sub-blocks per worker. let par_log = if log_d >= PARALLEL_FLOOR_LOG_D { - (log2_strict_usize(parallel::num_threads().next_power_of_two()) + LOG_SUBS_PER_WORKER) - .min(log_d - MIN_SUB_LOG) + (log2_strict_usize(workers.next_power_of_two()) + LOG_SUBS_PER_WORKER).min(log_d - MIN_SUB_LOG) } else { 0 }; @@ -294,17 +360,40 @@ impl AdditiveNttF64 { }; let deep_start = log_d.saturating_sub(deep).max(par_log).max(start); + // Residues a gathered group takes, as a log: several only for a lone worker, within its L3 budget. + let residues = |layer: usize, g: usize| { + if workers == 1 { + let room = (LONE_GATHER_WORDS / (num_ntts << g)).max(1).ilog2() as usize; + room.min(MAX_LOG_RESIDUES).min(log_d - layer - g) + } else { + 0 + } + }; + // A buffer this large is evicted before the next pass reads it back. let stream = data.len() >= STREAM_MIN_WORDS; + // A lane-major message is read in place by a first gathered pass that takes whole cache lines of it, + // eight residues of each lane at a time. Any other plan reads it transposed into the first replica. + let mut msg = msg; + if let Some(Message::Lanes(m)) = msg { + let first_g = (deep_start - start).min(fit2); + if !(LANE_GATHER && start < deep_start && residues(start, first_g) >= 3) { + let msg_len = data.len() >> start; + // SAFETY: the lane-major message is one replica of words, disjoint from the codeword. + let lanes = unsafe { std::slice::from_raw_parts(m.0.cast_const(), msg_len) }; + transpose_lane_major(&mut data[..msg_len], lanes, num_ntts, log_d - start); + msg = Some(Message::Rows(SendPtr(data.as_mut_ptr()))); + } + } + // Phase 1: gathered passes, each at most one L2 group of layers. // // Only the first one reads the message. - let mut msg = msg; let mut layer = start; while layer < deep_start { let g = (deep_start - layer).min(fit2); - self.gathered_pass(data, log_d, num_ntts, layer, g, msg.take(), stream); + self.gathered_pass(data, log_d, num_ntts, layer, g, residues(layer, g), msg.take(), stream); layer += g; } @@ -312,7 +401,9 @@ impl AdditiveNttF64 { // // message is block 0 of the buffer -> a task could overwrite it while others still copy it // no layer left to run -> no deep task would write the buffer at all - if let Some(m) = msg.take_if(|m| std::ptr::eq(m.0, data.as_mut_ptr()) || deep_start == log_d) { + if let Some(Message::Rows(m)) = + msg.take_if(|m| matches!(m, Message::Rows(r) if std::ptr::eq(r.0, data.as_mut_ptr())) || deep_start == log_d) + { replicate(data, m, data.len() >> start); } @@ -329,7 +420,7 @@ impl AdditiveNttF64 { // A separate message: build the sub-block in scratch, then write it out once. // // Streamed out whole, the codeword is written without ever being read. - Some(m) => with_scratch(sub_len, |scratch| { + Some(Message::Rows(m)) => with_scratch(sub_len, |scratch| { // A sub-block sits at the same offset in every replica. // // sub-block at offset off of its replica <- message words [off, off + sub_len) @@ -356,6 +447,8 @@ impl AdditiveNttF64 { f(sub_idx << log_sub, sub); } } + // A lane-major message is read by a gathered pass, or transposed into place before any pass. + Some(Message::Lanes(_)) => unreachable!("a lane-major message reaches no deep pass"), }); } @@ -389,10 +482,16 @@ impl AdditiveNttF64 { /// - The residue is below `step`, so that shift drops it. /// - The scratch transform therefore takes the same twiddles as the full one. /// + /// # Several residues a group + /// + /// - A group may take `2^log_r` adjacent residues: their rows sit side by side, and share every twiddle. + /// - So the group is one transform with `2^log_r` times the lanes, and each of its rows one contiguous run. + /// /// # With a message /// /// - Every block's rows come from the message. - /// - One task takes one residue across all blocks, block 0 last. + /// - One task takes one residue group across all blocks, block 0 last. + /// - A lane-major message gives each lane's word of eight adjacent residues as one cache line. #[allow(clippy::too_many_arguments)] fn gathered_pass( &self, @@ -401,45 +500,58 @@ impl AdditiveNttF64 { num_ntts: usize, layer: usize, g: usize, - msg: Option>, + log_r: usize, + msg: Option, stream: bool, ) { - // A group is 2^g rows, `step` rows apart. + // A group is 2^g rows, `step` rows apart, each `2^log_r` adjacent rows wide. let log_step = log_d - layer - g; let (rows, step) = (1usize << g, 1usize << log_step); - // With a message, a task is one residue across every block. - // Without one, a task is one (block, residue) pair. - let n_tasks = if msg.is_some() { step } else { step << layer }; + assert!(log_r <= log_step, "a group's residues are inside one step"); + let wide = num_ntts << log_r; + // A lane-major message is read eight residues of a lane at a time. + assert!( + !matches!(msg, Some(Message::Lanes(_))) || log_r >= 3, + "a lane-major message is read whole lines at a time" + ); + // With a message, a task is one residue group across every block. + // Without one, a task is one (block, residue group) pair. + let n_tasks = (if msg.is_some() { step } else { step << layer }) >> log_r; let base = SendPtr(data.as_mut_ptr()); parallel::for_each_chunk(n_tasks, |lo, hi| { - // One L2-resident scratch of 2^g rows serves every group of the task. - with_scratch(rows * num_ntts, |scratch| { + // One scratch of 2^g wide rows serves every group of the task. + with_scratch(rows * wide, |scratch| { // One fence at the end of the task covers all its streaming stores. let stream = stream.then(Stream::new); let mut group = |block: usize, r: usize| { - // Row `i` of the group, as a word offset inside its block. - let row = |i: usize| (r + (i << log_step)) * num_ntts; + // Row `i` of the group, as a row index and a word offset inside its block. + let row_idx = |i: usize| r + (i << log_step); + let row = |i: usize| row_idx(i) * num_ntts; let block_off = (block << (log_d - layer)) * num_ntts; // Gather: the scattered rows become one contiguous 2^g-row buffer. - for (i, dst) in scratch.chunks_exact_mut(num_ntts).enumerate() { + for (i, dst) in scratch.chunks_exact_mut(wide).enumerate() { // SAFETY: // - The message and the codeword both cover every row addressed here. // - Tasks own disjoint residues, so no other task writes these rows. // - Block 0 is written last, after every read of the message in it. - let src = unsafe { + unsafe { match msg { - Some(m) => std::slice::from_raw_parts(m.add(row(i)), num_ntts), - None => base.slice(block_off + row(i), num_ntts), + Some(Message::Rows(m)) => { + dst.copy_from_slice(std::slice::from_raw_parts(m.add(row(i)), wide)); + } + Some(Message::Lanes(m)) => { + gather_lanes(dst, m, 1 << (log_d - layer), num_ntts, row_idx(i)); + } + None => dst.copy_from_slice(base.slice(block_off + row(i), wide)), } - }; - dst.copy_from_slice(src); + } } // Transform: a (layer + g)-layer domain whose sub-block index is the global block. - self.run_layers(scratch, layer + g, num_ntts, layer, layer + g, layer, block); + self.run_layers(scratch, layer + g, wide, layer, layer + g, layer, block); // Scatter: every row returns to its place. - for (i, src) in scratch.chunks_exact(num_ntts).enumerate() { + for (i, src) in scratch.chunks_exact(wide).enumerate() { // SAFETY: this group alone owns these rows of the codeword. - let dst = unsafe { base.slice(block_off + row(i), num_ntts) }; + let dst = unsafe { base.slice(block_off + row(i), wide) }; match &stream { Some(s) => s.copy(dst, src), None => dst.copy_from_slice(src), @@ -447,6 +559,8 @@ impl AdditiveNttF64 { } }; for t in lo..hi { + // The group's first residue, in the task's step-indexed order. + let t = t << log_r; if msg.is_some() { // Block 0 may be the message itself, so it is transformed last. for block in (0..1usize << layer).rev() { @@ -734,6 +848,92 @@ pub fn transpose_lane_major(out: &mut [F64], msg: &[F64], n_lanes: usize, log_ro }); } +/// Whether a gathered pass reads a lane-major message in place; elsewhere it is transposed first. +/// +/// Only the AVX-512 build turns each eight lanes of eight residues around in registers. +const LANE_GATHER: bool = cfg!(all(target_arch = "x86_64", target_feature = "avx512f")); + +/// Gather `dst.len() / n` adjacent rows of a lane-major message, from row `row`, into row-major order. +/// +/// ```text +/// dst[j * n + (n - 1 - l)] = msg[l * msg_rows + row + j] lanes reversed, as in transpose_lane_major +/// ``` +/// +/// - Eight rows of one lane are one cache line, read whole, when `row` is a multiple of eight. +/// - Eight lanes of those eight rows are turned around in registers. +/// +/// # Safety +/// +/// - The message holds `n` lanes of `msg_rows` words each. +/// - `row + dst.len() / n <= msg_rows`, and `dst.len() / n` is a multiple of eight. +unsafe fn gather_lanes(dst: &mut [F64], msg: SendPtr, msg_rows: usize, n: usize, row: usize) { + let residues = dst.len() / n; + debug_assert!(residues.is_multiple_of(8) && row + residues <= msg_rows); + // Message lanes `0 .. whole` go eight at a time; the rest, the lowest codeword lanes, one at a time. + let whole = if LANE_GATHER { n - n % 8 } else { 0 }; + #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] + for q in (0..residues).step_by(8) { + for l0 in (0..whole).step_by(8) { + // SAFETY: line `k` is lane `l0 + 7 - k`, rows `row + q ..+ 8`, inside the message by the caller. + let lines: [__m512i; 8] = std::array::from_fn(|k| unsafe { + _mm512_loadu_si512(msg.add((l0 + 7 - k) * msg_rows + row + q).cast_const().cast()) + }); + // Column `j` is residue `q + j` across lanes `l0 + 7` down to `l0`: codeword lanes `n - 8 - l0 ..`. + for (j, col) in transpose_8x8_avx512(lines).into_iter().enumerate() { + let at = (q + j) * n + n - 8 - l0; + // SAFETY: eight words of residue `q + j`'s row, inside `dst`. + unsafe { _mm512_storeu_si512(dst.as_mut_ptr().add(at).cast(), col) }; + } + } + } + for l in whole..n { + for j in 0..residues { + // SAFETY: inside the message by the caller. + dst[j * n + n - 1 - l] = unsafe { *msg.add(l * msg_rows + row + j) }; + } + } +} + +/// Transpose eight rows of eight words: word `k` of output `j` is word `j` of input `k`. +#[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] +#[inline(always)] +fn transpose_8x8_avx512(r: [__m512i; 8]) -> [__m512i; 8] { + // SAFETY: the arm is compiled only with AVX-512F; registers only. + unsafe { + // Pairs of rows, 128-bit lane c holding words 2c and 2c + 1 of both. + let t: [__m512i; 8] = std::array::from_fn(|i| { + let (a, b) = (r[i & !1], r[i | 1]); + if i % 2 == 0 { + _mm512_unpacklo_epi64(a, b) + } else { + _mm512_unpackhi_epi64(a, b) + } + }); + // Four rows: u[c] holds words c and c + 4 of rows 0..4 (u[0..4]) or rows 4..8 (u[4..8]). + let u = [ + _mm512_shuffle_i64x2::<0x88>(t[0], t[2]), + _mm512_shuffle_i64x2::<0x88>(t[1], t[3]), + _mm512_shuffle_i64x2::<0xDD>(t[0], t[2]), + _mm512_shuffle_i64x2::<0xDD>(t[1], t[3]), + _mm512_shuffle_i64x2::<0x88>(t[4], t[6]), + _mm512_shuffle_i64x2::<0x88>(t[5], t[7]), + _mm512_shuffle_i64x2::<0xDD>(t[4], t[6]), + _mm512_shuffle_i64x2::<0xDD>(t[5], t[7]), + ]; + // Eight rows: words c from u[c] and u[4 + c], words c + 4 from their other halves. + [ + _mm512_shuffle_i64x2::<0x88>(u[0], u[4]), + _mm512_shuffle_i64x2::<0x88>(u[1], u[5]), + _mm512_shuffle_i64x2::<0x88>(u[2], u[6]), + _mm512_shuffle_i64x2::<0x88>(u[3], u[7]), + _mm512_shuffle_i64x2::<0xDD>(u[0], u[4]), + _mm512_shuffle_i64x2::<0xDD>(u[1], u[5]), + _mm512_shuffle_i64x2::<0xDD>(u[2], u[6]), + _mm512_shuffle_i64x2::<0xDD>(u[3], u[7]), + ] + } +} + /// Words one L2-resident unit of work may span: a gathered row group or a deep sub-block. /// /// # Why this value @@ -771,6 +971,21 @@ const LOG_SUBS_PER_WORKER: usize = 2; /// - That costs extra L3 traffic, but saves a whole sweep of DRAM. const L3_WORDS: usize = 1 << 18; +/// Words a gathered task may span when the pool has one worker. +/// +/// # Why this value +/// +/// - 2^22 words is 32 MiB, a whole L3 without stacked cache: a lone worker shares its L3 with no one. +/// - A gathered row group's layers cost no more from L3 than from L2, the butterflies being the bound. +/// - So the group can take several residues, and with them longer runs in and out of memory. +const LONE_GATHER_WORDS: usize = 1 << 22; + +/// Most residues one gathered group takes, as a log. +/// +/// - Eight residues make a row-major group's runs whole cache lines, for any lane count. +/// - Sixty-four make a lane-major message's reads eight lines of each lane. +const MAX_LOG_RESIDUES: usize = 6; + /// Buffers of at least this many words get streaming stores for data a pass does not read back. /// /// # Why this value @@ -1348,6 +1563,7 @@ mod tests { // (12, 2048, 1) two gathered passes, with streaming stores // // A non-zero start is the commit path, which enters at the rate layer. + // A lone worker's plan takes several residues per gathered group. let mut rng = Rng::new(0xC0FFEE); for (log_d, lanes, start_layer) in [ (7usize, 3usize, 0usize), @@ -1364,11 +1580,15 @@ mod tests { // Reference: one butterfly at a time. let mut want = original.clone(); forward_scalar_from_layer(&ntt, &mut want, lanes, start_layer); - // Under test: the pass-planning driver. - let mut got = original; - ntt.transform(&mut got, lanes, start_layer, None, None); - - assert_eq!(got, want, "log_d={log_d}, lanes={lanes}, start_layer={start_layer}"); + // Under test: the pass-planning driver, for the pool and for one worker. + for workers in [parallel::num_threads(), 1] { + let mut got = original.clone(); + ntt.transform_for(&mut got, lanes, start_layer, None, None, workers); + assert_eq!( + got, want, + "log_d={log_d}, lanes={lanes}, start_layer={start_layer}, workers={workers}" + ); + } } } @@ -1408,35 +1628,38 @@ mod tests { forward_scalar_from_layer(&ntt, &mut want, lanes, log_inv_rate); // Under test: only the first replica holds the message, the rest is zero. - let mut got = vec![F64::ZERO; msg_len << log_inv_rate]; - got[..msg_len].copy_from_slice(&msg); - let blocks = Mutex::new(Vec::new()); - ntt.encode_interleaved_in_place_with(&mut got, lanes, log_inv_rate, &|row, rows| { - blocks.lock().unwrap().push((row, rows.to_vec())); - }); - assert_eq!(got, want, "log_d={log_d}, lanes={lanes}, rate={log_inv_rate}"); - - let mut blocks = blocks.into_inner().unwrap(); - blocks.sort_by_key(|b| b.0); - let mut next = 0; - for (row, rows) in blocks { - assert_eq!(row, next, "blocks tile the rows, log_d={log_d}"); - assert_eq!( - rows[..], - want[row * lanes..][..rows.len()], - "a handed-over block is final" - ); - next += rows.len() / lanes; + for workers in [parallel::num_threads(), 1] { + let mut got = vec![F64::ZERO; msg_len << log_inv_rate]; + got[..msg_len].copy_from_slice(&msg); + let blocks = Mutex::new(Vec::new()); + let first = Message::Rows(SendPtr(got.as_mut_ptr())); + let sink = |row: usize, rows: &[F64]| blocks.lock().unwrap().push((row, rows.to_vec())); + ntt.transform_for(&mut got, lanes, log_inv_rate, Some(first), Some(&sink), workers); + assert_eq!(got, want, "log_d={log_d}, lanes={lanes}, rate={log_inv_rate}, workers={workers}"); + + let mut blocks = blocks.into_inner().unwrap(); + blocks.sort_by_key(|b| b.0); + let mut next = 0; + for (row, rows) in blocks { + assert_eq!(row, next, "blocks tile the rows, log_d={log_d}"); + assert_eq!( + rows[..], + want[row * lanes..][..rows.len()], + "a handed-over block is final" + ); + next += rows.len() / lanes; + } + assert_eq!(next, 1 << log_d, "every row handed over, log_d={log_d}"); } - assert_eq!(next, 1 << log_d, "every row handed over, log_d={log_d}"); } } /// Every lane of the codeword must be exactly the single-lane RS codeword of /// that lane's contiguous message block, which is what makes a commitment over /// `n_lanes` lanes equal to the `2^log_batch_size`-lane one with a zero tail. - /// The shapes cover the transposing fused first pass, its fallback, and lane - /// counts that are not powers of two (the padding-free commit's whole point). + /// The shapes cover the transposing fallback, the lone worker's first pass reading + /// the lanes in place (the last three, with and without lanes past a whole eight), + /// and lane counts that are not powers of two (the padding-free commit's whole point). #[test] fn lane_major_msg_encode_matches_per_lane_reference() { let mut rng = Rng::new(0x1A2E); @@ -1447,29 +1670,39 @@ mod tests { (9, 1, 5), (12, 2, 37), (14, 1, 64), + (13, 1, 41), + (13, 1, 64), + (16, 1, 5), ] { let log_d = log_rows + log_inv_rate; let ntt = AdditiveNttF64::standard(log_d); let rows = 1usize << log_rows; let msg: Vec = (0..rows * n_lanes).map(|_| F64(rng.next_u64())).collect(); - let mut got = vec![F64::ZERO; msg.len() << log_inv_rate]; - transpose_lane_major(&mut got[..msg.len()], &msg, n_lanes, log_rows); - ntt.encode_interleaved_in_place(&mut got, n_lanes, log_inv_rate); - let block_len = 1usize << log_d; - for lane in 0..n_lanes { - // Lane `lane` encodes message block `n_lanes - 1 - lane`. - let block = n_lanes - 1 - lane; - let mut want = vec![F64::ZERO; block_len]; - replicate_rows(&mut want, &msg[block * rows..(block + 1) * rows]); - ntt.transform(&mut want, 1, log_inv_rate, None, None); - for pos in 0..block_len { - assert_eq!( - got[pos * n_lanes + lane], - want[pos], - "lane {lane} pos {pos} at log_rows={log_rows}, rate={log_inv_rate}, n_lanes={n_lanes}" - ); + let want: Vec> = (0..n_lanes) + .map(|lane| { + // Lane `lane` encodes message block `n_lanes - 1 - lane`. + let block = n_lanes - 1 - lane; + let mut want = vec![F64::ZERO; block_len]; + replicate_rows(&mut want, &msg[block * rows..(block + 1) * rows]); + ntt.transform(&mut want, 1, log_inv_rate, None, None); + want + }) + .collect(); + for workers in [parallel::num_threads(), 1] { + let mut got = vec![F64::ZERO; msg.len() << log_inv_rate]; + let lanes = Message::Lanes(SendPtr(msg.as_ptr().cast_mut())); + ntt.transform_for(&mut got, n_lanes, log_inv_rate, Some(lanes), Some(&|_, _| {}), workers); + for (lane, want) in want.iter().enumerate() { + for (pos, want) in want.iter().enumerate() { + assert_eq!( + got[pos * n_lanes + lane], + *want, + "lane {lane} pos {pos} at log_rows={log_rows}, rate={log_inv_rate}, n_lanes={n_lanes}, \ + workers={workers}" + ); + } } } } diff --git a/crates/pcs/src/whir/commit.rs b/crates/pcs/src/whir/commit.rs index 9b2fc4e33..e4386a67c 100644 --- a/crates/pcs/src/whir/commit.rs +++ b/crates/pcs/src/whir/commit.rs @@ -54,17 +54,14 @@ pub fn commit(message: &[F64], log_n: usize, log_batch_size: usize, log_inv_rate let codeword_len = n_positions * n_lanes; // SAFETY: every codeword element is written before it is read. - // `transpose_lane_major` covers every word of the message region (its tiles are - // asserted to), and `encode_interleaved_in_place` writes every other replica from - // it before transforming that region in place. + // `encode_lane_major_with` writes every word of an uninitialized codeword before it reads it. let mut codeword = unsafe { primitives::uninit_vec::(codeword_len) }; // Leaves are hashed as the encode finishes each block of rows. let tree = MerkleBuilder::new(n_positions, n_lanes, 1usize << log_batch_size); tracing::info_span!("NTT", kind = "base encode", log_domain = k_code, lanes = n_lanes).in_scope(|| { - crate::ntt::transpose_lane_major(&mut codeword[..message.len()], message, n_lanes, log_rows); let ntt = AdditiveNttF64::standard(k_code); - ntt.encode_interleaved_in_place_with(&mut codeword, n_lanes, log_inv_rate, &|row, rows| { + ntt.encode_lane_major_with(&mut codeword, message, n_lanes, log_inv_rate, &|row, rows| { tree.absorb(row, rows); }); }); From 3c395bb76cf94ea05e9e7729ac0472d5e62d336a Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Sun, 4 Oct 2026 23:00:02 +0100 Subject: [PATCH 03/18] pcs: lane-major gather reads each lane's run in order; masked last vector for 2-7 spare lanes - gather_lanes walks one lane octet's lines front to back (8 lines per lane per row group), which halves the lane-major read time at 1T. - Rows ending in 2 to 7 lanes past a whole vector (fibonacci's 35 lanes) take one masked vector in separate #[inline(never)] sweeps; folding the mask into the common loop cost the whole-vector shapes about 5%. Byte-identical: fibonacci 2M 87a8a891167e3b9a, leanxmss 400 744c5cfc65b1157d. cargo test --release -p pcs passes. This and the previous two commits, 1T, base / this, A B A B: leanxmss 400: Commit 606, 608 -> 484, 484 ms; L1-L4 378, 378 -> 360, 350; Prove 5480, 5430 -> 5370, 5320 ms. fibonacci 2M: Commit 1150, 1150 -> 835, 830 ms; L1-L4 795, 797 -> 731, 730; Prove 9180, 9250 -> 8900, 8850 ms. 32T leanxmss 400 (--repeat 3): Commit 68.4, 67.5 -> 66.5, 67.3 ms; proving 0.556, 0.568 -> 0.559, 0.578 s (+-2 to 4% spreads): unchanged. Probe (1T, xmss shape, commit alone): transpose 50 + gathered in 39 + out 50 ms became gathered in 45 + out 22 ms. --- crates/pcs/src/ntt/additive_ntt_f64.rs | 171 ++++++++++++++++++++----- 1 file changed, 137 insertions(+), 34 deletions(-) diff --git a/crates/pcs/src/ntt/additive_ntt_f64.rs b/crates/pcs/src/ntt/additive_ntt_f64.rs index 4da83aa8d..8f1e21e4e 100644 --- a/crates/pcs/src/ntt/additive_ntt_f64.rs +++ b/crates/pcs/src/ntt/additive_ntt_f64.rs @@ -401,8 +401,8 @@ impl AdditiveNttF64 { // // message is block 0 of the buffer -> a task could overwrite it while others still copy it // no layer left to run -> no deep task would write the buffer at all - if let Some(Message::Rows(m)) = - msg.take_if(|m| matches!(m, Message::Rows(r) if std::ptr::eq(r.0, data.as_mut_ptr())) || deep_start == log_d) + if let Some(Message::Rows(m)) = msg + .take_if(|m| matches!(m, Message::Rows(r) if std::ptr::eq(r.0, data.as_mut_ptr())) || deep_start == log_d) { replicate(data, m, data.len() >> start); } @@ -744,6 +744,9 @@ fn butterfly_interleaved_fused_3layer(block: &mut [F64], t: &[F64; 7], eighth: u { // SAFETY: the arm is compiled only with AVX-512F; registers only. let tw = t.map(|x| unsafe { _mm512_set1_epi64(x.0 as i64) }); + if num_ntts % 8 >= 2 { + return butterfly_interleaved_fused_3layer_masked(block, &tw, eighth, num_ntts); + } let vectors = num_ntts / 8; let done = 8 * vectors; fused_rows::<8>(block, eighth, num_ntts, |rows| { @@ -760,6 +763,27 @@ fn butterfly_interleaved_fused_3layer(block: &mut [F64], t: &[F64; 7], eighth: u fused_rows::<8>(block, eighth, num_ntts, |rows| radix8_butterflies(rows, t)); } +/// [`butterfly_interleaved_fused_3layer`] for rows that end in two to seven lanes past a whole vector. +/// +/// - Those lanes take one masked vector, which costs less than two scalar butterflies each. +/// - It is a function of its own, so the whole-vector loop of the common shapes keeps its registers. +#[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] +#[inline(never)] +fn butterfly_interleaved_fused_3layer_masked(block: &mut [F64], tw: &[__m512i; 7], eighth: usize, num_ntts: usize) { + let (vectors, mask) = (num_ntts / 8, (1u8 << (num_ntts % 8)) - 1); + fused_rows::<8>(block, eighth, num_ntts, |rows| { + let rows: [*mut F64; 8] = std::array::from_fn(|i| rows[i].as_mut_ptr()); + // SAFETY: the target features are enabled at compile time; each row has `8 * vectors` words and the + // masked lanes after them. + unsafe { + radix8_avx512(rows, tw, vectors); + masked_column_avx512(rows.map(|row| row.add(8 * vectors)), mask, |r| { + radix8_regs_avx512(r, tw); + }); + } + }); +} + /// The twelve butterflies of one radix-8 row group, with its seven twiddles breadth-first. #[inline(always)] fn radix8_butterflies(rows: &mut [&mut [F64]; 8], t: &[F64; 7]) { @@ -872,8 +896,9 @@ unsafe fn gather_lanes(dst: &mut [F64], msg: SendPtr, msg_rows: usize, n: u // Message lanes `0 .. whole` go eight at a time; the rest, the lowest codeword lanes, one at a time. let whole = if LANE_GATHER { n - n % 8 } else { 0 }; #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - for q in (0..residues).step_by(8) { - for l0 in (0..whole).step_by(8) { + for l0 in (0..whole).step_by(8) { + // Each lane's lines in order, so a lane's run of `residues` words is read front to back. + for q in (0..residues).step_by(8) { // SAFETY: line `k` is lane `l0 + 7 - k`, rows `row + q ..+ 8`, inside the message by the caller. let lines: [__m512i; 8] = std::array::from_fn(|k| unsafe { _mm512_loadu_si512(msg.add((l0 + 7 - k) * msg_rows + row + q).cast_const().cast()) @@ -1062,6 +1087,9 @@ fn butterfly_interleaved_fused_2layer( { // SAFETY: the arm is compiled only with AVX-512F; registers only. let tw = [t_outer, t_inner_a, t_inner_b].map(|x| unsafe { _mm512_set1_epi64(x.0 as i64) }); + if num_ntts % 8 >= 2 { + return butterfly_interleaved_fused_2layer_masked(block, &tw, quarter, num_ntts); + } let vectors = num_ntts / 8; let done = 8 * vectors; fused_rows::<4>(block, quarter, num_ntts, |rows| { @@ -1089,8 +1117,32 @@ fn butterfly_interleaved_fused_2layer( }); } +/// [`butterfly_interleaved_fused_2layer`] for rows that end in two to seven lanes past a whole vector. +/// +/// As [`butterfly_interleaved_fused_3layer_masked`]. +#[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] +#[inline(never)] +fn butterfly_interleaved_fused_2layer_masked(block: &mut [F64], tw: &[__m512i; 3], quarter: usize, num_ntts: usize) { + let (vectors, mask) = (num_ntts / 8, (1u8 << (num_ntts % 8)) - 1); + fused_rows::<4>(block, quarter, num_ntts, |rows| { + let rows: [*mut F64; 4] = std::array::from_fn(|i| rows[i].as_mut_ptr()); + // SAFETY: the target features are enabled at compile time; each row has `8 * vectors` words and the + // masked lanes after them. + unsafe { + radix4_avx512(rows, tw, vectors); + masked_column_avx512(rows.map(|row| row.add(8 * vectors)), mask, |r| { + radix4_regs_avx512(r, tw); + }); + } + }); +} + #[inline] fn butterfly_interleaved_block(block: &mut [F64], twiddle: F64, block_size_half: usize, num_ntts: usize) { + #[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] + if num_ntts % 8 >= 2 { + return butterfly_interleaved_block_masked(block, twiddle, block_size_half, num_ntts); + } let half_offset = block_size_half * num_ntts; let (top, bot) = block.split_at_mut(half_offset); for r in 0..block_size_half { @@ -1099,6 +1151,29 @@ fn butterfly_interleaved_block(block: &mut [F64], twiddle: F64, block_size_half: } } +/// [`butterfly_interleaved_block`] for rows that end in two to seven lanes past a whole vector. +/// +/// As [`butterfly_interleaved_fused_3layer_masked`]. +#[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] +#[inline(never)] +fn butterfly_interleaved_block_masked(block: &mut [F64], twiddle: F64, block_size_half: usize, num_ntts: usize) { + let (vectors, mask) = (num_ntts / 8, (1u8 << (num_ntts % 8)) - 1); + let half_offset = block_size_half * num_ntts; + let base = block.as_mut_ptr(); + // SAFETY: the arm is compiled only with AVX-512F; registers only. + let tw = unsafe { _mm512_set1_epi64(twiddle.0 as i64) }; + for r in 0..block_size_half { + // SAFETY: rows `r` of the two halves, `num_ntts` words each, inside the block and disjoint. + unsafe { + let rows = [base.add(r * num_ntts), base.add(half_offset + r * num_ntts)]; + columns_avx512(rows, vectors, |p| butterfly_regs_avx512(p, 0, 1, tw)); + masked_column_avx512(rows.map(|row| row.add(8 * vectors)), mask, |p| { + butterfly_regs_avx512(p, 0, 1, tw); + }); + } + } +} + /// Butterfly all `num_ntts` lanes of one (top row, bottom row) pair with a /// shared twiddle: new_u = u + v*t; new_v = v + new_u. /// @@ -1399,6 +1474,58 @@ unsafe fn columns_avx512(rows: [*mut F64; N], vectors: usize, bo } } +/// [`columns_avx512`] for one column of fewer than eight lanes, the ones `mask` selects. +/// +/// # Safety +/// +/// - Requires AVX-512F. +/// - Each row pointer must address the lanes `mask` selects; the others are neither read nor written. +#[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] +#[inline(always)] +unsafe fn masked_column_avx512(rows: [*mut F64; N], mask: __mmask8, body: impl Fn(&mut [__m512i; N])) { + // SAFETY: masked-off lanes are neither read nor written, the rest are the caller's. + unsafe { + let mut r: [__m512i; N] = std::array::from_fn(|i| _mm512_maskz_loadu_epi64(mask, rows[i].cast())); + body(&mut r); + for (i, row) in r.iter().enumerate() { + _mm512_mask_storeu_epi64(rows[i].cast(), mask, *row); + } + } +} + +/// The twelve butterflies of a radix-8 row group held in registers, its seven twiddles breadth-first. +#[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] +#[inline(always)] +fn radix8_regs_avx512(r: &mut [__m512i; 8], t: &[__m512i; 7]) { + // Layer L: rows 4 apart, one twiddle for the whole block. + butterfly_regs_avx512(r, 0, 4, t[0]); + butterfly_regs_avx512(r, 1, 5, t[0]); + butterfly_regs_avx512(r, 2, 6, t[0]); + butterfly_regs_avx512(r, 3, 7, t[0]); + // Layer L+1: rows 2 apart, one twiddle per half. + butterfly_regs_avx512(r, 0, 2, t[1]); + butterfly_regs_avx512(r, 1, 3, t[1]); + butterfly_regs_avx512(r, 4, 6, t[2]); + butterfly_regs_avx512(r, 5, 7, t[2]); + // Layer L+2: adjacent rows, one twiddle per quarter. + butterfly_regs_avx512(r, 0, 1, t[3]); + butterfly_regs_avx512(r, 2, 3, t[4]); + butterfly_regs_avx512(r, 4, 5, t[5]); + butterfly_regs_avx512(r, 6, 7, t[6]); +} + +/// The four butterflies of a radix-4 row group held in registers: one twiddle for the block, then one per half. +#[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] +#[inline(always)] +fn radix4_regs_avx512(r: &mut [__m512i; 4], t: &[__m512i; 3]) { + // Layer L: rows 2 apart, one twiddle for the block. + butterfly_regs_avx512(r, 0, 2, t[0]); + butterfly_regs_avx512(r, 1, 3, t[0]); + // Layer L+1: adjacent rows, one twiddle per half. + butterfly_regs_avx512(r, 0, 1, t[1]); + butterfly_regs_avx512(r, 2, 3, t[2]); +} + /// A radix-8 row group with the eight rows in registers. /// /// Each row is loaded and stored once for its twelve butterflies, not three times. @@ -1412,25 +1539,7 @@ unsafe fn columns_avx512(rows: [*mut F64; N], vectors: usize, bo #[target_feature(enable = "vpclmulqdq", enable = "avx512f", enable = "avx2")] unsafe fn radix8_avx512(rows: [*mut F64; 8], t: &[__m512i; 7], vectors: usize) { // SAFETY: forwarded from the caller. - unsafe { - columns_avx512(rows, vectors, |r| { - // Layer L: rows 4 apart, one twiddle for the whole block. - butterfly_regs_avx512(r, 0, 4, t[0]); - butterfly_regs_avx512(r, 1, 5, t[0]); - butterfly_regs_avx512(r, 2, 6, t[0]); - butterfly_regs_avx512(r, 3, 7, t[0]); - // Layer L+1: rows 2 apart, one twiddle per half. - butterfly_regs_avx512(r, 0, 2, t[1]); - butterfly_regs_avx512(r, 1, 3, t[1]); - butterfly_regs_avx512(r, 4, 6, t[2]); - butterfly_regs_avx512(r, 5, 7, t[2]); - // Layer L+2: adjacent rows, one twiddle per quarter. - butterfly_regs_avx512(r, 0, 1, t[3]); - butterfly_regs_avx512(r, 2, 3, t[4]); - butterfly_regs_avx512(r, 4, 5, t[5]); - butterfly_regs_avx512(r, 6, 7, t[6]); - }); - } + unsafe { columns_avx512(rows, vectors, |r| radix8_regs_avx512(r, t)) }; } /// A radix-4 row group with the four rows in registers. @@ -1444,16 +1553,7 @@ unsafe fn radix8_avx512(rows: [*mut F64; 8], t: &[__m512i; 7], vectors: usize) { #[target_feature(enable = "vpclmulqdq", enable = "avx512f", enable = "avx2")] unsafe fn radix4_avx512(rows: [*mut F64; 4], t: &[__m512i; 3], vectors: usize) { // SAFETY: forwarded from the caller. - unsafe { - columns_avx512(rows, vectors, |r| { - // Layer L: rows 2 apart, one twiddle for the block. - butterfly_regs_avx512(r, 0, 2, t[0]); - butterfly_regs_avx512(r, 1, 3, t[0]); - // Layer L+1: adjacent rows, one twiddle per half. - butterfly_regs_avx512(r, 0, 1, t[1]); - butterfly_regs_avx512(r, 2, 3, t[2]); - }); - } + unsafe { columns_avx512(rows, vectors, |r| radix4_regs_avx512(r, t)) }; } /// Two F64 butterflies with a shared twiddle, NEON-resident end to end. @@ -1635,7 +1735,10 @@ mod tests { let first = Message::Rows(SendPtr(got.as_mut_ptr())); let sink = |row: usize, rows: &[F64]| blocks.lock().unwrap().push((row, rows.to_vec())); ntt.transform_for(&mut got, lanes, log_inv_rate, Some(first), Some(&sink), workers); - assert_eq!(got, want, "log_d={log_d}, lanes={lanes}, rate={log_inv_rate}, workers={workers}"); + assert_eq!( + got, want, + "log_d={log_d}, lanes={lanes}, rate={log_inv_rate}, workers={workers}" + ); let mut blocks = blocks.into_inner().unwrap(); blocks.sort_by_key(|b| b.0); From bcc61577ed76b1d486c30bc94dd641610d638681 Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Sun, 4 Oct 2026 23:01:13 +0100 Subject: [PATCH 04/18] pcs: a lone worker takes the L3 deep pass With one worker the deep pass takes fit3 rows whenever the plan has a gathered pass: its L3 is its own, a deep layer costs about the same from L3 as from L2, and the smaller gathered groups read the lane-major stack faster. 32T planning is unchanged. Probe (1T, encode + Merkle alone): leanxmss base commit 482 -> 459 ms, its L1 extension commit 236 -> 231 ms, fibonacci's L1 493 -> 472 ms (fibonacci's base commit already took fit3). Byte-identical: fibonacci 2M 87a8a891167e3b9a, leanxmss 400 744c5cfc65b1157d. cargo test --release -p pcs passes. All four commits, 1T, base / this, A B A B: leanxmss 400: Commit 604, 607 -> 464, 523 ms (second run noisy); L1-L4 378, 377 -> 342, 344 ms; Prove 5470, 5460 -> 5260, 5400 ms. fibonacci 2M: Commit 1170, 1160 -> 837, 898 ms; L1-L4 794, 809 -> 706, 709 ms; Prove 10100 (disturbed), 9370 -> 8830, 8890 ms. 32T leanxmss 400: Commit 66.6, 71.4 -> 66.8, 68.7 ms; proving 0.558, 0.58 -> 0.572, 0.629 s (last run +-21.7%, load 6.8): no change beyond noise. pcs throughput bench, 1T, PCS_LOG_N=22 (load 7 to 11): commit 47.9, 63.6 -> 43.0, 42.8 ms; encode alone 29.7, 33.7 -> 25.6, 24.7 ms. --- crates/pcs/src/ntt/additive_ntt_f64.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/crates/pcs/src/ntt/additive_ntt_f64.rs b/crates/pcs/src/ntt/additive_ntt_f64.rs index 8f1e21e4e..a73f771a7 100644 --- a/crates/pcs/src/ntt/additive_ntt_f64.rs +++ b/crates/pcs/src/ntt/additive_ntt_f64.rs @@ -352,8 +352,11 @@ impl AdditiveNttF64 { // L3 deep pass: 1 gathered sweep (9 layers), then 12 deep <- one sweep fewer // // 20 layers: 1 gathered sweep either way, so the deep pass stays in L2 + // + // A lone worker spills to an L3 it shares with no one, so it takes the L3 deep pass whenever it has one: + // the layers it moves out of the gathered pass shrink the gathered groups, and their scattered reads. let gathered_sweeps = |deep: usize| (log_d - start).saturating_sub(deep).div_ceil(fit2); - let deep = if gathered_sweeps(fit3) < gathered_sweeps(fit2) { + let deep = if workers == 1 || gathered_sweeps(fit3) < gathered_sweeps(fit2) { fit3 } else { fit2 From 4b9f431c4d6a053caf9f397a25b6c5182cc3b90a Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Sun, 4 Oct 2026 23:14:41 +0100 Subject: [PATCH 05/18] leanvm_core: word-level witnesses for the ALU, memory, shift and clock circuits; HASH's unrolled ALU, LOAD, STORE, LD, SD and SHIFT (`ClassSpec::witness`) and every clock circuit (`Clock::witness`) now write their packed witness by word arithmetic instead of the 64-lane walk of their gate list. Each generator computes the circuit's values natively and writes its rows in the order the circuit makes its products through `rv::circuits::Products`, which keeps each table's current word in a register and stores it once full. ALU's 192 interleaved logic rows go through `Products::push_interleaved3` (BMI2 `pdep`, with a portable spread elsewhere). DIV, EXT, and MUL off AVX2 still walk. HASH's generator was already word arithmetic, but it ORed each carry run into the tables at a slot read from a table the circuit build recorded. It is now unrolled by a macro: every run's place is a constant, so the rows are written with no branch and no read, and the recorded runs (`CARRIES`) are gone. A run starts above bit 0 only on round 0's first `c + d` of a column G, whose `c` is still the IV. `class_flock::Prepared::build` still asserts that every output port of every instance is what the interpreter computed. Tests: `semantics::tests::word_witness_is_the_walk` checks ALU, LOAD, STORE, LD, SD and SHIFT on a grid of edge instances (every legal flag word; zero, one, all ones and the sign bits in every operand; every shift amount from 0 to 64 and all ones, with and without high bits; every byte offset of a cell, so every misalignment of every width; cells with every byte's sign bit set or clear) plus 4096 random ones (half decoder-shaped, half edge-biased words in every port, flags included). Each instance is compared alone with the bit-by-bit walk, then the whole batch's packed tables and lincheck stripes with the 64-lane walk. The clock's test covers every table's clock circuit and one with all 32 slots on 4096 rows (padding, first and last cycle, all ones, random; previous timestamps around the row's, live bit flipped, zero, all ones, random); HASH's covers all-zero, all-ones and sign-bit blocks and edge counters. Mutations (a dropped carry, a wrong fill bit, a wrong sign bit, the clock's top carry, a short HASH carry run) each fail their test. The portable spread is tested with target-cpu=x86-64-v2. Byte-identical proofs: serialized proof digests equal to the base's for leanxmss 400, fibonacci 2M and leanda 1 blob, at 1 and 32 threads. Measured on a Ryzen 9 7950X3D, shared host, base (77077c6b) and this commit alternated, `--tracing`, loads 3.4 to 12.6: One thread (LEANVM_NUM_THREADS=1, taskset -c 2): leanxmss 400 Build flock witnesses 390, 398 -> 209, 212 ms proving 5.448, 5.475 -> 5.286, 5.298 s fibonacci 2M Build flock witnesses 571, 577 -> 297, 309 ms proving 9.191, 9.225 -> 8.961, 8.921 s leanda 1 blob Build flock witnesses 426, 425 -> 262, 261 ms proving 5.633, 5.670 -> 5.486, 5.508 s 32 threads: leanxmss 400 Build flock witnesses 47.4, 49.4 -> 47.2, 47.0 ms proving 0.558, 0.561 -> 0.588, 0.576 s fibonacci 2M Build flock witnesses 71.2, 71.4 -> 70.8, 69.1 ms proving 0.923, 0.923 -> 0.931, 0.947 s leanda 1 blob Build flock witnesses 52.4, 51.5 -> 50.2, 49.3 ms proving 0.601, 0.587 -> 0.600, 0.600 s At 32 threads the span and the proving totals stay within the run to run spread; `--repeat 3` runs at load under 10: leanxmss 0.557 to 0.566 s for the base and 0.559 to 0.601 s for this commit, fibonacci 0.906 to 0.942 s and 0.913 to 0.932 s. --- AGENTS.md | 2 +- crates/leanvm_core/src/class_flock.rs | 39 ++-- crates/leanvm_core/src/rv/circuits.rs | 98 ++++++++++ crates/leanvm_core/src/rv/semantics/alu.rs | 100 +++++++++- crates/leanvm_core/src/rv/semantics/hash.rs | 164 +++++++--------- crates/leanvm_core/src/rv/semantics/memory.rs | 185 +++++++++++++++++- crates/leanvm_core/src/rv/semantics/mod.rs | 65 +++++- crates/leanvm_core/src/rv/semantics/shift.rs | 69 ++++++- crates/leanvm_core/src/tables/clock.rs | 138 +++++++++++++ crates/leanvm_core/src/tables/spec.rs | 14 +- 10 files changed, 744 insertions(+), 130 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index f0ecf7cd2..aec51a0f0 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -135,7 +135,7 @@ Addition with carries, comparisons, shifts, AND/OR, multiplication and division - **Circuits are gate lists over word ports** (`flock::circuit`): inputs, outputs, the constant, then products in the order they are made. A port bit with no gate is an empty row, hence zero, which is what makes a one-bit output such as `taken` a 0 or 1 field element: give such an output a word to itself and never put a free wire on its spare bits. What Rust and Python must agree on is the port layout and the ORDER PRODUCTS ARE MADE IN; XOR order is free. `alu_is_its_reference` pins a circuit to its reference function, and the end-to-end tests pin the Python mirror. - **One reduction per circuit, one opening for all**: zerocheck plus lincheck per packed witness, every class circuit in table order then every clock circuit, after the exit claims, each leaving a claim on its own witness; `pcs::stack_open` takes one ring-switched region per witness, then one for the bytecode's multiplicity column, and ring-switches all their claims as one family under one map challenge. - **Batch floors.** Flock needs eight instances and a zerocheck cube of `2^13` bits for each of a table's two circuits (`class_flock::n_blocks_log`); padding rows supply them, as honest instances on zero registers. -- **Witness generation is the generic walk of the gate list, bit by bit**, except where a class sets `ClassSpec::witness`: HASH's is word arithmetic, the same tables from the native sums' carries. A word-level generator for the other circuits is a known follow-up, and each one needs a test pinning it to the walk. +- **Witness generation is word arithmetic wherever a circuit has a generator, and the generic walk of the gate list, bit by bit, elsewhere**: ALU, LOAD, STORE, LD, SD, SHIFT, MULH and HASH set `ClassSpec::witness`, MUL sets `batch_witness` on AVX2 only, and every clock circuit's generator is `Clock::witness`; DIV, EXT, and MUL off AVX2 walk. Each generator computes the circuit's values natively and writes its rows in the order the circuit makes its products (all but the multipliers' through `rv::circuits::Products`), so a change to a circuit's gate list must change its generator too. Each is pinned to the walk by a test that covers edge words (zero, all ones, the sign bits, every shift amount, every misalignment) and random instances: `semantics::tests::word_witness_is_the_walk` for ALU, the memory classes and SHIFT, and their own for MUL, MULH, HASH and the clock. A new generator needs such a test. HASH's generator is unrolled, so every carry run's place is a constant. - **Circuit sizes are structure, not measurements**: `k_log` and `clock_k_log` per class are pinned in its `ClassSpec` and asserted against the built circuits; product counts are in `doc/leanvm` Annex C. A product can be committed straight at an output port's position (`flock::circuit::Builder::and_output`), which saves the copy; the clock circuit's step uses it. CI compares every PR with its base, never with a history. The base is the merge commit's first parent; a workflow builds and runs both on one runner and saves both as an artifact, and a `workflow_run` workflow on the default branch (`main`, the only branch GitHub runs those from) posts them on the PR as one comment, edited on every run and deleted once nothing changed. The posting workflows never run the PR's code: they take only the PR number from a run's artifact, a fork's included, post nothing unless GitHub's API says that PR's head is the commit the run measured, and check every name and number before it reaches the comment. **Exact counts**: `counts.yml` runs `cargo leanvm bench --cycles-only`, every program at the sizes above without a proof (`cycles`, the program's own instructions; `proven-rows`, table heights after padding; `committed`, witness words) and the two aggregation trees above, one case per tree and kind of node (`aggregate-leanxmss-400-2to1-first` and `-node`, `aggregate-leanxmss-400-4to1-first` and `-node`, from `counted_trees()`: `rows`, the circuit's own; `proven-rows` and `committed` as for a program), puts the tables in each run's summary from the same pass (`--markdown-file "$GITHUB_STEP_SUMMARY"`, which appends what `--markdown` prints), saves the artifact `counts` and fails on any increase; `counts-comment.yml` posts the counts that changed. A tree's name says its shape, `Nto1` for N proofs combined into 1 at every node; `-first` is a first-level node (the RISC-V verifier in rows over N leaf proofs) and `-node` a higher node (the recursion verifier in rows over N child proofs). **Proving**: `bench.yml` runs one job per testbed and benchmark: each case of `proven()` in `bins/leanvm/src/tracked.rs` (at the sizes a 16 GB runner holds: leanXMSS and leanSPHINCS at a quarter, no leanDA), each of `proven_trees()` (the two aggregation trees above over leanXMSS at a quarter, `aggregate-leanxmss-100-2to1` and `aggregate-leanxmss-100-4to1`: the leaf proven once, then one first-level node and one node over copies of it, the two benchmarks `-first` and `-node`, each node's stages the children of a `Prove` span around it and the node's peak memory including the first-level node's), and each `benches/` target above with every benchmark it prints (flock's two at `FLOCK_N_LOG=18`, pcs's `throughput` at its default size). Each job builds both sides and runs them in turns, base, PR, PR, base, for `ROUNDS` rounds: `cargo leanvm bench --only ` gives the proving time, the proof's size in bytes (exact: the inputs are seeded and grinding takes the smallest nonce), the time of each of the proof's stages (`stage.`, the `Prove` span's direct children, `--tracing`'s first level, timed by a `tracing` layer that enables spans only) and the process's peak resident memory (`peak-memory`, that case's alone since each job proves one), and the verifying time, and a target's `--json` its own measures: `latency` for a whole operation (flock's proving times with the witness excluded, pcs's commit and open), `per-op` in nanoseconds per call for `kernels`, a hashed block of `hash_throughput` and one NTT encode. `bench-comment.yml` posts a size that changed and a time that moved the same way in every round with a median ratio at least 1% from one, every result collapsed beneath; it gates nothing. Run in turns on one runner, the comparison does not depend on which machine the job landed on, which a history across GitHub's runners does (one runner label spans CPU models of different speeds). The testbeds are GitHub-hosted: `x86-64` (built for `target-cpu=haswell`) and `arm64` (`ubuntu-24.04-arm`, `target-cpu=neoverse-n2` plus `+aes,+sha3`, since the NEON field kernels are chosen at compile time on those features and a generic aarch64 target enables neither; the job fails early if the CPU lacks them), each ISA pinned because a cached native build dies of SIGILL on an older CPU. A runner has 4 vCPUs, so a change to how the prover scales across many cores shows there only in part. A benchmark's name is what base and PR are matched by, so renaming one or changing its input shows it as new, as does one the base cannot run (it cannot build the target, or its first base run fails or prints no JSON, and the job carries on with the PR's runs alone). diff --git a/crates/leanvm_core/src/class_flock.rs b/crates/leanvm_core/src/class_flock.rs index 314fe1463..fb09b02fd 100644 --- a/crates/leanvm_core/src/class_flock.rs +++ b/crates/leanvm_core/src/class_flock.rs @@ -169,25 +169,13 @@ impl Prepared { *word = port.value(row, &entries[row.index as usize], &slots); } }; - // A class with a word-level witness skips the walk of its gate list; the others - // walk it 64 instances at a time. - let witness = spec.witness.filter(|_| part == Part::Class); - let batch_witness = spec.batch_witness.filter(|_| part == Part::Class); - let (z, a, b, z_lincheck) = batch_witness.map_or_else( - || { - witness.map_or_else( - || circuit.generate_witness_from(rows, &rows[0], n_blocks_log, input_words), - |witness| { - circuit.generate_witness_with(rows, &rows[0], n_blocks_log, |row, z, az, bz| { - let mut words = [0u64; MAX_INPUT_WORDS]; - let words = &mut words[..n_inputs]; - input_words(row, words); - witness(words, z, az, bz); - }) - }, - ) - }, - |batch| { + // A clock circuit, and a class with a word-level witness, skip the walk of the gate list; + // the others walk it 64 instances at a time. + let (z, a, b, z_lincheck) = match (part, spec.batch_witness, spec.witness) { + (Part::Clock, _, _) => circuit.generate_witness_with(rows, &rows[0], n_blocks_log, |row, z, az, bz| { + Clock::witness(&slots, row.ts, &row.prev()[..slots.len()], z, az, bz); + }), + (Part::Class, Some(batch), _) => { // Eight rows share a native arithmetic call before their byte stripe is packed. circuit.generate_witness_batched(rows, &rows[0], n_blocks_log, |rows, z, az, bz| { let mut words = [[0u64; MAX_INPUT_WORDS]; 8]; @@ -197,8 +185,17 @@ impl Prepared { let inputs = std::array::from_fn(|i| &words[i][..n_inputs]); batch(&inputs, z, az, bz); }) - }, - ); + } + (Part::Class, None, Some(witness)) => { + circuit.generate_witness_with(rows, &rows[0], n_blocks_log, |row, z, az, bz| { + let mut words = [0u64; MAX_INPUT_WORDS]; + let words = &mut words[..n_inputs]; + input_words(row, words); + witness(words, z, az, bz); + }) + } + (Part::Class, None, None) => circuit.generate_witness_from(rows, &rows[0], n_blocks_log, input_words), + }; assert_eq!(window.len(), z.len(), "the committed column is the wrong size"); let stride = 1 << stride_log(spec, part); // `F64` is `repr(transparent)` over `u64`, and the packing is bit `i` at diff --git a/crates/leanvm_core/src/rv/circuits.rs b/crates/leanvm_core/src/rv/circuits.rs index f1e3fef35..ab39db8dd 100644 --- a/crates/leanvm_core/src/rv/circuits.rs +++ b/crates/leanvm_core/src/rv/circuits.rs @@ -169,6 +169,104 @@ impl WordGadgets for Builder { } } +/// The rows of `z`, `A·z` and `B·z` from a word on, written in order: each row `A·z = a`, `B·z = b`, `z = a·b`. +/// +/// A word is stored once it is full, so the tables are written, never read. +pub(crate) struct Products<'a> { + tables: [&'a mut [u64]; 3], + word: usize, + /// Bits of `pending` taken. + used: u32, + /// The rows of each table's current word so far. + pending: [u64; 3], +} + +impl<'a> Products<'a> { + pub(crate) const fn new(tables: [&'a mut [u64]; 3], word: usize) -> Self { + Self { + tables, + word, + used: 0, + pending: [0; 3], + } + } + + /// The next `bits` rows, at most 64, their `A·z` and `B·z` the low bits of `a` and `b`, which hold no others. + #[inline(always)] + pub(crate) fn push(&mut self, a: u64, b: u64, bits: u32) { + let rows = [a & b, a, b]; + let at = self.used; + for (pending, row) in self.pending.iter_mut().zip(rows) { + *pending |= row << at; + } + self.used += bits; + if self.used >= 64 { + for ((table, pending), row) in self.tables.iter_mut().zip(&mut self.pending).zip(rows) { + table[self.word] = *pending; + // `(row >> 1) >> (63 - at)` is `row >> (64 - at)`, the rows that did not fit, with no overflowing shift at `at = 0`. + *pending = (row >> 1) >> (63 - at); + } + self.word += 1; + self.used -= 64; + } + } + + /// The next 192 rows, three per bit of three words: row `3i + k` has `A·z` bit `i` of `a[k]` and `B·z` bit `i` of `b[k]`. + #[inline(always)] + pub(crate) fn push_interleaved3(&mut self, a: [u64; 3], b: [u64; 3]) { + for (a, b) in interleave3(a).into_iter().zip(interleave3(b)) { + self.push(a, b, 64); + } + } + + /// The tables, the last word stored. + pub(crate) fn finish(mut self) -> [&'a mut [u64]; 3] { + if self.used > 0 { + for (table, pending) in self.tables.iter_mut().zip(self.pending) { + table[self.word] = pending; + } + } + self.tables + } +} + +/// The 192 bits `x0_0, x1_0, x2_0, x0_1, ...` of three words, as three words. +#[inline(always)] +fn interleave3([x, y, w]: [u64; 3]) -> [u64; 3] { + // Word `j` starts at bit 64j of the sequence, which is bit `64j / 3` of the word at `64j mod 3`. + [ + spread3(x) | spread3(y) << 1 | spread3(w) << 2, + spread3(y >> 21) | spread3(w >> 21) << 1 | spread3(x >> 22) << 2, + spread3(w >> 42) | spread3(x >> 43) << 1 | spread3(y >> 43) << 2, + ] +} + +/// The low 22 bits of `x` at every third bit, from bit 0. +#[cfg_attr( + not(all(target_arch = "x86_64", target_feature = "bmi2")), + expect(clippy::missing_const_for_fn, reason = "BMI2's bit deposit is a runtime intrinsic.") +)] +#[inline(always)] +fn spread3(x: u64) -> u64 { + #[cfg(all(target_arch = "x86_64", target_feature = "bmi2"))] + { + // SAFETY: BMI2 is enabled at compile time. + unsafe { std::arch::x86_64::_pdep_u64(x, 0x9249_2492_4924_9249) } + } + // Without it, bits 0 to 20 spread by halving strides, and bit 21 goes to bit 63. + #[cfg(not(all(target_arch = "x86_64", target_feature = "bmi2")))] + { + let top = (x >> 21 & 1) << 63; + let mut x = x & 0x1f_ffff; + x = (x | x << 32) & 0x001f_0000_0000_ffff; + x = (x | x << 16) & 0x001f_0000_ff00_00ff; + x = (x | x << 8) & 0x100f_00f0_0f00_f00f; + x = (x | x << 4) & 0x10c3_0c30_c30c_30c3; + x = (x | x << 2) & 0x1249_2492_4924_9249; + x | top + } +} + #[cfg(test)] mod tests { use crate::rv::entry::Class; diff --git a/crates/leanvm_core/src/rv/semantics/alu.rs b/crates/leanvm_core/src/rv/semantics/alu.rs index 9e8a7a245..38545cb61 100644 --- a/crates/leanvm_core/src/rv/semantics/alu.rs +++ b/crates/leanvm_core/src/rv/semantics/alu.rs @@ -1,7 +1,7 @@ //! The ALU: sums, differences, comparisons, bitwise logic, branches and jumps. use super::{InstructionClass, sext32}; -use crate::rv::circuits::{ClassCircuit, Word, WordGadgets}; +use crate::rv::circuits::{ClassCircuit, Products, Word, WordGadgets}; use crate::rv::entry::Class; use flock::circuit::{Builder, Circuit}; @@ -247,10 +247,99 @@ impl ClassCircuit for Alu { } } +impl Alu { + /// One instance of the circuit's witness by word arithmetic: what the walk of [`Alu::circuit`] writes, into zeroed buffers. + /// + /// Inputs `v1`, `v2`, `imm` and the flags' 15 bits, the outputs `out` and `taken`, the constant at bit 384, then the products: + /// + /// ```text + /// adder 64 A·z = v1 ^ c, B·z = (b ^ sub) ^ c, c the carries of v1 + (b ^ sub) + sub + /// any 63 A·z = the OR of diff's bits below, B·z = diff's bit + /// sext32 32 A·z = word, B·z = sum_31 ^ sum_i + /// none 64 A·z = none, B·z = the sign-extended sum + /// logic 192 per bit: v1 * b, and_or * that, or_xor * diff + /// comparisons 2 SEL_LT * lt, SEL_LTU * ltu + /// bit 0 1 !CLEAR_BIT0 * out_0 + /// branches 6 each condition's flag * whether it holds + /// ``` + pub(crate) fn witness(inputs: &[u64], z: &mut [u64], az: &mut [u64], bz: &mut [u64]) { + const FLAG_BITS: u64 = (1 << 15) - 1; + let (v1, v2, imm, flags) = (inputs[0], inputs[1], inputs[2], inputs[3] & FLAG_BITS); + let flag = |bit: u64| u64::from(flags & bit != 0); + let all = |bit: u64| u64::from(bit != 0).wrapping_neg(); + let b = v2 ^ imm; + let sub = flag(Self::SUB); + + // The adder: the carry into each bit, and out of the top. + let y = b ^ all(sub); + let (partial, o1) = v1.overflowing_add(y); + let (sum, o2) = partial.overflowing_add(sub); + let carries = sum ^ v1 ^ y; + let ltu = u64::from(!(o1 | o2)); + let lt = ltu ^ (v1 ^ b) >> 63; + let diff = v1 ^ b; + let ne = u64::from(diff != 0); + + let word = all(flag(Self::WORD)); + let sext_diff = ((sum >> 31 & 1).wrapping_neg() ^ sum) >> 32; + let extended = sum ^ (word & sext_diff) << 32; + let none = 1 + ^ flag(Self::SEL_LT) + ^ flag(Self::SEL_LTU) + ^ flag(Self::SEL_AND) + ^ flag(Self::SEL_OR) + ^ flag(Self::SEL_XOR); + let and_or = all(flag(Self::SEL_AND) ^ flag(Self::SEL_OR)); + let or_xor = all(flag(Self::SEL_OR) ^ flag(Self::SEL_XOR)); + let both = v1 & b; + let mut out = all(none) & extended ^ and_or & both ^ or_xor & diff; + let (lt_term, ltu_term) = (flag(Self::SEL_LT) & lt, flag(Self::SEL_LTU) & ltu); + let out0 = (out ^ lt_term ^ ltu_term) & 1; + let keep_bit0 = flag(Self::CLEAR_BIT0) ^ 1; + out = out & !1 | keep_bit0 & out0; + + let conditions = [ + (Self::BR_EQ, ne ^ 1), + (Self::BR_NE, ne), + (Self::BR_LT, lt), + (Self::BR_GE, lt ^ 1), + (Self::BR_LTU, ltu), + (Self::BR_GEU, ltu ^ 1), + ]; + let taken = conditions + .iter() + .fold(flag(Self::ALWAYS), |acc, &(when, holds)| acc ^ flag(when) & holds); + + // The ports. + for (i, bits) in [u64::MAX, u64::MAX, u64::MAX, FLAG_BITS].into_iter().enumerate() { + (z[i], az[i], bz[i]) = (inputs[i] & bits, inputs[i] & bits, bits); + } + (z[4], az[4], bz[4]) = (out, out, u64::MAX); + (z[5], az[5], bz[5]) = (taken, taken, 1); + + // The constant, then the products in the order the circuit makes them. + let mut rows = Products::new([z, az, bz], 6); + rows.push(1, 1, 1); + rows.push(v1 ^ carries, y ^ carries, 64); + let below = diff.isolate_lowest_one().wrapping_neg(); + rows.push(below & (u64::MAX >> 1), diff >> 1, 63); + rows.push(word >> 32, sext_diff, 32); + rows.push(all(none), extended, 64); + rows.push_interleaved3([v1, and_or, or_xor], [b, both, diff]); + rows.push(flag(Self::SEL_LT), lt, 1); + rows.push(flag(Self::SEL_LTU), ltu, 1); + rows.push(keep_bit0, out0, 1); + for (when, holds) in conditions { + rows.push(flag(when), holds, 1); + } + rows.finish(); + } +} + #[cfg(test)] mod tests { use super::*; - use crate::rv::semantics::tests::{circuit_matches_reference, edge_word}; + use crate::rv::semantics::tests::{EDGES, circuit_matches_reference, edge_word, grid, word_witness_is_the_walk}; use fiat_shamir::transcript::{ProverState, VerifierState}; use proptest::prelude::*; use proptest::sample::select; @@ -288,6 +377,13 @@ mod tests { circuit_matches_reference::(4096); } + #[test] + fn the_word_witness_is_the_gate_walk() { + // Every legal flag word on edge operands: `b` an edge word or its complement, equal to `v1` on the diagonal. + let edges = grid(&[&EDGES, &EDGES, &[0, u64::MAX], Alu::LEGAL]); + word_witness_is_the_walk::(Alu::witness, edges); + } + #[test] fn flock_proves_honest_alu_instances_and_refuses_a_flipped_bit() { // Fixture: 16 instances cycling through the legal words. diff --git a/crates/leanvm_core/src/rv/semantics/hash.rs b/crates/leanvm_core/src/rv/semantics/hash.rs index b3dd479b6..1d8589f16 100644 --- a/crates/leanvm_core/src/rv/semantics/hash.rs +++ b/crates/leanvm_core/src/rv/semantics/hash.rs @@ -1,11 +1,10 @@ //! The BLAKE2s compression instruction, and its access to its block. use super::InstructionClass; -use crate::rv::circuits::{ClassCircuit, Word, WordGadgets}; +use crate::rv::circuits::{ClassCircuit, Products, Word, WordGadgets}; use crate::rv::entry::Class; use flock::circuit::{Builder, Circuit, Wire}; use primitives::hash::{G_LANES, IV, SIGMA}; -use std::sync::OnceLock; /// One BLAKE2s compression instance, `blake2s rs1, rs2`: the finalization word, the counter, the block. /// @@ -122,22 +121,6 @@ impl ClassCircuit for Hash { let h: Vec = (0..8).map(|i| half(&c, 2 + i / 2, i)).collect(); let m: Vec = (0..16).map(|i| half(&c, 6 + i / 2, i)).collect(); - // Each addition records where its products went. - // - // Its products are its top bits, from the first bit where both operands exist. - // Once one product is made, the carry is a wire, so every later bit has one too. - let mut carries = Vec::with_capacity(SIGMA.len() * 8 * 6); - let mut add32 = |c: &mut Builder, x: &[Wire], y: &[Wire]| -> Word { - let slot = c.next_slot(); - let sum = c.add_wrapping(x, y); - let products = (c.next_slot() - slot) as u32; - carries.push(Carries { - slot: slot as u32, - low: 31 - products, - }); - sum - }; - // The working vector: h, the IV, with the counter and the finalization word XORed in. let mut v = h.clone(); v.extend(IV[..4].iter().map(|&x| literal(&c, x))); @@ -150,11 +133,11 @@ impl ClassCircuit for Hash { for round in &SIGMA { for (g, &[a, b, cc, d]) in G_LANES.iter().enumerate() { for (x, r1, r2) in [(&m[round[2 * g]], 16, 12), (&m[round[2 * g + 1]], 8, 7)] { - let ab = add32(&mut c, &v[a], &v[b]); - v[a] = add32(&mut c, &ab, x); + let ab = c.add_wrapping(&v[a], &v[b]); + v[a] = c.add_wrapping(&ab, x); let da = c.xor_word(&v[d], &v[a]); v[d] = rotr(&da, r1); - v[cc] = add32(&mut c, &v[cc], &v[d]); + v[cc] = c.add_wrapping(&v[cc], &v[d]); let bc = c.xor_word(&v[b], &v[cc]); v[b] = rotr(&bc, r2); } @@ -169,8 +152,6 @@ impl ClassCircuit for Hash { c.output(i / 2, 32 * (i % 2) + bit, wire); } } - // Publish the carry runs for the word witness; every build records the same runs. - let _ = CARRIES.set(carries); c.finish() } } @@ -178,22 +159,6 @@ impl ClassCircuit for Hash { /// The input ports in bits: `t`, `f0`, then the four words of `h` and the eight of `m`. const INPUT_BITS: [usize; 14] = [64, 32, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64]; -/// The hash circuit's carry runs, one per 32-bit addition, in the order the circuit makes them. -/// -/// Building the circuit records them, and the word witness reads them. -static CARRIES: OnceLock> = OnceLock::new(); - -/// Where one 32-bit addition put its carry products. -#[derive(Clone, Copy, Debug)] -struct Carries { - /// The slot of the first product. - slot: u32, - /// The lowest bit with a product. - /// - /// The bits below are structural zeros: a literal operand's low zero bits, with no carry yet. - low: u32, -} - /// One instance of the circuit's witness, by word arithmetic instead of the gate walk. /// /// It writes the same `z`, `A*z` and `B*z` the walk writes, into zeroed buffers. @@ -203,9 +168,11 @@ struct Carries { /// - words 0 to 13 are the inputs: `z` and `A*z` hold the word, `B*z` its wired bits; /// - words 14 to 17 are the outputs: `z` and `A*z` hold the word, `B*z` all ones; /// - bit 1152 is the constant, one in all three; -/// - from bit 1153, each 32-bit addition has a run of carry products. +/// - from bit 1153, each 32-bit addition has a run of carry products, one after another. /// -/// An addition `x + y` has carries `c = (x + y) ^ x ^ y`. +/// An addition `x + y` has carries `c = (x + y) ^ x ^ y`. Its run covers bits `low..31`, the carry out of bit 31 being no product; +/// `low` is zero but where a literal operand's low zero bits leave no carry yet, round 0's first `c + d` of a column G, whose `c` is +/// still the IV. /// /// At bit `i` of its run: /// @@ -217,18 +184,6 @@ struct Carries { pub fn blake2s_witness(inputs: &[u64], z: &mut [u64], az: &mut [u64], bz: &mut [u64]) { assert_eq!(inputs.len(), INPUT_BITS.len()); - // The carry runs, recorded by building the circuit, once. - let carries = CARRIES.get().unwrap_or_else(|| { - Hash::circuit(); - CARRIES.get().expect("building the circuit records its carry runs") - }); - - // Input ports: the word, masked to the port's width. - for (i, &bits) in INPUT_BITS.iter().enumerate() { - let wired = u64::MAX >> (64 - bits); - (z[i], az[i], bz[i]) = (inputs[i] & wired, inputs[i] & wired, wired); - } - // The working vector, as the circuit starts it. let half = |w: u64, i: usize| (w >> (32 * (i % 2))) as u32; let h: [u32; 8] = std::array::from_fn(|i| half(inputs[2 + i / 2], i)); @@ -241,34 +196,40 @@ pub fn blake2s_witness(inputs: &[u64], z: &mut [u64], az: &mut [u64], bz: &mut [ v[13] ^= (t >> 32) as u32; v[14] ^= f0; - // Each addition writes its run of carries, in the order the circuit made them. - let mut runs = carries.iter(); - let mut add = |x: u32, y: u32| -> u32 { - let Carries { slot, low } = *runs.next().expect("one run per addition"); - let sum = x.wrapping_add(y); - let carry = sum ^ x ^ y; - - // The run covers bits low..31: the carry out of bit 31 is no product. - let mask = (1u64 << (31 - low)) - 1; - let left = u64::from(x ^ carry) >> low & mask; - let right = u64::from(y ^ carry) >> low & mask; - or_run(z, slot, left & right); - or_run(az, slot, left); - or_run(bz, slot, right); - sum - }; + // The constant, right after the ports, then each addition's run of carries, in the order the circuit made them. + let n_ports = INPUT_BITS.len() + 4; + let mut rows = Products::new([&mut z[n_ports..], &mut az[n_ports..], &mut bz[n_ports..]], 0); + rows.push(1, 1, 1); // Ten rounds of eight G's, the working vector updated as the circuit updates it. - for round in &SIGMA { - for (g, &[a, b, c, d]) in G_LANES.iter().enumerate() { - for (x, r1, r2) in [(m[round[2 * g]], 16, 12), (m[round[2 * g + 1]], 8, 7)] { - let ab = add(v[a], v[b]); - v[a] = add(ab, x); - v[d] = (v[d] ^ v[a]).rotate_right(r1); - v[c] = add(v[c], v[d]); - v[b] = (v[b] ^ v[c]).rotate_right(r2); - } - } + // + // Unrolled, every run's place is a constant, so the rows are written with no branch. + macro_rules! round { + ($r:literal) => { + round!(@g $r 0); round!(@g $r 1); round!(@g $r 2); round!(@g $r 3); + round!(@g $r 4); round!(@g $r 5); round!(@g $r 6); round!(@g $r 7); + }; + (@g $r:literal $g:literal) => { + let low = if $r == 0 && $g < 4 { IV[$g].trailing_zeros() } else { 0 }; + g(&mut v, &mut rows, G_LANES[$g], [m[SIGMA[$r][2 * $g]], m[SIGMA[$r][2 * $g + 1]]], low); + }; + } + round!(0); + round!(1); + round!(2); + round!(3); + round!(4); + round!(5); + round!(6); + round!(7); + round!(8); + round!(9); + rows.finish(); + + // Input ports: the word, masked to the port's width. + for (i, &bits) in INPUT_BITS.iter().enumerate() { + let wired = u64::MAX >> (64 - bits); + (z[i], az[i], bz[i]) = (inputs[i] & wired, inputs[i] & wired, wired); } // Output ports: the new chaining value, every bit copied out. @@ -278,22 +239,31 @@ pub fn blake2s_witness(inputs: &[u64], z: &mut [u64], az: &mut [u64], bz: &mut [ let out = word(2 * i) | word(2 * i + 1) << 32; (z[n_in + i], az[n_in + i], bz[n_in + i]) = (out, out, u64::MAX); } +} - // The constant wire, right after the ports. - let one = 64 * (n_in + 4); - for buf in [z, az, bz] { - buf[one / 64] |= 1 << (one % 64); +/// One G on the working vector: its two halves, each three additions, the first `c + d`'s run from bit `low`. +#[inline(always)] +fn g(v: &mut [u32; 16], rows: &mut Products<'_>, [a, b, c, d]: [usize; 4], [x0, x1]: [u32; 2], low: u32) { + for (x, r1, r2, low) in [(x0, 16, 12, low), (x1, 8, 7, 0)] { + let ab = add(rows, v[a], v[b], 0); + v[a] = add(rows, ab, x, 0); + v[d] = (v[d] ^ v[a]).rotate_right(r1); + v[c] = add(rows, v[c], v[d], low); + v[b] = (v[b] ^ v[c]).rotate_right(r2); } } -/// OR a run of at most 32 bits into `buf`, from bit `slot`. +/// `x + y`, its run of carry rows from bit `low`: `A*z = x ^ c = sum ^ y`, `B*z = y ^ c = sum ^ x`. #[inline(always)] -const fn or_run(buf: &mut [u64], slot: u32, bits: u64) { - let (word, shift) = (slot as usize / 64, slot % 64); - buf[word] |= bits << shift; - - // The run's spill into the next word: (x >> 1) >> (63 - s) is x >> (64 - s), with no overflow at s = 0. - buf[word + 1] |= (bits >> 1) >> (63 - shift); +fn add(rows: &mut Products<'_>, x: u32, y: u32, low: u32) -> u32 { + let sum = x.wrapping_add(y); + let mask = (1u64 << (31 - low)) - 1; + rows.push( + u64::from(sum ^ y) >> low & mask, + u64::from(sum ^ x) >> low & mask, + 31 - low, + ); + sum } #[cfg(test)] @@ -334,7 +304,7 @@ mod tests { // Invariant: the word-level witness writes the tables the gate walk writes. // // The walk is the reference: it reads the circuit itself, slot by slot. - let n_log = 5; + let n_log = 8; let mut state = 0xA7u64; let mut next = || { state ^= state << 13; @@ -346,6 +316,7 @@ mod tests { // Fixture: counters at the edges of their halves, then random ones. // // Finalization words: the legal two, and any word, whose high half the port drops. + // Chaining values and messages: all zero, all ones, every half's sign bit alone or clear, then random words. let rows: Vec<[u64; 14]> = (0..1 << n_log) .map(|i| { let t = [0, u32::MAX as u64, 1 << 32, u64::MAX] @@ -353,9 +324,14 @@ mod tests { .copied() .unwrap_or_else(&mut next); let flags = [0, Hash::FINAL, next()][i % 3]; - std::array::from_fn(|k| match k { - 0 => t, - 1 => flags, + let block = i % 8; + std::array::from_fn(|k| match (k, block) { + (0, _) => t, + (1, _) => flags, + (_, 0) => 0, + (_, 1) => u64::MAX, + (_, 2) => 0x8000_0000_8000_0000, + (_, 3) => 0x7fff_ffff_7fff_ffff, _ => next(), }) }) diff --git a/crates/leanvm_core/src/rv/semantics/memory.rs b/crates/leanvm_core/src/rv/semantics/memory.rs index 8e0815e78..027100abc 100644 --- a/crates/leanvm_core/src/rv/semantics/memory.rs +++ b/crates/leanvm_core/src/rv/semantics/memory.rs @@ -3,7 +3,7 @@ //! A double word is its whole cell, so `ld` and `sd` are classes of their own, with no byte to select: [`Ld`] and [`Sd`]. use super::InstructionClass; -use crate::rv::circuits::{ClassCircuit, WordGadgets}; +use crate::rv::circuits::{ClassCircuit, Products, WordGadgets}; use crate::rv::entry::Class; use flock::circuit::{Builder, Circuit, Wire}; @@ -356,10 +356,165 @@ impl ClassCircuit for Sd { } } +/// All ones if `bit` is set, else zero. +const fn all(bit: u64) -> u64 { + (bit & 1).wrapping_neg() +} + +/// The address `v1 + imm` and the rows of its adder, which makes no carry out of the top bit: +/// `A·z = v1 ^ c`, `B·z = imm ^ c` on bits 0 to 62, `c` the carries. +fn address_rows(v1: u64, imm: u64, rows: &mut Products<'_>) -> u64 { + let address = v1.wrapping_add(imm); + let carries = address ^ v1 ^ imm; + rows.push((v1 ^ carries) & (u64::MAX >> 1), (imm ^ carries) & (u64::MAX >> 1), 63); + address +} + +/// The bus address of a load or a store of `2^log_width` bytes, `ge2` and `ge4` the width's thresholds, and its two rows. +fn bus_rows(address: u64, ge2: u64, ge4: u64, rows: &mut Products<'_>) -> u64 { + let (a0, a1) = (address & 1, address >> 1 & 1); + rows.push(a0, ge2, 1); + rows.push(a1, ge4, 1); + address & !7 | a0 & ge2 | (a1 & ge4) << 1 +} + +impl Load { + /// One instance of the circuit's witness by word arithmetic: what the walk of [`Load::circuit`] writes, into zeroed buffers. + /// + /// After the adder and the bus address, the cell's right shift by `8 * (address & 7)` bits, three byte stages of `A·z` the stage's + /// address bit and `B·z` the shifted and kept words' difference, the sign by width, the extension, and the muxes of bytes 1 to 3. + pub(crate) fn witness(inputs: &[u64], z: &mut [u64], az: &mut [u64], bz: &mut [u64]) { + let (v1, imm, flags, cell) = (inputs[0], inputs[1], inputs[2] & 7, inputs[3]); + let (ge2, ge4) = ((flags ^ flags >> 1) & 1, flags >> 1 & 1); + let mut rows = Products::new([&mut z[6..], &mut az[6..], &mut bz[6..]], 0); + rows.push(1, 1, 1); + let address = address_rows(v1, imm, &mut rows); + let bus = bus_rows(address, ge2, ge4, &mut rows); + + // The cell, right by whole bytes; the last stage makes only the low 32 bits. + let mut value = cell; + for (stage, bits) in [(0, 64), (1, 64), (2, 32)] { + let (by, s) = (8 << stage, address >> stage & 1); + let mask = u64::MAX >> (64 - bits); + let shifted = value >> by; + rows.push(all(s) & mask, (shifted ^ value) & mask, bits); + value = (if s == 1 { shifted } else { value }) & mask; + } + + let (w1, w2, w4) = (ge2 ^ 1, ge2 ^ ge4, ge4); + let mut sign = 0; + for (width, bit) in [(w1, 7), (w2, 15), (w4, 31)] { + let top = value >> bit & 1; + rows.push(width, top, 1); + sign ^= width & top; + } + let signed = flags >> 2; + rows.push(signed, sign, 1); + let extension = all(signed & sign); + rows.push(all(ge2) & 0xff, (value >> 8 ^ extension) & 0xff, 8); + rows.push(all(ge4) & 0xffff, (value >> 16 ^ extension) & 0xffff, 16); + rows.finish(); + let out = value & 0xff + | (if ge2 == 1 { value } else { extension }) & 0xff00 + | (if ge4 == 1 { value } else { extension }) & 0xffff_0000 + | extension & !0xffff_ffff; + + let ports = [ + (v1, u64::MAX), + (imm, u64::MAX), + (flags, 7), + (cell, u64::MAX), + (bus, !4), + (out, u64::MAX), + ]; + for (i, (word, wired)) in ports.into_iter().enumerate() { + (z[i], az[i], bz[i]) = (word, word, wired); + } + } +} + +impl Store { + /// One instance of the circuit's witness by word arithmetic: what the walk of [`Store::circuit`] writes, into zeroed buffers. + /// + /// After the adder and the bus address, the low half of `v2`'s left shift by `8 * (address & 7)` bits (three byte stages, + /// each making the bits a shifted or a kept bit can reach), the four ORs of the address's low bits with the thresholds, + /// then per byte whether it is written and its eight muxes. + pub(crate) fn witness(inputs: &[u64], z: &mut [u64], az: &mut [u64], bz: &mut [u64]) { + let (v1, v2, imm, flags, cell) = (inputs[0], inputs[1], inputs[2], inputs[3] & 3, inputs[4]); + let (ge2, ge4) = ((flags ^ flags >> 1) & 1, flags >> 1); + let mut rows = Products::new([&mut z[7..], &mut az[7..], &mut bz[7..]], 0); + rows.push(1, 1, 1); + let address = address_rows(v1, imm, &mut rows); + let bus = bus_rows(address, ge2, ge4, &mut rows); + + // The low half of `v2`, left by whole bytes: stage `k` makes the bits below `32 + 8 * (2^(k + 1) - 1)`, at most 64. + let mut value = v2 & 0xffff_ffff; + for (stage, bits) in [(0, 40), (1, 56), (2, 64)] { + let (by, s) = (8 << stage, address >> stage & 1); + let mask = u64::MAX >> (64 - bits); + let shifted = value << by; + rows.push(all(s) & mask, (shifted ^ value) & mask, bits); + value = if s == 1 { shifted } else { value }; + } + + // Bit `k` of a written byte's index is the address's, wherever the width does not span it. + let (a0, a1, a2) = (address & 1, address >> 1 & 1, address >> 2 & 1); + let mut spans = [[0; 2]; 3]; + for (k, (a, threshold)) in [(a0, ge2), (a1, ge4)].into_iter().enumerate() { + rows.push(a ^ 1, threshold, 1); + rows.push(a, threshold, 1); + spans[k] = [a ^ 1 | threshold, a | threshold]; + } + spans[2] = [a2 ^ 1, a2]; + let mut new = 0; + for j in 0..8 { + let low = spans[0][j & 1] & spans[1][j >> 1 & 1]; + rows.push(spans[0][j & 1], spans[1][j >> 1 & 1], 1); + rows.push(low, spans[2][j >> 2], 1); + let written = all(low & spans[2][j >> 2]); + let byte = 0xff << (8 * j); + rows.push(written & 0xff, (value ^ cell) >> (8 * j) & 0xff, 8); + new |= (written & value | !written & cell) & byte; + } + rows.finish(); + + let ports = [ + (v1, u64::MAX), + (v2, u64::MAX), + (imm, u64::MAX), + (flags, 3), + (cell, u64::MAX), + (bus, !4), + (new, u64::MAX), + ]; + for (i, (word, wired)) in ports.into_iter().enumerate() { + (z[i], az[i], bz[i]) = (word, word, wired); + } + } +} + +impl Ld { + /// One instance of the circuit's witness by word arithmetic, [`Sd`]'s too: what the walk of [`Ld::circuit`] writes, into zeroed buffers. + /// + /// The ports `v1`, `imm` and the address, the constant, then the adder's rows. + pub(crate) fn witness(inputs: &[u64], z: &mut [u64], az: &mut [u64], bz: &mut [u64]) { + let (v1, imm) = (inputs[0], inputs[1]); + let mut rows = Products::new([&mut z[3..], &mut az[3..], &mut bz[3..]], 0); + rows.push(1, 1, 1); + let address = address_rows(v1, imm, &mut rows); + rows.finish(); + for (i, word) in [v1, imm, address].into_iter().enumerate() { + (z[i], az[i], bz[i]) = (word, word, u64::MAX); + } + } +} + #[cfg(test)] mod tests { use super::*; - use crate::rv::semantics::tests::{circuit_matches_reference, edge_word, run}; + use crate::rv::semantics::tests::{ + EDGES, circuit_matches_reference, edge_word, grid, run, word_witness_is_the_walk, + }; use proptest::prelude::*; use proptest::sample::select; use proptest::strategy::BoxedStrategy; @@ -494,6 +649,32 @@ mod tests { circuit_matches_reference::(4096); } + /// Offsets that put an edge base at every byte of its cell, aligned or not, and carry through the whole adder. + const OFFSETS: [u64; 11] = [0, 1, 2, 3, 4, 5, 6, 7, u64::MAX, u64::MAX - 6, 1 << 63]; + + /// Cells and stored values: zero, all ones, every byte's sign bit set or clear, and two words with distinct bytes. + const BYTES: [u64; 6] = [ + 0, + u64::MAX, + 0x8080_8080_8080_8080, + 0x7f7f_7f7f_7f7f_7f7f, + 0x0123_4567_89ab_cdef, + 0xfedc_ba98_7654_3210, + ]; + + #[test] + fn the_word_witnesses_are_the_gate_walk() { + // Every legal flag word at every byte of a cell, signed or not, so each width's misalignments and sign bits are walked. + word_witness_is_the_walk::(Load::witness, grid(&[&EDGES, &OFFSETS, Load::LEGAL, &BYTES])); + word_witness_is_the_walk::( + Store::witness, + grid(&[&EDGES, &BYTES, &OFFSETS, Store::LEGAL, &BYTES[..3]]), + ); + let addresses = [grid(&[&EDGES, &OFFSETS]), grid(&[&EDGES, &EDGES])].concat(); + word_witness_is_the_walk::(Ld::witness, addresses.clone()); + word_witness_is_the_walk::(Ld::witness, addresses); + } + proptest! { #[test] fn a_misaligned_store_names_no_cell(store in any::(), offset in 1u64..8) { diff --git a/crates/leanvm_core/src/rv/semantics/mod.rs b/crates/leanvm_core/src/rv/semantics/mod.rs index 33b883feb..28ed09e63 100644 --- a/crates/leanvm_core/src/rv/semantics/mod.rs +++ b/crates/leanvm_core/src/rv/semantics/mod.rs @@ -77,14 +77,16 @@ const fn sext32(x: u64) -> u64 { #[cfg(test)] pub(super) mod tests { use super::*; + use crate::tables::InstanceWitness; use flock::circuit::Circuit; use proptest::prelude::*; use proptest::sample::select; + use proptest::strategy::ValueTree; use proptest::test_runner::{Config, TestRunner}; use std::fmt::Debug; /// The words arithmetic most often gets wrong: zero, one, all ones, and the sign bits. - const EDGES: [u64; 8] = [ + pub(super) const EDGES: [u64; 8] = [ 0, 1, u64::MAX, @@ -160,4 +162,65 @@ pub(super) mod tests { }) .unwrap(); } + + /// Every combination of one word per input port, the first port varying fastest. + pub(super) fn grid(ports: &[&[u64]]) -> Vec> { + ports.iter().fold(vec![vec![]], |rows, port| { + port.iter() + .flat_map(|&word| { + rows.iter().map(move |row| { + let mut row = row.clone(); + row.push(word); + row + }) + }) + .collect() + }) + } + + /// Check a class's word-level witness: on every instance, the tables it writes are the ones the walk of the circuit's gate list writes. + /// + /// The instances are `edges`, then 4096 random ones: every other one an instance the decoder can make, the rest edge-biased words in every input port, flags included. + /// Each instance is checked alone against the bit-by-bit walk, then the whole batch's packed tables and lincheck stripes against the 64-lane walk. + pub(super) fn word_witness_is_the_walk( + witness: InstanceWitness, + edges: Vec>, + ) { + let circuit = C::circuit(); + let n_in = circuit.n_input_words(); + let mut runner = TestRunner::deterministic(); + let random = (0..4096).map(|i| { + if i % 2 == 0 { + any::().new_tree(&mut runner).unwrap().current().input_words() + } else { + (0..n_in) + .map(|_| edge_word().new_tree(&mut runner).unwrap().current()) + .collect() + } + }); + let rows: Vec> = edges.into_iter().chain(random).collect(); + + let words = 1 << (circuit.k_log() - 6); + let tables = || [vec![0u64; words], vec![0; words], vec![0; words]]; + for row in &rows { + assert_eq!(row.len(), n_in, "an edge instance has the wrong number of input words"); + let [mut z, mut az, mut bz] = tables(); + circuit.witness_instance(row, &mut z, &mut az, &mut bz); + let walk = [z, az, bz]; + let [mut z, mut az, mut bz] = tables(); + witness(row, &mut z, &mut az, &mut bz); + for (name, (walk, word)) in ["z", "A·z", "B·z"].into_iter().zip(walk.iter().zip([z, az, bz])) { + let at = walk.iter().zip(&word).position(|(a, b)| a != b); + assert!(at.is_none(), "{name} word {at:?} on inputs {row:#x?}"); + } + } + + let n_log = rows.len().next_power_of_two().trailing_zeros() as usize; + let walk = circuit.generate_witness_from(&rows, &rows[0], n_log, |row, words| words.copy_from_slice(row)); + let words = circuit.generate_witness_with(&rows, &rows[0], n_log, |row, z, az, bz| witness(row, z, az, bz)); + assert!(walk.0[..] == words.0[..], "z"); + assert!(walk.1[..] == words.1[..], "A·z"); + assert!(walk.2[..] == words.2[..], "B·z"); + assert!(walk.3[..] == words.3[..], "lincheck stripes"); + } } diff --git a/crates/leanvm_core/src/rv/semantics/shift.rs b/crates/leanvm_core/src/rv/semantics/shift.rs index a5c2a4137..e9cf968f0 100644 --- a/crates/leanvm_core/src/rv/semantics/shift.rs +++ b/crates/leanvm_core/src/rv/semantics/shift.rs @@ -1,7 +1,7 @@ //! The shifter: logical and arithmetic shifts, on 64 or 32 bits. use super::{InstructionClass, sext32}; -use crate::rv::circuits::{ClassCircuit, Word, WordGadgets}; +use crate::rv::circuits::{ClassCircuit, Products, Word, WordGadgets}; use crate::rv::entry::Class; use flock::circuit::{Builder, Circuit}; @@ -119,10 +119,66 @@ impl ClassCircuit for Shift { } } +impl Shift { + /// One instance of the circuit's witness by word arithmetic: what the walk of [`Shift::circuit`] writes, into zeroed buffers. + /// + /// Every mux `s ? x : y` is a row `A·z = s`, `B·z = x ^ y`. In order: the amount's top bit, the low sign, the word's high + /// half, the fill, the reversal, the six stages, the reversal back, then the sign extension. + pub(crate) fn witness(inputs: &[u64], z: &mut [u64], az: &mut [u64], bz: &mut [u64]) { + let all = |bit: u64| (bit & 1).wrapping_neg(); + let (v1, v2, imm, flags) = (inputs[0], inputs[1], inputs[2], inputs[3] & 7); + let (right, arith, word) = (flags & 1, flags >> 1 & 1, flags >> 2); + let mut rows = Products::new([&mut z[5..], &mut az[5..], &mut bz[5..]], 0); + rows.push(1, 1, 1); + + let amount = v2 ^ imm; + rows.push(word ^ 1, amount >> 5 & 1, 1); + let amount = amount & if word == 1 { 31 } else { 63 }; + + let low_sign = arith & v1 >> 31; + rows.push(arith, v1 >> 31 & 1, 1); + let high = (all(low_sign) ^ v1) >> 32; + rows.push(all(word) >> 32, high, 32); + let x = v1 ^ (all(word) & high) << 32; + let fill = arith & x >> 63; + rows.push(arith, x >> 63, 1); + + let reverse_unless = |y: u64, rows: &mut Products<'_>| { + rows.push(all(right), y ^ y.reverse_bits(), 64); + if right == 1 { y } else { y.reverse_bits() } + }; + let mut y = reverse_unless(x, &mut rows); + for stage in 0..6 { + let by = 1 << stage; + let shifted = y >> by | all(fill) << (64 - by); + let bit = amount >> stage & 1; + rows.push(all(bit), shifted ^ y, 64); + y = if bit == 1 { shifted } else { y }; + } + let y = reverse_unless(y, &mut rows); + + let extension = (all(y >> 31) ^ y) >> 32; + rows.push(all(word) >> 32, extension, 32); + rows.finish(); + let out = y ^ (all(word) & extension) << 32; + + let ports = [ + (v1, u64::MAX), + (v2, u64::MAX), + (imm, u64::MAX), + (flags, 7), + (out, u64::MAX), + ]; + for (i, (port, wired)) in ports.into_iter().enumerate() { + (z[i], az[i], bz[i]) = (port, port, wired); + } + } +} + #[cfg(test)] mod tests { use super::*; - use crate::rv::semantics::tests::{circuit_matches_reference, edge_word}; + use crate::rv::semantics::tests::{EDGES, circuit_matches_reference, edge_word, grid, word_witness_is_the_walk}; use proptest::prelude::*; use proptest::sample::select; use proptest::strategy::BoxedStrategy; @@ -156,4 +212,13 @@ mod tests { // Legal flags and edge-biased operands pin the gate list to the reference function. circuit_matches_reference::(4096); } + + #[test] + fn the_word_witness_is_the_gate_walk() { + // Every legal flag word on edge operands and a word of distinct bits, by every amount and the next ones up, with and without high bits above the amount, which the shifter drops. + let operands: Vec = EDGES.into_iter().chain([0x0123_4567_89ab_cdef]).collect(); + let amounts: Vec = (0..=64).chain([u64::MAX]).collect(); + let edges = grid(&[&operands, &amounts, &[0, !63], Shift::LEGAL]); + word_witness_is_the_walk::(Shift::witness, edges); + } } diff --git a/crates/leanvm_core/src/tables/clock.rs b/crates/leanvm_core/src/tables/clock.rs index a6d0cb3c4..76d91a490 100644 --- a/crates/leanvm_core/src/tables/clock.rs +++ b/crates/leanvm_core/src/tables/clock.rs @@ -1,5 +1,6 @@ //! Timestamp encoding and the circuit that orders a row's memory accesses. +use crate::rv::circuits::Products; use flock::circuit::{Builder, Circuit}; /// A row timestamp, including its live bit and access slot. @@ -131,11 +132,98 @@ impl Clock { let carries = ((cycles + live) ^ cycles) & ((1 << (Self::CLOCK_BITS - Self::SLOT_BITS as usize)) - 1); carries << Self::SLOT_BITS | u64::from(fail) << Self::FAIL_BIT } + + /// One instance of the clock circuit's witness by word arithmetic: what the walk of [`Self::circuit`] on `slots` writes, into zeroed buffers. + /// + /// Its input ports are the row's clock `ts`, then each access's previous timestamp `prev`. + /// + /// ```text + /// words 0..=n inputs z = A·z = the word's 41 bits, B·z = those bits set + /// word n + 1 step bit 5 a copy of live, bits 6..41 its carries' products, bit 41 a copy of fail + /// bit 64(n + 2) constant z = A·z = B·z = 1 + /// then products per access its compare's carries, then its OR into disagree; + /// then the AND of the in-order bits, late, fail + /// ``` + /// + /// Access `i` compares by the sum `x + y`, `x = (ts ^ slot)` and `y = !prev` over 40 bits, whose carries `c = (x + y) ^ x ^ y` give each product: + /// + /// ```text + /// slot bits, above the slot's lowest one A·z = y_b B·z = c_b + /// bits 5..40 A·z = x_b ^ c_b B·z = y_b ^ c_b + /// ``` + pub fn witness(slots: &[u32], ts: u64, prev: &[u64], z: &mut [u64], az: &mut [u64], bz: &mut [u64]) { + let n = slots.len(); + assert_eq!(prev.len(), n); + const LOW: u64 = (1 << Clock::LIVE_BIT) - 1; + const READ: u64 = (1 << Clock::CLOCK_BITS) - 1; + const CYCLES: u32 = Clock::LIVE_BIT - Clock::SLOT_BITS; + let run = |bits: u32| (1u64 << bits) - 1; + + for (i, &word) in std::iter::once(&ts).chain(prev).enumerate() { + (z[i], az[i], bz[i]) = (word & READ, word & READ, READ); + } + let ts = ts & READ; + let live = ts >> Self::LIVE_BIT; + + // The constant, then the products, start word `n + 2`. + let mut products = Products::new([z, az, bz], n + 2); + products.push(1, 1, 1); + + let (mut in_order, mut disagree) = (0u32, 0); + for (i, (&prev, &slot)) in prev.iter().zip(slots).enumerate() { + let prev = prev & READ; + let x = ts & LOW & !(Self::CYCLE - 1) | u64::from(slot); + let y = !prev & LOW; + let c = (x + y) ^ x ^ y; + if slot != 0 { + let first = slot.trailing_zeros() + 1; + let bits = Self::SLOT_BITS - first; + products.push(y >> first & run(bits), c >> first & run(bits), bits); + } + let cycles = |w: u64| w >> Self::SLOT_BITS & run(CYCLES); + products.push(cycles(x ^ c), cycles(y ^ c), CYCLES); + in_order |= ((c >> Self::LIVE_BIT) as u32) << i; + + let differs = prev >> Self::LIVE_BIT ^ live; + if i > 0 { + products.push(disagree, differs, 1); + } + disagree |= differs; + } + let mut ordered = u64::from(in_order & 1); + for i in 1..n { + let next = u64::from(in_order >> i & 1); + products.push(ordered, next, 1); + ordered &= next; + } + let unordered = ordered ^ 1; + products.push(live, unordered, 1); + let late = live & unordered; + products.push(late, disagree, 1); + let fail = late | disagree; + + // The step is the carries of `cycle + live` from bit 5, and fail at bit 41. + // Bits 5 and 41 are copies (`B·z = 1`); bit `j + 1` between is the product of `ts_j` and the carry at bit `j`. + let cycle = ts >> Self::SLOT_BITS & run(CYCLES); + let carries = ((cycle + live) ^ cycle) & run(CYCLES + 1); + let step = carries << Self::SLOT_BITS | fail << Self::FAIL_BIT; + let copies = live << Self::SLOT_BITS | fail << Self::FAIL_BIT; + let ts_bits = cycle << (Self::SLOT_BITS + 1); + let carry_bits = (carries & run(CYCLES)) << (Self::SLOT_BITS + 1); + let [z, az, bz] = products.finish(); + (z[n + 1], az[n + 1], bz[n + 1]) = ( + step, + ts_bits | copies, + carry_bits | 1 << Self::SLOT_BITS | 1 << Self::FAIL_BIT, + ); + } } #[cfg(test)] mod tests { use super::*; + use crate::tables::ClassSpec; + use primitives::test_util::Rng; #[test] fn the_clock_circuit_is_its_reference() { @@ -184,4 +272,54 @@ mod tests { "an honest row advances one cycle" ); } + + #[test] + fn the_word_witness_is_the_gate_walk() { + // Invariant: the word-level clock witness writes the tables the 64-lane walk of the clock circuit's gate list writes. + // + // Fixture state: every table's clock circuit, and one with an access in each of the 32 slots. + // Rows: padding rows (clock zero), the first and the last cycle, random cycles, all ones and random words; each previous timestamp at `ts ^ slot`, around it, with the live bit flipped, a padding tuple, a seed, zero, all ones, or a random word. + let mut rng = Rng::new(0xC10C); + let n_log = 12; + let every_slot = (0..Clock::CYCLE as u32).collect(); + for slots in ClassSpec::ALL.iter().map(|spec| spec.slots()).chain([every_slot]) { + let circuit = Clock::circuit(&slots); + let rows: Vec> = (0..1 << n_log) + .map(|r| { + let ts = match r % 6 { + 0 => 0, + 1 => Clock::CLOCK_START, + 2 => Clock::SEED_CLOCK | (Clock::MAX_CYCLES * Clock::CYCLE), + 3 => Clock::SEED_CLOCK | (rng.next_u64() % Clock::MAX_CYCLES * Clock::CYCLE), + 4 => u64::MAX, + _ => rng.next_u64(), + }; + let prev = slots.iter().map(|&slot| { + let at = ts ^ u64::from(slot); + match rng.next_u64() % 10 { + 0 => at, + 1 => at.wrapping_sub(1), + 2 => at.wrapping_add(1), + 3 => at.wrapping_sub(Clock::CYCLE), + 4 => at ^ Clock::SEED_CLOCK, + 5 => u64::from(slot), + 6 => Clock::SEED_CLOCK, + 7 => 0, + 8 => u64::MAX, + _ => rng.next_u64(), + } + }); + std::iter::once(ts).chain(prev).collect() + }) + .collect(); + let walk = circuit.generate_witness_from(&rows, &rows[0], n_log, |row, words| words.copy_from_slice(row)); + let words = circuit.generate_witness_with(&rows, &rows[0], n_log, |row, z, az, bz| { + Clock::witness(&slots, row[0], &row[1..], z, az, bz); + }); + assert!(walk.0[..] == words.0[..], "z, slots {slots:?}"); + assert!(walk.1[..] == words.1[..], "A·z, slots {slots:?}"); + assert!(walk.2[..] == words.2[..], "B·z, slots {slots:?}"); + assert!(walk.3[..] == words.3[..], "lincheck stripes, slots {slots:?}"); + } + } } diff --git a/crates/leanvm_core/src/tables/spec.rs b/crates/leanvm_core/src/tables/spec.rs index 1b2f2598e..fca3dd48a 100644 --- a/crates/leanvm_core/src/tables/spec.rs +++ b/crates/leanvm_core/src/tables/spec.rs @@ -2,7 +2,7 @@ use super::Word; use super::clock::Clock; -use crate::rv::{Class, Ext, Hash, Mul, Mulh}; +use crate::rv::{Alu, Class, Ext, Hash, Ld, Load, Mul, Mulh, Shift, Store}; use crate::{class_flock, rv}; use std::ops::Range; @@ -95,7 +95,7 @@ impl ClassSpec { reads_rd: false, ram: Ram::None, copies: false, - witness: None, + witness: Some(Alu::witness), batch_witness: None, k_log: 10, ports: &[Word::V1, Word::V2, Word::Imm, Word::Flags, Word::Out, Word::Taken], @@ -113,7 +113,7 @@ impl ClassSpec { reads_rd: false, ram: Ram::Read, copies: false, - witness: None, + witness: Some(Load::witness), batch_witness: None, k_log: 10, ports: &[ @@ -138,7 +138,7 @@ impl ClassSpec { reads_rd: false, ram: Ram::Write, copies: false, - witness: None, + witness: Some(Store::witness), batch_witness: None, k_log: 10, ports: &[ @@ -164,7 +164,7 @@ impl ClassSpec { reads_rd: false, ram: Ram::Read, copies: true, - witness: None, + witness: Some(Ld::witness), batch_witness: None, k_log: 8, ports: &[Word::V1, Word::Imm, Word::Address], @@ -182,7 +182,7 @@ impl ClassSpec { reads_rd: false, ram: Ram::Write, copies: true, - witness: None, + witness: Some(Ld::witness), batch_witness: None, k_log: 8, ports: &[Word::V1, Word::Imm, Word::Address], @@ -200,7 +200,7 @@ impl ClassSpec { reads_rd: false, ram: Ram::None, copies: false, - witness: None, + witness: Some(Shift::witness), batch_witness: None, k_log: 10, ports: &[Word::V1, Word::V2, Word::Imm, Word::Flags, Word::Out], From 8afa3a58debd86cd75f7a0d32973e5e47884ce2e Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Sun, 4 Oct 2026 23:42:36 +0100 Subject: [PATCH 06/18] leanvm_core, flock: generate each packed witness's z straight into its committed column The generators now take the destination for z (`Circuit::generate_witness_{from,with,batched}_into`, the allocating forms wrapping them through `witness::with_z`), and the driver shows each instance's z to a `check` closure while its group is still in the worker's scratch. `Prepared::build` passes the class's window of the committed stack as z and checks every output port against the interpreter there, so the second pass that reread z, checked the ports and copied z into the window is gone, and `Prepared` no longer keeps its own copy of z: `Prepared::prove` reads it from the stack. Byte-identical proofs: serialized proof digests equal to the base's for leanxmss 400, fibonacci 2M and leanda 1 blob, at 1 and 32 threads. Tests: cargo test --release -p flock, the word-witness pins and leanvm_core's suite. Measured on a Ryzen 9 7950X3D, shared host at load 3 to 12, the previous commit and this one alternated. 32 threads, --repeat 3, proving: leanxmss 400 0.559, 0.564 -> 0.545, 0.545 s fibonacci 2M 0.927, 0.926 -> 0.879, 0.882 s leanda 1 blob 0.602, 0.607 -> 0.576, 0.587 s 32 threads, --tracing: leanxmss 400 Build flock witnesses 49.4, 46.6 -> 26.6, 27.7 ms Flock reductions 122, 122 -> 117, 134 ms fibonacci 2M Build flock witnesses 70.1, 69.5 -> 38.5, 39.3 ms Flock reductions 159, 166 -> 168, 164 ms leanda 1 blob Build flock witnesses 50.3, 50.7 -> 31.6, 30.2 ms Flock reductions 116, 113 -> 118, 114 ms One thread, --tracing (LEANVM_NUM_THREADS=1, taskset -c 2): leanxmss 400 Build flock witnesses 211 -> 188 ms proving 5.265 -> 5.256 s fibonacci 2M Build flock witnesses 290 -> 274 ms proving 8.910 -> 8.857 s leanda 1 blob Build flock witnesses 261 -> 242 ms proving 5.518 -> 5.512 s Peak memory: leanxmss 3.00 -> 2.75 GiB (32 threads), 2.91 -> 2.66 GiB (one); leanda 3.22 -> 2.95, 3.13 -> 2.86 GiB; fibonacci 4.70 -> 4.81, 4.67 -> 4.79 GiB. --- crates/flock/src/circuit.rs | 62 +++++++++++++++- crates/flock/src/hash.rs | 20 +++-- crates/flock/src/witness.rs | 48 +++++++++--- crates/leanvm_core/src/class_flock.rs | 103 ++++++++++++++------------ crates/leanvm_core/src/cpu/program.rs | 3 +- 5 files changed, 166 insertions(+), 70 deletions(-) diff --git a/crates/flock/src/circuit.rs b/crates/flock/src/circuit.rs index 051106743..564e6f25e 100644 --- a/crates/flock/src/circuit.rs +++ b/crates/flock/src/circuit.rs @@ -24,7 +24,9 @@ use crate::lincheck::LincheckCircuit; use crate::reduction::Block; -use crate::witness::{GroupTables, drive_witness_batched, drive_witness_groups, drive_witness_packed_and_lincheck}; +use crate::witness::{ + GroupTables, drive_witness_batched, drive_witness_groups, drive_witness_packed_and_lincheck, with_z, +}; use primitives::bits::transpose_64x64; use primitives::field::F192; @@ -331,6 +333,22 @@ impl Circuit { n_blocks_log: usize, input_words: impl Fn(&S, &mut [u64]) + Sync, ) -> (Vec, Vec, Vec, Vec) { + with_z(n_blocks_log, self.k_log, |z| { + self.generate_witness_from_into(z, rows, padding, n_blocks_log, input_words, |_, _| {}) + }) + } + + /// [`Self::generate_witness_from`] with `z` written into the caller's buffer, and `check(row, z)` shown each + /// instance's `z` words while they are in cache. Returns `(a, b, z_lincheck)`. + pub fn generate_witness_from_into( + &self, + z: &mut [u64], + rows: &[S], + padding: &S, + n_blocks_log: usize, + input_words: impl Fn(&S, &mut [u64]) + Sync, + check: impl Fn(&S, &[u64]) + Sync, + ) -> (Vec, Vec, Vec) { assert!(rows.len() <= 1 << n_blocks_log, "more rows than instances"); // A batch below 64 instances is walked in full and stored in part. let lanes = LANES.min(1 << n_blocks_log); @@ -340,6 +358,7 @@ impl Circuit { // Slot groups past the useful bits hold only zeros, so they skip the transpose. let live_words = self.useful_bits.div_ceil(64); drive_witness_groups( + z, n_blocks_log, self.k_log, lanes, @@ -390,6 +409,7 @@ impl Circuit { } } }, + |i, z| check(rows.get(i).unwrap_or(padding), z), ) } @@ -402,7 +422,25 @@ impl Circuit { n_blocks_log: usize, instance: impl Fn(&S, &mut [u64], &mut [u64], &mut [u64]) + Sync, ) -> (Vec, Vec, Vec, Vec) { - drive_witness_packed_and_lincheck(rows, Some(padding), n_blocks_log, self.k_log, instance) + with_z(n_blocks_log, self.k_log, |z| { + self.generate_witness_with_into(z, rows, padding, n_blocks_log, instance, |_, _| {}) + }) + } + + /// [`Self::generate_witness_with`] with `z` written into the caller's buffer, and `check(row, z)` shown each + /// instance's `z` words while they are in cache. Returns `(a, b, z_lincheck)`. + pub fn generate_witness_with_into( + &self, + z: &mut [u64], + rows: &[S], + padding: &S, + n_blocks_log: usize, + instance: impl Fn(&S, &mut [u64], &mut [u64], &mut [u64]) + Sync, + check: impl Fn(&S, &[u64]) + Sync, + ) -> (Vec, Vec, Vec) { + drive_witness_packed_and_lincheck(z, rows, Some(padding), n_blocks_log, self.k_log, instance, |i, z| { + check(rows.get(i).unwrap_or(padding), z); + }) } /// Build native witnesses eight instances at a time. @@ -415,8 +453,26 @@ impl Circuit { n_blocks_log: usize, batch: impl Fn([&S; 8], &mut [u64], &mut [u64], &mut [u64]) + Sync, ) -> (Vec, Vec, Vec, Vec) { + with_z(n_blocks_log, self.k_log, |z| { + self.generate_witness_batched_into(z, rows, padding, n_blocks_log, batch, |_, _| {}) + }) + } + + /// [`Self::generate_witness_batched`] with `z` written into the caller's buffer, and `check(row, z)` shown each + /// instance's `z` words while they are in cache. Returns `(a, b, z_lincheck)`. + pub fn generate_witness_batched_into( + &self, + z: &mut [u64], + rows: &[S], + padding: &S, + n_blocks_log: usize, + batch: impl Fn([&S; 8], &mut [u64], &mut [u64], &mut [u64]) + Sync, + check: impl Fn(&S, &[u64]) + Sync, + ) -> (Vec, Vec, Vec) { // Eight adjacent instances occupy one lincheck byte stripe. - drive_witness_batched(rows, padding, n_blocks_log, self.k_log, batch) + drive_witness_batched(z, rows, padding, n_blocks_log, self.k_log, batch, |i, z| { + check(rows.get(i).unwrap_or(padding), z); + }) } /// The matrix-vector products `(A_0 w, B_0 w)`, by one forward walk. diff --git a/crates/flock/src/hash.rs b/crates/flock/src/hash.rs index 095f8a389..5abf6840d 100644 --- a/crates/flock/src/hash.rs +++ b/crates/flock/src/hash.rs @@ -91,7 +91,7 @@ use crate::reduction; use crate::reduction::Block; use crate::verifier::FlockError; use crate::witness::{ - BitRecord, add_carry_parts, add3_fused_parts, drive_witness_packed_and_lincheck, or_bit_at, + BitRecord, add_carry_parts, add3_fused_parts, drive_witness_packed_and_lincheck, or_bit_at, with_z, write_lin_word_ab_packed, }; use fiat_shamir::transcript::{ProverState, VerifierState}; @@ -669,13 +669,17 @@ pub fn generate_witness_with_ab_packed_and_lincheck( n_blocks_log: usize, ) -> (Vec, Vec, Vec, Vec) { let padding = padding_block(); - drive_witness_packed_and_lincheck( - blocks, - Some(&padding), - n_blocks_log, - K_LOG, - |&(ref h, ref m, t, f0, f1), z, a, b| build_block_witness_ab_packed_into(h, m, t, f0, f1, z, a, b), - ) + with_z(n_blocks_log, K_LOG, |z| { + drive_witness_packed_and_lincheck( + z, + blocks, + Some(&padding), + n_blocks_log, + K_LOG, + |&(ref h, ref m, t, f0, f1), z, a, b| build_block_witness_ab_packed_into(h, m, t, f0, f1, z, a, b), + |_, _| {}, + ) + }) } // --------------------------------------------------------------------------- diff --git a/crates/flock/src/witness.rs b/crates/flock/src/witness.rs index e2b67fa0f..9be2c2441 100644 --- a/crates/flock/src/witness.rs +++ b/crates/flock/src/witness.rs @@ -136,21 +136,24 @@ pub(crate) struct GroupTables<'a> { pub stripes: &'a mut [u8], } -/// The four witness tables of `2^n_blocks_log` instances, built `group` instances at a time. +/// The four witness tables of `2^n_blocks_log` instances, built `group` instances at a time, `z` into the caller's buffer. /// /// - The fill closure writes every word and byte of the group starting at the instance it is given. /// - Each worker keeps one scratch state, built once and reused across its groups. +/// - `check(i, z)` sees instance `i`'s `z` words while its group is still in cache. /// /// A group builds in its worker's buffers, which stay in cache, then streams out. /// /// Building in place instead would fetch every output line before writing it. pub(crate) fn drive_witness_groups( + z: &mut [u64], n_blocks_log: usize, k_log: usize, group: usize, init: I, fill: F, -) -> (Vec, Vec, Vec, Vec) + check: impl Fn(usize, &[u64]) + Sync, +) -> (Vec, Vec, Vec) where St: Send, I: Fn() -> St + Sync, @@ -168,10 +171,10 @@ where ); let total_words = n_total * (k / 64); + assert_eq!(z.len(), total_words, "z holds every instance's words"); // SAFETY: group `g` publishes chunk `g` of every table in full below, and the chunk counts match. - let (mut z, mut a, mut b, mut z_lincheck) = unsafe { + let (mut a, mut b, mut z_lincheck) = unsafe { ( - primitives::uninit_vec::(total_words), primitives::uninit_vec::(total_words), primitives::uninit_vec::(total_words), primitives::uninit_vec::((n_total / 8) * k), @@ -181,7 +184,7 @@ where // A group's share: its packed words in each table, and one stripe per 8 instances. let group_words = group * (k / 64); let group_bytes = (group / 8) * k; - let z_chunks = Chunks::new(&mut z, group_words); + let z_chunks = Chunks::new(z, group_words); let a_chunks = Chunks::new(&mut a, group_words); let b_chunks = Chunks::new(&mut b, group_words); let stripe_chunks = Chunks::new(&mut z_lincheck, group_bytes); @@ -205,6 +208,9 @@ where stripes, }, ); + for (i, z) in z_grp.chunks_exact(k / 64).enumerate() { + check(g * group + i, z); + } let stream = Stream::new(); // SAFETY: each group `g` takes chunk `g` of each table exactly once, and @@ -219,13 +225,25 @@ where |(), ()| (), ); + (a, b, z_lincheck) +} + +/// The four tables of `2^n_blocks_log` instances of `2^k_log` bits, `z` allocated here, from a driver that writes the other three. +pub(crate) fn with_z( + n_blocks_log: usize, + k_log: usize, + drive: impl FnOnce(&mut [u64]) -> (Vec, Vec, Vec), +) -> (Vec, Vec, Vec, Vec) { + // SAFETY: every driver publishes all of `z`. + let mut z = unsafe { primitives::uninit_vec::((1usize << n_blocks_log) * ((1usize << k_log) / 64)) }; + let (a, b, z_lincheck) = drive(&mut z); (z, a, b, z_lincheck) } /// Drive the parallel chunked witness build for `n_blocks` instances padded -/// to `2^n_blocks_log` slots, one instance at a time. Returns `(z, a, b, z_lincheck)`: -/// the three bit-packed `u64` tables (`K / 64` words per instance) and the lincheck -/// byte stripe. +/// to `2^n_blocks_log` slots, one instance at a time, `z` into the caller's buffer. +/// Returns `(a, b, z_lincheck)`: the bit-packed `u64` tables (`K / 64` words per instance) +/// and the lincheck byte stripe. /// /// `per_block(initial, z_u64, a_u64, b_u64)` populates one block's worth of /// `(z, a, b)` data: 3 zero-initialized `u64`-buffers of length `K / 64`. @@ -239,12 +257,14 @@ where /// that pin a constant wire need this so the constant column is all-ones /// across *every* batched instance (see `lincheck's `LincheckCircuit::const_pin_col``). pub(crate) fn drive_witness_packed_and_lincheck( + z: &mut [u64], initial_states: &[S], padding: Option<&S>, n_blocks_log: usize, k_log: usize, per_block: F, -) -> (Vec, Vec, Vec, Vec) + check: impl Fn(usize, &[u64]) + Sync, +) -> (Vec, Vec, Vec) where F: Fn(&S, &mut [u64], &mut [u64], &mut [u64]) + Sync, { @@ -257,6 +277,7 @@ where // Eight blocks per group, the lincheck stripe of one group being their bit transpose. drive_witness_groups( + z, n_blocks_log, k_log, 8, @@ -285,20 +306,24 @@ where bit_transpose_64bytes(rows.as_flattened().try_into().expect("64 bytes"), out); } }, + check, ) } -/// Build native witnesses eight instances at a time, then pack their byte stripe. +/// Build native witnesses eight instances at a time, then pack their byte stripe; `z` into the caller's buffer. pub(crate) fn drive_witness_batched( + z: &mut [u64], rows: &[S], padding: &S, n_blocks_log: usize, k_log: usize, batch: impl Fn([&S; 8], &mut [u64], &mut [u64], &mut [u64]) + Sync, -) -> (Vec, Vec, Vec, Vec) { + check: impl Fn(usize, &[u64]) + Sync, +) -> (Vec, Vec, Vec) { assert!(rows.len() <= 1 << n_blocks_log, "more rows than instances"); let words = (1usize << k_log) / 64; drive_witness_groups( + z, n_blocks_log, k_log, 8, @@ -317,5 +342,6 @@ pub(crate) fn drive_witness_batched( bit_transpose_64bytes(bits.as_flattened().try_into().expect("eight word lanes"), out); } }, + check, ) } diff --git a/crates/leanvm_core/src/class_flock.rs b/crates/leanvm_core/src/class_flock.rs index fb09b02fd..d44f3765c 100644 --- a/crates/leanvm_core/src/class_flock.rs +++ b/crates/leanvm_core/src/class_flock.rs @@ -136,19 +136,18 @@ pub const fn n_blocks_log(spec: &ClassSpec, n_rows: usize) -> usize { if natural > floor { natural } else { floor } } -/// The flock-native tables of one class's batch, kept from the pass that wrote its -/// committed column so the reduction needs no second witness pass. +/// The flock-native tables of one class's batch but `z`, which is its committed column, +/// kept from the pass that wrote that column so the reduction needs no second witness pass. pub(crate) struct Prepared { flock: usize, n_blocks_log: usize, - z: Vec, a: Vec, b: Vec, z_lincheck: Vec, } impl Prepared { - /// Build packed witness `f`'s batch, one instance per row of its table, and write it + /// Build packed witness `f`'s batch, one instance per row of its table, its `z` straight /// into `window`, its committed column. pub(crate) fn build(f: usize, rows: &[Row], entries: &[Entry], window: &mut [F64]) -> Self { let (t, part) = flock(f); @@ -169,70 +168,80 @@ impl Prepared { *word = port.value(row, &entries[row.index as usize], &slots); } }; + // What the circuit computed is what the interpreter did, or the bus would carry one and flock prove the other. + let check = |row: &Row, z: &[u64]| { + for (k, &port) in ports.iter().enumerate().skip(n_inputs) { + let expected = port.value(row, &entries[row.index as usize], &slots); + assert_eq!( + z[k], expected, + "{}'s {part:?} circuit disagrees with the interpreter on {port:?}", + spec.name + ); + } + }; + // SAFETY: `F64` is `repr(transparent)` over `u64`, and the packing is bit `i` at position `i` on both sides. + let z = unsafe { std::slice::from_raw_parts_mut(window.as_mut_ptr().cast::(), window.len()) }; // A clock circuit, and a class with a word-level witness, skip the walk of the gate list; // the others walk it 64 instances at a time. - let (z, a, b, z_lincheck) = match (part, spec.batch_witness, spec.witness) { - (Part::Clock, _, _) => circuit.generate_witness_with(rows, &rows[0], n_blocks_log, |row, z, az, bz| { - Clock::witness(&slots, row.ts, &row.prev()[..slots.len()], z, az, bz); - }), + let (a, b, z_lincheck) = match (part, spec.batch_witness, spec.witness) { + (Part::Clock, _, _) => circuit.generate_witness_with_into( + z, + rows, + &rows[0], + n_blocks_log, + |row, z, az, bz| Clock::witness(&slots, row.ts, &row.prev()[..slots.len()], z, az, bz), + check, + ), (Part::Class, Some(batch), _) => { // Eight rows share a native arithmetic call before their byte stripe is packed. - circuit.generate_witness_batched(rows, &rows[0], n_blocks_log, |rows, z, az, bz| { - let mut words = [[0u64; MAX_INPUT_WORDS]; 8]; - for (row, words) in rows.into_iter().zip(&mut words) { - input_words(row, &mut words[..n_inputs]); - } - let inputs = std::array::from_fn(|i| &words[i][..n_inputs]); - batch(&inputs, z, az, bz); - }) + circuit.generate_witness_batched_into( + z, + rows, + &rows[0], + n_blocks_log, + |rows, z, az, bz| { + let mut words = [[0u64; MAX_INPUT_WORDS]; 8]; + for (row, words) in rows.into_iter().zip(&mut words) { + input_words(row, &mut words[..n_inputs]); + } + let inputs = std::array::from_fn(|i| &words[i][..n_inputs]); + batch(&inputs, z, az, bz); + }, + check, + ) } - (Part::Class, None, Some(witness)) => { - circuit.generate_witness_with(rows, &rows[0], n_blocks_log, |row, z, az, bz| { + (Part::Class, None, Some(witness)) => circuit.generate_witness_with_into( + z, + rows, + &rows[0], + n_blocks_log, + |row, z, az, bz| { let mut words = [0u64; MAX_INPUT_WORDS]; let words = &mut words[..n_inputs]; input_words(row, words); witness(words, z, az, bz); - }) + }, + check, + ), + (Part::Class, None, None) => { + circuit.generate_witness_from_into(z, rows, &rows[0], n_blocks_log, input_words, check) } - (Part::Class, None, None) => circuit.generate_witness_from(rows, &rows[0], n_blocks_log, input_words), }; - assert_eq!(window.len(), z.len(), "the committed column is the wrong size"); - let stride = 1 << stride_log(spec, part); - // `F64` is `repr(transparent)` over `u64`, and the packing is bit `i` at - // position `i` on both sides. - const BATCH: usize = 1 << 10; - parallel::chunks_mut_zip(window, &z, stride * BATCH, |batch, dst, src| { - for (j, src) in src.chunks_exact(stride).enumerate() { - // What the circuit computed is what the interpreter did, or the bus - // would carry one and flock prove the other. - let row = &rows[batch * BATCH + j]; - for (k, &port) in ports.iter().enumerate().skip(n_inputs) { - let expected = port.value(row, &entries[row.index as usize], &slots); - assert_eq!( - src[k], expected, - "{}'s {part:?} circuit disagrees with the interpreter on {port:?}", - spec.name - ); - } - } - for (d, &s) in dst.iter_mut().zip(src) { - *d = F64(s); - } - }); Self { flock: f, n_blocks_log, - z, a, b, z_lincheck, } } - /// Flock's zerocheck then lincheck, leaving the one claim on the committed column. - pub(crate) fn prove(&self, ps: &mut ProverState) -> SliceClaim { + /// Flock's zerocheck then lincheck on `z`, the committed column, leaving the one claim on it. + pub(crate) fn prove(&self, z: &[F64], ps: &mut ProverState) -> SliceClaim { + // SAFETY: `F64` is `repr(transparent)` over `u64`. + let z = unsafe { std::slice::from_raw_parts(z.as_ptr().cast::(), z.len()) }; let block = circuit(self.flock).block(); - let stage = block.prove_zerocheck(self.n_blocks_log, &self.z, &self.a, &self.b, ps); + let stage = block.prove_zerocheck(self.n_blocks_log, z, &self.a, &self.b, ps); block.prove_lincheck(self.n_blocks_log, stage, &self.z_lincheck, ps) } } diff --git a/crates/leanvm_core/src/cpu/program.rs b/crates/leanvm_core/src/cpu/program.rs index 6ab8c216f..2b199a7d4 100644 --- a/crates/leanvm_core/src/cpu/program.rs +++ b/crates/leanvm_core/src/cpu/program.rs @@ -268,7 +268,8 @@ impl Program { .enumerate() .map(|(f, prepared)| { let window = w.layout.witness_window(f); - let reduced = prepared.prove(&mut ps); + let z = &w.q[window.offset..window.offset + (1 << window.n_vars)]; + let reduced = prepared.prove(z, &mut ps); flock::reduction::ring_switch_open(window.n_vars, window.offset, &reduced) }) .collect() From b084b241d002055071c133366bfd47261bc98a82 Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Mon, 5 Oct 2026 05:38:25 +0100 Subject: [PATCH 07/18] pcs: the wide commit NTT plan runs for up to two workers The plan with wide gathered groups (up to 64 adjacent residues), the L3 deep pass and the lane-major stack read in place now runs whenever the transform is planned for at most FEW_WORKERS = 2 workers, not one alone. The tests plan for 1, 2 and 3 workers, so both sides of the gate are pinned to the scalar reference whatever the pool's size. Gate, measured on a Ryzen 9 7950X3D (workers pinned to the 32 MiB L3 CCD), Commit span of the base encode, one-worker gate / the plan at every count, two runs each: 2 workers (2 cores): leanxmss 400 307, 287 -> 281, 245 ms; leanda 305, 329 -> 263, 255 ms; fibonacci 2M 549 -> 467, 511 ms 4 workers (2 cores + SMT): leanxmss 220, 216 -> 232, 234 ms; fibonacci 406, 408 -> 455, 447 ms; leanda 235, 235 -> 238, 264 ms 4 workers (4 cores): leanxmss 148, 147 -> 164, 180 ms; leanda 157, 182 -> 193, 166 ms 8 workers (8 cores): leanxmss 92.5, 87.2 -> 132, 133 ms; leanda 116, 128 -> 141, 141 ms 32 workers, pcs throughput bench PCS_LOG_N=24 rate 1: commit 24.4, 24.0 -> 30.8, 30.9 ms; encode alone 16.5, 16.3 -> 24.6, 24.6 ms Three workers were not measured and keep the narrow plan. This branch against upstream 77077c6b, Commit span (base encode), runs base -> this: 1 worker: leanxmss 601, 613, 613 -> 533, 472, 458 ms; fibonacci 1160, 1150, 1150 -> 846, 836, 850 ms; leanda 646, 663, 648 -> 484, 493, 497 ms 2 workers: leanxmss 279, 312, 312 -> 326, 260, 262 ms; fibonacci 550, 540, 540 -> 452, 461, 457 ms; leanda 335, 334, 336 -> 266, 281, 276 ms 4 workers (2 cores + SMT, the narrow plan): leanxmss 260, 259 -> 238, 238 ms; fibonacci 462, 473 -> 434, 429 ms; leanda 286, 257 -> 262, 260 ms 8 workers: leanxmss 92.7, 91.0 -> 90.3, 90.7 ms; fibonacci 171, 171 -> 163, 164 ms; leanda 98.7, 98.4 -> 99.2, 92.4 ms 32 workers: leanxmss 67.0, 67.3 -> 65.6, 64.9 ms; fibonacci 125, 120 -> 121, 119 ms; leanda 69.9, 71.3 -> 69.2, 69.9 ms Byte-identical proofs: leanxmss 400, fibonacci 2M and leanda 1 hash to the same digest on base and branch at every worker count. cargo test --release -p pcs passes, natively (AVX-512), with target-cpu=znver3 (the AVX2 arm) and on aarch64 under qemu. --- crates/pcs/src/ntt/additive_ntt_f64.rs | 47 ++++++++++++++++---------- 1 file changed, 29 insertions(+), 18 deletions(-) diff --git a/crates/pcs/src/ntt/additive_ntt_f64.rs b/crates/pcs/src/ntt/additive_ntt_f64.rs index a73f771a7..9cafa8390 100644 --- a/crates/pcs/src/ntt/additive_ntt_f64.rs +++ b/crates/pcs/src/ntt/additive_ntt_f64.rs @@ -291,12 +291,12 @@ impl AdditiveNttF64 { /// - Each deep sub-block is final, and hands its rows to `on_rows` while they are in L2. /// - Without a split deep pass, the rows go over in parallel blocks at the end. /// - /// # One worker + /// # Few workers /// - /// - A lone worker has its core's whole L3, so a gathered group takes up to `2^MAX_LOG_RESIDUES` residues. + /// - Up to [`FEW_WORKERS`] workers, a gathered group takes up to `2^MAX_LOG_RESIDUES` residues. /// - Their rows are adjacent, so each copy in or out is one long run, and a row group has that many times /// the lanes. - /// - Several workers keep one residue a group: their scratch then stays in their share of L2. + /// - More workers keep one residue a group: their scratch then stays in their share of L2. fn transform( &self, data: &mut [F64], @@ -353,20 +353,21 @@ impl AdditiveNttF64 { // // 20 layers: 1 gathered sweep either way, so the deep pass stays in L2 // - // A lone worker spills to an L3 it shares with no one, so it takes the L3 deep pass whenever it has one: + // Few workers share L3 with few others, so they take the L3 deep pass whenever there is one: // the layers it moves out of the gathered pass shrink the gathered groups, and their scattered reads. + let few = workers <= FEW_WORKERS; let gathered_sweeps = |deep: usize| (log_d - start).saturating_sub(deep).div_ceil(fit2); - let deep = if workers == 1 || gathered_sweeps(fit3) < gathered_sweeps(fit2) { + let deep = if few || gathered_sweeps(fit3) < gathered_sweeps(fit2) { fit3 } else { fit2 }; let deep_start = log_d.saturating_sub(deep).max(par_log).max(start); - // Residues a gathered group takes, as a log: several only for a lone worker, within its L3 budget. + // Residues a gathered group takes, as a log: several only for few workers, within their L3 budget. let residues = |layer: usize, g: usize| { - if workers == 1 { - let room = (LONE_GATHER_WORDS / (num_ntts << g)).max(1).ilog2() as usize; + if few { + let room = (WIDE_GATHER_WORDS / (num_ntts << g)).max(1).ilog2() as usize; room.min(MAX_LOG_RESIDUES).min(log_d - layer - g) } else { 0 @@ -999,14 +1000,24 @@ const LOG_SUBS_PER_WORKER: usize = 2; /// - That costs extra L3 traffic, but saves a whole sweep of DRAM. const L3_WORDS: usize = 1 << 18; -/// Words a gathered task may span when the pool has one worker. +/// Most workers whose transform takes wide gathered groups and the L3 deep pass. /// /// # Why this value /// -/// - 2^22 words is 32 MiB, a whole L3 without stacked cache: a lone worker shares its L3 with no one. -/// - A gathered row group's layers cost no more from L3 than from L2, the butterflies being the bound. +/// - A wide group spills from L2 into L3, and in return copies long runs and reads a lane-major message in place. +/// - With one or two workers the runs save more than the spill costs. +/// - From four workers on, SMT siblings or not, their spills share L3 and memory with more of them, and the narrow +/// plan is faster. +const FEW_WORKERS: usize = 2; + +/// Words a gathered task may span when few workers run. +/// +/// # Why this value +/// +/// - 2^22 words is 32 MiB, a whole L3 without stacked cache. +/// - With few workers, a gathered row group's layers cost no more from L3 than from L2, the butterflies being the bound. /// - So the group can take several residues, and with them longer runs in and out of memory. -const LONE_GATHER_WORDS: usize = 1 << 22; +const WIDE_GATHER_WORDS: usize = 1 << 22; /// Most residues one gathered group takes, as a log. /// @@ -1666,7 +1677,7 @@ mod tests { // (12, 2048, 1) two gathered passes, with streaming stores // // A non-zero start is the commit path, which enters at the rate layer. - // A lone worker's plan takes several residues per gathered group. + // Up to `FEW_WORKERS` workers, the plan takes several residues per gathered group. let mut rng = Rng::new(0xC0FFEE); for (log_d, lanes, start_layer) in [ (7usize, 3usize, 0usize), @@ -1683,8 +1694,8 @@ mod tests { // Reference: one butterfly at a time. let mut want = original.clone(); forward_scalar_from_layer(&ntt, &mut want, lanes, start_layer); - // Under test: the pass-planning driver, for the pool and for one worker. - for workers in [parallel::num_threads(), 1] { + // Under test: the pass-planning driver, planned for workers on both sides of `FEW_WORKERS`. + for workers in [1, FEW_WORKERS, FEW_WORKERS + 1] { let mut got = original.clone(); ntt.transform_for(&mut got, lanes, start_layer, None, None, workers); assert_eq!( @@ -1731,7 +1742,7 @@ mod tests { forward_scalar_from_layer(&ntt, &mut want, lanes, log_inv_rate); // Under test: only the first replica holds the message, the rest is zero. - for workers in [parallel::num_threads(), 1] { + for workers in [1, FEW_WORKERS, FEW_WORKERS + 1] { let mut got = vec![F64::ZERO; msg_len << log_inv_rate]; got[..msg_len].copy_from_slice(&msg); let blocks = Mutex::new(Vec::new()); @@ -1763,7 +1774,7 @@ mod tests { /// Every lane of the codeword must be exactly the single-lane RS codeword of /// that lane's contiguous message block, which is what makes a commitment over /// `n_lanes` lanes equal to the `2^log_batch_size`-lane one with a zero tail. - /// The shapes cover the transposing fallback, the lone worker's first pass reading + /// The shapes cover the transposing fallback, the few-workers first pass reading /// the lanes in place (the last three, with and without lanes past a whole eight), /// and lane counts that are not powers of two (the padding-free commit's whole point). #[test] @@ -1796,7 +1807,7 @@ mod tests { want }) .collect(); - for workers in [parallel::num_threads(), 1] { + for workers in [1, FEW_WORKERS, FEW_WORKERS + 1] { let mut got = vec![F64::ZERO; msg.len() << log_inv_rate]; let lanes = Message::Lanes(SendPtr(msg.as_ptr().cast_mut())); ntt.transform_for(&mut got, n_lanes, log_inv_rate, Some(lanes), Some(&|_, _| {}), workers); From f5739f4b38154e08581fec2616d80422fa7ce64e Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Mon, 5 Oct 2026 20:13:05 +0100 Subject: [PATCH 08/18] leanvm_core: word-level witness for EXT's clock circuit, its flag bits and limb addresses --- AGENTS.md | 2 +- crates/leanvm_core/src/class_flock.rs | 23 +++-- crates/leanvm_core/src/rv/semantics/ext.rs | 58 ++++++++++++ crates/leanvm_core/src/tables/clock.rs | 104 +++++++++++++++------ crates/leanvm_core/src/tables/spec.rs | 10 ++ 5 files changed, 160 insertions(+), 37 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index fa2f89688..c95fa730d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -135,7 +135,7 @@ Addition with carries, comparisons, shifts, AND/OR, multiplication and division - **Circuits are gate lists over word ports** (`flock::circuit`): inputs, outputs, the constant, then products in the order they are made. A port bit with no gate is an empty row, hence zero, which is what makes a one-bit output such as `taken` a 0 or 1 field element: give such an output a word to itself and never put a free wire on its spare bits. What Rust and Python must agree on is the port layout and the ORDER PRODUCTS ARE MADE IN; XOR order is free. `alu_is_its_reference` pins a circuit to its reference function, and the end-to-end tests pin the Python mirror. - **One batched reduction for all circuits, one opening for all**: one zerocheck and one lincheck over every packed witness under shared challenges (`flock::reduction::{prove, verify_deferred}`, `doc/leanvm` §flock:batch), every class circuit in table order then every clock circuit, after the exit claims, each circuit leaving a claim on its own witness. Both sumchecks are front-loaded like the table sumcheck: every circuit takes part from the first round and is done after its own rounds (the zerocheck binding each circuit's lowest variable first, lincheck its top inner coordinate, as for one circuit), and a done circuit adds only to the coefficient the round's claim fixes. Circuit `f`'s zerocheck takes `λ^f` and its lincheck `α^{4f}`; the zerocheck sends every circuit's `v_a, v_b, v_c` and checks their `λ`-combination. Lincheck sends each circuit's 64 slices, then the value of its matrices' form, and checks the batch's final claim against those values and each circuit's closed terms (pin and `C`); each value is left as that circuit's deferred matrix claim, so the verifier's core builds no circuit. A batch of one is how the recursion machine proves its hash rows. The zerocheck prover is a per-circuit state machine (`zerocheck::CircuitProver`) stepped one round at a time, so all twenty-one circuits' folded tables are alive at once. `pcs::stack_open` takes one ring-switched region per witness, then one for the bytecode's multiplicity column, and ring-switches all their claims as one family under one map challenge; claims whose points are prefixes of one another (the circuits of one block size) share one pass of the Frobenius products (`stack_open::PrefixGroup`, also the rows' grouping, by wire), each region closing once. - **Batch floors.** Flock needs eight instances and a zerocheck cube of `2^13` bits for each of a table's circuits (`class_flock::n_blocks_log`); padding rows supply them, as honest instances on zero registers. -- **Witness generation is word arithmetic wherever a circuit has a generator, and the generic walk of the gate list, bit by bit, elsewhere**: ALU, LOAD, STORE, LD, SD, SHIFT, MULH and HASH set `ClassSpec::witness`, MUL sets `batch_witness` on AVX2 only, and every clock circuit's generator is `Clock::witness` but EXT's, which has ports of its own; DIV, MUL off AVX2 and EXT's clock circuit walk. Each generator computes the circuit's values natively and writes its rows in the order the circuit makes its products (all but the multipliers' through `rv::circuits::Products`), so a change to a circuit's gate list must change its generator too. Each is pinned to the walk by a test that covers edge words (zero, all ones, the sign bits, every shift amount, every misalignment) and random instances: `semantics::tests::word_witness_is_the_walk` for ALU, the memory classes and SHIFT, and their own for MUL, MULH, HASH and the clock. A new generator needs such a test. HASH's generator is unrolled, so every carry run's place is a constant. +- **Witness generation is word arithmetic wherever a circuit has a generator, and the generic walk of the gate list, bit by bit, elsewhere**: ALU, LOAD, STORE, LD, SD, SHIFT, MULH and HASH set `ClassSpec::witness`, MUL sets `batch_witness` on AVX2 only, and every clock circuit has one (`ClassSpec::clock_witness`: `Clock::witness`, and for EXT's, whose ports past the clock's split the flags and add 8 and 16 to the pointers, `Ext::clock_witness` on top of `Clock::witness_with`); DIV and MUL off AVX2 walk. Each generator computes the circuit's values natively and writes its rows in the order the circuit makes its products (all but the multipliers' through `rv::circuits::Products`), so a change to a circuit's gate list must change its generator too. Each is pinned to the walk by a test that covers edge words (zero, all ones, the sign bits, every shift amount, every misalignment) and random instances: `semantics::tests::word_witness_is_the_walk` for ALU, the memory classes and SHIFT, and their own for MUL, MULH, HASH and the clock circuits (`clock::tests::the_word_witness_is_the_gate_walk`, every table's, EXT's included). A new generator needs such a test. HASH's generator is unrolled, so every carry run's place is a constant. - **Circuit sizes are structure, not measurements**: `k_log` and `clock_k_log` per class are pinned in its `ClassSpec` and asserted against the built circuits; product counts are in `doc/leanvm` Annex C. A product can be committed straight at an output port's position (`flock::circuit::Builder::and_output`), which saves the copy; the clock circuit's step uses it. CI compares every PR with its base, never with a history. The base is the merge commit's first parent; a workflow builds and runs both on one runner and saves both as an artifact, and a `workflow_run` workflow on the default branch (`main`, the only branch GitHub runs those from) posts them on the PR as one comment, edited on every run and deleted once nothing changed. The posting workflows never run the PR's code: they take only the PR number from a run's artifact, a fork's included, post nothing unless GitHub's API says that PR's head is the commit the run measured, and check every name and number before it reaches the comment. **Exact counts**: `counts.yml` runs `cargo leanvm bench --cycles-only`, every program at the sizes above without a proof (`cycles`, the program's own instructions; `proven-rows`, table heights after padding; `committed`, witness words) and the two aggregation trees above, one case per tree and kind of node (`aggregate-leanxmss-400-2to1-first` and `-node`, `aggregate-leanxmss-400-4to1-first` and `-node`, from `counted_trees()`: `rows`, the circuit's own; `proven-rows` and `committed` as for a program), puts the tables in each run's summary from the same pass (`--markdown-file "$GITHUB_STEP_SUMMARY"`, which appends what `--markdown` prints), saves the artifact `counts` and fails on any increase; `counts-comment.yml` posts the counts that changed. A tree's name says its shape, `Nto1` for N proofs combined into 1 at every node; `-first` is a first-level node (the RISC-V verifier in rows over N leaf proofs) and `-node` a higher node (the recursion verifier in rows over N child proofs). **Proving**: `bench.yml` runs one job per testbed and benchmark: each case of `proven()` in `bins/leanvm/src/tracked.rs` (at the sizes a 16 GB runner holds: leanXMSS and leanSPHINCS at a quarter, no leanDA; leanXMSS also on one thread, `leanxmss-100-1thread`, which the CLI proves in a process of its own with `LEANVM_NUM_THREADS=1` since the pool is sized once per process), each of `proven_trees()` (the two aggregation trees above over leanXMSS at a quarter, `aggregate-leanxmss-100-2to1` and `aggregate-leanxmss-100-4to1`: the leaf proven once, then one first-level node and one node over copies of it, the two benchmarks `-first` and `-node`, each node's stages the children of a `Prove` span around it and the node's peak memory including the first-level node's; the 2-to-1 tree also on one thread, `aggregate-leanxmss-100-2to1-1thread`, proven whole, leaf included, in a process of its own like `leanxmss-100-1thread`), and each `benches/` target above with every benchmark it prints (flock's two at `FLOCK_N_LOG=18`, pcs's `throughput` at its default size). Each job builds both sides and runs them in turns, base, PR, PR, base, for `ROUNDS` rounds: `cargo leanvm bench --only ` gives the proving time, the proof's size in bytes (exact: the inputs are seeded and grinding takes the smallest nonce), the time of each of the proof's stages (`stage.`, the `Prove` span's direct children, `--tracing`'s first level, timed by a `tracing` layer that enables spans only) and the process's peak resident memory (`peak-memory`, that case's alone since each job proves one), and the verifying time, and a target's `--json` its own measures: `latency` for a whole operation (flock's proving times with the witness excluded, pcs's commit and open), `per-op` in nanoseconds per call for `kernels`, a hashed block of `hash_throughput` and one NTT encode. `bench-comment.yml` posts a size that changed and a time that moved the same way in every round with a median ratio at least 1% from one, every result collapsed beneath; it gates nothing. Run in turns on one runner, the comparison does not depend on which machine the job landed on, which a history across GitHub's runners does (one runner label spans CPU models of different speeds). The testbeds are GitHub-hosted: `x86-64` (built for `target-cpu=haswell`) and `arm64` (`ubuntu-24.04-arm`, `target-cpu=neoverse-n2` plus `+aes,+sha3`, since the NEON field kernels are chosen at compile time on those features and a generic aarch64 target enables neither; the job fails early if the CPU lacks them), each ISA pinned because a cached native build dies of SIGILL on an older CPU. A runner has 4 vCPUs, so a change to how the prover scales across many cores shows there only in part. A benchmark's name is what base and PR are matched by, so renaming one or changing its input shows it as new, as does one the base cannot run (it cannot build the target, or its first base run fails or prints no JSON, and the job carries on with the PR's runs alone). diff --git a/crates/leanvm_core/src/class_flock.rs b/crates/leanvm_core/src/class_flock.rs index d9cb8b3b3..9f12b70d4 100644 --- a/crates/leanvm_core/src/class_flock.rs +++ b/crates/leanvm_core/src/class_flock.rs @@ -29,8 +29,8 @@ use std::sync::OnceLock; /// plus its fixed-point dimensions), which floors the batch of a small circuit. pub const MIN_CUBE_LOG: usize = flock::zerocheck::MIN_LOG_N; -/// The most input ports a circuit with a word-level witness has: the hash's fourteen. -const MAX_INPUT_WORDS: usize = 14; +/// The most input ports a circuit with a word-level witness has: EXT's clock circuit's seventeen. +const MAX_INPUT_WORDS: usize = 17; /// The packed witnesses: every class circuit in table order (the tables that have one come first), then every /// table's clock circuit. @@ -222,10 +222,21 @@ impl Prepared { // A clock circuit, and a class with a word-level witness, skip the walk of the gate list; // the others walk it 64 instances at a time. let (a, b, z_lincheck) = match (part, spec.batch_witness, spec.witness) { - // The EXT clock circuit's extra ports (#400) have no generator: it walks. - (Part::Clock, _, _) if !spec.clock_inputs.is_empty() || !spec.clock_outputs.is_empty() => { - circuit.generate_witness_from_into(z, rows, &rows[0], n_blocks_log, input_words, check) - } + // A clock circuit with ports of its own (EXT's) takes its input words through its ports. + (Part::Clock, _, _) if !spec.clock_inputs.is_empty() => circuit.generate_witness_with_into( + z, + rows, + &rows[0], + n_blocks_log, + |row, z, az, bz| { + let mut words = [0u64; MAX_INPUT_WORDS]; + let words = &mut words[..n_inputs]; + input_words(row, words); + spec.clock_witness(&slots, words, z, az, bz); + }, + check, + ), + // The others read the row's clock and previous timestamps straight off it. (Part::Clock, _, _) => circuit.generate_witness_with_into( z, rows, diff --git a/crates/leanvm_core/src/rv/semantics/ext.rs b/crates/leanvm_core/src/rv/semantics/ext.rs index b3480cc5a..50ffeb8a9 100644 --- a/crates/leanvm_core/src/rv/semantics/ext.rs +++ b/crates/leanvm_core/src/rv/semantics/ext.rs @@ -149,6 +149,64 @@ impl Ext { } c.finish() } + + /// One instance of [`Self::clock_circuit`]'s witness by word arithmetic: what the walk of its gate list writes, + /// into zeroed buffers, from its input words (`ts`, each access's previous timestamp, `v1`, `v2`, `vd`, `flags`). + /// + /// The clock's words and products are [`Clock::witness`]'s; its own ports and products, `n` being the accesses: + /// + /// ```text + /// words n + 1 ..= n + 3 v1, v2, vd z = A·z = the word, B·z = all ones + /// word n + 4 flags z = A·z = its two bits, B·z = 0b11 + /// word n + 5 step the clock's + /// words n + 6, n + 7 accumulate, base, copies of the flag bits: z = A·z = the bit, B·z = 1 + /// words n + 8 ..= n + 13 p + 8, p + 16 for p = v1, v2, vd: copies, z = A·z = the sum, B·z = all ones, + /// but v2's products: A·z = !base on every bit, B·z = the sum + /// then products after the clock's, the same order: each incrementer's carries + /// ``` + /// + /// Adding `2^bit` to `p` carries `c_i = p_bit & .. & p_{i-1}` into bit `i > bit`, so the products of bits + /// `bit + 1 ..= 62` are `A·z = p_i`, `B·z = c_i`, and the carries are those of the native sum. + pub fn clock_witness(slots: &[u32], inputs: &[u64], z: &mut [u64], az: &mut [u64], bz: &mut [u64]) { + const BITS: [u32; 2] = [3, 4]; + let n = slots.len(); + let [pointers @ .., flags]: [u64; 4] = inputs[n + 1..].try_into().expect("the clock's inputs, then four"); + let base = flags >> 1 & 1; + let sums = pointers.map(|p| BITS.map(|bit| p.wrapping_add(1 << bit))); + let tables = [&mut *z, &mut *az, &mut *bz]; + Clock::witness_with( + slots, + inputs[0], + &inputs[1..=n], + [4, 2 + Self::OFFSET_LIMBS.len()], + tables, + |products| { + for (&p, sums) in pointers.iter().zip(&sums) { + for (bit, sum) in BITS.into_iter().zip(sums) { + let carries = sum ^ p ^ 1 << bit; + let rows = 62 - bit; + let run = (1 << rows) - 1; + products.push(p >> (bit + 1) & run, carries >> (bit + 1) & run, rows); + } + } + }, + ); + for (k, &p) in pointers.iter().enumerate() { + (z[n + 1 + k], az[n + 1 + k], bz[n + 1 + k]) = (p, p, !0); + } + (z[n + 4], az[n + 4], bz[n + 4]) = (flags & 3, flags & 3, 3); + (z[n + 6], az[n + 6], bz[n + 6]) = (flags & 1, flags & 1, 1); + (z[n + 7], az[n + 7], bz[n + 7]) = (base, base, 1); + let not_base = base.wrapping_sub(1); + for (k, &sum) in sums.as_flattened().iter().enumerate() { + let w = n + 8 + k; + (z[w], az[w], bz[w]) = if k / 2 == 1 { + (not_base & sum, not_base, sum) + } else { + (sum, sum, !0) + }; + } + } } /// `x + 2^bit` modulo `2^64`. diff --git a/crates/leanvm_core/src/tables/clock.rs b/crates/leanvm_core/src/tables/clock.rs index cfb746364..3ece4bc11 100644 --- a/crates/leanvm_core/src/tables/clock.rs +++ b/crates/leanvm_core/src/tables/clock.rs @@ -168,8 +168,25 @@ impl Clock { /// bits 5..40 A·z = x_b ^ c_b B·z = y_b ^ c_b /// ``` pub fn witness(slots: &[u32], ts: u64, prev: &[u64], z: &mut [u64], az: &mut [u64], bz: &mut [u64]) { + Self::witness_with(slots, ts, prev, [0, 0], [z, az, bz], |_| {}); + } + + /// [`Self::witness`] for a circuit of [`Self::builder`] with `extra` more input and output ports past its own: + /// the clock's words and products, then the products `more` pushes. + /// + /// The extra ports' words are the caller's to write: the step is word `n + 1 + extra[0]`, the constant the first + /// bit of the word `extra[1]` past it. + pub(crate) fn witness_with( + slots: &[u32], + ts: u64, + prev: &[u64], + extra: [usize; 2], + [z, az, bz]: [&mut [u64]; 3], + more: impl FnOnce(&mut Products), + ) { let n = slots.len(); assert_eq!(prev.len(), n); + let step_word = n + 1 + extra[0]; const LOW: u64 = (1 << Clock::LIVE_BIT) - 1; const READ: u64 = (1 << Clock::CLOCK_BITS) - 1; const CYCLES: u32 = Clock::LIVE_BIT - Clock::SLOT_BITS; @@ -181,8 +198,8 @@ impl Clock { let ts = ts & READ; let live = ts >> Self::LIVE_BIT; - // The constant, then the products, start word `n + 2`. - let mut products = Products::new([z, az, bz], n + 2); + // The constant, then the products, start past the output ports. + let mut products = Products::new([z, az, bz], step_word + 1 + extra[1]); products.push(1, 1, 1); let (mut in_order, mut disagree) = (0u32, 0); @@ -226,8 +243,9 @@ impl Clock { let copies = live << Self::SLOT_BITS | fail << Self::FAIL_BIT; let ts_bits = cycle << (Self::SLOT_BITS + 1); let carry_bits = (carries & run(CYCLES)) << (Self::SLOT_BITS + 1); + more(&mut products); let [z, az, bz] = products.finish(); - (z[n + 1], az[n + 1], bz[n + 1]) = ( + (z[step_word], az[step_word], bz[step_word]) = ( step, ts_bits | copies, carry_bits | 1 << Self::SLOT_BITS | 1 << Self::FAIL_BIT, @@ -238,7 +256,8 @@ impl Clock { #[cfg(test)] mod tests { use super::*; - use crate::tables::ClassSpec; + use crate::rv::Ext; + use crate::tables::{ClassSpec, Word}; use primitives::test_util::Rng; #[test] @@ -293,13 +312,15 @@ mod tests { fn the_word_witness_is_the_gate_walk() { // Invariant: the word-level clock witness writes the tables the 64-lane walk of the clock circuit's gate list writes. // - // Fixture state: every table's clock circuit, and one with an access in each of the 32 slots. + // Fixture state: every table's clock circuit (EXT's with its own ports), and one with an access in each of the 32 slots. // Rows: padding rows (clock zero), the first and the last cycle, random cycles, all ones and random words; each previous timestamp at `ts ^ slot`, around it, with the live bit flipped, a padding tuple, a seed, zero, all ones, or a random word. + // EXT's operands: pointers at zero, all ones, the sign bit, one limb or two below a wrap of 2^64 or of the sign bit, with every low bit set, aligned random or random; flags every legal word or a random word. let mut rng = Rng::new(0xC10C); let n_log = 12; - let every_slot = (0..Clock::CYCLE as u32).collect(); - for slots in ClassSpec::ALL.iter().map(|spec| spec.slots()).chain([every_slot]) { - let circuit = Clock::circuit(&slots); + for spec in ClassSpec::ALL.into_iter().map(Some).chain([None]) { + let slots = spec.map_or_else(|| (0..Clock::CYCLE as u32).collect(), |spec| spec.slots()); + let circuit = spec.map_or_else(|| Clock::circuit(&slots), |spec| spec.clock_circuit()); + let operands = spec.map_or(&[][..], |spec| spec.clock_inputs); let rows: Vec> = (0..1 << n_log) .map(|r| { let ts = match r % 6 { @@ -310,32 +331,55 @@ mod tests { 4 => u64::MAX, _ => rng.next_u64(), }; - let prev = slots.iter().map(|&slot| { - let at = ts ^ u64::from(slot); - match rng.next_u64() % 10 { - 0 => at, - 1 => at.wrapping_sub(1), - 2 => at.wrapping_add(1), - 3 => at.wrapping_sub(Clock::CYCLE), - 4 => at ^ Clock::SEED_CLOCK, - 5 => u64::from(slot), - 6 => Clock::SEED_CLOCK, - 7 => 0, - 8 => u64::MAX, + let prev: Vec = (slots.iter()) + .map(|&slot| { + let at = ts ^ u64::from(slot); + match rng.next_u64() % 10 { + 0 => at, + 1 => at.wrapping_sub(1), + 2 => at.wrapping_add(1), + 3 => at.wrapping_sub(Clock::CYCLE), + 4 => at ^ Clock::SEED_CLOCK, + 5 => u64::from(slot), + 6 => Clock::SEED_CLOCK, + 7 => 0, + 8 => u64::MAX, + _ => rng.next_u64(), + } + }) + .collect(); + let operands: Vec = (operands.iter()) + .map(|&word| match (word, rng.next_u64() % 12) { + (Word::Flags, k) if k < 8 => Ext::LEGAL[k as usize % Ext::LEGAL.len()], + (Word::Flags, _) => rng.next_u64(), + (_, 0) => 0, + (_, 1) => u64::MAX, + (_, 2) => 1 << 63, + (_, 3) => 8u64.wrapping_neg(), + (_, 4) => 16u64.wrapping_neg(), + (_, 5) => (1 << 63) - 8, + (_, 6) => (1 << 63) - 16, + (_, 7) => (1 << 63) - 1, + (_, 8) => 7, + (_, 9) => rng.next_u64() & !7, _ => rng.next_u64(), - } - }); - std::iter::once(ts).chain(prev).collect() + }) + .collect(); + [vec![ts], prev, operands].concat() }) .collect(); - let walk = circuit.generate_witness_from(&rows, &rows[0], n_log, |row, words| words.copy_from_slice(row)); - let words = circuit.generate_witness_with(&rows, &rows[0], n_log, |row, z, az, bz| { - Clock::witness(&slots, row[0], &row[1..], z, az, bz); + let walk = circuit.generate_witness_from(&rows, &rows[0], n_log, |row: &Vec, words| { + words.copy_from_slice(row) + }); + let words = circuit.generate_witness_with(&rows, &rows[0], n_log, |row: &Vec, z, az, bz| match spec { + Some(spec) => spec.clock_witness(&slots, row, z, az, bz), + None => Clock::witness(&slots, row[0], &row[1..], z, az, bz), }); - assert!(walk.0[..] == words.0[..], "z, slots {slots:?}"); - assert!(walk.1[..] == words.1[..], "A·z, slots {slots:?}"); - assert!(walk.2[..] == words.2[..], "B·z, slots {slots:?}"); - assert!(walk.3[..] == words.3[..], "lincheck stripes, slots {slots:?}"); + let name = spec.map_or("every slot", |spec| spec.name); + assert!(walk.0[..] == words.0[..], "z, {name}"); + assert!(walk.1[..] == words.1[..], "A·z, {name}"); + assert!(walk.2[..] == words.2[..], "B·z, {name}"); + assert!(walk.3[..] == words.3[..], "lincheck stripes, {name}"); } } } diff --git a/crates/leanvm_core/src/tables/spec.rs b/crates/leanvm_core/src/tables/spec.rs index 4b9e78e6c..8deca3e9c 100644 --- a/crates/leanvm_core/src/tables/spec.rs +++ b/crates/leanvm_core/src/tables/spec.rs @@ -473,6 +473,16 @@ impl ClassSpec { } } + /// One instance of [`Self::clock_circuit`]'s witness by word arithmetic, from its input words (the + /// [`Self::clock_ports`] before the step), into zeroed buffers: [`Clock::witness`], or EXT's + /// ([`Ext::clock_witness`]). + pub fn clock_witness(&self, slots: &[u32], inputs: &[u64], z: &mut [u64], az: &mut [u64], bz: &mut [u64]) { + match self.class { + Class::Ext => Ext::clock_witness(slots, inputs, z, az, bz), + _ => Clock::witness(slots, inputs[0], &inputs[1..], z, az, bz), + } + } + /// The clock circuit's port words: the clock, each access's previous timestamp and [`Self::clock_inputs`], then /// the step and [`Self::clock_outputs`]. pub fn clock_ports(&self) -> Vec { From ddbef89c28348f36f9ca3b9047cdee12071f724b Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Mon, 5 Oct 2026 20:15:15 +0100 Subject: [PATCH 09/18] leanvm_core: end the clock witness calls with a semicolon --- crates/leanvm_core/src/class_flock.rs | 2 +- crates/leanvm_core/src/tables/clock.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/leanvm_core/src/class_flock.rs b/crates/leanvm_core/src/class_flock.rs index 9f12b70d4..3fdc004ef 100644 --- a/crates/leanvm_core/src/class_flock.rs +++ b/crates/leanvm_core/src/class_flock.rs @@ -244,7 +244,7 @@ impl Prepared { n_blocks_log, |row, z, az, bz| { let row = view(row); - Clock::witness(&slots, row.row.ts, &row.prev()[..slots.len()], z, az, bz) + Clock::witness(&slots, row.row.ts, &row.prev()[..slots.len()], z, az, bz); }, check, ), diff --git a/crates/leanvm_core/src/tables/clock.rs b/crates/leanvm_core/src/tables/clock.rs index 3ece4bc11..f8998de51 100644 --- a/crates/leanvm_core/src/tables/clock.rs +++ b/crates/leanvm_core/src/tables/clock.rs @@ -369,7 +369,7 @@ mod tests { }) .collect(); let walk = circuit.generate_witness_from(&rows, &rows[0], n_log, |row: &Vec, words| { - words.copy_from_slice(row) + words.copy_from_slice(row); }); let words = circuit.generate_witness_with(&rows, &rows[0], n_log, |row: &Vec, z, az, bz| match spec { Some(spec) => spec.clock_witness(&slots, row, z, az, bz), From 21cc3b39b718b3598ac68c7ebeeb281a612424b3 Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Mon, 5 Oct 2026 22:07:05 +0100 Subject: [PATCH 10/18] pcs: keep the opening's weight from the first pass on small pools The stacked opening's lifted weight was filled twice per proof: once by the first pass (the Basis span) and again by the first lane fold (Sumcheck round 4), each fill paying the ring-switch map (apply_sliced_add), the point claims' eq tables and the scatter. On a pool of at most KEEP_WEIGHT_MAX_THREADS (4) workers the first pass now also streams each filled chunk into a dense buffer, and the first fold reads that back as Basis::Dense instead of filling it again. On a larger pool the refill, spread over the workers, costs less than writing the weight out and reading it back, so that path is unchanged. Proofs are byte-identical. Measured on a 7950X3D, cargo leanvm aggregate --n 100 --leaves 4 --arity0 2 --arity 2 --repeat 3 --cooldown 1, AGG1_NOTREE, pooled per-pass medians (min) over alternating base/branch rounds, loads 1 to 12: 1T (5 rounds a side, taskset -c 2) first-level node base 0.978 s (0.969) branch 0.960 s (0.954) -1.8% node base 1.046 s (1.035) branch 1.025 s (1.018) -2.0% 32T (3 rounds a side, gate off) first-level node base 0.148 s (0.144) branch 0.148 s (0.143) node base 0.175 s (0.169) branch 0.174 s (0.169) Spans at 1T (tracing, one pass): first-level node Basis 79.0 -> 89.1 ms, Sumcheck round 4 73.1 -> 43.9 ms, PCS open 289 -> 270 ms; node Basis 76.9 -> 89.2 ms, round 4 72.4 -> 43.8 ms, PCS open 292 -> 275 ms; leaf Basis 100 -> 113 ms, round 4 84.8 -> 47.6 ms. Crossover (PCS open, store vs refill, tracing): 2T -9 and -6 ms, 4T -3 and -3 ms, 32T +5 to +7 ms, which sets the gate at 4. Peak memory at 1T grows by the kept weight (24 bytes per committed word): aggregate 1.02 -> 1.31 GiB, leanxmss --n 100 0.83 -> 1.08 GiB. leanxmss --n 100 at 1T: 1.002 -> 0.986 s (pooled per-pass medians, 3 rounds a side). --- crates/pcs/src/stack_open.rs | 13 +-- crates/pcs/src/stack_open/basis.rs | 2 +- crates/pcs/src/whir.rs | 2 +- crates/pcs/src/whir/sumcheck.rs | 7 +- crates/pcs/src/whir/sumcheck/first_pass.rs | 101 +++++++++++++++++++-- 5 files changed, 105 insertions(+), 20 deletions(-) diff --git a/crates/pcs/src/stack_open.rs b/crates/pcs/src/stack_open.rs index 10b79346c..cd28f6f5f 100644 --- a/crates/pcs/src/stack_open.rs +++ b/crates/pcs/src/stack_open.rs @@ -57,7 +57,7 @@ use super::pack::PACKING_WIDTH; use super::ring_switch; use super::ring_switch::RsEqQuery; -use super::whir::{Basis, ProverConfig, ProverData, VerifierConfig, WhirError, recursive_verifier_with_basis_succinct}; +use super::whir::{ProverConfig, ProverData, VerifierConfig, WhirError, recursive_verifier_with_basis_succinct}; use crate::merkle::Hash; use basis::StackWeight; use fiat_shamir::transcript::{Challenger, Receiver, Transmitter}; @@ -267,17 +267,18 @@ pub fn open_batch_mixed_whir_stacked( .zip(lambdas_pd) .fold(F192::ZERO, |sum, (claim, &lambda)| sum + lambda * claim.value()); - // The lifted weight is never stored. + // The lifted weight is filled one chunk at a time: // // first pass: each chunk is filled, then feeds the first lane rounds' sums while hot // first fold: each chunk is filled again, then folded by those rounds' challenges // - // Filling costs less than writing the weight out and reading it back. + // On a small pool the first pass also writes each chunk out, and the first fold reads it back rather than filling + // it again. let lane_block = 1usize << (log_n - config.initial_k()); let weight = StackWeight::new(stack.len(), lane_block, point_claims, lambdas_pd, rings, &rs_outputs); let fill = |start: usize, dst: &mut [F192]| weight.fill(start, dst); - let initial = tracing::info_span!("Basis") - .in_scope(|| super::whir::initial_rounds(stack, lane_block, config.initial_k(), &Basis::Virtual(&fill))); + let (initial, basis) = tracing::info_span!("Basis") + .in_scope(|| super::whir::initial_rounds_virtual(stack, lane_block, config.initial_k(), &fill)); // 4. One WHIR over the full stack against the combined claim (the // stack is borrowed by the prover; no copy). @@ -285,7 +286,7 @@ pub fn open_batch_mixed_whir_stacked( config, log_n, stack, - Basis::Virtual(&fill), + basis, target, &prover_data.codeword, &prover_data.merkle_tree, diff --git a/crates/pcs/src/stack_open/basis.rs b/crates/pcs/src/stack_open/basis.rs index 1e6d5039b..b40985b8f 100644 --- a/crates/pcs/src/stack_open/basis.rs +++ b/crates/pcs/src/stack_open/basis.rs @@ -76,7 +76,7 @@ impl<'a> PointWeight<'a> { /// - each ring-switched region's combined weight, over that region; /// - each point claim's equality weight, over its own support. /// -/// It is never stored whole: the opening's first pass and its first fold each refill the chunks they read. +/// It is filled one chunk at a time, as the opening reads it. pub(super) struct StackWeight<'a> { /// Each point claim's weight, its high equality table built once. weights: Vec>, diff --git a/crates/pcs/src/whir.rs b/crates/pcs/src/whir.rs index 5b66b2827..fc7083235 100644 --- a/crates/pcs/src/whir.rs +++ b/crates/pcs/src/whir.rs @@ -51,7 +51,7 @@ pub use crate::whir_induce::*; pub use commit::{Commitment, ProverData, commit}; pub use prove::recursive_prover_with_basis; pub(crate) use prove::recursive_prover_with_prepared_basis; -pub(crate) use sumcheck::{Basis, INITIAL_BASIS_CHUNK, initial_rounds}; +pub(crate) use sumcheck::{INITIAL_BASIS_CHUNK, initial_rounds_virtual}; pub use verify::{WhirError, recursive_verifier_with_basis_succinct}; /// Mixed inner product `Σ_i b[i] · witness[i]` (E x K via `mul_base`). The diff --git a/crates/pcs/src/whir/sumcheck.rs b/crates/pcs/src/whir/sumcheck.rs index 7216ecfba..cb0b6f26e 100644 --- a/crates/pcs/src/whir/sumcheck.rs +++ b/crates/pcs/src/whir/sumcheck.rs @@ -18,7 +18,7 @@ use std::ops::Add; mod first_pass; -pub(crate) use first_pass::{InitialRounds, initial_rounds}; +pub(crate) use first_pass::{InitialRounds, initial_rounds, initial_rounds_virtual}; // =================================================================== // Tuning constants @@ -50,6 +50,9 @@ pub(crate) const INITIAL_BASIS_CHUNK: usize = 256; /// Elements a stored (dense) weight's lane fold stages in L1 before publishing them. const DENSE_STAGE: usize = 128; +/// Pool sizes up to which the first pass writes out a regenerated weight for the first fold. +const KEEP_WEIGHT_MAX_THREADS: usize = 4; + // =================================================================== // Stateful sumcheck over E with a two-phase (Base then Ext) witness // =================================================================== @@ -421,7 +424,7 @@ pub(crate) enum Basis<'a> { Dense(Vec), /// Regenerated when read, by chunks of the initial fill size. /// - /// Only the first pass and the first fold read it, so it is never stored. + /// Only the first pass and the first fold read it. Virtual(&'a BasisFill<'a>), } diff --git a/crates/pcs/src/whir/sumcheck/first_pass.rs b/crates/pcs/src/whir/sumcheck/first_pass.rs index ebbdbabe6..ed0aa70f7 100644 --- a/crates/pcs/src/whir/sumcheck/first_pass.rs +++ b/crates/pcs/src/whir/sumcheck/first_pass.rs @@ -10,13 +10,18 @@ //! precomputation of Bagad, Dao, Domb and Thaler, : //! the witness is in `K`, so every product is a mixed one). -use super::{Basis, FIRST_PASS_PAR_THRESHOLD, INITIAL_BASIS_CHUNK, PRECOMPUTED_ROUNDS, SumcheckMessage, window}; +use super::{ + Basis, BasisFill, FIRST_PASS_PAR_THRESHOLD, INITIAL_BASIS_CHUNK, KEEP_WEIGHT_MAX_THREADS, PRECOMPUTED_ROUNDS, + SumcheckMessage, window, +}; +use parallel::SendPtr; #[cfg(not(any( all(target_arch = "x86_64", target_feature = "pclmulqdq"), all(target_arch = "aarch64", target_feature = "aes") )))] use primitives::field::gf2_64::mul_wide; use primitives::field::{F64, F192}; +use primitives::stream::Stream; use std::ops::Range; /// Lanes per group, and points of `{0, 1, inf}^R`, `R` being [`PRECOMPUTED_ROUNDS`]. @@ -618,6 +623,46 @@ impl InitialRounds { /// The first pass: the grid sums of the first `min(PRECOMPUTED_ROUNDS, initial_k)` lane rounds. pub(crate) fn initial_rounds(f: &[F64], block: usize, initial_k: usize, b: &Basis<'_>) -> InitialRounds { + first_pass(f, block, initial_k, b, None) +} + +/// [`initial_rounds`] over a regenerated weight, and the weight the first fold then reads. +/// +/// On a small pool the pass also writes the weight out, so the fold reads it back rather than refilling it; on a +/// larger one the refill, spread over the workers, costs less than the memory traffic. +pub(crate) fn initial_rounds_virtual<'a>( + f: &[F64], + block: usize, + initial_k: usize, + fill: &'a BasisFill<'a>, +) -> (InitialRounds, Basis<'a>) { + if parallel::num_threads() <= KEEP_WEIGHT_MAX_THREADS { + let (rounds, kept) = initial_rounds_kept(f, block, initial_k, fill); + (rounds, Basis::Dense(kept)) + } else { + ( + first_pass(f, block, initial_k, &Basis::Virtual(fill), None), + Basis::Virtual(fill), + ) + } +} + +/// [`initial_rounds`] over a regenerated weight, which it also writes out whole. +fn initial_rounds_kept(f: &[F64], block: usize, initial_k: usize, fill: &BasisFill<'_>) -> (InitialRounds, Vec) { + // SAFETY: the first pass fills every window of every lane of `f` once, and keeps each in its slots. + let mut kept = unsafe { primitives::uninit_vec(f.len()) }; + let rounds = first_pass( + f, + block, + initial_k, + &Basis::Virtual(fill), + Some(SendPtr(kept.as_mut_ptr())), + ); + (rounds, kept) +} + +/// [`initial_rounds`], writing each window of the weight it reads to `keep` when given. +fn first_pass(f: &[F64], block: usize, initial_k: usize, b: &Basis<'_>, keep: Option>) -> InitialRounds { if let Basis::Dense(b) = b { assert_eq!(b.len(), f.len()); } @@ -627,7 +672,7 @@ pub(crate) fn initial_rounds(f: &[F64], block: usize, initial_k: usize, b: &Basi let n_lanes = f.len() / block; let whole = n_lanes >> rounds << rounds; let mut grid = if whole > 0 { - grid_pass(rounds, f, block, b, 0..whole) + grid_pass(rounds, f, block, b, 0..whole, keep) } else { vec![F192::ZERO; 3usize.pow(rounds as u32)] }; @@ -635,7 +680,7 @@ pub(crate) fn initial_rounds(f: &[F64], block: usize, initial_k: usize, b: &Basi // lane bit is clear: at 1 the group has nothing, and at 0 and at inf it has those sums. if whole < n_lanes { let digits = (n_lanes - whole).next_power_of_two().ilog2() as usize; - let tail = grid_pass(digits, f, block, b, whole..n_lanes); + let tail = grid_pass(digits, f, block, b, whole..n_lanes, keep); let low = tail.len(); for &high in &LANE_IN_GRID[..1 << (rounds - digits)] { let base = low * 2 * high; @@ -649,18 +694,31 @@ pub(crate) fn initial_rounds(f: &[F64], block: usize, initial_k: usize, b: &Basi /// The grid sums over `lanes`, `rounds` digits at a time: whole groups, or one group whose /// missing lanes are zero. -fn grid_pass(rounds: usize, f: &[F64], block: usize, b: &Basis<'_>, lanes: Range) -> Vec { +fn grid_pass( + rounds: usize, + f: &[F64], + block: usize, + b: &Basis<'_>, + lanes: Range, + keep: Option>, +) -> Vec { match rounds { - 0 => grid_pass_with::<0>(f, block, b, lanes), - 1 => grid_pass_with::<1>(f, block, b, lanes), - 2 => grid_pass_with::<2>(f, block, b, lanes), - 3 => grid_pass_with::<3>(f, block, b, lanes), - 4 => grid_pass_with::<4>(f, block, b, lanes), + 0 => grid_pass_with::<0>(f, block, b, lanes, keep), + 1 => grid_pass_with::<1>(f, block, b, lanes, keep), + 2 => grid_pass_with::<2>(f, block, b, lanes, keep), + 3 => grid_pass_with::<3>(f, block, b, lanes, keep), + 4 => grid_pass_with::<4>(f, block, b, lanes, keep), _ => unreachable!("at most PRECOMPUTED_ROUNDS digits"), } } -fn grid_pass_with(f: &[F64], block: usize, b: &Basis<'_>, lanes: Range) -> Vec { +fn grid_pass_with( + f: &[F64], + block: usize, + b: &Basis<'_>, + lanes: Range, + keep: Option>, +) -> Vec { let group = 1 << R; let points = 3usize.pow(R as u32); // A regenerated weight is filled one aligned chunk at a time, or one whole block below that. @@ -689,6 +747,10 @@ fn grid_pass_with(f: &[F64], block: usize, b: &Basis<'_>, lanes: let at = lane * block + x0; fs[l] = &f[at..at + chunk]; bs[l] = window(b, raw, at, chunk); + if let Some(keep) = keep { + // SAFETY: `keep` holds `f.len()` weights, and no other task fills the window at `at`. + Stream::new().copy(unsafe { keep.slice(at, chunk) }, bs[l]); + } } } for x in (0..chunk).step_by(ROW) { @@ -870,8 +932,27 @@ impl WeightFold { #[cfg(test)] mod tests { use super::*; + use crate::whir::INITIAL_FOLDING_FACTOR; use primitives::test_util::Rng; + #[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 + // and blocks below one fill chunk included. + let mut rng = Rng::new(0x6EE9); + for block in [1, 16, INITIAL_BASIS_CHUNK, 4 * INITIAL_BASIS_CHUNK] { + for lanes in [1, 3, 5, GROUP - 1, GROUP, GROUP + 1, 37] { + let f: Vec = (0..block * lanes).map(|_| F64(rng.next_u64())).collect(); + let weight = rng.ext_vec(f.len()); + let fill = |start: usize, out: &mut [F192]| out.copy_from_slice(&weight[start..start + out.len()]); + let (rounds, kept) = initial_rounds_kept(&f, block, INITIAL_FOLDING_FACTOR, &fill); + let expected = initial_rounds(&f, block, INITIAL_FOLDING_FACTOR, &Basis::Dense(weight.clone())); + assert_eq!(rounds.grid, expected.grid, "block={block}, lanes={lanes}"); + assert_eq!(kept, weight, "block={block}, lanes={lanes}"); + } + } + } + /// Every arm this target compiles accumulates the products of the definition, not only the dispatched one. #[cfg(any( all(target_arch = "x86_64", target_feature = "pclmulqdq"), From b4af91633e9596a526bb9fe010ce67b8015bf908 Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Mon, 5 Oct 2026 22:17:23 +0100 Subject: [PATCH 11/18] pcs: a deep NTT task takes whole sub-blocks for at least 64 rows A high-rate encode of a short message (the last WHIR level: 4 message rows, rate 2^13) cut its deep pass into one task per 4-row sub-block, so the Merkle builder got 4 leaves per call, short of a hash batch, and hashed them one by one. A task now runs 2^MIN_TASK_LOG rows of contiguous sub-blocks and hands them over at once. The rows, the codeword and the tree are unchanged. Measured on the 7950X3D (aggregate --n 100 --leaves 4, 2 alternating rounds, --repeat 3, load under 10), upstream f774828d vs this: - level-4 extension commit at 1T: 21.3 ms -> 3.4 ms (span tree) - 1T first-level node 972 -> 955 ms, node 1039 -> 1019 ms, leanXMSS-100 leaf 1011 -> 994 ms - 32T first-level node 152 -> 145 ms, node 174 -> 170 ms Proof digests unchanged (leaf, first-level node, node). --- crates/pcs/src/ntt/additive_ntt_f64.rs | 83 ++++++++++++++++---------- 1 file changed, 50 insertions(+), 33 deletions(-) diff --git a/crates/pcs/src/ntt/additive_ntt_f64.rs b/crates/pcs/src/ntt/additive_ntt_f64.rs index 311434030..e177eea39 100644 --- a/crates/pcs/src/ntt/additive_ntt_f64.rs +++ b/crates/pcs/src/ntt/additive_ntt_f64.rs @@ -232,7 +232,7 @@ impl AdditiveNttF64 { /// /// # Finished rows /// - /// - Each deep sub-block is final, and hands its rows to `on_rows` while they are in L2. + /// - Each deep task's sub-blocks are final, and hand their rows to `on_rows` while they are in L2. /// - Without a split deep pass, the rows go over in parallel blocks at the end. fn transform( &self, @@ -307,44 +307,54 @@ impl AdditiveNttF64 { replicate(data, m, data.len() >> start); } - // A lone sub-block runs on this thread, so its rows go over in parallel afterwards. - let fuse = 0 < deep_start && deep_start < log_d; + // A deep task takes whole sub-blocks, enough of them for `2^MIN_TASK_LOG` rows. + let log_sub = log_d - deep_start; + let log_group = MIN_TASK_LOG.saturating_sub(log_sub).min(deep_start); + + // A lone task runs on this thread, so its rows go over in parallel afterwards. + let fuse = log_group < deep_start && deep_start < log_d; let deep_rows = on_rows.filter(|_| fuse); - // Phase 3: the deep pass, one task per contiguous sub-block. + // Phase 3: the deep pass, one task per run of contiguous sub-blocks. if deep_start < log_d { - let log_sub = log_d - deep_start; let sub_len = num_ntts << log_sub; let block_len = data.len() >> start; - parallel::chunks_mut(data, sub_len, |sub_idx, sub| match msg { - // A separate message: build the sub-block in scratch, then write it out once. - // - // Streamed out whole, the codeword is written without ever being read. - Some(m) => with_scratch(sub_len, |scratch| { - // A sub-block sits at the same offset in every replica. + parallel::chunks_mut(data, sub_len << log_group, |task_idx, task| { + let first_sub = task_idx << log_group; + match msg { + // A separate message: build the sub-blocks in scratch, then write them out once. // - // sub-block at offset off of its replica <- message words [off, off + sub_len) - let off = (sub_idx * sub_len) % block_len; - // SAFETY: - // - The external message is valid for one whole replica of words. - // - It is disjoint from the codeword. - // - This sub-block ends inside its replica, so the read stays in bounds. - scratch.copy_from_slice(unsafe { std::slice::from_raw_parts(m.add(off), sub_len) }); - self.run_layers(scratch, log_d, num_ntts, deep_start, log_d, deep_start, sub_idx); - if let Some(f) = deep_rows { - f(sub_idx << log_sub, scratch); - } - if stream { - Stream::new().copy(sub, scratch); - } else { - sub.copy_from_slice(scratch); - } - }), - // The replicas are already in place: run the layers where they are. - None => { - self.run_layers(sub, log_d, num_ntts, deep_start, log_d, deep_start, sub_idx); - if let Some(f) = deep_rows { - f(sub_idx << log_sub, sub); + // Streamed out whole, the codeword is written without ever being read. + Some(m) => with_scratch(task.len(), |scratch| { + for (i, sub) in scratch.chunks_exact_mut(sub_len).enumerate() { + // A sub-block sits at the same offset in every replica. + // + // sub-block at offset off of its replica <- message words [off, off + sub_len) + let off = ((first_sub + i) * sub_len) % block_len; + // SAFETY: + // - The external message is valid for one whole replica of words. + // - It is disjoint from the codeword. + // - This sub-block ends inside its replica, so the read stays in bounds. + sub.copy_from_slice(unsafe { std::slice::from_raw_parts(m.add(off), sub_len) }); + self.run_layers(sub, log_d, num_ntts, deep_start, log_d, deep_start, first_sub + i); + } + if let Some(f) = deep_rows { + f(first_sub << log_sub, scratch); + } + if stream { + Stream::new().copy(task, scratch); + } else { + task.copy_from_slice(scratch); + } + }), + // The replicas are already in place: run the layers where they are. + None => { + for (i, sub) in task.chunks_exact_mut(sub_len).enumerate() { + self.run_layers(sub, log_d, num_ntts, deep_start, log_d, deep_start, first_sub + i); + } + if let Some(f) = deep_rows { + f(first_sub << log_sub, task); + } } } }); @@ -717,6 +727,13 @@ const LOG_SUBS_PER_WORKER: usize = 0; #[cfg(all(target_arch = "aarch64", target_os = "macos"))] const LOG_SUBS_PER_WORKER: usize = 2; +/// Fewest rows a deep task covers, as a log. +/// +/// - A high-rate encode of a short message leaves sub-blocks of a few rows. +/// - One task each would hand the Merkle tree a few leaves at a time, short of a hash batch, which it hashes one by one. +/// - So a task takes enough contiguous sub-blocks for whole batches, which also spreads its fixed cost. +const MIN_TASK_LOG: usize = 6; + /// Words a deep sub-block may span when that saves a whole sweep. /// /// # Why this value From 7e8753654f745acc97ca6cbcf8c6913b24e7efd9 Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Mon, 5 Oct 2026 22:19:06 +0100 Subject: [PATCH 12/18] rec: a tree proves and reads with the fixed columns it built once Tree::new already builds each circuit's fixed columns. Proving a node (prove_rows) and reading a child (read, which prove_node calls for every child and verify calls for the root) built them again through TableArgument::of, a serial pass over every slot of every table. They now take the tree's columns (Circuit::prove_with, Circuit::verify_to_raw_with); Circuit::prove and verify_to_raw still build their own. Measured on the 7950X3D, upstream f774828d vs this, aggregate --n 100: - sampled profile at 1T: FixedColumns::of is 0.51% of a first-level node proof (about 5 ms) and 1.19% of a node proof (about 12 ms: the prover's build and one per child read) - 32T, 2 alternating rounds: node 174/180 -> 167/170 ms; first-level node 152/149 -> 145/152 ms (within noise) - 1T end to end is within the run-to-run noise of these rounds (host load 4 to 13) Proof digests unchanged (leaf, first-level node, node). --- crates/leanvm_core/src/rec/proof.rs | 48 +++++++++++++++++++++----- crates/leanvm_core/src/rec/tree/mod.rs | 11 ++++-- 2 files changed, 47 insertions(+), 12 deletions(-) diff --git a/crates/leanvm_core/src/rec/proof.rs b/crates/leanvm_core/src/rec/proof.rs index bf33d4514..47cabdaf7 100644 --- a/crates/leanvm_core/src/rec/proof.rs +++ b/crates/leanvm_core/src/rec/proof.rs @@ -17,6 +17,7 @@ use crate::pcs::{Rate, RingSwitch, SliceClaim, StackClaim}; use crate::{constraints, pcs, witness}; use fiat_shamir::transcript::{Challenger, ProofTranscript, ProverState, RawProof, VerifierState}; use primitives::field::{F64, F192}; +use std::borrow::Cow; /// The transcript's public input for a statement: the hash of its words' limbs, in order. pub fn statement_seed(statement: &[Limbs]) -> [F64; 4] { @@ -146,11 +147,10 @@ impl<'a> TableArgument<'a> { Self { layout, blocks } } - /// The argument of a circuit and its statement, the public rows reading the statement. - fn of(circuit: &Circuit, statement: &[Limbs], layout: &'a RecLayout) -> Self { - let fixed = FixedColumns::of(circuit, &layout.taus); + /// The argument of a circuit's fixed columns and its statement, the public rows reading the statement. + fn of(fixed: &FixedColumns, statement: &[Limbs], layout: &'a RecLayout) -> Self { let public = fixed.public_values(statement); - Self::new(layout, BusBlocks::new(&fixed, public, layout)) + Self::new(layout, BusBlocks::new(fixed, public, layout)) } /// Prove the bus, then every owned table's summand at the bus's point. @@ -211,8 +211,19 @@ impl Circuit { /// # Panics /// /// Panics if the assignment is not one of the circuit, or its bus does not balance. - #[tracing::instrument(name = "Prove recursion", skip_all)] pub fn prove(&self, a: &Assignment, iv: [F64; 4], rate: Rate) -> Result { + self.prove_with(a, iv, rate, None) + } + + /// [`Self::prove`], with the circuit's fixed columns at its heights if the caller has them already. + #[tracing::instrument(name = "Prove recursion", skip_all)] + pub(crate) fn prove_with( + &self, + a: &Assignment, + iv: [F64; 4], + rate: Rate, + fixed: Option<&FixedColumns>, + ) -> Result { assert_eq!(a.statement.len(), self.statement_len, "the assignment's statement"); assert!( Table::ALL @@ -226,7 +237,8 @@ impl Circuit { let w = crate::stage!("Build witness", || RecWitness::build(&layout, a)); let committed = crate::stage!("Commit", || pcs::commit(&mut ps, &w.q, layout.shape, log_inv_rate)); - let slots = TableArgument::of(self, &a.statement, &layout).prove(&w, &mut ps); + let fixed = fixed.map_or_else(|| Cow::Owned(FixedColumns::of(self, &layout.taus)), Cow::Borrowed); + let slots = TableArgument::of(&fixed, &a.statement, &layout).prove(&w, &mut ps); let RecWitness { q, ports, batch } = w; drop(ports); @@ -264,7 +276,19 @@ impl Circuit { rate: Rate, proof: &ProofTranscript, ) -> Result { - self.verify_seeded(statement, iv, statement_seed(statement), rate, proof) + self.verify_seeded(statement, iv, statement_seed(statement), rate, proof, None) + } + + /// [`Self::verify_to_raw`], with the circuit's fixed columns at its heights. + pub(crate) fn verify_to_raw_with( + &self, + statement: &[Limbs], + iv: [F64; 4], + rate: Rate, + proof: &ProofTranscript, + fixed: &FixedColumns, + ) -> Result { + self.verify_seeded(statement, iv, statement_seed(statement), rate, proof, Some(fixed)) } /// Verify a proof whose transcript absorbed the given public input in place of the statement. @@ -275,6 +299,7 @@ impl Circuit { public_input: [F64; 4], rate: Rate, proof: &ProofTranscript, + fixed: Option<&FixedColumns>, ) -> Result { if statement.len() != self.statement_len { return Err(RecError::StatementLength { @@ -285,7 +310,8 @@ impl Circuit { let layout = RecLayout::new(self)?; let mut vs = VerifierState::new(iv, proof, public_input); let root = pcs::read_commitment(&mut vs)?; - let slots = TableArgument::of(self, statement, &layout).verify(&mut vs)?; + let fixed = fixed.map_or_else(|| Cow::Owned(FixedColumns::of(self, &layout.taus)), Cow::Borrowed); + let slots = TableArgument::of(&fixed, statement, &layout).verify(&mut vs)?; let hash_claim = HashBatch::verify(&layout, &mut vs)?; let ring = layout.hash_window().ring(hash_claim); pcs::verify( @@ -475,7 +501,11 @@ mod tests { } // Under the honest seed the bus accepts the forged words; seeded with them, the proof is refused. - assert!(circuit.verify_seeded(&forged, IV, seed, Rate::MIN, &proof).is_ok()); + assert!( + circuit + .verify_seeded(&forged, IV, seed, Rate::MIN, &proof, None) + .is_ok() + ); assert!(matches!( verify_run(&circuit, &forged, &proof), Err(RecError::Bus(BusError::Gkr(_))) diff --git a/crates/leanvm_core/src/rec/tree/mod.rs b/crates/leanvm_core/src/rec/tree/mod.rs index c074f2b27..d4d271da8 100644 --- a/crates/leanvm_core/src/rec/tree/mod.rs +++ b/crates/leanvm_core/src/rec/tree/mod.rs @@ -568,8 +568,13 @@ impl<'p> Tree<'p> { /// Verify a tree proof's recursion proof, short of its claims, returning it as its verifier read it. fn read(&self, p: &TreeProof) -> Result { let limbs: Vec<[u64; 4]> = p.words.iter().map(|w| [w.c0, w.c1, w.c2, 0]).collect(); - self.circuit(p.kind) - .verify_to_raw(&limbs, self.design.iv, self.design.rate, &p.proof) + self.circuit(p.kind).verify_to_raw_with( + &limbs, + self.design.iv, + self.design.rate, + &p.proof, + &self.columns[p.kind as usize], + ) } /// Prove a circuit's rows: its reduction, then its recursion proof. @@ -587,7 +592,7 @@ impl<'p> Tree<'p> { return Err(TreeError::Unsatisfied(first)); } let proof = (self.circuit(kind)) - .prove(&assignment, d.iv, d.rate) + .prove_with(&assignment, d.iv, d.rate, Some(&self.columns[kind as usize])) .map_err(|_| TreeError::TooLarge)?; let words = (assignment.statement().iter()) .map(|l| F192::new(l[0], l[1], l[2])) From af9403dce78ac19ece74dcec8e4e5ac6da4a12bb Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Mon, 5 Oct 2026 22:19:06 +0100 Subject: [PATCH 13/18] rec: the hash rows' packed witness is generated straight into the stack On top of #463's generate_witness_with_into: HashBatch writes z into the stack's packed-witness window and copies each live row's 18 port words into the port buffers while the row is in cache (the padding rows' ports are one fill afterwards), so the 2^15-instance z (64 MiB) is neither allocated nor copied into the stack, and the ports are no longer gathered by 18 strided passes over it. The flock reduction reads z from the stack. Measured on the 7950X3D, #463 (ddbef89c) vs this, aggregate --n 100, 1T: - Build witness span 43.3/43.4 -> 38.0/38.6 ms (first-level node and node, 2 traced runs each) - end to end at 1T within noise (load 2 to 12); 32T first-level node 155/152 -> 139/153 ms, node 172/188 -> 166/185 ms Proof digests unchanged (leaf, first-level node, node). --- crates/leanvm_core/src/rec/proof.rs | 72 ++++++++++++++++++----------- 1 file changed, 44 insertions(+), 28 deletions(-) diff --git a/crates/leanvm_core/src/rec/proof.rs b/crates/leanvm_core/src/rec/proof.rs index bf33d4514..532a99e7d 100644 --- a/crates/leanvm_core/src/rec/proof.rs +++ b/crates/leanvm_core/src/rec/proof.rs @@ -16,6 +16,7 @@ use crate::constraints::{Columns, ConstraintError}; use crate::pcs::{Rate, RingSwitch, SliceClaim, StackClaim}; use crate::{constraints, pcs, witness}; use fiat_shamir::transcript::{Challenger, ProofTranscript, ProverState, RawProof, VerifierState}; +use parallel::SendPtr; use primitives::field::{F64, F192}; /// The transcript's public input for a statement: the hash of its words' limbs, in order. @@ -24,10 +25,9 @@ pub fn statement_seed(statement: &[Limbs]) -> [F64; 4] { chain(&limbs).map(F64) } -/// The hash table's flock batch, one instance per row. +/// The hash table's flock batch, one instance per row; its `z` is the stack's packed-witness window. struct HashBatch { tau: usize, - z: Vec, a: Vec, b: Vec, z_lincheck: Vec, @@ -47,28 +47,49 @@ pub(crate) struct TableArgument<'a> { } impl HashBatch { - /// Every hash row's packed witness, at the hash table's height. - fn build(hash: &[Compression], tau: usize) -> Self { - let (z, a, b, z_lincheck) = - HashFlock::circuit().generate_witness_with(hash, &Compression::PADDING, tau, |row, z, az, bz| { - crate::rv::circuits::blake2s_witness(row.inputs(), z, az, bz); - }); - Self { - tau, + /// Every hash row's packed witness, at the hash table's height, written into `z`, and each row's ports into `ports`. + fn build(hash: &[Compression], tau: usize, z: &mut [F64], ports: &mut [&mut [F64]]) -> Self { + // SAFETY: `F64` is `repr(transparent)` over `u64`, and the packing is bit `i` at position `i` on both sides. + let z = unsafe { std::slice::from_raw_parts_mut(z.as_mut_ptr().cast::(), z.len()) }; + let port_ptrs: Vec> = ports.iter_mut().map(|p| SendPtr(p.as_mut_ptr())).collect(); + // A live row's ports are copied while its words are in cache; the padding rows' after. + let ports_of = |row: &Compression, z: &[u64]| { + let offset = std::ptr::from_ref(row).addr().wrapping_sub(hash.as_ptr().addr()); + let j = offset / size_of::(); + if j < hash.len() && std::ptr::eq(&hash[j], row) { + for (ptr, &word) in port_ptrs.iter().zip(z) { + // SAFETY: row `j` is visited once, and each port buffer holds every row. + unsafe { *ptr.0.add(j) = F64(word) }; + } + } + }; + let (a, b, z_lincheck) = HashFlock::circuit().generate_witness_with_into( z, - a, - b, - z_lincheck, + hash, + &Compression::PADDING, + tau, + |row, z, az, bz| crate::rv::circuits::blake2s_witness(row.inputs(), z, az, bz), + ports_of, + ); + if hash.len() < 1 << tau { + let padding = &z[hash.len() << HashFlock::stride_log()..]; + for (port, &word) in ports.iter_mut().zip(padding) { + port[hash.len()..].fill(F64(word)); + } } + Self { tau, a, b, z_lincheck } } /// Prove every hash row: zerocheck, then lincheck, a batch of one circuit, to one ring-switched claim on the packed witness. - fn prove(self, layout: &RecLayout, ps: &mut ProverState) -> RingSwitch { + fn prove(self, layout: &RecLayout, q: &[F64], ps: &mut ProverState) -> RingSwitch { let window = layout.hash_window(); + let z = &q[window.offset..window.offset + (1 << window.n_vars)]; + // SAFETY: `F64` is `repr(transparent)` over `u64`. + let z = unsafe { std::slice::from_raw_parts(z.as_ptr().cast::(), z.len()) }; let instance = flock::reduction::Instance { block: HashFlock::circuit().block(), n_blocks_log: self.tau, - z: &self.z, + z, a: &self.a, b: &self.b, z_lincheck: &self.z_lincheck, @@ -109,18 +130,13 @@ impl RecWitness { table.fill(a, &mut windows[RecLayout::columns(table)]); } - let batch = HashBatch::build(&a.hash, layout.tau(Table::Hash)); - parallel::chunks_mut_zip(windows[RecLayout::HASH_WITNESS], &batch.z, 1 << 16, |_, dst, src| { - for (d, &s) in dst.iter_mut().zip(src) { - *d = F64(s); - } - }); - let stride_log = HashFlock::stride_log(); - for port in 0..HashFlock::N_PORTS { - for (j, cell) in windows[ports_at + port].iter_mut().enumerate() { - *cell = F64(batch.z[(j << stride_log) + port]); - } - } + let (head, rest) = windows.split_at_mut(ports_at); + let batch = HashBatch::build( + &a.hash, + layout.tau(Table::Hash), + head[RecLayout::HASH_WITNESS], + &mut rest[..HashFlock::N_PORTS], + ); drop(windows); Self { q, ports, batch } } @@ -230,7 +246,7 @@ impl Circuit { let RecWitness { q, ports, batch } = w; drop(ports); - let ring = crate::stage!("Flock reduction", || batch.prove(&layout, &mut ps)); + let ring = crate::stage!("Flock reduction", || batch.prove(&layout, &q, &mut ps)); crate::stage!("PCS open", || pcs::open(&mut ps, &committed, &q, &slots, &[ring])); Ok(ps.into_proof()) } From 0d2c7e0f1cb0e7f51ca07b32625194bfb41dc680 Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Mon, 5 Oct 2026 23:24:23 +0100 Subject: [PATCH 14/18] rec::tree::reduce: dense reduction over the live prefix, factored weights, fused rounds The node's dense reduction was a degree-two sumcheck over the whole fixed polynomial (2^22 entries) against a full weight table, each round a products pass and two fold passes. Prover only; the transcript, the verifier and every proof byte are unchanged. - Live prefix. The weight table is zero past the blocks its terms touch (about 1.45M of the node's 4.2M entries), so every round folds and sums that prefix alone. Where the prefix's last entry pairs with one past it, that entry is the polynomial's value on its block at the challenges so far, an `mle_eval_par` of the table as given; those blocks partition the rest of the table, so the final value is still exact. - Factored weights. Terms of one block sharing their low coordinates are one group (each of the node's blocks has two terms at one point, so one group). While every block has at most two groups, the weight table is never written out: a round sums each block's low table against the table's run, then scales by the high entry, until the low tables are bound (10 rounds); then it is written out once at that level and folded from there. Otherwise (the first level's bytecode, 18 groups) it is written out at the start, in one pass with the first message. - Fused rounds. Each round folds the table (and the weights once written) tile by tile and builds the next round's message from the folded tile in L1, with `mul4`, `mul_base8` and unreduced `mul_unreduced4` sums reduced once per message. - The matrix reduction's row and column rounds fold and build the next message in one pass the same way (`FoldTable::fold_pair`). - `round(i)` is now `message()`, the message the last bind built. A new test reduces claims whose blocks leave a prefix odd at several rounds, with factored and written-out parts, and checks the verifier accepts and the values are the polynomials'. Measured on the 7950X3D, `aggregate --n 100 --leaves 4 --arity0 2 --arity 2 --repeat 3`, AGG1_NOTREE, alternating base (f774828d) and branch runs; "proving" means of 3 passes, runs with 1-minute load above 20 or a disturbed pass dropped (loads in parentheses). 1T (taskset -c 2): - first-level node: base 0.973 s (9.1), 0.991 s (6.5), 0.972 s (9.9), median 0.973 s; branch 0.971 s (13.7), 0.958 s (10.7), median 0.965 s - node: base 1.042 s (9.1), 1.052 s (6.5), 1.039 s (9.9), median 1.042 s; branch 0.993 s (13.7), 0.975 s (10.7), and per-pass medians 0.990 s, 0.980 s of two runs with one disturbed pass each, median 0.985 s (-5.5%) 32T: - first-level node: base 0.149 s (9.6), 0.153 s (11.6), 0.160 s (13.7), median 0.153 s; branch 0.149 s (10.8), 0.149 s (10.1), 0.158 s (13.7), median 0.149 s - node: base 0.176 s (9.6), 0.203 s (11.6), 0.192 s (13.7), median 0.192 s; branch 0.156 s (10.8), 0.155 s (10.1), 0.169 s (13.7), median 0.156 s Spans (--tracing, one pass), base then branch: - 1T first-level: Reduce claims 25.7 to 13.2 ms (Dense reduction 17.2 to 5.25 ms, Matrix reduction 8.55 to 7.90 ms) - 1T node: Reduce claims 85.4 to 25.4 ms (Dense reduction 70.7 to 12.0 ms, Matrix reduction 14.7 to 13.3 ms) - 32T first-level: Reduce claims 8.99 to 6.33 ms (Dense reduction 2.61 to 0.99 ms) - 32T node: Reduce claims 35.1 to 12.9 ms (Dense reduction 25.6 to 4.53 ms, Matrix reduction 9.47 to 8.35 ms) Digests equal base at --n 100 (leaf 1ddc4572a71ec2af, first 6cc4ab4bf8ee46b7, node d0d2b1643bb63a90) and --n 400 (leaf 9311cba19de1b111, first cac5d72e17c200f1, node bc70f47ff4643d4b). --- .../leanvm_core/src/rec/tree/reduce/dense.rs | 463 ++++++++++++++---- .../leanvm_core/src/rec/tree/reduce/matrix.rs | 96 ++-- crates/leanvm_core/src/rec/tree/reduce/mod.rs | 151 ++++-- crates/leanvm_core/src/rec/tree/tests.rs | 75 ++- 4 files changed, 605 insertions(+), 180 deletions(-) diff --git a/crates/leanvm_core/src/rec/tree/reduce/dense.rs b/crates/leanvm_core/src/rec/tree/reduce/dense.rs index e2193889d..106250a29 100644 --- a/crates/leanvm_core/src/rec/tree/reduce/dense.rs +++ b/crates/leanvm_core/src/rec/tree/reduce/dense.rs @@ -9,13 +9,17 @@ //! lowest variable first. //! A polynomial of fewer variables is bound early, and its share then waits on the rest: each later round multiplies it by its challenge. -use super::{DenseTables, FoldTable, PAR_LEN, ReduceError, products_par, table_of}; +use super::{DenseTables, Entry, Msg, ReduceError, TILE, ZERO, xor}; use crate::arith::{Arith, Verifier}; use crate::rec::tree::claims::{DenseClaim, DensePoly, DenseTerm}; use fiat_shamir::transcript::{Challenger, ProverState, Transmitter}; -use primitives::field::{F64, F192}; -use primitives::multilinear::eq_table; +use parallel::Chunks; +use primitives::field::{F64, F192, F192Unreduced, mul_unreduced4}; +use primitives::multilinear::{eq_table, eq_table_seeded, mle_eval_par}; +use primitives::uninit_vec; +use std::cmp::Ordering; use std::collections::BTreeMap; +use std::ops::Range; /// Each dense polynomial's variables: the bytecode table's, the image's, the fixed polynomial's. #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -35,24 +39,38 @@ pub(crate) struct DenseReduced { pub(crate) struct DenseProver<'a> { /// The reduced polynomials, in order. parts: Vec>, + /// The challenges so far. + point: Vec, + /// The next round's message. + message: [F192; 2], } /// One reduced polynomial in the prover: its variables, its table, and its weight table `Xi_j`. +/// +/// The weight table is zero past the prefix its terms touch, so each round folds and sums that prefix alone. Where the +/// prefix's last entry pairs with one past it, that entry is the polynomial's value on its block at the challenges so +/// far, evaluated from the table as given. +/// +/// When each of its blocks is a sum of few low tables times high ones, the weight table stays factored while the low +/// tables have variables left: a round sums each low table against the table, then scales by the high one. Past that, +/// or from the start otherwise, it is a table. struct Part<'a> { /// The polynomial's variables. n_vars: usize, - /// Its values, folded. - table: Table<'a>, - /// Its weight table, folded. - weights: FoldTable, -} - -/// A polynomial's table: its values in `K` until the first round binds a variable. -enum Table<'a> { - /// The polynomial as given. - Base(&'a [F64]), - /// The polynomial folded by at least one round. - Ext(FoldTable), + /// Its values as given. + base: &'a [F64], + /// The length of the prefix its terms touch, at least one. + live: usize, + /// The rounds before its weight table is written out. + factored: usize, + /// Its weight table's blocks, folded, until it is written out. + blocks: Vec, + /// Its weight table once written out, folded. + weights: Vec, + /// Its values, folded from the first round on. + table: Vec, + /// The buffers the next fold writes. + spare: [Vec; 2], } /// One term of a weight table: `c eq(p, .)` on one aligned block of the table. @@ -210,7 +228,7 @@ impl<'a> DenseProver<'a> { let theta = ps.sample(); let mut prover = Self::new(vars, tables, claims, theta); for i in 0..prover.rounds() { - ps.add_scalars(&prover.round(i)); + ps.add_scalars(&prover.message()); let r = ps.sample(); prover.bind(i, r); } @@ -231,20 +249,23 @@ impl<'a> DenseProver<'a> { power *= theta; } } - let parts = (DensePoly::ALL.into_iter().zip(placed)) + let mut message = ZERO; + let parts: Vec = (DensePoly::ALL.into_iter().zip(placed)) .filter(|&(p, _)| reduced[p as usize]) .map(|(p, terms)| { let n_vars = vars.0[p as usize]; - let table = &tables.0[p as usize]; - assert_eq!(table.len(), 1 << n_vars, "a dense table has its variables"); - Part { - n_vars, - table: Table::Base(table), - weights: FoldTable::new(weight_table(n_vars, &terms)), - } + let base = &tables.0[p as usize]; + assert_eq!(base.len(), 1 << n_vars, "a dense table has its variables"); + let (part, first) = Part::new(n_vars, base, &terms); + message = xor(message, first); + part }) .collect(); - Self { parts } + Self { + point: Vec::with_capacity(parts.iter().map(|p| p.n_vars).max().unwrap_or(0)), + parts, + message: message.map(F192Unreduced::reduce), + } } /// The number of rounds. @@ -252,86 +273,360 @@ impl<'a> DenseProver<'a> { self.parts.iter().map(|p| p.n_vars).max().unwrap_or(0) } - /// Round `i`'s message: `h(0)` and the leading coefficient, the claim fixing the linear one. + /// The next round's message: `h(0)` and the leading coefficient, the claim fixing the linear one. /// /// A polynomial bound already contributes a linear term only, which the claim accounts for. - pub(crate) fn round(&self, i: usize) -> [F192; 2] { - (self.parts.iter().filter(|p| p.n_vars > i)).fold([F192::ZERO; 2], |[c0, c2], p| { - let w = &p.weights.values; - let [a, b] = match &p.table { - Table::Base(t) => products_par(w, t, F192::mul_base), - Table::Ext(t) => products_par(w, &t.values, |x, y| x * y), - }; - [c0 + a, c2 + b] - }) + pub(crate) const fn message(&self) -> [F192; 2] { + self.message } - /// Bind round `i`'s variable to `r`. + /// Bind round `i`'s variable to `r`, and build the next round's message from the folded pairs. pub(crate) fn bind(&mut self, i: usize, r: F192) { - for p in self.parts.iter_mut().filter(|p| p.n_vars > i) { - p.weights.fold(r, true); - match &mut p.table { - Table::Base(t) => { - let t = *t; - let folded = table_of(t.len() / 2, |k| { - r.mul_base(t[2 * k] + t[2 * k + 1]) + F192::from(t[2 * k]) - }); - p.table = Table::Ext(FoldTable::new(folded)); - } - Table::Ext(t) => t.fold(r, true), - } - } + assert_eq!(i, self.point.len(), "rounds come in order"); + self.point.push(r); + let point = &self.point; + let message = (self.parts.iter_mut().filter(|p| p.n_vars > i)).fold(ZERO, |acc, p| xor(acc, p.bind(point))); + self.message = message.map(F192Unreduced::reduce); } /// Each reduced polynomial's value at its prefix of the point. pub(crate) fn finals(&self) -> Vec { (self.parts.iter()) - .map(|p| match &p.table { - Table::Base(t) => F192::from(t[0]), - Table::Ext(t) => t.values[0], - }) + .map(|p| p.table.first().copied().unwrap_or_else(|| F192::from(p.base[0]))) .collect() } } -/// `Xi(x) = sum_t coef_t eq(p_t, x)` over `2^n` entries, each block's terms added in one pass. -fn weight_table(n: usize, terms: &[Placed]) -> Vec { - let mut table = table_of(1 << n, |_| F192::ZERO); - let mut blocks: BTreeMap<(usize, usize), Vec<&Placed>> = BTreeMap::new(); - for t in terms { - blocks.entry((t.offset, t.point.len())).or_default().push(t); +impl<'a> Part<'a> { + /// The part of the polynomial `base` of `n_vars` variables, weighed by `terms`, and its first round's message. + fn new(n_vars: usize, base: &'a [F64], terms: &[Placed]) -> (Self, Msg) { + let blocks = Block::all(terms); + let factored = if blocks.iter().all(|b| b.groups.len() <= FACTORED_GROUPS) { + blocks + .iter() + .map(|b| b.row.trailing_zeros() as usize) + .min() + .unwrap_or(0) + } else { + 0 + }; + let live = terms + .iter() + .map(|t| t.offset + (1 << t.point.len())) + .max() + .unwrap_or(0) + .max(1); + let mut part = Self { + n_vars, + base, + live, + factored, + blocks: Vec::new(), + weights: Vec::new(), + table: Vec::new(), + spare: [Vec::new(), Vec::new()], + }; + let len = part.len_at(0); + let (source, mut weights) = if factored > 0 { + (Source::Blocks(&blocks), Vec::new()) + } else { + // SAFETY: the pass writes every entry before any is read. + (Source::Write(&blocks), unsafe { uninit_vec(len) }) + }; + let wc = Chunks::new(&mut weights, PAR_LEN); + let message = chunked(len, |c| { + let (from, to) = (c * PAR_LEN, len.min((c + 1) * PAR_LEN)); + let touching = source.touching(from, to); + // SAFETY: each chunk index is taken once, and the table outlives the dispatch. + let wo: &mut [F192] = if wc.count() == 0 { &mut [] } else { unsafe { wc.get(c) } }; + (from..to).step_by(TILE).fold(ZERO, |m, a| { + let b = to.min(a + TILE); + let wo = wo.get_mut(a - from..b - from).unwrap_or_default(); + xor(m, source.message(&touching, a, &base[a..b], wo, F192::ZERO)) + }) + }); + part.blocks = blocks; + part.weights = weights; + (part, message) } - for block in blocks.values() { - add_eqs(&mut table, block); + + /// The entries kept after `i` rounds: the pairs covering the live prefix, or the value once bound. + const fn len_at(&self, i: usize) -> usize { + if i < self.n_vars { + 2 * self.live.div_ceil(2 << i) + } else { + 1 + } + } + + /// Bind the next variable to the last challenge of `point`; the next round's message, zero after the last. + fn bind(&mut self, point: &[F192]) -> Msg { + let i = point.len() - 1; + let r = point[i]; + if i < self.factored { + self.blocks.iter_mut().for_each(|b| b.fold(r)); + } + let half = self.live.div_ceil(2 << i); + let len = self.len_at(i + 1); + let extra = (len > half).then(|| { + let size = 2 << i; + mle_eval_par(&self.base[half * size..(half + 1) * size], point) + }); + let source = match (i + 1).cmp(&self.factored) { + Ordering::Less => Source::Blocks(&self.blocks), + Ordering::Equal => Source::Write(&self.blocks), + Ordering::Greater => Source::Fold(&self.weights), + }; + let w_len = if matches!(source, Source::Blocks(_)) || len == 1 { + 0 + } else { + len + }; + let [w_spare, t_spare] = &mut self.spare; + let mut w = buffer(w_spare, w_len); + let mut t = buffer(t_spare, len); + let message = if i == 0 { + fold_round(source, self.base, r, half, extra, &mut w, &mut t) + } else { + fold_round(source, &self.table, r, half, extra, &mut w, &mut t) + }; + if i + 1 == self.factored { + self.blocks = Vec::new(); + } + self.spare = [ + std::mem::replace(&mut self.weights, w), + std::mem::replace(&mut self.table, t), + ]; + message } - table } -/// The variables of the low eq tables a weight table's pass keeps in L1. -const LOW_VARS: usize = 10; +/// A buffer of `len` entries the caller writes before reading, reusing `spare`'s allocation. +fn buffer(spare: &mut Vec, len: usize) -> Vec { + let mut v = std::mem::take(spare); + if v.len() < len { + // SAFETY: the fold writes every entry before any is read. + return unsafe { uninit_vec(len) }; + } + v.truncate(len); + v +} + +/// Below this many entries a pass runs on the calling thread; a task of the pool takes this many. +const PAR_LEN: usize = 1 << 12; + +/// The sum of `task(c)` over the chunks of `n` entries, `PAR_LEN` each, across the pool. +fn chunked(n: usize, task: impl Fn(usize) -> Msg + Sync) -> Msg { + match n.div_ceil(PAR_LEN) { + 0 => ZERO, + 1 => task(0), + count => parallel::map_reduce(count, || ZERO, task, xor), + } +} + +/// Where a pass finds the weights of the entries it reaches. +#[derive(Clone, Copy)] +enum Source<'w> { + /// The blocks, at the pass's level: the weights stay factored. + Blocks(&'w [Block]), + /// The blocks, at the pass's level: the pass writes the weights out. + Write(&'w [Block]), + /// The table one level up: the pass folds it. + Fold(&'w [F192]), +} + +impl Source<'_> { + /// The blocks reaching entries `from..to`. + fn touching(&self, from: usize, to: usize) -> Vec<&Block> { + match *self { + Self::Blocks(blocks) | Self::Write(blocks) => (blocks.iter()) + .filter(|b| b.start < to && from < b.start + b.len) + .collect(), + Self::Fold(_) => Vec::new(), + } + } + + /// The message of the table's entries `a..a + t.len()`, `touching` the blocks reaching them. Unless the weights stay + /// factored, the pass writes them to `w`: from the blocks, or folding the table one level up by `r`. + fn message(&self, touching: &[&Block], a: usize, t: &[T], w: &mut [F192], r: F192) -> Msg { + match *self { + Self::Blocks(_) => touching.iter().fold(ZERO, |m, b| xor(m, b.message(a, t))), + Self::Write(_) => { + let mut acc = [F192Unreduced::ZERO; TILE]; + let acc = &mut acc[..w.len()]; + for b in touching { + b.add_to(a, acc); + } + for (w, x) in w.iter_mut().zip(acc) { + *w = x.reduce(); + } + T::dot(w, t) + } + Self::Fold(table) => { + F192::fold_into(&table[2 * a..2 * (a + w.len())], r, w); + T::dot(w, t) + } + } + } +} -/// `table[offset + x] += sum_t coef_t eq(p_t, x)` for terms of one block. +/// One round's fold of a part's table `t` into `t_out`, its weights as `source` says into `w_out`, and the next +/// round's message. /// -/// Each term's eq table is the tensor of a low one and a high one, the coefficient in the high one. -fn add_eqs(table: &mut [F192], terms: &[&Placed]) { - let k = terms[0].point.len(); - let l = k.min(LOW_VARS); - let lows: Vec> = terms.iter().map(|t| eq_table(&t.point[..l])).collect(); - let highs: Vec> = (terms.iter()) - .map(|t| eq_table(&t.point[l..]).into_iter().map(|e| e * t.coef).collect()) - .collect(); - let row = |h: usize, dst: &mut [F192]| { - for (low, high) in lows.iter().zip(&highs) { - for (slot, &e) in dst.iter_mut().zip(low) { - *slot += high[h] * e; +/// The first `half` entries fold from pairs. Past them, when `t_out` is one longer, are the weight's zero and the +/// table's `extra`. The message is over the folded pairs, none when one entry is left. +fn fold_round( + source: Source, + t: &[T], + r: F192, + half: usize, + extra: Option, + w_out: &mut [F192], + t_out: &mut [F192], +) -> Msg { + if t_out.len() == 1 { + t_out[0] = T::fold(t[0], t[1], r); + return ZERO; + } + let pairs = half & !1; + let tc = Chunks::new(&mut t_out[..pairs], PAR_LEN); + let w_pairs = pairs.min(w_out.len()); + let wc = Chunks::new(&mut w_out[..w_pairs], PAR_LEN); + let mut message = chunked(pairs, |c| { + let (from, to) = (c * PAR_LEN, pairs.min((c + 1) * PAR_LEN)); + let touching = source.touching(from, to); + // SAFETY: each chunk index is taken once, and both tables outlive the dispatch. + let (tables, wo): (&mut [F192], &mut [F192]) = + unsafe { (tc.get(c), if wc.count() == 0 { &mut [] } else { wc.get(c) }) }; + (from..to).step_by(TILE).fold(ZERO, |m, a| { + let b = to.min(a + TILE); + let tile = &mut tables[a - from..b - from]; + T::fold_into(&t[2 * a..2 * b], r, tile); + let wo = wo.get_mut(a - from..b - from).unwrap_or_default(); + xor(m, source.message(&touching, a, tile, wo, r)) + }) + }); + if let Some(extra) = extra { + let k = half - 1; + t_out[k..].copy_from_slice(&[T::fold(t[2 * k], t[2 * k + 1], r), extra]); + let w = &mut w_out[k..]; + w[0] = match source { + Source::Blocks(_) => unreachable!("a factored level's prefix is whole pairs"), + Source::Write(blocks) => { + let mut acc = [F192Unreduced::ZERO]; + blocks.iter().for_each(|b| b.add_to(k, &mut acc)); + acc[0].reduce() } + Source::Fold(table) => F192::fold(table[2 * k], table[2 * k + 1], r), + }; + w[1] = F192::ZERO; + message = xor(message, F192::dot(w, &t_out[k..])); + } + message +} + +/// The variables of a block's low tables, which a pass keeps in L1. +const LOW_VARS: usize = 10; + +/// The most groups a block may have for its part's weights to stay factored. +const FACTORED_GROUPS: usize = 2; + +/// One aligned block of a weight table at the current level, `sum_g low_g ⊗ high_g`: terms sharing their low +/// coordinates share the low table, against the sum of their high ones scaled by their coefficients. +struct Block { + /// The block's first entry. + start: usize, + /// Its entries. + len: usize, + /// The entries of its low tables, one row of the block. + row: usize, + /// Each group's low and high tables. + groups: Vec<(Vec, Vec)>, +} + +impl Block { + /// The blocks of a weight table's terms, at the first level. + fn all(terms: &[Placed]) -> Vec { + let mut by_block: BTreeMap<(usize, usize), Vec<&Placed>> = BTreeMap::new(); + for t in terms { + by_block.entry((t.offset, t.point.len())).or_default().push(t); } - }; - let offset = terms[0].offset; - let dst = &mut table[offset..offset + (1 << k)]; - if dst.len() < PAR_LEN { - dst.chunks_mut(1 << l).enumerate().for_each(|(h, d)| row(h, d)); - } else { - parallel::chunks_mut(dst, 1 << l, row); + by_block.values().map(|terms| Self::new(terms)).collect() + } + + /// The block of these terms, which share their offset and length. + fn new(terms: &[&Placed]) -> Self { + let k = terms[0].point.len(); + let l = k.min(LOW_VARS); + let mut groups: Vec<(&[F192], Vec)> = Vec::new(); + for t in terms { + let (low, high) = t.point.split_at(l); + let high = eq_table_seeded(high, t.coef); + match groups.iter_mut().find(|(p, _)| *p == low) { + Some((_, sum)) => sum.iter_mut().zip(high).for_each(|(s, e)| *s += e), + None => groups.push((low, high)), + } + } + Self { + start: terms[0].offset, + len: 1 << k, + row: 1 << l, + groups: groups.into_iter().map(|(p, high)| (eq_table(p), high)).collect(), + } + } + + /// Bind the lowest variable, one of the low tables', to `r`. + fn fold(&mut self, r: F192) { + for (low, _) in &mut self.groups { + let mut folded = vec![F192::ZERO; low.len() / 2]; + F192::fold_into(low, r, &mut folded); + *low = folded; + } + (self.start, self.len, self.row) = (self.start / 2, self.len / 2, self.row / 2); + } + + /// `f(h, x, run)` for each run of entries `a..a + n` in the block, `run` in row `h` from column `x`. + fn runs(&self, a: usize, n: usize, mut f: impl FnMut(usize, usize, Range)) { + let (mut lo, hi) = (a.max(self.start), (a + n).min(self.start + self.len)); + while lo < hi { + let (h, x) = ((lo - self.start) / self.row, (lo - self.start) % self.row); + let end = hi.min(lo - x + self.row); + f(h, x, lo..end); + lo = end; + } + } + + /// `acc[y] += Xi(a + y)` over the block. + fn add_to(&self, a: usize, acc: &mut [F192Unreduced]) { + self.runs(a, acc.len(), |h, x, run| { + let acc = &mut acc[run.start - a..run.end - a]; + for (low, high) in &self.groups { + let (e, low) = (high[h], &low[x..x + acc.len()]); + let (quads, rest) = low.as_chunks::<4>(); + let (acc_quads, acc_rest) = acc.as_chunks_mut::<4>(); + for (s, &q) in acc_quads.iter_mut().zip(quads) { + for (s, p) in s.iter_mut().zip(mul_unreduced4([e; 4], q)) { + *s ^= p; + } + } + for (s, &l) in acc_rest.iter_mut().zip(rest) { + *s ^= e.mul_unreduced(l); + } + } + }); + } + + /// The message of the table's entries `a..a + t.len()` against the block's weights: each row's run is summed + /// against the low tables, then scaled by the high ones. A run is whole pairs while a row has two entries. + fn message(&self, a: usize, t: &[T]) -> Msg { + let mut m = ZERO; + self.runs(a, t.len(), |h, x, run| { + let t = &t[run.start - a..run.end - a]; + for (low, high) in &self.groups { + let [c0, c2] = T::dot(&low[x..x + t.len()], t).map(F192Unreduced::reduce); + m = xor(m, [high[h].mul_unreduced(c0), high[h].mul_unreduced(c2)]); + } + }); + m } } diff --git a/crates/leanvm_core/src/rec/tree/reduce/matrix.rs b/crates/leanvm_core/src/rec/tree/reduce/matrix.rs index 586ae77d2..3b7145e01 100644 --- a/crates/leanvm_core/src/rec/tree/reduce/matrix.rs +++ b/crates/leanvm_core/src/rec/tree/reduce/matrix.rs @@ -7,7 +7,7 @@ //! //! A circuit of fewer variables waits on the rest in each phase, as the dense reduction's smaller polynomials do. -use super::{FoldTable, ReduceError, products}; +use super::{FoldTable, Msg, ReduceError, ZERO, xor}; use crate::arith::{Arith, Verifier}; use crate::class_flock; use crate::rec::tree::claims::{Coefficient, ColWeight, MatrixClaim, RowWeight}; @@ -15,7 +15,8 @@ use crate::rec::verifier::SkipDomain; use fiat_shamir::transcript::{Challenger, ProverState, Transmitter}; use flock::lincheck::{LincheckCircuit, build_quirky_eq_table}; use flock::zerocheck::K_SKIP; -use primitives::field::F192; +use parallel::Chunks; +use primitives::field::{F192, F192Unreduced}; use primitives::multilinear::eq_table; /// What the matrix reduction leaves: the row and column points, and each circuit's `A` and `B` at its prefixes of them. @@ -32,12 +33,16 @@ pub(crate) struct MatrixReduced { pub(crate) struct MatrixProver { /// Each claim's tables. claims: Vec, + /// The next round's message. + message: [F192; 2], } /// The matrix reduction's prover, in its column phase: per circuit, `A(r, .)`, `Theta_A`, `B(r, .)` and `Theta_B`. pub(crate) struct ColumnPhase { /// Each circuit's four tables. circuits: Vec<[FoldTable; 4]>, + /// The next round's message. + message: [F192; 2], } /// One claim's tables in the row phase. @@ -230,7 +235,7 @@ impl MatrixProver { let mut rows = Self::new(claims, theta); let r: Vec = (0..class_flock::max_k_log()) .map(|i| { - ps.add_scalars(&rows.round(i)); + ps.add_scalars(&rows.message()); let x = ps.sample(); rows.bind(i, x); x @@ -238,7 +243,7 @@ impl MatrixProver { .collect(); let mut cols = rows.columns(claims, &r); for i in 0..class_flock::max_k_log() { - ps.add_scalars(&cols.round(i)); + ps.add_scalars(&cols.message()); let x = ps.sample(); cols.bind(i, x); } @@ -250,7 +255,7 @@ impl MatrixProver { let powers: Vec = std::iter::successors(Some(F192::ONE), |&p| Some(p * theta)) .take(claims.len()) .collect(); - let claims = parallel::map_collect(claims.len(), |i| { + let mut claims = parallel::map_collect(claims.len(), |i| { let c = &claims[i]; let u: Vec = c.row.table().into_iter().map(|x| powers[i] * x).collect(); let (ra, rb) = class_flock::circuit(c.circuit).row_values(&c.col.table()); @@ -262,40 +267,32 @@ impl MatrixProver { g: FoldTable::new(g), } }); - Self { claims } + let message = each_mut( + &mut claims, + |_, t| if t.k > 0 { FoldTable::message(&t.u, &t.g) } else { ZERO }, + ); + Self { claims, message } } - /// Row round `i`'s message: `h(0)` and the leading coefficient. - pub(crate) fn round(&self, i: usize) -> [F192; 2] { - let add = |[a, b]: [F192; 2], [c, d]: [F192; 2]| [a + c, b + d]; - parallel::map_reduce( - self.claims.len(), - || [F192::ZERO; 2], - |c| { - let t = &self.claims[c]; - if t.k > i { - products(&t.u.values, &t.g.values) - } else { - [F192::ZERO; 2] - } - }, - add, - ) + /// The next row round's message: `h(0)` and the leading coefficient. + pub(crate) const fn message(&self) -> [F192; 2] { + self.message } - /// Bind row round `i`'s variable to `x`. + /// Bind row round `i`'s variable to `x`, and build the next round's message from the folded pairs. pub(crate) fn bind(&mut self, i: usize, x: F192) { - parallel::for_each_mut(&mut self.claims, |_, t| { + self.message = each_mut(&mut self.claims, |_, t| { if t.k > i { - t.u.fold(x, false); - t.g.fold(x, false); + FoldTable::fold_pair(&mut t.u, &mut t.g, x) + } else { + ZERO } }); } /// The column phase at the row point `r`: per circuit, two backward walks for `A(r, .)` and `B(r, .)`, and the claims' column weights. pub(crate) fn columns(&self, claims: &[MatrixClaim], r: &[F192]) -> ColumnPhase { - let circuits = parallel::map_collect(class_flock::N_FLOCKS, |f| { + let mut circuits = parallel::map_collect(class_flock::N_FLOCKS, |f| { let k = k_of(f); let circuit = class_flock::circuit(f); let eq = eq_table(&r[..k]); @@ -315,33 +312,40 @@ impl MatrixProver { } [at, wa, bt, wb].map(FoldTable::new) }); - ColumnPhase { circuits } + let message = each_mut(&mut circuits, |f, [at, wa, bt, wb]| { + if k_of(f) > 0 { + xor(FoldTable::message(wa, at), FoldTable::message(wb, bt)) + } else { + ZERO + } + }); + ColumnPhase { circuits, message } } } +/// The sum of `f(i, item)` over the items, each on a task of the pool, reduced. +fn each_mut(items: &mut [T], f: impl Fn(usize, &mut T) -> Msg + Sync) -> [F192; 2] { + let chunks = Chunks::new(items, 1); + let task = |i: usize| { + // SAFETY: each item is taken once, and the items outlive the dispatch. + f(i, &mut unsafe { chunks.get(i) }[0]) + }; + parallel::map_reduce(chunks.count(), || ZERO, task, xor).map(F192Unreduced::reduce) +} + impl ColumnPhase { - /// Column round `i`'s message. - pub(crate) fn round(&self, i: usize) -> [F192; 2] { - let add = |[a, b]: [F192; 2], [c, d]: [F192; 2]| [a + c, b + d]; - parallel::map_reduce( - self.circuits.len(), - || [F192::ZERO; 2], - |f| { - if k_of(f) <= i { - return [F192::ZERO; 2]; - } - let [at, wa, bt, wb] = &self.circuits[f]; - add(products(&wa.values, &at.values), products(&wb.values, &bt.values)) - }, - add, - ) + /// The next column round's message. + pub(crate) const fn message(&self) -> [F192; 2] { + self.message } - /// Bind column round `i`'s variable to `x`. + /// Bind column round `i`'s variable to `x`, and build the next round's message from the folded pairs. pub(crate) fn bind(&mut self, i: usize, x: F192) { - parallel::for_each_mut(&mut self.circuits, |f, tables| { + self.message = each_mut(&mut self.circuits, |f, [at, wa, bt, wb]| { if k_of(f) > i { - tables.iter_mut().for_each(|t| t.fold(x, false)); + xor(FoldTable::fold_pair(wa, at, x), FoldTable::fold_pair(wb, bt, x)) + } else { + ZERO } }); } diff --git a/crates/leanvm_core/src/rec/tree/reduce/mod.rs b/crates/leanvm_core/src/rec/tree/reduce/mod.rs index 726bf0121..7efc82afc 100644 --- a/crates/leanvm_core/src/rec/tree/reduce/mod.rs +++ b/crates/leanvm_core/src/rec/tree/reduce/mod.rs @@ -12,8 +12,7 @@ use super::claims::{DensePoly, NodeClaims}; use crate::arith::Verifier; use crate::rec::circuit::{Limbs, digest_limbs}; use fiat_shamir::transcript::{ProofTranscript, ProverState, TranscriptError, Transmitter}; -use primitives::field::{F64, F192}; -use std::ops::Add; +use primitives::field::{F64, F192, F192Unreduced, mul_base8, mul_unreduced4, mul4}; use thiserror::Error; mod dense; @@ -93,43 +92,97 @@ impl NodeClaims { } } -/// The table of `f(k)` for `k < n`, written across the pool. -fn table_of(n: usize, f: impl Fn(usize) -> F192 + Sync) -> Vec { - if n < PAR_LEN { - return (0..n).map(f).collect(); - } - parallel::map_collect(n, f) +/// A round's message, `h(0)` and the leading coefficient, unreduced. +type Msg = [F192Unreduced; 2]; + +/// The zero message. +const ZERO: Msg = [F192Unreduced::ZERO; 2]; + +/// The sum of two messages. +fn xor([a, b]: Msg, [c, d]: Msg) -> Msg { + [a ^ c, b ^ d] } -/// Below this many entries a pass runs on the calling thread. -const PAR_LEN: usize = 1 << 12; +/// The entries a pass keeps in L1 at once. +const TILE: usize = 64; + +/// An entry of a table a round folds: in `K` before the first round, in `E` after. +trait Entry: Copy + Sync { + /// `a + r (a + b)`. + fn fold(a: Self, b: Self, r: F192) -> F192; -/// `h(0)` and the leading coefficient of `sum_k u(X, k) g(X, k)` over the lowest variable, on the calling thread. -fn products(u: &[F192], g: &[F192]) -> [F192; 2] { - (0..u.len() / 2).fold([F192::ZERO; 2], |[c0, c2], k| { - let (u0, u1, g0, g1) = (u[2 * k], u[2 * k + 1], g[2 * k], g[2 * k + 1]); - [c0 + u0 * g0, c2 + (u0 + u1) * (g0 + g1)] - }) + /// The folds of `t`'s pairs into `out`. + fn fold_into(t: &[Self], r: F192, out: &mut [F192]); + + /// The message of `sum_k w(X, k) t(X, k)` over the lowest variable: `sum_k w(2k) t(2k)` and + /// `sum_k (w(2k) + w(2k + 1)) (t(2k) + t(2k + 1))`. + fn dot(w: &[F192], t: &[Self]) -> Msg; } -/// The same sums across the pool, `g` of either field, with its product by an element of `E`. -fn products_par(u: &[F192], g: &[G], times: impl Fn(F192, G) -> F192 + Sync) -> [F192; 2] -where - G: Copy + Sync + Add, -{ - let pairs = u.len() / 2; - let task = |i: usize| { - (i * PAR_LEN..pairs.min((i + 1) * PAR_LEN)).fold([F192::ZERO; 2], |[c0, c2], k| { - let (u0, u1, g0, g1) = (u[2 * k], u[2 * k + 1], g[2 * k], g[2 * k + 1]); - [c0 + times(u0, g0), c2 + times(u0 + u1, g0 + g1)] +impl Entry for F192 { + fn fold(a: Self, b: Self, r: F192) -> F192 { + a + r * (a + b) + } + + fn fold_into(t: &[Self], r: F192, out: &mut [F192]) { + let (t8, t_rest) = t.as_chunks::<8>(); + let (out4, out_rest) = out.as_chunks_mut::<4>(); + for (o, x) in out4.iter_mut().zip(t8) { + let d = mul4([r; 4], [x[0] + x[1], x[2] + x[3], x[4] + x[5], x[6] + x[7]]); + *o = [x[0] + d[0], x[2] + d[1], x[4] + d[2], x[6] + d[3]]; + } + for (o, x) in out_rest.iter_mut().zip(t_rest.as_chunks::<2>().0) { + *o = Self::fold(x[0], x[1], r); + } + } + + fn dot(w: &[F192], t: &[Self]) -> Msg { + let (w4, w_rest) = w.as_chunks::<4>(); + let (t4, t_rest) = t.as_chunks::<4>(); + let mut m = ZERO; + for (a, b) in w4.iter().zip(t4) { + let p = mul_unreduced4( + [a[0], a[0] + a[1], a[2], a[2] + a[3]], + [b[0], b[0] + b[1], b[2], b[2] + b[3]], + ); + m = xor(m, [p[0] ^ p[2], p[1] ^ p[3]]); + } + for (a, b) in w_rest.as_chunks::<2>().0.iter().zip(t_rest.as_chunks::<2>().0) { + m = xor(m, [a[0].mul_unreduced(b[0]), (a[0] + a[1]).mul_unreduced(b[0] + b[1])]); + } + m + } +} + +impl Entry for F64 { + fn fold(a: Self, b: Self, r: F192) -> F192 { + F192::from(a) + r.mul_base(a + b) + } + + fn fold_into(t: &[Self], r: F192, out: &mut [F192]) { + let (t16, t_rest) = t.as_chunks::<16>(); + let (out8, out_rest) = out.as_chunks_mut::<8>(); + for (o, x) in out8.iter_mut().zip(t16) { + let d = mul_base8(r, std::array::from_fn(|k| x[2 * k] + x[2 * k + 1])); + *o = std::array::from_fn(|k| d[k] + F192::from(x[2 * k])); + } + for (o, x) in out_rest.iter_mut().zip(t_rest.as_chunks::<2>().0) { + *o = Self::fold(x[0], x[1], r); + } + } + + fn dot(w: &[F192], t: &[Self]) -> Msg { + let pairs = w.as_chunks::<2>().0.iter().zip(t.as_chunks::<2>().0); + pairs.fold(ZERO, |m, (a, b)| { + xor( + m, + [ + a[0].mul_base_unreduced(b[0]), + (a[0] + a[1]).mul_base_unreduced(b[0] + b[1]), + ], + ) }) - }; - let chunks = pairs.div_ceil(PAR_LEN); - if chunks <= 1 { - return task(0); } - let add = |[a, b]: [F192; 2], [c, d]: [F192; 2]| [a + c, b + d]; - parallel::map_reduce(chunks, || [F192::ZERO; 2], task, add) } /// A table the prover folds one variable at a time, lowest first, into a second buffer it then swaps with. @@ -150,19 +203,25 @@ impl FoldTable { } } - /// Bind the lowest variable to `r`, on the calling thread or across the pool. - fn fold(&mut self, r: F192, par: bool) { - let n = self.values.len() / 2; - self.spare.resize(n, F192::ZERO); - let t = &self.values; - let at = |k: usize| t[2 * k] + r * (t[2 * k] + t[2 * k + 1]); - if par && n >= PAR_LEN { - parallel::fill(&mut self.spare, at); - } else { - for (k, slot) in self.spare.iter_mut().enumerate() { - *slot = at(k); - } - } - std::mem::swap(&mut self.values, &mut self.spare); + /// The message of `sum_k u(X, k) g(X, k)` over the lowest variable, on the calling thread. + fn message(u: &Self, g: &Self) -> Msg { + F192::dot(&u.values, &g.values) + } + + /// Bind the lowest variable of `u` and `g` to `r`, and the next round's message from the folded pairs. + fn fold_pair(u: &mut Self, g: &mut Self, r: F192) -> Msg { + let n = u.values.len() / 2; + u.spare.resize(n, F192::ZERO); + g.spare.resize(n, F192::ZERO); + let outs = u.spare.chunks_mut(TILE).zip(g.spare.chunks_mut(TILE)); + let ins = u.values.chunks(2 * TILE).zip(g.values.chunks(2 * TILE)); + let m = (outs.zip(ins)).fold(ZERO, |m, ((uo, go), (ui, gi))| { + F192::fold_into(ui, r, uo); + F192::fold_into(gi, r, go); + xor(m, F192::dot(uo, go)) + }); + std::mem::swap(&mut u.values, &mut u.spare); + std::mem::swap(&mut g.values, &mut g.spare); + m } } diff --git a/crates/leanvm_core/src/rec/tree/tests.rs b/crates/leanvm_core/src/rec/tree/tests.rs index 230bdf50b..3127b22a4 100644 --- a/crates/leanvm_core/src/rec/tree/tests.rs +++ b/crates/leanvm_core/src/rec/tree/tests.rs @@ -406,7 +406,7 @@ fn forged_reduction(vars: &DenseVars, tables: &DenseTables, claims: &NodeClaims< let mut dense = DenseProver::new(vars, tables, &claims.dense, theta); let point: Vec = (0..dense.rounds()) .map(|i| { - ps.add_scalars(&dense.round(i)); + ps.add_scalars(&dense.message()); let r = ps.sample(); dense.bind(i, r); r @@ -428,7 +428,7 @@ fn forged_reduction(vars: &DenseVars, tables: &DenseTables, claims: &NodeClaims< let mut rows = MatrixProver::new(&claims.matrices, theta); let r: Vec = (0..crate::class_flock::max_k_log()) .map(|i| { - ps.add_scalars(&rows.round(i)); + ps.add_scalars(&rows.message()); let x = ps.sample(); rows.bind(i, x); x @@ -437,7 +437,7 @@ fn forged_reduction(vars: &DenseVars, tables: &DenseTables, claims: &NodeClaims< let mut cols = rows.columns(&claims.matrices, &r); let s: Vec = (0..crate::class_flock::max_k_log()) .map(|i| { - ps.add_scalars(&cols.round(i)); + ps.add_scalars(&cols.message()); let x = ps.sample(); cols.bind(i, x); x @@ -614,7 +614,7 @@ fn the_dense_reduction_reduces_to_the_polynomials() { }); let point: Vec = (0..p.rounds()) .map(|i| { - let [c0, c2] = p.round(i); + let [c0, c2] = p.message(); ps.add_scalars(&[c0, c2]); let r = ps.sample(); claim = c0 + (claim + c2) * r + c2 * r * r; @@ -665,6 +665,73 @@ fn the_dense_reduction_reduces_to_the_polynomials() { assert!(falsified.count() > 0, "a false claim reduced to true values"); } +// A claim on `poly` whose terms each weigh one block: `(n_low, bits)` names the `2^n_low` entries at `bits << n_low`. +fn block_claim( + rng: &mut Rng, + tables: &DenseTables, + vars: &DenseVars, + poly: DensePoly, + blocks: &[(usize, usize)], +) -> DenseClaim { + let n = vars.0[poly as usize]; + let low = rng.ext_vec(n); + let mut terms = Vec::new(); + for &(n_low, bits) in blocks { + let mut point = low[..n_low].to_vec(); + point.extend((0..n - n_low).map(|i| F192::new((bits >> i & 1) as u64, 0, 0))); + terms.push(DenseTerm { + n_low, + bits: Bits { + value: bits, + len: n - n_low, + }, + top: None, + scale: Some(rng.ext()), + value: mle_eval(&tables.0[poly as usize], &point), + }); + } + DenseClaim { poly, low, terms } +} + +#[test] +fn the_dense_reduction_reduces_claims_on_a_prefix() { + // Weights zero past a prefix whose length is odd at several rounds, so the prover evaluates the table past it. The + // blocks of the image and the fixed polynomial have one or two low points, so their weights stay factored for a few + // rounds; the bytecode's has three, written out at once. + let mut rng = Rng::new(23); + let vars = DenseVars([3, 5, 14]); + let tables = DenseTables(vars.0.map(|n| (0..1 << n).map(|_| F64(rng.next_u64())).collect())); + let mut claims: Vec> = (0..3) + .map(|_| { + let point = rng.ext_vec(3); + let value = mle_eval(&tables.0[DensePoly::Bytecode as usize], &point); + DenseClaim::at(DensePoly::Bytecode, point, None, value) + }) + .collect(); + claims.push(block_claim(&mut rng, &tables, &vars, DensePoly::Image, &[(1, 2)])); + claims.push(block_claim( + &mut rng, + &tables, + &vars, + DensePoly::Fixed, + &[(2, 0xAAA), (5, 7), (2, 0xAAA)], + )); + claims.push(block_claim(&mut rng, &tables, &vars, DensePoly::Fixed, &[(5, 7)])); + let mut ps = ProverState::from_label(LABEL); + DenseProver::prove(&mut ps, &vars, &tables, &claims); + let proof = ps.into_proof(); + let mut vs = VerifierState::from_label(LABEL, &proof); + let reduced = vars.verify(&mut vs, &claims).expect("an honest reduction"); + for poly in DensePoly::ALL { + let point = &reduced.point[..vars.0[poly as usize]]; + assert_eq!( + reduced.values[poly as usize], + Some(mle_eval(&tables.0[poly as usize], point)), + "{poly:?}" + ); + } +} + // Each circuit's matrices at random points: lincheck claims and both carried claims, on three circuits. // // The lincheck claims share their points as one batch leaves them: one alpha and one skip point, the other coordinates From 5fe18270cfab91bdb45a1cc7a61ae6eb7162102c Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Tue, 6 Oct 2026 00:09:29 +0100 Subject: [PATCH 15/18] leanvm_core: rustfmt the merge --- crates/leanvm_core/src/class_flock.rs | 1 - crates/leanvm_core/src/rv/semantics/alu.rs | 9 ++++++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/crates/leanvm_core/src/class_flock.rs b/crates/leanvm_core/src/class_flock.rs index bc8d2ee97..84d900bc9 100644 --- a/crates/leanvm_core/src/class_flock.rs +++ b/crates/leanvm_core/src/class_flock.rs @@ -267,7 +267,6 @@ impl Prepared { }, check, ) - } (Part::Class, None, Some(witness)) => circuit.generate_witness_with_into( z, diff --git a/crates/leanvm_core/src/rv/semantics/alu.rs b/crates/leanvm_core/src/rv/semantics/alu.rs index 81a41be9f..33f33fbaf 100644 --- a/crates/leanvm_core/src/rv/semantics/alu.rs +++ b/crates/leanvm_core/src/rv/semantics/alu.rs @@ -415,7 +415,14 @@ mod tests { #[test] fn the_word_witness_is_the_gate_walk() { // Every legal flag word on edge operands: `b` an edge word or its complement, equal to `v1` on the diagonal, `dt` and `pc4` each none or all of their bits. - let edges = grid(&[&EDGES, &EDGES, &[0, u64::MAX], Alu::LEGAL, &[0, u64::MAX], &[0, u64::MAX]]); + let edges = grid(&[ + &EDGES, + &EDGES, + &[0, u64::MAX], + Alu::LEGAL, + &[0, u64::MAX], + &[0, u64::MAX], + ]); word_witness_is_the_walk::(Alu::witness, edges); } From 7a88b9664d32223c3df176571dfb8971c3e32cdb Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Tue, 6 Oct 2026 01:49:50 +0100 Subject: [PATCH 16/18] flock: sum a batch's identical padding tail once in the zerocheck A batch padded to a power of two ends in copies of one padding instance: the recursion machine's HASH table (about 20k to 23k live rows of 2^15) and every RISC-V table's fill rows. The zerocheck binds the lowest variables first, so while it binds bits inside an instance (or inside a group of instances), every group past the live ones has the same per-round polynomial and the same folded tables. The prover now sums the head (the live instances, rounded up to the kernel's unit of work) as before, sums the cube's last group once, and adds it weighted by the tail's eq mass, sum_{g >= head} eq(r_groups, g), which is one plus the mass below the head and costs O(log) products. Field arithmetic is exact, so every message is unchanged and the proof bytes are identical. - `PaddingSpec::live_blocks` and `PaddingSpec::tail` split a kernel's cube at the tail; a group is a block, or enough whole blocks for the kernel's smallest cube, so small circuits and rounds that bind instance bits split too. - Round 1 (univariate skip), the paired bit passes, the storing bit pass and the paired table passes take the split; the storing and table passes write the last group in place and copy it over the tail, so the stored tables stay whole and the single-round tail is untouched. - `reduction::Instance::live` carries the count; `HashBatch` sets it to its row count, and `class_flock::prove_reductions` detects each RISC-V batch's tail from its packed `z` (`reduction::live_instances`). The verifier and the proof are unchanged. - Test: `an_identical_tail_changes_no_word` pins the proof stream and claims of a batch told of its tail to the one not told, with random live instances and identical tails starting at block 0, mid step, mid task and past every step, over shapes reaching every kernel (both table-pass arities included). Measured on the 7950X3D, `aggregate --n 100 --leaves 4 --arity0 2 --arity 2 --repeat 3 --cooldown 1` with AGG1_NOTREE=1, proving time (mean of 3 passes), base f774828d: | node kind, threads | base runs | branch runs | median change | | ------------------- | ---------------- | ----------------------- | ------------- | | first-level, 1T | 0.971, 0.974 s | 0.907, 0.900, 0.902 s | -7.2% | | node, 1T | 1.041, 1.041 s | 0.952, 0.952, 0.951 s | -8.5% | | first-level, 32T | 0.150, 0.147 s | 0.144, 0.143 s | -3% (noise) | | node, 32T | 0.174, 0.175 s | 0.173, 0.171 s | -1.5% (noise) | Spreads 0.1% to 2.2% at 1T, 1.4% to 4.7% at 32T; loads 2.7 to 4.9 at start. An earlier pair of 1T rounds at loads 1 to 5 gave base 0.973, 0.977 / 1.038, 1.043 s. Span tree, 1T, one tracing pass each: first-level Flock reduction 249 -> 170 ms (Round 1 91.8 -> 57.9, Rounds 149 -> 104, Lincheck 8.6 unchanged), node 237 -> 147 ms (Round 1 79.4 -> 50.4, Rounds 149 -> 88.1). HASH holds 22,871 (first-level) and 19,665 (node) live rows of 32,768. Leaf, `leanxmss --n 100 --repeat 3 --cooldown 1`: 1T base 1.009, 1.041, 0.999 s, branch 0.989, 0.973, 0.971 s (clean pair 0.999 -> 0.971, -2.8%); 32T base 0.141, 0.145 s, branch 0.144, 0.143 s. Its Flock reductions span 250 -> 222 ms at 1T. Digests unchanged at n = 100 (leaf 1ddc4572a71ec2af, first 6cc4ab4bf8ee46b7, node d0d2b1643bb63a90) and n = 400 (leaf 9311cba19de1b111, first cac5d72e17c200f1, node bc70f47ff4643d4b). --- crates/flock/benches/arithmetic_batch.rs | 1 + crates/flock/src/arith.rs | 2 + crates/flock/src/hash.rs | 1 + crates/flock/src/reduction.rs | 17 ++ crates/flock/src/reduction_tests.rs | 1 + crates/flock/src/zerocheck.rs | 142 ++++++++++++- crates/flock/src/zerocheck/multilinear.rs | 198 ++++++++++++++---- .../zerocheck/univariate_skip_optimized.rs | 42 +++- crates/leanvm_core/src/class_flock.rs | 22 +- crates/leanvm_core/src/rec/proof.rs | 4 + crates/leanvm_core/src/rec/verifier/tests.rs | 1 + crates/leanvm_core/src/rv/semantics/alu.rs | 1 + 12 files changed, 371 insertions(+), 61 deletions(-) diff --git a/crates/flock/benches/arithmetic_batch.rs b/crates/flock/benches/arithmetic_batch.rs index a805358b7..134cb0e76 100644 --- a/crates/flock/benches/arithmetic_batch.rs +++ b/crates/flock/benches/arithmetic_batch.rs @@ -108,6 +108,7 @@ fn bench(op: U64Op, quiet: bool) -> (usize, Timing) { let instance = [flock::reduction::Instance { block, n_blocks_log: n_log, + live: 1 << n_log, z: &z_packed, a: &a_packed, b: &b_packed, diff --git a/crates/flock/src/arith.rs b/crates/flock/src/arith.rs index e7d6396fc..1de3e2158 100644 --- a/crates/flock/src/arith.rs +++ b/crates/flock/src/arith.rs @@ -255,6 +255,7 @@ mod tests { let instance = Instance { block, n_blocks_log: n_log, + live: 1 << n_log, z: &z, a: &a, b: &b, @@ -303,6 +304,7 @@ mod tests { .map(|((z, a, b, z_lincheck), &(block, n_blocks_log))| Instance { block, n_blocks_log, + live: 1 << n_blocks_log, z, a, b, diff --git a/crates/flock/src/hash.rs b/crates/flock/src/hash.rs index 0939ccc26..86eb6cab4 100644 --- a/crates/flock/src/hash.rs +++ b/crates/flock/src/hash.rs @@ -736,6 +736,7 @@ impl Blake2sSetup { Instance { block: BLOCK, n_blocks_log: self.n_blocks_log, + live: 1 << self.n_blocks_log, z, a, b, diff --git a/crates/flock/src/reduction.rs b/crates/flock/src/reduction.rs index 0d81f5e8f..58ef395e1 100644 --- a/crates/flock/src/reduction.rs +++ b/crates/flock/src/reduction.rs @@ -31,6 +31,18 @@ pub const fn min_n_blocks_log(n_blocks: usize) -> usize { n.next_power_of_two().trailing_zeros() as usize } +/// The instances before a packed witness's identical tail: from the returned index on, every +/// instance's `z` is the last one's, and so are its `A·z` and `B·z`. +pub fn live_instances(z: &[u64], k_log: usize) -> usize { + assert!(k_log >= 6, "an instance is whole words"); + let words = 1usize << (k_log - 6); + let Some(last) = z.rchunks_exact(words).next() else { + return 0; + }; + let copies = z.rchunks_exact(words).skip(1).take_while(|w| *w == last).count(); + z.len() / words - 1 - copies +} + /// A circuit as the reduction sees it: `2^k_log` witness bits per instance, of /// which `[useful_bits, 2^k_log)` are zero padding the prover skips. #[derive(Clone, Copy)] @@ -63,10 +75,14 @@ pub struct ReductionReplay { /// One circuit's batch as the prover holds it: the packed `z`, `A·z` and `B·z` of /// `2^n_blocks_log` instances, and `z` again in the lincheck stripe layout. +/// +/// The instances from `live` on are copies of one instance, which the zerocheck sums once while it binds bits inside an instance. +/// Nothing of the proof depends on `live`: `1 << n_blocks_log` claims no such tail. #[derive(Clone, Copy)] pub struct Instance<'a> { pub block: Block<'a>, pub n_blocks_log: usize, + pub live: usize, pub z: &'a [u64], pub a: &'a [u64], pub b: &'a [u64], @@ -131,6 +147,7 @@ pub fn prove_zerocheck(instances: &[Instance<'_>], ps: &mut ProverState) -> Zero padding: PaddingSpec { k_log: i.block.k_log, useful_bits_per_block: i.block.useful_bits, + live_blocks: i.live, }, } }) diff --git a/crates/flock/src/reduction_tests.rs b/crates/flock/src/reduction_tests.rs index 9b7766c12..0004632f4 100644 --- a/crates/flock/src/reduction_tests.rs +++ b/crates/flock/src/reduction_tests.rs @@ -61,6 +61,7 @@ fn prove(n: usize, tamper: Option) -> (usize, ProofTranscript) { let instance = reduction::Instance { block: BLOCK, n_blocks_log: n_log, + live: 1 << n_log, z: &z, a: &a, b: &b, diff --git a/crates/flock/src/zerocheck.rs b/crates/flock/src/zerocheck.rs index 635d2aeb4..b0a48da4b 100644 --- a/crates/flock/src/zerocheck.rs +++ b/crates/flock/src/zerocheck.rs @@ -73,19 +73,23 @@ fn equality_tail(m: usize, mut sample_vec: impl FnMut(usize) -> Vec) -> Ve .collect() } -/// Where the zero padding of a batched witness lies, so the zerocheck can skip it. +/// Where the padding of a batched witness lies, so the zerocheck can skip it. /// /// The witness is `2^(m - k_log)` blocks of `2^k_log` bits. /// /// Each block holds its data first and zero padding after it. /// /// A chunk of zero bits adds nothing to a round message, so skipping it leaves the output unchanged. +/// +/// The blocks from `live_blocks` on are copies of one padding block: while the rounds bind variables inside a block, they sum one copy, weighted by the tail's eq mass. #[derive(Clone, Copy, Debug)] pub struct PaddingSpec { /// Log of the bits in one block. pub k_log: usize, /// Bits at the start of each block that carry data; the rest are zero. pub useful_bits_per_block: usize, + /// Blocks before the identical tail; at least the block count when there is none. + pub live_blocks: usize, } impl PaddingSpec { @@ -94,8 +98,84 @@ impl PaddingSpec { Self { k_log: m, useful_bits_per_block: 1usize << m, + live_blocks: 1, + } + } + + /// The same blocks, with no tail known. + pub(crate) const fn without_tail(self) -> Self { + Self { + live_blocks: usize::MAX, + ..self + } + } + + /// Split a kernel's cube of `2^m` bits at the tail, or `None` when it saves nothing. + /// + /// - The head is a whole number of groups and of `2^step_log`-bit steps, the kernel's unit of work. + /// - A group is a block, or enough whole blocks for the kernel's smallest cube, `2^min_group_log` bits. + /// - `r` are the kernel's eq challenges, the last of which index the groups. + /// - The cube's last group stands for every group past the head: they are all copies of it. + pub(crate) fn tail(&self, m: usize, min_group_log: usize, step_log: usize, r: &[F192]) -> Option { + let group_log = self.k_log.max(min_group_log); + if group_log >= m { + return None; + } + let n_groups_log = m - group_log; + let live_groups = self.live_blocks.div_ceil(1 << (group_log - self.k_log)); + if live_groups >= 1 << n_groups_log { + return None; + } + let first = live_groups.next_multiple_of(1 << step_log.saturating_sub(group_log)); + if first >= 1 << n_groups_log { + return None; + } + let head = first << group_log; + Some(Tail { + head, + group_log, + r_inner: r.len() - n_groups_log, + weight: eq_mass_from(&r[r.len() - n_groups_log..], first), + }) + } +} + +/// A kernel's cube split by [`PaddingSpec::tail`]: the head it sums as is, then its last group once for the rest. +#[derive(Clone, Copy, Debug)] +pub(crate) struct Tail { + /// Bits in the head. + pub head: usize, + /// Log of the bits in a group. + pub group_log: usize, + /// The eq challenges inside a group, the first `r_inner` of the kernel's. + pub r_inner: usize, + /// The eq mass of the groups past the head. + pub weight: F192, +} + +impl Tail { + /// The last group's bytes of a packed witness. + pub fn group<'a>(&self, packed: &'a [u8]) -> &'a [u8] { + &packed[packed.len() - (1 << self.group_log) / 8..] + } +} + +/// `sum_{x >= from} eq(r, x)`, `x` read low bit first. +/// +/// The whole cube's mass is one, so this is one plus the mass below `from`. +fn eq_mass_from(r: &[F192], from: usize) -> F192 { + let mut below = F192::ZERO; + // The eq factor of the bits above the current one, set to `from`'s. + let mut above = F192::ONE; + for (k, &r_k) in r.iter().enumerate().rev() { + if from >> k & 1 == 1 { + below += above * (F192::ONE + r_k); + above *= r_k; + } else { + above *= F192::ONE + r_k; } } + F192::ONE + below } // --------------------------------------------------------------------------- @@ -310,7 +390,7 @@ impl<'a> CircuitProver<'a> { /// A round on the stored tables: the first of a paired pass, or a single round. fn table_round(&mut self, j: usize) -> (Option, F192, F192) { - let r = self.r; + let (r, padding) = (self.r, self.padding); let n_mlv = r.len(); let tb = self.tables.as_mut().expect("the tables are stored"); @@ -320,11 +400,14 @@ impl<'a> CircuitProver<'a> { if paired { let [a, b, c] = &tb.t; let [an, bn, cn] = &mut tb.nxt; + // A level-`j` value covers `2^(K_SKIP + j)` bits of the witness. let pair = fold_and_round_pair_into( [a, b, c], [&mut an[..n_out], &mut bn[..n_out], &mut cn[..n_out]], &tb.pending, &r[j + 1..], + &padding, + K_SKIP + j, ); tb.swap_in(n_out); tb.pending.clear(); @@ -905,6 +988,7 @@ mod tests { let padding = PaddingSpec { k_log: 9, useful_bits_per_block: 400, + live_blocks: usize::MAX, }; let inputs: Vec> = (circuits.iter()) .map(|c| { @@ -988,4 +1072,58 @@ mod tests { } vs.finish().unwrap(); } + + /// **Telling the prover of an identical tail changes no word of the proof.** + /// + /// Each circuit's blocks from `live` on are copies of one padding block, and the tails start at a block, inside a kernel's step or task, past every step (so no kernel skips), and at block 0. + /// Shapes: BLAKE2s blocks deep enough for table passes folding one and two pending challenges, blocks smaller than a round-1 window over a table pass, and blocks the size of a window. + #[test] + fn an_identical_tail_changes_no_word() { + let shapes = [(23usize, 14usize, 16_000usize), (21, 9, 400), (19, 13, 7_000)]; + let mut rng = Rng::new(0x7A11_5EED); + for (trial, frac) in [(0, 0.27), (1, 0.6), (2, 1.0 - 1e-9), (3, 0.0), (4, 0.51)] { + let lives: Vec = (shapes.iter()) + .map(|&(m, k_log, _)| ((frac * (1usize << (m - k_log)) as f64) as usize) | usize::from(trial == 4)) + .collect(); + let witnesses: Vec<[Vec; 3]> = (shapes.iter().zip(&lives)) + .map(|(&(m, k_log, useful), &live)| { + let block = 1usize << k_log; + let pad = [rng.bits(block), rng.bits(block)]; + let [a, b]: [Vec; 2] = std::array::from_fn(|w| { + (0..1usize << m) + .map(|i| { + let (blk, off) = (i / block, i % block); + let bit = if blk < live { + rng.next_u64() & 1 == 1 + } else { + pad[w][off] + }; + off < useful && bit + }) + .collect() + }); + let c: Vec = a.iter().zip(&b).map(|(x, y)| x & y).collect(); + [pack_bits(&a), pack_bits(&b), pack_bits(&c)] + }) + .collect(); + let run = |tail: bool| { + let inputs: Vec> = (shapes.iter().zip(&lives).zip(&witnesses)) + .map(|((&(m, k_log, useful), &live), [a, b, c])| ZerocheckInput { + bits: PackedWitness { a, b }, + c, + m, + padding: PaddingSpec { + k_log, + useful_bits_per_block: useful, + live_blocks: if tail { live } else { usize::MAX }, + }, + }) + .collect(); + let mut ps = ProverState::from_label(b"flock-test-v0"); + let claims = prove(&inputs, &mut ps); + (claims, ps.into_proof().stream) + }; + assert_eq!(run(true), run(false), "trial {trial}, live {lives:?}"); + } + } } diff --git a/crates/flock/src/zerocheck/multilinear.rs b/crates/flock/src/zerocheck/multilinear.rs index 67544521d..29d255e2c 100644 --- a/crates/flock/src/zerocheck/multilinear.rs +++ b/crates/flock/src/zerocheck/multilinear.rs @@ -324,14 +324,15 @@ impl RoundPair { /// /// The linear `c` term reaches `G(1)` only. pub fn bit_round_pair(bits: PackedWitness<'_>, fold: &BitFold, r_eq: &[F192], padding: &PaddingSpec) -> RoundPair { - match fold.n_chunks() { + let sums = match fold.n_chunks() { 8 => bit_round_pair_kernel::<8>(bits, fold, r_eq, padding), 16 => bit_round_pair_kernel::<16>(bits, fold, r_eq, padding), 32 => bit_round_pair_kernel::<32>(bits, fold, r_eq, padding), 64 => bit_round_pair_kernel::<64>(bits, fold, r_eq, padding), 128 => bit_round_pair_kernel::<128>(bits, fold, r_eq, padding), n => panic!("no bit-round kernel for {n}-byte rows"), - } + }; + RoundPair::from_sums(sums, r_eq[0]) } /// One round straight from the packed bits, storing the folded `(a, b, c)` tables for the rounds that follow. @@ -343,17 +344,24 @@ pub fn bit_round_materialize( r_eq: &[F192], padding: &PaddingSpec, ) -> ((F192, F192), [Vec; 3]) { - match fold.n_chunks() { - 8 => bit_round_store_kernel::<8>(bits, fold, r_eq, padding), - 16 => bit_round_store_kernel::<16>(bits, fold, r_eq, padding), - 32 => bit_round_store_kernel::<32>(bits, fold, r_eq, padding), - 64 => bit_round_store_kernel::<64>(bits, fold, r_eq, padding), - 128 => bit_round_store_kernel::<128>(bits, fold, r_eq, padding), + let n_pos = bits.a.len() / fold.n_chunks(); + let mut out: [Vec; 3] = std::array::from_fn(|_| { + // SAFETY: the kernel writes every slot, padding and tail included. + unsafe { primitives::uninit_vec(n_pos) } + }); + let outs = out.each_mut().map(|o| &mut o[..]); + let message = match fold.n_chunks() { + 8 => bit_round_store_kernel::<8>(bits, fold, r_eq, padding, outs), + 16 => bit_round_store_kernel::<16>(bits, fold, r_eq, padding, outs), + 32 => bit_round_store_kernel::<32>(bits, fold, r_eq, padding, outs), + 64 => bit_round_store_kernel::<64>(bits, fold, r_eq, padding, outs), + 128 => bit_round_store_kernel::<128>(bits, fold, r_eq, padding, outs), n => panic!("no bit-round kernel for {n}-byte rows"), - } + }; + (message, out) } -/// The two-round pass, for rows of `CHUNKS` bytes. +/// The two-round pass, for rows of `CHUNKS` bytes: the eight sums [`RoundPair::from_sums`] reads. /// /// Positions group in quads `4k + u + 2v`: `u` is round `t`'s variable and `v` round `t + 1`'s. /// @@ -379,14 +387,14 @@ fn bit_round_pair_kernel( fold: &BitFold, r_eq: &[F192], padding: &PaddingSpec, -) -> RoundPair { +) -> [F192; 8] { let rows = bits.rows::(); let n_quads = rows[0].len() / 4; assert!(n_quads >= 1, "two rounds need four positions"); assert_eq!(r_eq.len(), n_quads.trailing_zeros() as usize + 1); // `r_eq[0]` weights round `t`'s split by `v`; the rest weight the quads. - let (r_v, r_quad) = (r_eq[0], &r_eq[1..]); + let r_quad = &r_eq[1..]; let SplitEq { low: eq_lo, high: eq_hi, @@ -397,7 +405,11 @@ fn bit_round_pair_kernel( // A quad covers 2^6 skip bits times its 4 * 2^t bound rows. let quad_log = (32 * CHUNKS).trailing_zeros() as usize; let (quad_in_block_mask, live_quads) = padding_pairs(padding, quad_log - 1); - let live = |quad: usize| (quad & quad_in_block_mask) < live_quads; + // The quads before the tail, in whole folded blocks of sixteen. + let m = quad_log + n_quads.trailing_zeros() as usize; + let tail = padding.tail(m, quad_log, quad_log + 4, r_eq); + let head_quads = tail.map_or(n_quads, |t| t.head >> quad_log); + let live = |quad: usize| quad < head_quads && (quad & quad_in_block_mask) < live_quads; #[cfg(all( target_arch = "x86_64", @@ -409,8 +421,8 @@ fn bit_round_pair_kernel( ))] let eq_planes = (lo_size >= BLOCK / 4).then(|| planar::planes(&eq_lo)); - let sums = parallel::map_reduce( - eq_hi.len(), + let mut sums = parallel::map_reduce( + head_quads.div_ceil(lo_size), || [F192::ZERO; 8], |hi| { #[cfg(all( @@ -432,7 +444,7 @@ fn bit_round_pair_kernel( for lo_first in (0..lo_size).step_by(BLOCK / 4) { let n = (lo_size - lo_first).min(BLOCK / 4); let quad_first = hi * lo_size + lo_first; - // A block wholly in padding folds to zero. + // A block wholly in padding folds to zero, and the tail is summed apart. if !(quad_first..quad_first + n).any(live) { continue; } @@ -456,18 +468,31 @@ fn bit_round_pair_kernel( |x, y| std::array::from_fn(|i| x[i] + y[i]), ); - RoundPair::from_sums(sums, r_v) + if let Some(tail) = tail { + let group = PackedWitness { + a: tail.group(bits.a), + b: tail.group(bits.b), + }; + let group_sums = bit_round_pair_kernel::(group, fold, &r_eq[..tail.r_inner], &padding.without_tail()); + for (s, g) in sums.iter_mut().zip(group_sums) { + *s += tail.weight * g; + } + } + sums } -/// The storing single-round pass, for rows of `CHUNKS` bytes. +/// The storing single-round pass, for rows of `CHUNKS` bytes, writing the folded tables to `outs`. /// /// Positions pair up as `(2k, 2k + 1)`, the low index bit being the variable this round binds. +/// +/// The identical tail's tables are copies of its last group's. fn bit_round_store_kernel( bits: PackedWitness<'_>, fold: &BitFold, r_eq: &[F192], padding: &PaddingSpec, -) -> ((F192, F192), [Vec; 3]) { + mut outs: [&mut [F192]; 3], +) -> (F192, F192) { let rows = bits.rows::(); let n_pos = rows[0].len(); assert!(n_pos >= 2, "a round needs two positions"); @@ -484,16 +509,16 @@ fn bit_round_store_kernel( let position_log = (8 * CHUNKS).trailing_zeros() as usize; let (pair_in_block_mask, live_pairs) = padding_pairs(padding, position_log); let live = |pair: usize| (pair & pair_in_block_mask) < live_pairs; + // The pairs before the tail, in whole folded blocks of thirty-two. + let m = position_log + n_pos.trailing_zeros() as usize; + let tail = padding.tail(m, position_log + 1, position_log + 6, r_eq); + let head_pairs = tail.map_or(n_pos / 2, |t| t.head >> (position_log + 1)); - let mut out: [Vec; 3] = std::array::from_fn(|_| { - // SAFETY: every slot is written below, padding included. - unsafe { primitives::uninit_vec(n_pos) } - }); - let [out_a, out_b, out_c] = &mut out; - let chunks = [out_a, out_b, out_c].map(|o| Chunks::new(o, 2 * lo_size)); + assert!(outs.iter().all(|o| o.len() == n_pos), "one output per position"); + let chunks = outs.each_mut().map(|o| Chunks::new(o, 2 * lo_size)); - let message = parallel::map_reduce( - eq_hi.len(), + let mut message = parallel::map_reduce( + head_pairs.div_ceil(lo_size), || (F192::ZERO, F192::ZERO), |hi| { // SAFETY: task `hi` takes chunk `hi` of each output once, and the buffers outlive the dispatch. @@ -506,6 +531,10 @@ fn bit_round_store_kernel( for lo_first in (0..lo_size).step_by(BLOCK / 2) { let n = (lo_size - lo_first).min(BLOCK / 2); let pair_first = hi * lo_size + lo_first; + // The tail is copied after the dispatch. + if pair_first >= head_pairs { + break; + } let (o_first, o_len) = (2 * lo_first, 2 * n); // A block wholly in padding folds to zero. if !(pair_first..pair_first + n).any(live) { @@ -572,7 +601,25 @@ fn bit_round_store_kernel( }, |(s1, si), (t1, ti)| (s1 + t1, si + ti), ); - (message, out) + if let Some(tail) = tail { + // The last group is stored in place, then copied over the rest of the tail. + let group_len = 1 << (tail.group_log - position_log); + let group = PackedWitness { + a: tail.group(bits.a), + b: tail.group(bits.b), + }; + let (g1, g_inf) = bit_round_store_kernel::( + group, + fold, + &r_eq[..tail.r_inner], + &padding.without_tail(), + outs.each_mut().map(|o| &mut o[n_pos - group_len..]), + ); + message.0 += tail.weight * g1; + message.1 += tail.weight * g_inf; + copy_group(outs, 2 * head_pairs, group_len); + } + message } /// Rounds `t` and `t + 1` from the stored tables, folding the challenges still pending on them first. @@ -581,6 +628,7 @@ fn bit_round_store_kernel( /// - `rhos` are those variables' challenges, lowest first; each output folds `2^rhos.len()` inputs. /// - `outs` receive the level-`t` tables, `ins.len() >> rhos.len()` values each. /// - `r_eq` are the eq challenges of the variables round `t` does not bind. +/// - `padding` is the witness's, and each output covers `2^out_log` of its bits. /// /// ```text /// two pending: read level t - 2 (n) -> write level t (n / 4) + rounds t, t + 1 @@ -588,21 +636,33 @@ fn bit_round_store_kernel( /// ``` /// /// The rounds are built from each quad of folded values while they are in registers, as in the bit pass. -pub fn fold_and_round_pair_into(ins: [&[F192]; 3], outs: [&mut [F192]; 3], rhos: &[F192], r_eq: &[F192]) -> RoundPair { - match *rhos { - [rho] => fold_and_round_pair_kernel::<1>(ins, outs, [rho, F192::ZERO], r_eq), - [rho_0, rho_1] => fold_and_round_pair_kernel::<2>(ins, outs, [rho_0, rho_1], r_eq), +pub fn fold_and_round_pair_into( + ins: [&[F192]; 3], + outs: [&mut [F192]; 3], + rhos: &[F192], + r_eq: &[F192], + padding: &PaddingSpec, + out_log: usize, +) -> RoundPair { + let sums = match *rhos { + [rho] => fold_and_round_pair_kernel::<1>(ins, outs, [rho, F192::ZERO], r_eq, padding, out_log), + [rho_0, rho_1] => fold_and_round_pair_kernel::<2>(ins, outs, [rho_0, rho_1], r_eq, padding, out_log), _ => panic!("one or two pending challenges"), - } + }; + RoundPair::from_sums(sums, r_eq[0]) } -/// The paired pass for `K` pending challenges, `rhos[..K]`. +/// The paired pass for `K` pending challenges, `rhos[..K]`: the eight sums [`RoundPair::from_sums`] reads. +/// +/// The identical tail's outputs are copies of its last group's. fn fold_and_round_pair_kernel( ins: [&[F192]; 3], - outs: [&mut [F192]; 3], + mut outs: [&mut [F192]; 3], rhos: [F192; 2], r_eq: &[F192], -) -> RoundPair { + padding: &PaddingSpec, + out_log: usize, +) -> [F192; 8] { let n_out = ins[0].len() >> K; assert!(ins.iter().all(|t| t.len() == n_out << K), "a, b, c have one length"); assert!( @@ -616,7 +676,7 @@ fn fold_and_round_pair_kernel( // `r_eq[0]` weighs round `t`'s split by `v`; the rest weigh the quads. // // Up to 2^7 high eq indices, one task each: enough tasks for every worker at every table size. - let (r_v, r_quad) = (r_eq[0], &r_eq[1..]); + let r_quad = &r_eq[1..]; let SplitEq { low: eq_lo, high: eq_hi, @@ -624,10 +684,15 @@ fn fold_and_round_pair_kernel( } = SplitEq::with_high_vars(r_quad, EQ_HIGH_VARS); let lo_size = eq_lo.len(); + // The quads before the tail, in the pairs the outputs are published in. + let quad_log = out_log + 2; + let m = quad_log + n_quads.trailing_zeros() as usize; + let tail = padding.tail(m, quad_log, quad_log + 1, r_eq); + let head_quads = tail.map_or(n_quads, |t| t.head >> quad_log); + // One task per high eq index: `lo_size` quads, `4 * lo_size` outputs of each table. let (chunk_in, chunk_out) = ((4 * lo_size) << K, 4 * lo_size); - let [out_a, out_b, out_c] = outs; - let chunks = [out_a, out_b, out_c].map(|o| Chunks::new(o, chunk_out)); + let chunks = outs.each_mut().map(|o| Chunks::new(o, chunk_out)); let rho = |j: usize| (rhos[j], rhos[j], rhos[j], rhos[j]); // Four outputs of one table, each folded from its `2^K` inputs. @@ -666,8 +731,8 @@ fn fold_and_round_pair_kernel( [lo[0] + f.0, lo[1] + f.1, lo[2] + f.2, lo[3] + f.3] }; - let sums = parallel::map_reduce( - eq_hi.len(), + let mut sums = parallel::map_reduce( + head_quads.div_ceil(lo_size), || [F192::ZERO; 8], |hi| { // SAFETY: task `hi` takes chunk `hi` of each output once, and the buffers outlive the dispatch. @@ -677,7 +742,8 @@ fn fold_and_round_pair_kernel( let mut acc = [F192Unreduced::ZERO; 8]; // Two quads of a table are eight outputs, three whole cache lines, published at once. let mut staged = [[F192::ZERO; 8]; 3]; - for q in 0..lo_size { + let n_q = lo_size.min(head_quads - hi * lo_size); + for q in 0..n_q { let [a, b, c] = ins.map(|t| fold_quad(&t[(4 * q) << K..(4 * (q + 1)) << K])); let [lo, hi] = quad_pair_terms(a, b, c); @@ -699,7 +765,7 @@ fn fold_and_round_pair_kernel( for (out, stage) in outs.iter_mut().zip(&staged) { stream.copy(&mut out[4 * (q - 1)..4 * (q + 1)], stage); } - } else if q + 1 == lo_size { + } else if q + 1 == n_q { for (out, stage) in outs.iter_mut().zip(&staged) { out[4 * q..4 * q + 4].copy_from_slice(&stage[..4]); } @@ -709,7 +775,41 @@ fn fold_and_round_pair_kernel( }, |x, y| std::array::from_fn(|i| x[i] + y[i]), ); - RoundPair::from_sums(sums, r_v) + + if let Some(tail) = tail { + // The last group is folded in place, then copied over the rest of the tail. + let group_len = 1 << (tail.group_log - out_log); + let group_ins = ins.map(|t| &t[t.len() - (group_len << K)..]); + let group_sums = fold_and_round_pair_kernel::( + group_ins, + outs.each_mut().map(|o| &mut o[n_out - group_len..]), + rhos, + &r_eq[..tail.r_inner], + &padding.without_tail(), + out_log, + ); + for (s, g) in sums.iter_mut().zip(group_sums) { + *s += tail.weight * g; + } + copy_group(outs, head_quads * 4, group_len); + } + sums +} + +/// Copy each table's last `group_len` values over its values from `head` on. +fn copy_group>(tables: [T; 3], head: usize, group_len: usize) { + // Whole groups per copy task, about 2^12 values. + let task = group_len.max(1 << 12); + for mut t in tables { + let tail = &mut t.as_mut()[head..]; + let (body, group) = tail.split_at_mut(tail.len() - group_len); + let group = &*group; + parallel::chunks_mut(body, task, |_, dst| { + for d in dst.chunks_exact_mut(group_len) { + d.copy_from_slice(group); + } + }); + } } // --------------------------------------------------------------------------- @@ -959,7 +1059,14 @@ mod tests { // The pass under test. let mut outs: [Vec; 3] = std::array::from_fn(|_| vec![F192::ZERO; 1 << log_out]); let [oa, ob, oc] = &mut outs; - let pair = fold_and_round_pair_into([&tables[0], &tables[1], &tables[2]], [oa, ob, oc], &rhos, &r_eq); + let pair = fold_and_round_pair_into( + [&tables[0], &tables[1], &tables[2]], + [oa, ob, oc], + &rhos, + &r_eq, + &PaddingSpec::dense(log_out), + 0, + ); assert_eq!(outs, [level_a, level_b, level_c], "tables, k={k}, log_out={log_out}"); assert_eq!(pair.first, first, "round t, k={k}, log_out={log_out}"); assert_eq!(pair.second(rho_t), second, "round t + 1, k={k}, log_out={log_out}"); @@ -1142,6 +1249,7 @@ mod tests { let padding = PaddingSpec { k_log, useful_bits_per_block: useful, + live_blocks: usize::MAX, }; let lagrange = skip_lagrange_weights(K_SKIP, rng.ext()); let r_rest = rng.ext_vec(m - K_SKIP); diff --git a/crates/flock/src/zerocheck/univariate_skip_optimized.rs b/crates/flock/src/zerocheck/univariate_skip_optimized.rs index fb9446537..c2c96931a 100644 --- a/crates/flock/src/zerocheck/univariate_skip_optimized.rs +++ b/crates/flock/src/zerocheck/univariate_skip_optimized.rs @@ -1005,7 +1005,7 @@ fn accumulate_x_outer( ); } -/// Process one outer value. +/// Process one outer value, up to `n_windows` windows of the cube. #[inline] #[expect( clippy::too_many_arguments, @@ -1014,6 +1014,7 @@ fn accumulate_x_outer( fn process_one_x_hi( x_hi: usize, big_lo_size: usize, + n_windows: usize, n_lo_and_inner: usize, within_outer_mask: usize, b_med_counts: &[u8], @@ -1028,8 +1029,9 @@ fn process_one_x_hi( state.partials = Convert::new(); let n_lo = n_lo_and_inner - N_INNER; + let lo_len = big_lo_size.min(n_windows - x_hi * big_lo_size); - for (x_outer_lo, &eq_lo_val) in eq_lo_scaled.iter().enumerate().take(big_lo_size) { + for (x_outer_lo, &eq_lo_val) in eq_lo_scaled.iter().enumerate().take(lo_len) { let x_outer = x_outer_lo | (x_hi << n_lo); let within_hash_outer = x_outer & within_outer_mask; let n_b_med = b_med_counts[within_hash_outer] as usize; @@ -1117,7 +1119,8 @@ fn build_b_med_counts(padding: &PaddingSpec) -> (usize, Vec) { /// /// Skips 512-bit b_med sub-windows that fall entirely in the zero padding of /// every witness block per `padding`, which is byte-identical to the dense -/// path when those bits are honestly zero. +/// path when those bits are honestly zero. The identical tail of blocks is +/// summed once, its last group weighted by the tail's eq mass. #[expect( clippy::too_many_arguments, reason = "The proof kernel keeps its independent inputs explicit." @@ -1132,6 +1135,8 @@ pub fn round1_shift_reduce_extract_c_packed_padded( inv_table: &InvNttTableByteSingleGf8, padding: &PaddingSpec, ) -> (Vec, Vec) { + // The bits of one `x_outer` window, the smallest cube. + const WINDOW_LOG: usize = K_SKIP + N_INNER; assert_eq!(k_skip, K_SKIP, "optimized variant is k_skip=6 only"); assert!( m >= k_skip + N_INNER, @@ -1145,9 +1150,12 @@ pub fn round1_shift_reduce_extract_c_packed_padded( assert_eq!(r_rest.len(), m - k_skip); assert_eq!(inv_table.k, k_skip); + let tail = padding.tail(m, WINDOW_LOG, WINDOW_LOG, r_rest); + let n_windows = tail.map_or(1 << (m - WINDOW_LOG), |t| t.head >> WINDOW_LOG); + let eq = SplitEq::with_high_vars(&r_rest[N_INNER..], EQ_HIGH_VARS); let big_lo_size = eq.low.len(); - let hi_size = eq.high.len(); + let hi_size = n_windows.div_ceil(big_lo_size); let n_lo_and_inner = eq.low_log() + N_INNER; let d_inv_val = d_inv(); @@ -1166,6 +1174,7 @@ pub fn round1_shift_reduce_extract_c_packed_padded( process_one_x_hi( x_hi, big_lo_size, + n_windows, n_lo_and_inner, within_outer_mask, &b_med_counts, @@ -1188,8 +1197,28 @@ pub fn round1_shift_reduce_extract_c_packed_padded( ) .into_results(); - let res_c_lifted = ntt_extend_vec(&res_c_s, inv_table); - (res_ab.to_vec(), res_c_lifted) + let mut res_c_lifted = ntt_extend_vec(&res_c_s, inv_table); + let mut res_ab = res_ab.to_vec(); + if let Some(tail) = tail { + let [a, b, c] = [a_packed, b_packed, c_packed].map(|p| tail.group(p)); + let (group_ab, group_c) = round1_shift_reduce_extract_c_packed_padded( + a, + b, + c, + tail.group_log, + k_skip, + &r_rest[..tail.r_inner], + inv_table, + &padding.without_tail(), + ); + for (x, y) in res_ab.iter_mut().zip(group_ab) { + *x += tail.weight * y; + } + for (x, y) in res_c_lifted.iter_mut().zip(group_c) { + *x += tail.weight * y; + } + } + (res_ab, res_c_lifted) } #[cfg(test)] @@ -1466,6 +1495,7 @@ mod tests { let padding = PaddingSpec { k_log, useful_bits_per_block: useful_bits, + live_blocks: usize::MAX, }; let (padded_ab, padded_c) = round1_shift_reduce_extract_c_packed_padded(&a_p, &b_p, &c_p, m, K_SKIP, &r, &table, &padding); diff --git a/crates/leanvm_core/src/class_flock.rs b/crates/leanvm_core/src/class_flock.rs index d1db75afb..13d857524 100644 --- a/crates/leanvm_core/src/class_flock.rs +++ b/crates/leanvm_core/src/class_flock.rs @@ -20,7 +20,7 @@ use ::pcs::stack_open::SliceClaim; use fiat_shamir::transcript::{ProverState, VerifierState}; use flock::circuit::Circuit; use flock::lincheck::MatrixClaim; -use flock::reduction::{Instance, ReductionReplay, Shape, min_n_blocks_log}; +use flock::reduction::{Instance, ReductionReplay, Shape, live_instances, min_n_blocks_log}; use flock::verifier::FlockError; use primitives::field::F64; use std::sync::OnceLock; @@ -272,15 +272,21 @@ impl Prepared { /// Flock's batched zerocheck then lincheck over every packed witness, every class /// circuit then every clock circuit, leaving one claim on each committed column. +/// +/// A table's padding rows are alike, so its batch ends in an identical tail the zerocheck sums once. pub(crate) fn prove_reductions(batches: &[Prepared], ps: &mut ProverState) -> Vec { let instances: Vec> = (batches.iter()) - .map(|p| Instance { - block: circuit(p.flock).block(), - n_blocks_log: p.n_blocks_log, - z: &p.z, - a: &p.a, - b: &p.b, - z_lincheck: &p.z_lincheck, + .map(|p| { + let block = circuit(p.flock).block(); + Instance { + block, + n_blocks_log: p.n_blocks_log, + live: live_instances(&p.z, block.k_log), + z: &p.z, + a: &p.a, + b: &p.b, + z_lincheck: &p.z_lincheck, + } }) .collect(); flock::reduction::prove(&instances, ps) diff --git a/crates/leanvm_core/src/rec/proof.rs b/crates/leanvm_core/src/rec/proof.rs index bf33d4514..b68618bbf 100644 --- a/crates/leanvm_core/src/rec/proof.rs +++ b/crates/leanvm_core/src/rec/proof.rs @@ -27,6 +27,8 @@ pub fn statement_seed(statement: &[Limbs]) -> [F64; 4] { /// The hash table's flock batch, one instance per row. struct HashBatch { tau: usize, + /// The rows, before the padding rows: every instance from here on is the padding compression's. + live: usize, z: Vec, a: Vec, b: Vec, @@ -55,6 +57,7 @@ impl HashBatch { }); Self { tau, + live: hash.len(), z, a, b, @@ -68,6 +71,7 @@ impl HashBatch { let instance = flock::reduction::Instance { block: HashFlock::circuit().block(), n_blocks_log: self.tau, + live: self.live, z: &self.z, a: &self.a, b: &self.b, diff --git a/crates/leanvm_core/src/rec/verifier/tests.rs b/crates/leanvm_core/src/rec/verifier/tests.rs index 1e4912a92..ee4d68038 100644 --- a/crates/leanvm_core/src/rec/verifier/tests.rs +++ b/crates/leanvm_core/src/rec/verifier/tests.rs @@ -253,6 +253,7 @@ fn prove_reductions(batches: &[Batch]) -> ProofTranscript { Instance { block: class_flock::circuit(batch.f).block(), n_blocks_log: batch.n_blocks_log, + live: 1 << batch.n_blocks_log, z, a, b, diff --git a/crates/leanvm_core/src/rv/semantics/alu.rs b/crates/leanvm_core/src/rv/semantics/alu.rs index 9fed80040..641e6331c 100644 --- a/crates/leanvm_core/src/rv/semantics/alu.rs +++ b/crates/leanvm_core/src/rv/semantics/alu.rs @@ -300,6 +300,7 @@ mod tests { let instance = flock::reduction::Instance { block, n_blocks_log: n_log, + live: 1 << n_log, z: &z, a: &a, b: &b, From eb99e2437d690d96eadea4c1ad5aae3b62df6237 Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Tue, 6 Oct 2026 11:24:51 +0100 Subject: [PATCH 17/18] rec: rustfmt the hash witness merge --- crates/leanvm_core/src/rec/proof.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/crates/leanvm_core/src/rec/proof.rs b/crates/leanvm_core/src/rec/proof.rs index d3bffd429..fba34b014 100644 --- a/crates/leanvm_core/src/rec/proof.rs +++ b/crates/leanvm_core/src/rec/proof.rs @@ -132,8 +132,9 @@ impl RecWitness { let tau = layout.tau(Table::Hash); // A port is not in the stack: its values are a buffer of their own, written by the batch. - let mut ports: Vec]>> = - (0..HashFlock::N_PORTS).map(|_| Box::new_uninit_slice(1 << tau)).collect(); + let mut ports: Vec]>> = (0..HashFlock::N_PORTS) + .map(|_| Box::new_uninit_slice(1 << tau)) + .collect(); let batch = HashBatch::build(&a.hash, tau, windows[RecLayout::HASH_WITNESS], &mut ports); drop(windows); // SAFETY: the windows tile the stack up to its zeroed tail, and each was filled above. From dc24a7524856a3437661f3283a5e037ae371b25e Mon Sep 17 00:00:00 2001 From: Kevaundray Wedderburn Date: Wed, 7 Oct 2026 08:41:22 +0100 Subject: [PATCH 18/18] pcs: take PR 464's commit NTT plan out; keep ntt-tasks on upstream's NTT PR 464 (the commit NTT plan for few workers) is open on its own, so this branch no longer carries it. Its three files go back to upstream's versions: pcs/src/ntt/additive_ntt_f64.rs, pcs/src/ntt.rs and pcs/src/whir/commit.rs. That also drops the AVX-512 register kernels, including the register version of upstream's transposed butterfly, which belonged to PR 464's register path. The deep-task sizing of kw/agg1-ntt-tasks goes back onto upstream's AdditiveNttF64::transform as it was first written: a deep task takes 2^MIN_TASK_LOG (64) rows of contiguous sub-blocks and hands them to the row sink at once. The codeword and the tree are unchanged. Nothing else in this branch used PR 464. --- crates/pcs/src/ntt.rs | 2 +- crates/pcs/src/ntt/additive_ntt_f64.rs | 902 ++++++------------------- crates/pcs/src/whir/commit.rs | 6 +- 3 files changed, 211 insertions(+), 699 deletions(-) diff --git a/crates/pcs/src/ntt.rs b/crates/pcs/src/ntt.rs index 4986252e3..4ae5080da 100644 --- a/crates/pcs/src/ntt.rs +++ b/crates/pcs/src/ntt.rs @@ -3,4 +3,4 @@ pub(crate) mod additive_ntt_f64; pub use additive_ntt_f64::AdditiveNttF64; -pub(crate) use additive_ntt_f64::RowSink; +pub(crate) use additive_ntt_f64::{RowSink, transpose_lane_major}; diff --git a/crates/pcs/src/ntt/additive_ntt_f64.rs b/crates/pcs/src/ntt/additive_ntt_f64.rs index 0ec9cc058..e2e1f657b 100644 --- a/crates/pcs/src/ntt/additive_ntt_f64.rs +++ b/crates/pcs/src/ntt/additive_ntt_f64.rs @@ -63,52 +63,9 @@ fn span_get(basis: &[F64], idx: usize) -> F64 { acc } -/// The twiddles of consecutive blocks of one layer, one addition a block. -/// -/// A twiddle is `Σ_j bit_j(block) · row[offset + j]`, so from block `b` to `b + 1` it changes by the -/// entries under the bits that flip: `b`'s trailing ones and the zero above them, a prefix of the row. -struct TwiddleWalk { - /// The current block's twiddle. - t: F64, - /// `steps[k] = Σ_{j <= k} row[offset + j]`, the change past `k` trailing ones. - steps: [F64; 64], -} - -impl TwiddleWalk { - /// The walk over a layer of `2^layer` blocks, from block `first`. - fn new(row: &[F64], offset: usize, layer: usize, first: usize) -> Self { - let mut walk = Self { - t: span_get(&row[offset..offset + layer], first), - steps: [F64::ZERO; 64], - }; - let mut acc = F64::ZERO; - for (step, &entry) in walk.steps.iter_mut().zip(&row[offset..offset + layer]) { - acc += entry; - *step = acc; - } - walk - } - - /// Move from block `block` to `block + 1`. - #[inline] - fn step(&mut self, block: usize) { - // Past the layer's last block the index has `layer` trailing ones, which reads a zero step. - self.t += self.steps[block.trailing_ones() as usize]; - } -} - /// Receives a finished block of codeword rows, as `(first_row, rows)`. pub(crate) type RowSink<'a> = dyn Fn(usize, &[F64]) + Sync + 'a; -/// Where an encode's message is read from, while the codeword does not hold it yet. -#[derive(Clone, Copy)] -enum Message { - /// Row-major, `num_ntts` words a row; it may be the codeword's own first replica. - Rows(SendPtr), - /// Lane-major, in its own buffer: block `b` of the rows is codeword lane `num_ntts - 1 - b`. - Lanes(SendPtr), -} - /// Additive NTT over F_{2^64} with the standard polynomial-basis subspace /// `{1, x, x², …}`: the F_2-subspace is `{0, 1, …, 2^ℓ−1}` under the natural /// integer encoding, exactly as in the extension-field version (whose domain already @@ -147,6 +104,31 @@ impl AdditiveNttF64 { span_get(&v[1..], block) } + /// The seven twiddles a radix-8 group needs, breadth-first: layer `layer`, + /// then `layer + 1` (one per half), then `layer + 2` (one per quarter). + /// + /// `span_get` is F_2-linear in the block index, so the six deeper twiddles are + /// the block's own contribution plus a fixed correction per sub-block index: + /// one scan of the three basis rows replaces seven. + pub(crate) fn twiddles_radix8(&self, layer: usize, block: usize) -> [F64; 7] { + let l = self.log_domain_size(); + let (v0, v1, v2) = ( + &self.evals[l - layer - 1], + &self.evals[l - layer - 2], + &self.evals[l - layer - 3], + ); + let (mut t0, mut a, mut c) = (F64::ZERO, F64::ZERO, F64::ZERO); + for j in 0..layer { + if (block >> j) & 1 == 1 { + t0 += v0[1 + j]; + a += v1[2 + j]; + c += v2[3 + j]; + } + } + let (d, e0, e1) = (v1[1], v2[1], v2[2]); + [t0, a, a + d, c, c + e0, c + e1, c + e0 + e1] + } + /// RS-encode a message already stored in the codeword's first replica. /// /// # Overview @@ -162,66 +144,44 @@ impl AdditiveNttF64 { /// - The first pass therefore reads its rows from the message itself. /// - No pass fills the replicas only to read them back. /// - /// # Lane-major callers + /// # Why the transpose stays separate /// - /// - A lane-major message goes to the crate's `encode_lane_major_with` instead, which reads it in place. + /// - A lane-major caller transposes its message into place first. + /// - The transpose wants one long contiguous run per lane. + /// - The first pass wants hundreds of scattered rows at once. pub fn encode_interleaved_in_place(&self, data: &mut [F64], num_ntts: usize, log_inv_rate: usize) { // The message is the buffer's own first replica. - let msg = Message::Rows(SendPtr(data.as_mut_ptr())); + let msg = SendPtr(data.as_mut_ptr()); self.transform(data, num_ntts, log_inv_rate, Some(msg), None); } - /// RS-encode a message held in a buffer of its own, handing `on_rows` every finished block of rows. - /// - /// - The result equals encoding in place. - /// - Every codeword word is written before it is read, so the codeword may start uninitialized. - /// - `on_rows(first_row, rows)` is called from pool tasks, each row handed over exactly once, in aligned blocks of - /// one power-of-two size while their rows are still in cache. + /// Encode in place, handing `on_rows` every finished block of rows. /// - /// # Panics - /// - /// Panics unless the codeword is exactly `2^r` messages long. - pub(crate) fn encode_interleaved_with( + /// - It is called as `on_rows(first_row, rows)`, from pool tasks. + /// - Each row is handed over exactly once. + /// - Blocks are aligned, and all of one power-of-two size. + /// - A block is handed over while its rows are still in cache. + pub(crate) fn encode_interleaved_in_place_with( &self, data: &mut [F64], - msg: &[F64], num_ntts: usize, log_inv_rate: usize, on_rows: &RowSink<'_>, ) { - assert_eq!( - msg.len() << log_inv_rate, - data.len(), - "the codeword is 2^log_inv_rate messages" - ); - // Read-only from here on: the pointer only feeds the first pass's reads. - let msg = Message::Rows(SendPtr(msg.as_ptr().cast_mut())); + let msg = SendPtr(data.as_mut_ptr()); self.transform(data, num_ntts, log_inv_rate, Some(msg), Some(on_rows)); } - /// RS-encode a lane-major message, handing `on_rows` every finished block of rows. - /// - /// # Layout - /// - /// ```text - /// message, lane-major: [ block 0 | block 1 | ... | block n-1 ] each 2^(d - r) words - /// codeword lane t encodes message block n - 1 - t - /// ``` + /// RS-encode a message held in a buffer of its own, handing `on_rows` every finished block of rows. /// - /// - The lane order is [`transpose_lane_major`]'s: the result equals that transpose, then the in-place encode. + /// - The result equals encoding in place. /// - Every codeword word is written before it is read, so the codeword may start uninitialized. /// - The blocks are as for the in-place encode. /// - /// # Why the message is read in place - /// - /// - The first pass reads its rows from the message, so a row-major copy of it would be written only to be read once. - /// - Where the plan's first pass takes eight neighbouring rows at once, each lane gives it whole cache lines. - /// - Otherwise the message is transposed into the first replica, then encoded in place. - /// /// # Panics /// /// Panics unless the codeword is exactly `2^r` messages long. - pub(crate) fn encode_lane_major_with( + pub(crate) fn encode_interleaved_with( &self, data: &mut [F64], msg: &[F64], @@ -234,8 +194,8 @@ impl AdditiveNttF64 { data.len(), "the codeword is 2^log_inv_rate messages" ); - // Read-only from here on: the pointer only feeds the first pass's reads, or the transpose. - let msg = Message::Lanes(SendPtr(msg.as_ptr().cast_mut())); + // Read-only from here on: the pointer only feeds the first pass's reads. + let msg = SendPtr(msg.as_ptr().cast_mut()); self.transform(data, num_ntts, log_inv_rate, Some(msg), Some(on_rows)); } @@ -274,33 +234,13 @@ impl AdditiveNttF64 { /// /// - Each deep task's sub-blocks are final, and hand their rows to `on_rows` while they are in L2. /// - Without a split deep pass, the rows go over in parallel blocks at the end. - /// - /// # Few workers - /// - /// - Up to [`FEW_WORKERS`] workers, a gathered group takes up to `2^MAX_LOG_RESIDUES` residues. - /// - Their rows are adjacent, so each copy in or out is one long run, and a row group has that many times - /// the lanes. - /// - More workers keep one residue a group: their scratch then stays in their share of L2. fn transform( &self, data: &mut [F64], num_ntts: usize, start: usize, - msg: Option, + msg: Option>, on_rows: Option<&RowSink<'_>>, - ) { - self.transform_for(data, num_ntts, start, msg, on_rows, parallel::num_threads()); - } - - /// [`Self::transform`], planned for `workers` workers. - fn transform_for( - &self, - data: &mut [F64], - num_ntts: usize, - start: usize, - msg: Option, - on_rows: Option<&RowSink<'_>>, - workers: usize, ) { // The buffer is 2^log_d rows of `num_ntts` words. assert!(num_ntts > 0); @@ -322,7 +262,8 @@ impl AdditiveNttF64 { // From 2^12 rows on, cut at least 2^LOG_SUBS_PER_WORKER deep sub-blocks per worker. let par_log = if log_d >= PARALLEL_FLOOR_LOG_D { - (log2_strict_usize(workers.next_power_of_two()) + LOG_SUBS_PER_WORKER).min(log_d - MIN_SUB_LOG) + (log2_strict_usize(parallel::num_threads().next_power_of_two()) + LOG_SUBS_PER_WORKER) + .min(log_d - MIN_SUB_LOG) } else { 0 }; @@ -336,52 +277,25 @@ impl AdditiveNttF64 { // L3 deep pass: 1 gathered sweep (9 layers), then 12 deep <- one sweep fewer // // 20 layers: 1 gathered sweep either way, so the deep pass stays in L2 - // - // Few workers share L3 with few others, so they take the L3 deep pass whenever there is one: - // the layers it moves out of the gathered pass shrink the gathered groups, and their scattered reads. - let few = workers <= FEW_WORKERS; let gathered_sweeps = |deep: usize| (log_d - start).saturating_sub(deep).div_ceil(fit2); - let deep = if few || gathered_sweeps(fit3) < gathered_sweeps(fit2) { + let deep = if gathered_sweeps(fit3) < gathered_sweeps(fit2) { fit3 } else { fit2 }; let deep_start = log_d.saturating_sub(deep).max(par_log).max(start); - // Residues a gathered group takes, as a log: several only for few workers, within their L3 budget. - let residues = |layer: usize, g: usize| { - if few { - let room = (WIDE_GATHER_WORDS / (num_ntts << g)).max(1).ilog2() as usize; - room.min(MAX_LOG_RESIDUES).min(log_d - layer - g) - } else { - 0 - } - }; - // A buffer this large is evicted before the next pass reads it back. let stream = data.len() >= STREAM_MIN_WORDS; - // A lane-major message is read in place by a first gathered pass that takes whole cache lines of it, - // eight residues of each lane at a time. Any other plan reads it transposed into the first replica. - let mut msg = msg; - if let Some(Message::Lanes(m)) = msg { - let first_g = (deep_start - start).min(fit2); - if !(LANE_GATHER && start < deep_start && residues(start, first_g) >= 3) { - let msg_len = data.len() >> start; - // SAFETY: the lane-major message is one replica of words, disjoint from the codeword. - let lanes = unsafe { std::slice::from_raw_parts(m.0.cast_const(), msg_len) }; - transpose_lane_major(&mut data[..msg_len], lanes, num_ntts, log_d - start); - msg = Some(Message::Rows(SendPtr(data.as_mut_ptr()))); - } - } - // Phase 1: gathered passes, each at most one L2 group of layers. // // Only the first one reads the message. + let mut msg = msg; let mut layer = start; while layer < deep_start { let g = (deep_start - layer).min(fit2); - self.gathered_pass(data, log_d, num_ntts, layer, g, residues(layer, g), msg.take(), stream); + self.gathered_pass(data, log_d, num_ntts, layer, g, msg.take(), stream); layer += g; } @@ -389,9 +303,7 @@ impl AdditiveNttF64 { // // message is block 0 of the buffer -> a task could overwrite it while others still copy it // no layer left to run -> no deep task would write the buffer at all - if let Some(Message::Rows(m)) = msg - .take_if(|m| matches!(m, Message::Rows(r) if std::ptr::eq(r.0, data.as_mut_ptr())) || deep_start == log_d) - { + if let Some(m) = msg.take_if(|m| std::ptr::eq(m.0, data.as_mut_ptr()) || deep_start == log_d) { replicate(data, m, data.len() >> start); } @@ -413,7 +325,7 @@ impl AdditiveNttF64 { // A separate message: build the sub-blocks in scratch, then write them out once. // // Streamed out whole, the codeword is written without ever being read. - Some(Message::Rows(m)) => with_scratch(task.len(), |scratch| { + Some(m) => with_scratch(task.len(), |scratch| { for (i, sub) in scratch.chunks_exact_mut(sub_len).enumerate() { // A sub-block sits at the same offset in every replica. // @@ -444,8 +356,6 @@ impl AdditiveNttF64 { f(first_sub << log_sub, task); } } - // A lane-major message is read by a gathered pass, or transposed into place before any pass. - Some(Message::Lanes(_)) => unreachable!("a lane-major message reaches no deep pass"), } }); } @@ -480,16 +390,10 @@ impl AdditiveNttF64 { /// - The residue is below `step`, so that shift drops it. /// - The scratch transform therefore takes the same twiddles as the full one. /// - /// # Several residues a group - /// - /// - A group may take `2^log_r` adjacent residues: their rows sit side by side, and share every twiddle. - /// - So the group is one transform with `2^log_r` times the lanes, and each of its rows one contiguous run. - /// /// # With a message /// /// - Every block's rows come from the message. - /// - One task takes one residue group across all blocks, block 0 last. - /// - A lane-major message gives each lane's word of eight adjacent residues as one cache line. + /// - One task takes one residue across all blocks, block 0 last. #[allow(clippy::too_many_arguments)] fn gathered_pass( &self, @@ -498,58 +402,45 @@ impl AdditiveNttF64 { num_ntts: usize, layer: usize, g: usize, - log_r: usize, - msg: Option, + msg: Option>, stream: bool, ) { - // A group is 2^g rows, `step` rows apart, each `2^log_r` adjacent rows wide. + // A group is 2^g rows, `step` rows apart. let log_step = log_d - layer - g; let (rows, step) = (1usize << g, 1usize << log_step); - assert!(log_r <= log_step, "a group's residues are inside one step"); - let wide = num_ntts << log_r; - // A lane-major message is read eight residues of a lane at a time. - assert!( - !matches!(msg, Some(Message::Lanes(_))) || log_r >= 3, - "a lane-major message is read whole lines at a time" - ); - // With a message, a task is one residue group across every block. - // Without one, a task is one (block, residue group) pair. - let n_tasks = (if msg.is_some() { step } else { step << layer }) >> log_r; + // With a message, a task is one residue across every block. + // Without one, a task is one (block, residue) pair. + let n_tasks = if msg.is_some() { step } else { step << layer }; let base = SendPtr(data.as_mut_ptr()); parallel::for_each_chunk(n_tasks, |lo, hi| { - // One scratch of 2^g wide rows serves every group of the task. - with_scratch(rows * wide, |scratch| { + // One L2-resident scratch of 2^g rows serves every group of the task. + with_scratch(rows * num_ntts, |scratch| { // One fence at the end of the task covers all its streaming stores. let stream = stream.then(Stream::new); let mut group = |block: usize, r: usize| { - // Row `i` of the group, as a row index and a word offset inside its block. - let row_idx = |i: usize| r + (i << log_step); - let row = |i: usize| row_idx(i) * num_ntts; + // Row `i` of the group, as a word offset inside its block. + let row = |i: usize| (r + (i << log_step)) * num_ntts; let block_off = (block << (log_d - layer)) * num_ntts; // Gather: the scattered rows become one contiguous 2^g-row buffer. - for (i, dst) in scratch.chunks_exact_mut(wide).enumerate() { + for (i, dst) in scratch.chunks_exact_mut(num_ntts).enumerate() { // SAFETY: // - The message and the codeword both cover every row addressed here. // - Tasks own disjoint residues, so no other task writes these rows. // - Block 0 is written last, after every read of the message in it. - unsafe { + let src = unsafe { match msg { - Some(Message::Rows(m)) => { - dst.copy_from_slice(std::slice::from_raw_parts(m.add(row(i)), wide)); - } - Some(Message::Lanes(m)) => { - gather_lanes(dst, m, 1 << (log_d - layer), num_ntts, row_idx(i)); - } - None => dst.copy_from_slice(base.slice(block_off + row(i), wide)), + Some(m) => std::slice::from_raw_parts(m.add(row(i)), num_ntts), + None => base.slice(block_off + row(i), num_ntts), } - } + }; + dst.copy_from_slice(src); } // Transform: a (layer + g)-layer domain whose sub-block index is the global block. - self.run_layers(scratch, layer + g, wide, layer, layer + g, layer, block); + self.run_layers(scratch, layer + g, num_ntts, layer, layer + g, layer, block); // Scatter: every row returns to its place. - for (i, src) in scratch.chunks_exact(wide).enumerate() { + for (i, src) in scratch.chunks_exact(num_ntts).enumerate() { // SAFETY: this group alone owns these rows of the codeword. - let dst = unsafe { base.slice(block_off + row(i), wide) }; + let dst = unsafe { base.slice(block_off + row(i), num_ntts) }; match &stream { Some(s) => s.copy(dst, src), None => dst.copy_from_slice(src), @@ -557,8 +448,6 @@ impl AdditiveNttF64 { } }; for t in lo..hi { - // The group's first residue, in the task's step-indexed order. - let t = t << log_r; if msg.is_some() { // Block 0 may be the message itself, so it is transformed last. for block in (0..1usize << layer).rev() { @@ -577,12 +466,7 @@ impl AdditiveNttF64 { /// /// - The domain has `2^d` rows and splits into `2^o` equal sub-blocks. /// - The buffer is one of them; its index fixes the global block index, and so the twiddle, of each block. - /// - Three layers fuse into one radix-8 sweep, the one or two layers left over into a first sweep of their own. - /// - /// # Why the leftover goes first - /// - /// - A sweep costs a twiddle computation per block, and the first layers have the fewest, largest blocks. - /// - Last, a lone layer would pair two-row blocks, each with its own twiddle for one butterfly row. + /// - Three layers fuse into one radix-8 sweep where blocks are wide enough, then two, then one. #[allow(clippy::too_many_arguments)] fn run_layers( &self, @@ -600,69 +484,55 @@ impl AdditiveNttF64 { let num_blocks_in_buf = 1usize << (layer - outer_log); let block_size = 1usize << (log_d - layer); let block_elems = block_size * num_ntts; - // The global index of the buffer's first block, which picks its twiddle; the others follow it. - let first = sub_idx * num_blocks_in_buf; - // Blocks span at least the layers left, since the end layer is at most `log_d`. - let blocks = buf.chunks_exact_mut(block_elems).zip(first..); - - match (end_layer - layer) % 3 { - 0 => { - // A radix-8 sweep, its seven twiddles breadth-first. - let (mut t0, mut a, mut c) = ( - TwiddleWalk::new(self.row(layer), 1, layer, first), - TwiddleWalk::new(self.row(layer + 1), 2, layer, first), - TwiddleWalk::new(self.row(layer + 2), 3, layer, first), - ); - let (d, e0, e1) = (self.row(layer + 1)[1], self.row(layer + 2)[1], self.row(layer + 2)[2]); - for (block, global) in blocks { - let (a_t, c_t) = (a.t, c.t); - let t = [t0.t, a_t, a_t + d, c_t, c_t + e0, c_t + e1, c_t + e0 + e1]; - butterfly_interleaved_fused_3layer(block, &t, block_size >> 3, num_ntts); - t0.step(global); - a.step(global); - c.step(global); - } - layer += 3; + // A block's index in the whole domain, which picks its twiddle. + let global = |block_in_buf: usize| sub_idx * num_blocks_in_buf + block_in_buf; + + if layer + 2 < end_layer && block_size >= 8 { + // Three layers to go and blocks of at least 8 rows: one radix-8 sweep. + let eighth = block_size >> 3; + for block_in_buf in 0..num_blocks_in_buf { + let t = self.twiddles_radix8(layer, global(block_in_buf)); + let start = block_in_buf * block_elems; + butterfly_interleaved_fused_3layer(&mut buf[start..start + block_elems], &t, eighth, num_ntts); } - 2 => { - // A radix-4 sweep: one twiddle for the block, then one per half. - let (mut outer, mut inner) = ( - TwiddleWalk::new(self.row(layer), 1, layer, first), - TwiddleWalk::new(self.row(layer + 1), 2, layer, first), + layer += 3; + } else if layer + 1 < end_layer && block_size >= 4 { + // Two layers to go: one radix-4 sweep. + let quarter = block_size >> 2; + for block_in_buf in 0..num_blocks_in_buf { + let global_block = global(block_in_buf); + let t_outer = self.twiddle(layer, global_block); + let t_inner_a = self.twiddle(layer + 1, 2 * global_block); + let t_inner_b = self.twiddle(layer + 1, 2 * global_block + 1); + let start = block_in_buf * block_elems; + butterfly_interleaved_fused_2layer( + &mut buf[start..start + block_elems], + t_outer, + t_inner_a, + t_inner_b, + quarter, + num_ntts, ); - let d = self.row(layer + 1)[1]; - for (block, global) in blocks { - butterfly_interleaved_fused_2layer( - block, - outer.t, - inner.t, - inner.t + d, - block_size >> 2, - num_ntts, - ); - outer.step(global); - inner.step(global); - } - layer += 2; } - _ => { - // One layer: a plain butterfly sweep. - let mut walk = TwiddleWalk::new(self.row(layer), 1, layer, first); - for (block, global) in blocks { - butterfly_interleaved_block(block, walk.t, block_size >> 1, num_ntts); - walk.step(global); - } - layer += 1; + layer += 2; + } else { + // One layer: a plain butterfly sweep. + let block_size_half = block_size >> 1; + for block_in_buf in 0..num_blocks_in_buf { + let twiddle = self.twiddle(layer, global(block_in_buf)); + let start = block_in_buf * block_elems; + butterfly_interleaved_block( + &mut buf[start..start + block_elems], + twiddle, + block_size_half, + num_ntts, + ); } + layer += 1; } } } - /// The table row a layer's twiddles are subset sums of. - fn row(&self, layer: usize) -> &[F64] { - &self.evals[self.log_domain_size() - layer - 1] - } - /// Recover novel-basis coefficients from evaluations with a scalar inverse NTT. pub fn inverse_transform(&self, data: &mut [F64]) { let log_d = log2_strict_usize(data.len()); @@ -738,50 +608,9 @@ fn fused_rows( /// - The eight rows stay in L1 across all twelve butterflies. /// - The seven twiddles are breadth-first: one for layer L, two for L+1, four for L+2. fn butterfly_interleaved_fused_3layer(block: &mut [F64], t: &[F64; 7], eighth: usize, num_ntts: usize) { - #[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] - { - // SAFETY: the arm is compiled only with AVX-512F; registers only. - let tw = t.map(|x| unsafe { _mm512_set1_epi64(x.0 as i64) }); - if num_ntts % 8 >= 2 { - return butterfly_interleaved_fused_3layer_masked(block, &tw, eighth, num_ntts); - } - let vectors = num_ntts / 8; - let done = 8 * vectors; - fused_rows::<8>(block, eighth, num_ntts, |rows| { - // SAFETY: the target features are enabled at compile time, and each row has `num_ntts` words. - unsafe { radix8_avx512(std::array::from_fn(|i| rows[i].as_mut_ptr()), &tw, vectors) }; - // The lane no vector covers, if any. - if done < num_ntts { - let mut tail = rows.each_mut().map(|row| &mut row[done..]); - radix8_butterflies(&mut tail, t); - } - }); - } - #[cfg(not(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f")))] fused_rows::<8>(block, eighth, num_ntts, |rows| radix8_butterflies(rows, t)); } -/// [`butterfly_interleaved_fused_3layer`] for rows that end in two to seven lanes past a whole vector. -/// -/// - Those lanes take one masked vector, which costs less than two scalar butterflies each. -/// - It is a function of its own, so the whole-vector loop of the common shapes keeps its registers. -#[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] -#[inline(never)] -fn butterfly_interleaved_fused_3layer_masked(block: &mut [F64], tw: &[__m512i; 7], eighth: usize, num_ntts: usize) { - let (vectors, mask) = (num_ntts / 8, (1u8 << (num_ntts % 8)) - 1); - fused_rows::<8>(block, eighth, num_ntts, |rows| { - let rows: [*mut F64; 8] = std::array::from_fn(|i| rows[i].as_mut_ptr()); - // SAFETY: the target features are enabled at compile time; each row has `8 * vectors` words and the - // masked lanes after them. - unsafe { - radix8_avx512(rows, tw, vectors); - masked_column_avx512(rows.map(|row| row.add(8 * vectors)), mask, |r| { - radix8_regs_avx512(r, tw); - }); - } - }); -} - /// The twelve butterflies of one radix-8 row group, with its seven twiddles breadth-first. #[inline(always)] fn radix8_butterflies(rows: &mut [&mut [F64]; 8], t: &[F64; 7]) { @@ -870,93 +699,6 @@ pub(crate) fn transpose_lane_major(out: &mut [F64], msg: &[F64], n_lanes: usize, }); } -/// Whether a gathered pass reads a lane-major message in place; elsewhere it is transposed first. -/// -/// Only the AVX-512 build turns each eight lanes of eight residues around in registers. -const LANE_GATHER: bool = cfg!(all(target_arch = "x86_64", target_feature = "avx512f")); - -/// Gather `dst.len() / n` adjacent rows of a lane-major message, from row `row`, into row-major order. -/// -/// ```text -/// dst[j * n + (n - 1 - l)] = msg[l * msg_rows + row + j] lanes reversed, as in transpose_lane_major -/// ``` -/// -/// - Eight rows of one lane are one cache line, read whole, when `row` is a multiple of eight. -/// - Eight lanes of those eight rows are turned around in registers. -/// -/// # Safety -/// -/// - The message holds `n` lanes of `msg_rows` words each. -/// - `row + dst.len() / n <= msg_rows`, and `dst.len() / n` is a multiple of eight. -unsafe fn gather_lanes(dst: &mut [F64], msg: SendPtr, msg_rows: usize, n: usize, row: usize) { - let residues = dst.len() / n; - debug_assert!(residues.is_multiple_of(8) && row + residues <= msg_rows); - // Message lanes `0 .. whole` go eight at a time; the rest, the lowest codeword lanes, one at a time. - let whole = if LANE_GATHER { n - n % 8 } else { 0 }; - #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] - for l0 in (0..whole).step_by(8) { - // Each lane's lines in order, so a lane's run of `residues` words is read front to back. - for q in (0..residues).step_by(8) { - // SAFETY: line `k` is lane `l0 + 7 - k`, rows `row + q ..+ 8`, inside the message by the caller. - let lines: [__m512i; 8] = std::array::from_fn(|k| unsafe { - _mm512_loadu_si512(msg.add((l0 + 7 - k) * msg_rows + row + q).cast_const().cast()) - }); - // Column `j` is residue `q + j` across lanes `l0 + 7` down to `l0`: codeword lanes `n - 8 - l0 ..`. - for (j, col) in transpose_8x8_avx512(lines).into_iter().enumerate() { - let at = (q + j) * n + n - 8 - l0; - // SAFETY: eight words of residue `q + j`'s row, inside `dst`. - unsafe { _mm512_storeu_si512(dst.as_mut_ptr().add(at).cast(), col) }; - } - } - } - for l in whole..n { - for j in 0..residues { - // SAFETY: inside the message by the caller. - dst[j * n + n - 1 - l] = unsafe { *msg.add(l * msg_rows + row + j) }; - } - } -} - -/// Transpose eight rows of eight words: word `k` of output `j` is word `j` of input `k`. -#[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] -#[inline(always)] -fn transpose_8x8_avx512(r: [__m512i; 8]) -> [__m512i; 8] { - // SAFETY: the arm is compiled only with AVX-512F; registers only. - unsafe { - // Pairs of rows, 128-bit lane c holding words 2c and 2c + 1 of both. - let t: [__m512i; 8] = std::array::from_fn(|i| { - let (a, b) = (r[i & !1], r[i | 1]); - if i % 2 == 0 { - _mm512_unpacklo_epi64(a, b) - } else { - _mm512_unpackhi_epi64(a, b) - } - }); - // Four rows: u[c] holds words c and c + 4 of rows 0..4 (u[0..4]) or rows 4..8 (u[4..8]). - let u = [ - _mm512_shuffle_i64x2::<0x88>(t[0], t[2]), - _mm512_shuffle_i64x2::<0x88>(t[1], t[3]), - _mm512_shuffle_i64x2::<0xDD>(t[0], t[2]), - _mm512_shuffle_i64x2::<0xDD>(t[1], t[3]), - _mm512_shuffle_i64x2::<0x88>(t[4], t[6]), - _mm512_shuffle_i64x2::<0x88>(t[5], t[7]), - _mm512_shuffle_i64x2::<0xDD>(t[4], t[6]), - _mm512_shuffle_i64x2::<0xDD>(t[5], t[7]), - ]; - // Eight rows: words c from u[c] and u[4 + c], words c + 4 from their other halves. - [ - _mm512_shuffle_i64x2::<0x88>(u[0], u[4]), - _mm512_shuffle_i64x2::<0x88>(u[1], u[5]), - _mm512_shuffle_i64x2::<0x88>(u[2], u[6]), - _mm512_shuffle_i64x2::<0x88>(u[3], u[7]), - _mm512_shuffle_i64x2::<0xDD>(u[0], u[4]), - _mm512_shuffle_i64x2::<0xDD>(u[1], u[5]), - _mm512_shuffle_i64x2::<0xDD>(u[2], u[6]), - _mm512_shuffle_i64x2::<0xDD>(u[3], u[7]), - ] - } -} - /// Words one L2-resident unit of work may span: a gathered row group or a deep sub-block. /// /// # Why this value @@ -1001,31 +743,6 @@ const MIN_TASK_LOG: usize = 6; /// - That costs extra L3 traffic, but saves a whole sweep of DRAM. const L3_WORDS: usize = 1 << 18; -/// Most workers whose transform takes wide gathered groups and the L3 deep pass. -/// -/// # Why this value -/// -/// - A wide group spills from L2 into L3, and in return copies long runs and reads a lane-major message in place. -/// - With one or two workers the runs save more than the spill costs. -/// - From four workers on, SMT siblings or not, their spills share L3 and memory with more of them, and the narrow -/// plan is faster. -const FEW_WORKERS: usize = 2; - -/// Words a gathered task may span when few workers run. -/// -/// # Why this value -/// -/// - 2^22 words is 32 MiB, a whole L3 without stacked cache. -/// - With few workers, a gathered row group's layers cost no more from L3 than from L2, the butterflies being the bound. -/// - So the group can take several residues, and with them longer runs in and out of memory. -const WIDE_GATHER_WORDS: usize = 1 << 22; - -/// Most residues one gathered group takes, as a log. -/// -/// - Eight residues make a row-major group's runs whole cache lines, for any lane count. -/// - Sixty-four make a lane-major message's reads eight lines of each lane. -const MAX_LOG_RESIDUES: usize = 6; - /// Buffers of at least this many words get streaming stores for data a pass does not read back. /// /// # Why this value @@ -1098,29 +815,6 @@ fn butterfly_interleaved_fused_2layer( quarter: usize, num_ntts: usize, ) { - #[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] - { - // SAFETY: the arm is compiled only with AVX-512F; registers only. - let tw = [t_outer, t_inner_a, t_inner_b].map(|x| unsafe { _mm512_set1_epi64(x.0 as i64) }); - if num_ntts % 8 >= 2 { - return butterfly_interleaved_fused_2layer_masked(block, &tw, quarter, num_ntts); - } - let vectors = num_ntts / 8; - let done = 8 * vectors; - fused_rows::<4>(block, quarter, num_ntts, |rows| { - // SAFETY: the target features are enabled at compile time, and each row has `num_ntts` words. - unsafe { radix4_avx512(std::array::from_fn(|i| rows[i].as_mut_ptr()), &tw, vectors) }; - // The lane no vector covers, if any. - if done < num_ntts { - let [row_a, row_b, row_c, row_d] = rows.each_mut().map(|row| &mut row[done..]); - butterfly_lanes(row_a, row_c, t_outer); - butterfly_lanes(row_b, row_d, t_outer); - butterfly_lanes(row_a, row_b, t_inner_a); - butterfly_lanes(row_c, row_d, t_inner_b); - } - }); - } - #[cfg(not(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f")))] fused_rows::<4>(block, quarter, num_ntts, |rows| { let [row_a, row_b, row_c, row_d] = rows; // Layer L: rows 2 apart, one twiddle for the block. @@ -1132,32 +826,8 @@ fn butterfly_interleaved_fused_2layer( }); } -/// [`butterfly_interleaved_fused_2layer`] for rows that end in two to seven lanes past a whole vector. -/// -/// As [`butterfly_interleaved_fused_3layer_masked`]. -#[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] -#[inline(never)] -fn butterfly_interleaved_fused_2layer_masked(block: &mut [F64], tw: &[__m512i; 3], quarter: usize, num_ntts: usize) { - let (vectors, mask) = (num_ntts / 8, (1u8 << (num_ntts % 8)) - 1); - fused_rows::<4>(block, quarter, num_ntts, |rows| { - let rows: [*mut F64; 4] = std::array::from_fn(|i| rows[i].as_mut_ptr()); - // SAFETY: the target features are enabled at compile time; each row has `8 * vectors` words and the - // masked lanes after them. - unsafe { - radix4_avx512(rows, tw, vectors); - masked_column_avx512(rows.map(|row| row.add(8 * vectors)), mask, |r| { - radix4_regs_avx512(r, tw); - }); - } - }); -} - #[inline] fn butterfly_interleaved_block(block: &mut [F64], twiddle: F64, block_size_half: usize, num_ntts: usize) { - #[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] - if num_ntts % 8 >= 2 { - return butterfly_interleaved_block_masked(block, twiddle, block_size_half, num_ntts); - } let half_offset = block_size_half * num_ntts; let (top, bot) = block.split_at_mut(half_offset); for r in 0..block_size_half { @@ -1166,29 +836,6 @@ fn butterfly_interleaved_block(block: &mut [F64], twiddle: F64, block_size_half: } } -/// [`butterfly_interleaved_block`] for rows that end in two to seven lanes past a whole vector. -/// -/// As [`butterfly_interleaved_fused_3layer_masked`]. -#[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] -#[inline(never)] -fn butterfly_interleaved_block_masked(block: &mut [F64], twiddle: F64, block_size_half: usize, num_ntts: usize) { - let (vectors, mask) = (num_ntts / 8, (1u8 << (num_ntts % 8)) - 1); - let half_offset = block_size_half * num_ntts; - let base = block.as_mut_ptr(); - // SAFETY: the arm is compiled only with AVX-512F; registers only. - let tw = unsafe { _mm512_set1_epi64(twiddle.0 as i64) }; - for r in 0..block_size_half { - // SAFETY: rows `r` of the two halves, `num_ntts` words each, inside the block and disjoint. - unsafe { - let rows = [base.add(r * num_ntts), base.add(half_offset + r * num_ntts)]; - columns_avx512(rows, vectors, |p| butterfly_regs_avx512(p, 0, 1, tw)); - masked_column_avx512(rows.map(|row| row.add(8 * vectors)), mask, |p| { - butterfly_regs_avx512(p, 0, 1, tw); - }); - } - } -} - /// Butterfly all `num_ntts` lanes of one (top row, bottom row) pair with a /// shared twiddle: new_u = u + v*t; new_v = v + new_u. #[inline] @@ -1229,16 +876,16 @@ fn lane_butterflies(top: &mut [F64], bot: &mut [F64], tw #[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] let done = { let vectors = top.len() / 8; - // SAFETY: the target features are enabled at compile time, and both rows have the lanes covered. + // SAFETY: the target features are enabled at compile time and each + // iteration reads and writes exactly eight elements from both rows. unsafe { - let tw = _mm512_set1_epi64(twiddle.0 as i64); - columns_avx512([top.as_mut_ptr(), bot.as_mut_ptr()], vectors, |r| { - if TRANSPOSED { - transposed_butterfly_regs_avx512(r, 0, 1, tw); - } else { - butterfly_regs_avx512(r, 0, 1, tw); - } - }); + for i in 0..vectors { + butterfly_lanes_avx512::( + top.as_mut_ptr().add(8 * i), + bot.as_mut_ptr().add(8 * i), + twiddle.0, + ); + } } 8 * vectors }; @@ -1458,7 +1105,7 @@ unsafe fn butterfly_lanes_neon_8(top: *mut F64, bot: *mu } } -/// The products `v * t` of eight F64 lanes by one broadcast twiddle, reduced. +/// Eight F64 butterflies with a shared twiddle, one per 64-bit lane of an AVX-512 register. /// /// # Algorithm /// @@ -1470,38 +1117,46 @@ unsafe fn butterfly_lanes_neon_8(top: *mut F64, bot: *mu /// /// reduce: p = lo + hi * x^64, x^64 = x^4 + x^3 + x + 1 /// p mod f = lo ^ g(hi ^ spill), g(y) = y ^ y<<1 ^ y<<3 ^ y<<4 -/// spill = n ^ n>>1 ^ n>>3 with n = hi>>60, the bits g pushes past x^63 +/// spill = hi>>63 ^ hi>>61 ^ hi>>60, the bits g pushes past x^63 /// ``` /// /// - The unpacks stay inside 128-bit lanes, so no shuffle crosses lanes. -/// - The reduction takes no further carry-less multiply: the multiplier issues one every other cycle. -/// - Shifts and shuffles share two pipes with the unpacks, adds have four, so the left shifts of `g` are doublings. -/// - The spill is one table lookup of the top nibble where AVX-512BW has `vpshufb`. +/// - The reduction takes shifts and three-way XORs, `vpternlogq 0x96`, instead of two more carry-less multiplies. +/// - Zen 5 has one carry-less multiplier per core, so those two would dominate the kernel. +/// +/// # Safety +/// +/// - Requires VPCLMULQDQ and AVX-512F. +/// - Each pointer must address eight readable and writable words. #[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] -#[inline(always)] -fn mul_lanes_avx512(v: __m512i, tw: __m512i) -> __m512i { - const XOR3: i32 = 0x96; - // SAFETY: the arm is compiled only with VPCLMULQDQ and AVX-512F (and BW where it uses it); registers only. +#[inline] +#[target_feature(enable = "vpclmulqdq", enable = "avx512f", enable = "avx2")] +unsafe fn butterfly_lanes_avx512(top: *mut F64, bot: *mut F64, twiddle: u64) { + // SAFETY: + // - The caller supplies two valid eight-word rows. + // - This function's target features cover every intrinsic below. unsafe { - let even = _mm512_clmulepi64_epi128::<0x00>(v, tw); - let odd = _mm512_clmulepi64_epi128::<0x11>(v, tw); + // Load both rows and broadcast the twiddle to every lane. + let u = _mm512_loadu_si512(top.cast()); + let v = _mm512_loadu_si512(bot.cast()); + let tw = _mm512_set1_epi64(twiddle as i64); + // The row multiplied by the twiddle, and the row the product is added to. + let (m, acc) = if TRANSPOSED { + (_mm512_xor_si512(u, v), v) + } else { + (v, u) + }; + + // Products m * t: even lanes, then odd lanes, one 128-bit product per 128-bit lane. + let even = _mm512_clmulepi64_epi128::<0x00>(m, tw); + let odd = _mm512_clmulepi64_epi128::<0x11>(m, tw); + // Back to lane order: qword i of lo / hi is the low / high half of lane i's product. let lo = _mm512_unpacklo_epi64(even, odd); let hi = _mm512_unpackhi_epi64(even, odd); - #[cfg(target_feature = "avx512bw")] - let spill = { - const SPILL: [u8; 16] = { - let mut table = [0u8; 16]; - let mut n = 0; - while n < 16 { - table[n] = (n ^ (n >> 1) ^ (n >> 3)) as u8; - n += 1; - } - table - }; - let table = _mm512_broadcast_i32x4(_mm_loadu_si128(SPILL.as_ptr().cast())); - _mm512_shuffle_epi8(table, _mm512_srli_epi64::<60>(hi)) - }; - #[cfg(not(target_feature = "avx512bw"))] + + // Reduce modulo x^64 + x^4 + x^3 + x + 1 with shifts and three-way XORs. + const XOR3: i32 = 0x96; + // The bits of hi * (x^4 + x^3 + x + 1) that land past x^63. let spill = _mm512_ternarylogic_epi64::( _mm512_srli_epi64::<63>(hi), _mm512_srli_epi64::<61>(hi), @@ -1509,145 +1164,21 @@ fn mul_lanes_avx512(v: __m512i, tw: __m512i) -> __m512i { ); // Both hi and spill are multiplied by the same constant, so fold them first. let x = _mm512_xor_si512(hi, spill); - let x2 = _mm512_add_epi64(x, x); - let x8 = { - let x4 = _mm512_add_epi64(x2, x2); - _mm512_add_epi64(x4, x4) + // g(x) = x ^ x<<1 ^ x<<3 ^ x<<4, split across two three-way XORs. + let fx = _mm512_ternarylogic_epi64::(x, _mm512_slli_epi64::<1>(x), _mm512_slli_epi64::<3>(x)); + // acc + m * t, with the product's lo and g(x) folded in one step. + let sum = _mm512_ternarylogic_epi64::(acc, lo, _mm512_xor_si512(fx, _mm512_slli_epi64::<4>(x))); + // Forward: u' = u + v * t, then v' = v + u'. Transposed: u' = u + v, then v' = v + u' * t. + let (new_u, new_v) = if TRANSPOSED { + (m, sum) + } else { + (sum, _mm512_xor_si512(v, sum)) }; - let x16 = _mm512_add_epi64(x8, x8); - _mm512_ternarylogic_epi64::(_mm512_ternarylogic_epi64::(lo, x, x2), x8, x16) - } -} - -/// One butterfly of eight lanes held in registers: `u' = u + v * t`, `v' = v + u'`. -#[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] -#[inline(always)] -fn butterfly_regs_avx512(rows: &mut [__m512i], top: usize, bot: usize, tw: __m512i) { - // SAFETY: the arm is compiled only with AVX-512F; registers only. - unsafe { - let new_u = _mm512_xor_si512(rows[top], mul_lanes_avx512(rows[bot], tw)); - rows[bot] = _mm512_xor_si512(rows[bot], new_u); - rows[top] = new_u; - } -} - -/// One transposed butterfly of eight lanes held in registers: `u' = u + v`, `v' = v + u' * t`. -#[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] -#[inline(always)] -fn transposed_butterfly_regs_avx512(rows: &mut [__m512i], top: usize, bot: usize, tw: __m512i) { - // SAFETY: the arm is compiled only with AVX-512F; registers only. - unsafe { - let new_u = _mm512_xor_si512(rows[top], rows[bot]); - rows[bot] = _mm512_xor_si512(rows[bot], mul_lanes_avx512(new_u, tw)); - rows[top] = new_u; - } -} - -/// Run `body` on columns of eight lanes of `N` rows held in registers, each row loaded and stored once. -/// -/// The lanes past the last whole vector are the caller's. -/// -/// # Safety -/// -/// - Requires AVX-512F. -/// - Each row pointer must address `8 * vectors` readable and writable words. -#[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] -#[inline(always)] -unsafe fn columns_avx512(rows: [*mut F64; N], vectors: usize, body: impl Fn(&mut [__m512i; N])) { - // SAFETY: every access is inside the words the caller guarantees. - unsafe { - for c in 0..vectors { - let mut r: [__m512i; N] = std::array::from_fn(|i| _mm512_loadu_si512(rows[i].add(8 * c).cast())); - body(&mut r); - for (i, row) in r.iter().enumerate() { - _mm512_storeu_si512(rows[i].add(8 * c).cast(), *row); - } - } + _mm512_storeu_si512(top.cast(), new_u); + _mm512_storeu_si512(bot.cast(), new_v); } } -/// [`columns_avx512`] for one column of fewer than eight lanes, the ones `mask` selects. -/// -/// # Safety -/// -/// - Requires AVX-512F. -/// - Each row pointer must address the lanes `mask` selects; the others are neither read nor written. -#[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] -#[inline(always)] -unsafe fn masked_column_avx512(rows: [*mut F64; N], mask: __mmask8, body: impl Fn(&mut [__m512i; N])) { - // SAFETY: masked-off lanes are neither read nor written, the rest are the caller's. - unsafe { - let mut r: [__m512i; N] = std::array::from_fn(|i| _mm512_maskz_loadu_epi64(mask, rows[i].cast())); - body(&mut r); - for (i, row) in r.iter().enumerate() { - _mm512_mask_storeu_epi64(rows[i].cast(), mask, *row); - } - } -} - -/// The twelve butterflies of a radix-8 row group held in registers, its seven twiddles breadth-first. -#[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] -#[inline(always)] -fn radix8_regs_avx512(r: &mut [__m512i; 8], t: &[__m512i; 7]) { - // Layer L: rows 4 apart, one twiddle for the whole block. - butterfly_regs_avx512(r, 0, 4, t[0]); - butterfly_regs_avx512(r, 1, 5, t[0]); - butterfly_regs_avx512(r, 2, 6, t[0]); - butterfly_regs_avx512(r, 3, 7, t[0]); - // Layer L+1: rows 2 apart, one twiddle per half. - butterfly_regs_avx512(r, 0, 2, t[1]); - butterfly_regs_avx512(r, 1, 3, t[1]); - butterfly_regs_avx512(r, 4, 6, t[2]); - butterfly_regs_avx512(r, 5, 7, t[2]); - // Layer L+2: adjacent rows, one twiddle per quarter. - butterfly_regs_avx512(r, 0, 1, t[3]); - butterfly_regs_avx512(r, 2, 3, t[4]); - butterfly_regs_avx512(r, 4, 5, t[5]); - butterfly_regs_avx512(r, 6, 7, t[6]); -} - -/// The four butterflies of a radix-4 row group held in registers: one twiddle for the block, then one per half. -#[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] -#[inline(always)] -fn radix4_regs_avx512(r: &mut [__m512i; 4], t: &[__m512i; 3]) { - // Layer L: rows 2 apart, one twiddle for the block. - butterfly_regs_avx512(r, 0, 2, t[0]); - butterfly_regs_avx512(r, 1, 3, t[0]); - // Layer L+1: adjacent rows, one twiddle per half. - butterfly_regs_avx512(r, 0, 1, t[1]); - butterfly_regs_avx512(r, 2, 3, t[2]); -} - -/// A radix-8 row group with the eight rows in registers. -/// -/// Each row is loaded and stored once for its twelve butterflies, not three times. -/// -/// # Safety -/// -/// - Requires VPCLMULQDQ and AVX-512F. -/// - Each row pointer must address `8 * vectors` readable and writable words. -#[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] -#[inline] -#[target_feature(enable = "vpclmulqdq", enable = "avx512f", enable = "avx2")] -unsafe fn radix8_avx512(rows: [*mut F64; 8], t: &[__m512i; 7], vectors: usize) { - // SAFETY: forwarded from the caller. - unsafe { columns_avx512(rows, vectors, |r| radix8_regs_avx512(r, t)) }; -} - -/// A radix-4 row group with the four rows in registers. -/// -/// # Safety -/// -/// - Requires VPCLMULQDQ and AVX-512F. -/// - Each row pointer must address `8 * vectors` readable and writable words. -#[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] -#[inline] -#[target_feature(enable = "vpclmulqdq", enable = "avx512f", enable = "avx2")] -unsafe fn radix4_avx512(rows: [*mut F64; 4], t: &[__m512i; 3], vectors: usize) { - // SAFETY: forwarded from the caller. - unsafe { columns_avx512(rows, vectors, |r| radix4_regs_avx512(r, t)) }; -} - /// Two F64 butterflies with a shared twiddle, NEON-resident end to end. /// The two products issue as PMULL/PMULL2 on the loaded row (no lane /// extraction) and reduce through the all-PMULL lane-pair fold @@ -1760,7 +1291,6 @@ mod tests { // (12, 2048, 1) two gathered passes, with streaming stores // // A non-zero start is the commit path, which enters at the rate layer. - // Up to `FEW_WORKERS` workers, the plan takes several residues per gathered group. let mut rng = Rng::new(0xC0FFEE); for (log_d, lanes, start_layer) in [ (7usize, 3usize, 0usize), @@ -1777,15 +1307,11 @@ mod tests { // Reference: one butterfly at a time. let mut want = original.clone(); forward_scalar_from_layer(&ntt, &mut want, lanes, start_layer); - // Under test: the pass-planning driver, planned for workers on both sides of `FEW_WORKERS`. - for workers in [1, FEW_WORKERS, FEW_WORKERS + 1] { - let mut got = original.clone(); - ntt.transform_for(&mut got, lanes, start_layer, None, None, workers); - assert_eq!( - got, want, - "log_d={log_d}, lanes={lanes}, start_layer={start_layer}, workers={workers}" - ); - } + // Under test: the pass-planning driver. + let mut got = original; + ntt.transform(&mut got, lanes, start_layer, None, None); + + assert_eq!(got, want, "log_d={log_d}, lanes={lanes}, start_layer={start_layer}"); } } @@ -1825,41 +1351,35 @@ mod tests { forward_scalar_from_layer(&ntt, &mut want, lanes, log_inv_rate); // Under test: only the first replica holds the message, the rest is zero. - for workers in [1, FEW_WORKERS, FEW_WORKERS + 1] { - let mut got = vec![F64::ZERO; msg_len << log_inv_rate]; - got[..msg_len].copy_from_slice(&msg); - let blocks = Mutex::new(Vec::new()); - let first = Message::Rows(SendPtr(got.as_mut_ptr())); - let sink = |row: usize, rows: &[F64]| blocks.lock().unwrap().push((row, rows.to_vec())); - ntt.transform_for(&mut got, lanes, log_inv_rate, Some(first), Some(&sink), workers); + let mut got = vec![F64::ZERO; msg_len << log_inv_rate]; + got[..msg_len].copy_from_slice(&msg); + let blocks = Mutex::new(Vec::new()); + ntt.encode_interleaved_in_place_with(&mut got, lanes, log_inv_rate, &|row, rows| { + blocks.lock().unwrap().push((row, rows.to_vec())); + }); + assert_eq!(got, want, "log_d={log_d}, lanes={lanes}, rate={log_inv_rate}"); + + let mut blocks = blocks.into_inner().unwrap(); + blocks.sort_by_key(|b| b.0); + let mut next = 0; + for (row, rows) in blocks { + assert_eq!(row, next, "blocks tile the rows, log_d={log_d}"); assert_eq!( - got, want, - "log_d={log_d}, lanes={lanes}, rate={log_inv_rate}, workers={workers}" + rows[..], + want[row * lanes..][..rows.len()], + "a handed-over block is final" ); - - let mut blocks = blocks.into_inner().unwrap(); - blocks.sort_by_key(|b| b.0); - let mut next = 0; - for (row, rows) in blocks { - assert_eq!(row, next, "blocks tile the rows, log_d={log_d}"); - assert_eq!( - rows[..], - want[row * lanes..][..rows.len()], - "a handed-over block is final" - ); - next += rows.len() / lanes; - } - assert_eq!(next, 1 << log_d, "every row handed over, log_d={log_d}"); + next += rows.len() / lanes; } + assert_eq!(next, 1 << log_d, "every row handed over, log_d={log_d}"); } } /// Every lane of the codeword must be exactly the single-lane RS codeword of /// that lane's contiguous message block, which is what makes a commitment over /// `n_lanes` lanes equal to the `2^log_batch_size`-lane one with a zero tail. - /// The shapes cover the transposing fallback, the few-workers first pass reading - /// the lanes in place (the last three, with and without lanes past a whole eight), - /// and lane counts that are not powers of two (the padding-free commit's whole point). + /// The shapes cover the transposing fused first pass, its fallback, and lane + /// counts that are not powers of two (the padding-free commit's whole point). #[test] fn lane_major_msg_encode_matches_per_lane_reference() { let mut rng = Rng::new(0x1A2E); @@ -1870,39 +1390,29 @@ mod tests { (9, 1, 5), (12, 2, 37), (14, 1, 64), - (13, 1, 41), - (13, 1, 64), - (16, 1, 5), ] { let log_d = log_rows + log_inv_rate; let ntt = AdditiveNttF64::standard(log_d); let rows = 1usize << log_rows; let msg: Vec = (0..rows * n_lanes).map(|_| F64(rng.next_u64())).collect(); + let mut got = vec![F64::ZERO; msg.len() << log_inv_rate]; + transpose_lane_major(&mut got[..msg.len()], &msg, n_lanes, log_rows); + ntt.encode_interleaved_in_place(&mut got, n_lanes, log_inv_rate); + let block_len = 1usize << log_d; - let want: Vec> = (0..n_lanes) - .map(|lane| { - // Lane `lane` encodes message block `n_lanes - 1 - lane`. - let block = n_lanes - 1 - lane; - let mut want = vec![F64::ZERO; block_len]; - replicate_rows(&mut want, &msg[block * rows..(block + 1) * rows]); - ntt.transform(&mut want, 1, log_inv_rate, None, None); - want - }) - .collect(); - for workers in [1, FEW_WORKERS, FEW_WORKERS + 1] { - let mut got = vec![F64::ZERO; msg.len() << log_inv_rate]; - let lanes = Message::Lanes(SendPtr(msg.as_ptr().cast_mut())); - ntt.transform_for(&mut got, n_lanes, log_inv_rate, Some(lanes), Some(&|_, _| {}), workers); - for (lane, want) in want.iter().enumerate() { - for (pos, want) in want.iter().enumerate() { - assert_eq!( - got[pos * n_lanes + lane], - *want, - "lane {lane} pos {pos} at log_rows={log_rows}, rate={log_inv_rate}, n_lanes={n_lanes}, \ - workers={workers}" - ); - } + for lane in 0..n_lanes { + // Lane `lane` encodes message block `n_lanes - 1 - lane`. + let block = n_lanes - 1 - lane; + let mut want = vec![F64::ZERO; block_len]; + replicate_rows(&mut want, &msg[block * rows..(block + 1) * rows]); + ntt.transform(&mut want, 1, log_inv_rate, None, None); + for pos in 0..block_len { + assert_eq!( + got[pos * n_lanes + lane], + want[pos], + "lane {lane} pos {pos} at log_rows={log_rows}, rate={log_inv_rate}, n_lanes={n_lanes}" + ); } } } diff --git a/crates/pcs/src/whir/commit.rs b/crates/pcs/src/whir/commit.rs index 3ce37560f..ce082d401 100644 --- a/crates/pcs/src/whir/commit.rs +++ b/crates/pcs/src/whir/commit.rs @@ -59,10 +59,12 @@ pub fn commit(message: &[F64], log_n: usize, log_batch_size: usize, log_inv_rate let tree = MerkleBuilder::new(n_positions, n_lanes, 1usize << log_batch_size); tracing::info_span!("NTT", kind = "base encode", log_domain = k_code, lanes = n_lanes).in_scope(|| { // SAFETY: every codeword element is written before it is read. - // `encode_lane_major_with` writes every word of an uninitialized codeword before it reads it. + // The transpose covers every word of the message region (its tiles are asserted to). + // The encode writes every other replica from it before transforming that region in place. let codeword = unsafe { primitives::write_only(&mut codeword) }; + crate::ntt::transpose_lane_major(&mut codeword[..message.len()], message, n_lanes, log_rows); let ntt = AdditiveNttF64::standard(k_code); - ntt.encode_lane_major_with(codeword, message, n_lanes, log_inv_rate, &|row, rows| { + ntt.encode_interleaved_in_place_with(codeword, n_lanes, log_inv_rate, &|row, rows| { tree.absorb(row, rows); }); });