diff --git a/crates/leanvm_core/src/colval.rs b/crates/leanvm_core/src/colval.rs deleted file mode 100644 index 89cbcb39d..000000000 --- a/crates/leanvm_core/src/colval.rs +++ /dev/null @@ -1,105 +0,0 @@ -//! Column operations shared by the `F64` joining round and subsequent `F192` rounds. - -use primitives::field::{F64, F192, F192Unreduced}; -use std::ops::{Add, BitXor, BitXorAssign, Mul}; - -/// A value a constraint reads out of a column. -pub trait ColVal: Copy + Send + Sync + Add + Mul { - const ZERO: Self; - const ONE: Self; - - /// Where this column's products XOR-accumulate before the one reduction that - /// ends a form. - type Unreduced: Copy + Send + Sync + BitXor + BitXorAssign; - - /// Times an `E` value: an `η`-power, a bus coefficient, a machine word. - fn mul_e(self, e: F192) -> F192; - - /// The same product, left for the caller to accumulate. - fn mul_e_unreduced(self, e: F192) -> Self::Unreduced; - - /// An already-reduced `E` value as an accumulator term; exact, because - /// reducing a reduced value fixes it. - fn lift(e: F192) -> Self::Unreduced; - - fn reduce(acc: Self::Unreduced) -> F192; - - /// `Σ coeffs[i]·vals[i]`, unreduced. - fn dot_unreduced(coeffs: &[F192], vals: &[Self]) -> Self::Unreduced; - - /// `constant + Σ coeffs[i]·vals[i]`, one reduction for the whole slice. - #[inline(always)] - fn dot(coeffs: &[F192], vals: &[Self], constant: F192) -> F192 { - Self::reduce(Self::dot_unreduced(coeffs, vals) ^ Self::lift(constant)) - } -} - -impl ColVal for F64 { - const ZERO: Self = F64::ZERO; - const ONE: Self = F64::ONE; - - type Unreduced = F192Unreduced; - - #[inline(always)] - fn mul_e(self, e: F192) -> F192 { - e.mul_base(self) - } - - #[inline(always)] - fn mul_e_unreduced(self, e: F192) -> Self::Unreduced { - e.mul_base_unreduced(self) - } - - #[inline(always)] - fn lift(e: F192) -> Self::Unreduced { - e.into() - } - - #[inline(always)] - fn reduce(acc: Self::Unreduced) -> F192 { - acc.reduce() - } - - #[inline(always)] - fn dot_unreduced(coeffs: &[F192], vals: &[Self]) -> Self::Unreduced { - coeffs - .iter() - .zip(vals) - .fold(F192Unreduced::ZERO, |acc, (&w, &v)| acc ^ w.mul_base_unreduced(v)) - } -} - -impl ColVal for F192 { - const ZERO: Self = F192::ZERO; - const ONE: Self = F192::ONE; - - type Unreduced = F192Unreduced; - - #[inline(always)] - fn mul_e(self, e: F192) -> F192 { - self * e - } - - #[inline(always)] - fn mul_e_unreduced(self, e: F192) -> Self::Unreduced { - self.mul_unreduced(e) - } - - #[inline(always)] - fn lift(e: F192) -> Self::Unreduced { - e.into() - } - - #[inline(always)] - fn reduce(acc: Self::Unreduced) -> F192 { - acc.reduce() - } - - #[inline(always)] - fn dot_unreduced(coeffs: &[F192], vals: &[Self]) -> Self::Unreduced { - coeffs - .iter() - .zip(vals) - .fold(F192Unreduced::ZERO, |acc, (&w, &v)| acc ^ w.mul_unreduced(v)) - } -} diff --git a/crates/leanvm_core/src/constraints.rs b/crates/leanvm_core/src/constraints.rs index a57a2efc9..4c61f7fb1 100644 --- a/crates/leanvm_core/src/constraints.rs +++ b/crates/leanvm_core/src/constraints.rs @@ -32,11 +32,10 @@ //! verifier certifies that maximum with one hinted `g`-power (§recursion), so there //! are no rounds in which no table has joined. -use crate::PAR_THRESHOLD; -use crate::colval::ColVal; +use crate::leaf::BusForm; use crate::transcript::ProverState; -use primitives::field::{F64, F192, F192Unreduced, powers}; -use primitives::multilinear::{eq_table, fold_high_inplace, fold_high_k, poly_eval, shrink_eq_high}; +use primitives::field::{F64, F64x8, F192, F192x8, Prod8}; +use primitives::multilinear::{eq_table, interp, mle_eval, poly_eval, shrink_eq_high}; /// One table's involved columns' evaluations at its table-sumcheck point. #[derive(Clone, Debug, PartialEq, Eq)] @@ -45,17 +44,16 @@ pub struct Claims { pub evals: Vec, } -/// One table's row constraint: identity `i` is weighted by `pows[i]`; `true` selects its homogeneous quadratic part. -pub type Constraint<'a> = Box F192 + Sync + 'a>; -/// The same form over `K`-valued columns, for the round a table joins the batch. -pub type ConstraintK<'a> = Box F192 + Sync + 'a>; - /// One table's place in the shared batch. pub struct Air<'a> { pub tau: usize, - pub n_constraints: usize, - pub eval: Constraint<'a>, - pub eval_k: ConstraintK<'a>, + /// Its committed columns, `2^tau` values each, read in place. + pub cols: &'a [&'a [F64]], + /// Its summand, a degree-2 form over `cols`. + pub form: BusForm, + /// What the summand sums to against `eq(ζ, ·)`: prover-side only, for the line the table + /// contributes while it waits. The verifier derives the total. + pub sigma: F192, } /// Start of each table's disjoint range of `η`-powers. @@ -69,55 +67,336 @@ pub fn xi_offsets(n_constraints: impl Iterator) -> Vec { .collect() } -/// An active round: one endpoint evaluation and the quadratic coefficient. +/// A linear form over a table's columns, as `(column, coefficient)` terms. +type Lin = Vec<(usize, F192)>; + +/// Blocks of eight rows per task of the round passes. +const TASK_BLOCKS: usize = 128; + +/// Vectors per task of a fold. +const FOLD_CHUNK: usize = 1 << 10; + +/// A table's summand `P(v) = constant + lin(v) + Σ_g v[left_g] · right_g(v)`, its products grouped +/// by a shared column. Its homogeneous quadratic part is `Q(v) = Σ_g v[left_g] · right_g(v)`. /// -/// Generic twice over: in the column element, `K` before a table's columns are -/// folded and `E` after ([`ColVal`]), and in the container, `Vec` for the former -/// and `Vec` for the latter. `#[inline(always)]` matters here, on this and on -/// every `ColVal` method: this is the body of the constraint sumcheck's innermost -/// loop, and without it the generic stops inlining and costs measurable prover -/// time. Nothing is lifted into `E`, so a `K` round evaluates the identity and the -/// bus forms in 64-bit arithmetic, and its scratch is a third the size. +/// The `right_g` and `lin` are linear, so a joined table folds them as single `E` columns +/// instead of folding every column: its state is `[lin, left_0, right_0, left_1, right_1, ..]`. +struct Form { + constant: F192, + lin: Lin, + groups: Vec<(usize, Lin)>, + /// Every column a term reads, ascending. + used: Vec, +} + +impl Form { + /// The summand of `form`, its products grouped by a greedy vertex cover: the column with the most + /// remaining products becomes a group's left factor. + fn new(form: &BusForm) -> Self { + let lin: Lin = form + .coeffs + .iter() + .copied() + .enumerate() + .filter(|(_, c)| !c.is_zero()) + .collect(); + let mut edges: Vec<(usize, usize, F192)> = Vec::new(); + for &(a, b, q) in &form.prods { + let (a, b) = (a.min(b), a.max(b)); + match edges.iter_mut().find(|e| (e.0, e.1) == (a, b)) { + Some(e) => e.2 += q, + None => edges.push((a, b, q)), + } + } + edges.retain(|e| !e.2.is_zero()); + let n_cols = form.coeffs.len(); + let mut groups = Vec::new(); + while !edges.is_empty() { + let mut degree = vec![0usize; n_cols]; + for &(a, b, _) in &edges { + degree[a] += 1; + degree[b] += usize::from(a != b); + } + let left = (0..n_cols).max_by_key(|&c| degree[c]).unwrap(); + let mut right = Vec::new(); + edges.retain(|&(a, b, q)| { + let partner = if a == left { b } else { a }; + let mine = a == left || b == left; + if mine { + right.push((partner, q)); + } + !mine + }); + groups.push((left, right)); + } + let mut used: Vec = lin.iter().map(|t| t.0).collect(); + for (left, right) in &groups { + used.push(*left); + used.extend(right.iter().map(|t| t.0)); + } + used.sort_unstable(); + used.dedup(); + Self { + constant: form.constant, + lin, + groups, + used, + } + } + + /// Columns of the folded state: `lin`, then `left_g` and `right_g` per group. + fn width(&self) -> usize { + 1 + 2 * self.groups.len() + } + + /// The state of the first `rows` (at most eight) rows, unfolded, from the `K` columns. + fn small_state(&self, cols: &[&[F64]], rows: usize) -> State { + let dot = |terms: &Lin, i: usize| { + terms + .iter() + .fold(F192::ZERO, |acc, &(c, w)| acc + w.mul_base(cols[c][i])) + }; + let column = |value: &dyn Fn(usize) -> F192| { + F192x8::from_array(&std::array::from_fn(|i| if i < rows { value(i) } else { F192::ZERO })) + }; + let mut state = vec![column(&|i| dot(&self.lin, i))]; + for (left, right) in &self.groups { + state.push(column(&|i| F192::from(cols[*left][i]))); + state.push(column(&|i| dot(right, i))); + } + state + } +} + +fn load_k(col: &[F64], i: usize) -> F64x8 { + F64x8::from_array(col[i..i + 8].try_into().unwrap()) +} + +fn load_e(col: &[F192], i: usize) -> F192x8 { + F192x8::from_array(col[i..i + 8].try_into().unwrap()) +} + +fn splat(terms: &Lin) -> Vec<(usize, F192x8)> { + terms.iter().map(|&(c, w)| (c, F192x8::splat(w))).collect() +} + +/// A worker's scratch for [`load_rows`] over `n` columns: `end` and `sum`. +fn rows_scratch(n: usize) -> (Vec, Vec) { + let zero = F64x8::from_array(&[F64::ZERO; 8]); + (vec![zero; n], vec![zero; n]) +} + +/// Loads the vectors of rows `i` and `i + half` of every column in `read`: `end` takes the upper vector +/// with `at_one` and the lower without, `sum` their sum. #[inline(always)] -fn table_message + Sync>( - cols: &[C], - eval: &(impl Fn(&[F192], &[T], bool) -> F192 + Sync + ?Sized), - pows: &[F192], +fn load_rows( + cols: &[&[F64]], + read: &[usize], + i: usize, half: usize, - eqr: &[F192], at_one: bool, -) -> [F192; 2] { - let ncols = cols.len(); - // The X² coefficient is Q(hi + lo); linear and constant terms cannot contribute. - let summand = |i: usize, scratch: &mut [T]| -> [F192Unreduced; 2] { - let e = eqr[i]; - let (endpoint, slope) = scratch.split_at_mut(ncols); - for (ci, c) in cols.iter().enumerate() { - let (lo, hi) = (c[i], c[i + half]); - endpoint[ci] = if at_one { hi } else { lo }; - slope[ci] = lo + hi; - } - [ - e.mul_unreduced(eval(pows, endpoint, false)), - e.mul_unreduced(eval(pows, slope, true)), - ] + (end, sum): &mut (Vec, Vec), +) { + for &c in read { + let (lo, hi) = (load_k(cols[c], i), load_k(cols[c], i + half)); + (end[c], sum[c]) = (if at_one { hi } else { lo }, lo + hi); + } +} + +/// A joined table's folded state: per vector of eight rows, its `width` columns side by side, +/// `state[b * width + c]`. Below eight rows, one vector per column holds them all. +type State = Vec; + +/// The blocks task `t` covers out of `n`. +fn task_blocks(t: usize, n: usize) -> std::ops::Range { + t * TASK_BLOCKS..((t + 1) * TASK_BLOCKS).min(n) +} + +fn xor2(a: [Prod8; 2], b: [Prod8; 2]) -> [Prod8; 2] { + [a[0] ^ b[0], a[1] ^ b[1]] +} + +fn finish([endpoint, slope]: [Prod8; 2], constant: F192) -> [F192; 2] { + [endpoint.reduce() + constant, slope.reduce()] +} + +/// An active round of a table that has not joined: its message from the `K` columns, in place. +/// Per vector of rows, the linear forms are `E x K` mixed dots, and every product of the +/// summand is one mixed product against a reduced `right_g`. +fn message_k(form: &Form, cols: &[&[F64]], eqr: &[F192], half: usize, at_one: bool) -> [F192; 2] { + if half < 8 { + return message_e(form, &form.small_state(cols, 2 * half), eqr, half, at_one); + } + let lin = splat(&form.lin); + let rights: Vec<_> = form.groups.iter().map(|(_, right)| splat(right)).collect(); + let dot = |terms: &[(usize, F192x8)], vals: &[F64x8]| { + terms + .iter() + .fold(Prod8::ZERO, |acc, &(c, w)| acc ^ w.mul_base_unreduced(vals[c])) }; - let xor = |a: [F192Unreduced; 2], b: [F192Unreduced; 2]| [a[0] ^ b[0], a[1] ^ b[1]]; - let acc = if half >= PAR_THRESHOLD { - // The `2 * ncols` scratch is per-worker, not per-row: `map_reduce_with_state` - // creates it once and threads it through every row that worker claims. - parallel::map_reduce_with_state( - half, - || vec![T::ZERO; 2 * ncols], - || [F192Unreduced::ZERO; 2], - |scratch, acc, i| *acc = xor(*acc, summand(i, scratch)), - xor, - ) - } else { - let mut scratch = vec![T::ZERO; 2 * ncols]; - (0..half).fold([F192Unreduced::ZERO; 2], |acc, i| xor(acc, summand(i, &mut scratch))) + let blocks = half / 8; + let acc = parallel::map_reduce_with_state( + blocks.div_ceil(TASK_BLOCKS), + || rows_scratch(cols.len()), + || [Prod8::ZERO; 2], + |rows, acc, t| { + for b in task_blocks(t, blocks) { + let i = 8 * b; + load_rows(cols, &form.used, i, half, at_one, rows); + let (end, sum) = &*rows; + let (mut endpoint, mut slope) = (dot(&lin, end), Prod8::ZERO); + for ((left, _), right) in form.groups.iter().zip(&rights) { + endpoint ^= dot(right, end).reduce_lanes().mul_base_unreduced(end[*left]); + slope ^= dot(right, sum).reduce_lanes().mul_base_unreduced(sum[*left]); + } + let eq = load_e(eqr, i); + acc[0] ^= endpoint.reduce_lanes().mul_unreduced(eq); + acc[1] ^= slope.reduce_lanes().mul_unreduced(eq); + } + }, + xor2, + ); + finish(acc, form.constant) +} + +/// An active round of a joined table, from its folded state. +fn message_e(form: &Form, state: &State, eqr: &[F192], half: usize, at_one: bool) -> [F192; 2] { + if half < 8 { + // Fewer row pairs than a vector: `[Σ_i eqr[i]·end_i, Σ_i eqr[i]·Q(lo_i + hi_i)]` by lanes. + // `eqr` sums to one, so the form's constant seeds the endpoint. + let rows: Vec<[F192; 8]> = state.iter().map(|v| v.to_array()).collect(); + return (0..half).fold([form.constant, F192::ZERO], |acc, i| { + let (lo, hi) = (|c: usize| rows[c][i], |c: usize| rows[c][i + half]); + let end = |c| if at_one { hi(c) } else { lo(c) }; + let (mut endpoint, mut slope) = (end(0), F192::ZERO); + for g in (1..state.len()).step_by(2) { + endpoint += end(g) * end(g + 1); + slope += (lo(g) + hi(g)) * (lo(g + 1) + hi(g + 1)); + } + [acc[0] + eqr[i] * endpoint, acc[1] + eqr[i] * slope] + }); + } + let (blocks, width) = (half / 8, form.width()); + let (lows, highs) = state.split_at(blocks * width); + let acc = parallel::fold_reduce( + blocks.div_ceil(TASK_BLOCKS), + || [Prod8::ZERO; 2], + |acc, t| { + for b in task_blocks(t, blocks) { + let (lo, hi) = (&lows[b * width..][..width], &highs[b * width..][..width]); + let end = if at_one { hi } else { lo }; + let (mut endpoint, mut slope) = (Prod8::ZERO, Prod8::ZERO); + for g in (1..width).step_by(2) { + endpoint ^= end[g].mul_unreduced(end[g + 1]); + slope ^= (lo[g] + hi[g]).mul_unreduced(lo[g + 1] + hi[g + 1]); + } + let eq = load_e(eqr, 8 * b); + acc[0] ^= (endpoint.reduce_lanes() + end[0]).mul_unreduced(eq); + acc[1] ^= slope.reduce_lanes().mul_unreduced(eq); + } + }, + xor2, + ); + finish(acc, form.constant) +} + +/// The round a table joins, after its challenge `r`: the state `lo + r·(lo + hi)` of every +/// entry, built straight from the `K` columns. By linearity `lin` and `right_g` take the +/// terms `w·lo + (w·r)·(lo + hi)`, so no row pays an `E x E` product. +fn join(form: &Form, cols: &[&[F64]], r: F192, half: usize) -> State { + if half < 8 { + let mut state = form.small_state(cols, 2 * half); + fold_state(&mut state, r, half); + return state; + } + let (blocks, width) = (half / 8, form.width()); + let pair = |terms: &Lin| -> Vec<_> { + terms + .iter() + .map(|&(c, w)| (c, F192x8::splat(w), F192x8::splat(w * r))) + .collect() }; - acc.map(F192Unreduced::reduce) + let lin = pair(&form.lin); + let rights: Vec<_> = form.groups.iter().map(|(_, right)| pair(right)).collect(); + let r8 = F192x8::splat(r); + // SAFETY: every task writes all `width` entries of each of its blocks. + unsafe { + primitives::filled_vec(blocks * width, |state| { + parallel::chunks_mut(state, TASK_BLOCKS * width, |t, out| { + let mut rows = rows_scratch(cols.len()); + for (b, out) in task_blocks(t, blocks).zip(out.chunks_exact_mut(width)) { + load_rows(cols, &form.used, 8 * b, half, false, &mut rows); + let (lo, sum) = &rows; + let dot = |terms: &[(usize, F192x8, F192x8)]| { + terms + .iter() + .fold(Prod8::ZERO, |acc, &(c, w, wr)| { + acc ^ w.mul_base_unreduced(lo[c]) ^ wr.mul_base_unreduced(sum[c]) + }) + .reduce_lanes() + }; + out[0] = dot(&lin); + for (g, ((left, _), right)) in form.groups.iter().zip(&rights).enumerate() { + out[1 + 2 * g] = F192x8::from(lo[*left]) + r8.mul_base_unreduced(sum[*left]).reduce_lanes(); + out[2 + 2 * g] = dot(right); + } + } + }); + }) + } +} + +/// The state folded by `r`, `v[i] = v[i] + r·(v[i] + v[i + half])` in every column, the upper half dropped. +/// From eight rows up the halves are whole vectors, so the fold is elementwise. +fn fold_state(state: &mut State, r: F192, half: usize) { + if half < 8 { + for v in state.iter_mut() { + let mut rows = v.to_array(); + (0..half).for_each(|i| rows[i] = interp(rows[i], rows[i + half], r)); + *v = F192x8::from_array(&rows); + } + return; + } + let r8 = F192x8::splat(r); + let n = state.len() / 2; + let (lows, highs) = state.split_at_mut(n); + parallel::chunks_mut(lows, FOLD_CHUNK, |t, lows| { + for (lo, &hi) in lows.iter_mut().zip(&highs[t * FOLD_CHUNK..]) { + *lo = *lo + (*lo + hi) * r8; + } + }); + state.truncate(n); +} + +/// One table of the batch, and what the rounds have made of it. +struct Table<'a> { + tau: usize, + sigma: F192, + /// Its summand's factor: the challenges it sat out, then the eq factors of the rounds it was active. + weight: F192, + form: Form, + cols: &'a [&'a [F64]], + /// The folded state, once the table has joined: a few `E` columns, the batch's largest transients. + state: Option, +} + +impl Table<'_> { + /// The message of the round that binds variable `m`, while the table is active. + fn message(&mut self, eqr: &[F192], m: usize, at_one: bool) -> [F192; 2] { + let (form, half) = (&self.form, 1 << m); + match &self.state { + Some(state) => message_e(form, state, eqr, half, at_one), + None => message_k(form, self.cols, eqr, half, at_one), + } + } + + /// Binds variable `m` to `r`: the join of a table that has not joined, the fold of one that has. + fn bind(&mut self, r: F192, m: usize) { + match &mut self.state { + Some(state) => fold_state(state, r, 1 << m), + None => self.state = Some(join(&self.form, self.cols, r, 1 << m)), + } + } } fn round_polynomial([endpoint, quadratic]: [F192; 2], zeta: F192, claim: F192, waiting: F192) -> [F192; 4] { @@ -133,57 +412,44 @@ fn round_polynomial([endpoint, quadratic]: [F192; 2], zeta: F192, claim: F192, w } /// Prove that every table's batched constraint vanishes on all of its rows, as ONE -/// sumcheck over `max τ_t` variables. `cols[t]` holds table `t`'s involved columns -/// (`2^{τ_t}` values each, folded in place). Returns the per-table claims, in input -/// order, on the nested points `ρ[..τ_t]`. -pub fn prove( - airs: &[Air<'_>], - cols: &[Vec<&[F64]>], - xi: F192, - zeta: &[F192], - sigma: &[F192], - ps: &mut ProverState, -) -> Vec { +/// sumcheck over `max τ_t` variables. Returns the per-table claims, in input order, on the +/// nested points `ρ[..τ_t]`. +pub fn prove(airs: &[Air<'_>], zeta: &[F192], ps: &mut ProverState) -> Vec { let n = airs.iter().map(|a| a.tau).max().unwrap_or(0); debug_assert!(zeta.len() >= n, "the eq point must cover the tallest table"); - let offsets = xi_offsets(airs.iter().map(|a| a.n_constraints)); - let pows = powers(xi, airs.iter().map(|a| a.n_constraints).sum()); - // η^{offset_t}, already inside `pows`; the rounds then fold in the pre-join - // challenges and the eq factor, so `weights` is the whole per-table state. - let mut weights = vec![F192::ONE; airs.len()]; + let mut tables: Vec = airs + .iter() + .map(|air| Table { + tau: air.tau, + sigma: air.sigma, + weight: F192::ONE, + form: Form::new(&air.form), + cols: air.cols, + state: None, + }) + .collect(); // ONE eq table over the low (still free) variables serves every active table. let mut eqr = eq_table(&zeta[..n.saturating_sub(1)]); let mut chi = vec![F192::ZERO; n]; - // The folded tables are the batch's largest transients: one E-lifted copy of - // every column of every still-active table. - let mut folded: Vec>>> = (0..airs.len()).map(|_| None).collect(); - // `k`, the challenges drawn so far, common to every air that is still waiting. + // `k`, the challenges drawn so far, common to every table that is still waiting. let mut k = F192::ONE; - let mut claim = sigma.iter().copied().fold(F192::ZERO, |a, b| a + b); + let mut claim = airs.iter().fold(F192::ZERO, |acc, a| acc + a.sigma); for j in 0..n { let m = n - 1 - j; // the variable this round binds - // The waiting airs contribute the line `Y·k·Σσ`, whose slope `u` is all there + let at_one = zeta[m].is_zero(); + // The waiting tables contribute the line `Y·k·Σσ`, whose slope `u` is all there // is to it. It is NOT sent on its own: it folds into `h` below, and only `h` - // travels. `msg` is the joined airs' degree-2 cofactor, `h`'s multiplicand. - let waiting = airs + // travels. `msg` is the active tables' degree-2 cofactor, `h`'s multiplicand. + let waiting = tables .iter() - .zip(sigma) - .filter(|(a, _)| a.tau <= m) - .fold(F192::ZERO, |acc, (_, &s)| acc + s); + .filter(|t| t.tau <= m) + .fold(F192::ZERO, |acc, t| acc + t.sigma); let u = k * waiting; let mut msg = [F192::ZERO; 2]; - for (t, air) in airs.iter().enumerate() { - if air.tau > m { - let w = &pows[offsets[t]..offsets[t] + air.n_constraints]; - let p = if let Some(table) = &folded[t] { - table_message(table, &*air.eval, w, 1 << m, &eqr, zeta[m].is_zero()) - } else { - table_message(&cols[t], &*air.eval_k, w, 1 << m, &eqr, zeta[m].is_zero()) - }; - for i in 0..2 { - msg[i] += weights[t] * p[i]; - } - } + for t in tables.iter_mut().filter(|t| t.tau > m) { + let p = t.message(&eqr, m, at_one); + msg[0] += t.weight * p[0]; + msg[1] += t.weight * p[1]; } shrink_eq_high(&mut eqr); // The running claim recovers the missing endpoint of the quadratic cofactor. @@ -195,138 +461,148 @@ pub fn prove( chi[m] = rk; k *= rk; let eq_k = F192::ONE + zeta[m] + rk; - for (t, air) in airs.iter().enumerate() { - weights[t] *= if air.tau > m { eq_k } else { rk }; - if air.tau <= m { - continue; - } - if let Some(table) = &mut folded[t] { - if m >= PAR_THRESHOLD.trailing_zeros() as usize { - let cols = parallel::Chunks::new(table, 1); - parallel::for_each(cols.count(), |ci| { - // SAFETY: column `ci` is folded by exactly one task. - let col = unsafe { &mut cols.get(ci)[0] }; - fold_high_inplace(col, rk); - }); - } else { - table.iter_mut().for_each(|c| fold_high_inplace(c, rk)); - } + for t in &mut tables { + if t.tau > m { + t.weight *= eq_k; + t.bind(rk, m); } else { - // The round a table joins is the largest fold it ever does, so it - // fans out like the in-place ones below rather than running on the - // dispatcher alone. - folded[t] = Some(parallel::map_collect(cols[t].len(), |ci| fold_high_k(cols[t][ci], rk))); + t.weight *= rk; } } } - airs.iter() - .enumerate() - .map(|(t, air)| { - let evals: Vec = if let Some(table) = &folded[t] { - table.iter().map(|c| c[0]).collect() - } else { - cols[t].iter().map(|c| F192::from(c[0])).collect() - }; - ps.add_scalars(&evals); - Claims { - chi: chi[..air.tau].to_vec(), - evals, - } - }) - .collect() + // Each column's evaluation at its table's point, one task per column: the columns were never folded. + crate::stage!("Constraint evals", || { + tables + .iter() + .map(|t| { + let chi = &chi[..t.tau]; + let evals = parallel::map_collect(t.cols.len(), |c| mle_eval(t.cols[c], chi)); + ps.add_scalars(&evals); + Claims { + chi: chi.to_vec(), + evals, + } + }) + .collect() + }) } #[cfg(test)] mod tests { use super::*; - use primitives::field::F64; - fn synth_eval(pows: &[F192], v: &[T], quadratic: bool) -> F192 { - if quadratic { - return (v[0] * v[1]).mul_e(pows[0]); - } - (v[0] * v[1] + v[2]).mul_e(pows[0]) + (v[0] + v[3]).mul_e(pows[1]) + const COLS: usize = 6; + + /// Products sharing column 0, one of them named twice and either way round, a square, linear + /// terms and a constant. + fn synth(pows: &[F192]) -> BusForm { + let mut form = BusForm::new(COLS); + form.constant = F192::ONE; + (form.coeffs[2], form.coeffs[1], form.coeffs[4]) = (pows[0], pows[1], pows[2]); + form.prods = vec![ + (0, 1, pows[0]), + (3, 0, pows[1]), + (0, 4, pows[1]), + (5, 5, pows[2]), + (1, 0, pows[2]), + ]; + form + } + + fn splitmix(state: &mut u64) -> u64 { + *state = state.wrapping_add(0x9e37_79b9_7f4a_7c15); + let mut z = *state; + z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); + z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); + z ^ (z >> 31) + } + + fn random_e(state: &mut u64) -> F192 { + F192::new(splitmix(state), splitmix(state), splitmix(state)) + } + + fn random_columns(state: &mut u64, rows: usize) -> Vec> { + (0..COLS) + .map(|_| (0..rows).map(|_| F64(splitmix(state))).collect()) + .collect() } #[test] - fn round_coefficients_match_full_evaluations() { - fn check>(cols: &[Vec]) { - let weights = powers(F192::new(3, 5, 7), 3); - let eq = eq_table(&[F192::new(11, 13, 17), F192::new(19, 23, 29)]); - let eval = |p: &[F192], v: &[T], quadratic| { - synth_eval_attached(p, v, quadratic) + if quadratic { F192::ZERO } else { F192::ONE } - }; - let full_eval = |r| { - (0..4).fold(F192::ZERO, |sum, i| { - let v: Vec<_> = cols - .iter() - .map(|c| primitives::multilinear::interp(c[i].into(), c[i + 4].into(), r)) - .collect(); - sum + eq[i] * (synth_eval_attached(&weights, &v, false) + F192::ONE) - }) - }; - let waiting = F192::new(43, 47, 53); - for zeta in [F192::ZERO, F192::ONE, F192::new(59, 61, 67)] { - let message = table_message(cols, &eval, &weights, 4, &eq, zeta.is_zero()); - let claim = (F192::ONE + zeta) * full_eval(F192::ZERO) + zeta * full_eval(F192::ONE) + waiting; - let h = round_polynomial(message, zeta, claim, waiting); - for r in [F192::ZERO, F192::ONE, F192::new(31, 37, 41)] { - assert_eq!(poly_eval(&h, r), (F192::ONE + zeta + r) * full_eval(r) + r * waiting); - } - } + fn grouped_form_reproduces_the_summand() { + let mut s = 1; + let pows: Vec = (0..3).map(|_| random_e(&mut s)).collect(); + let form = Form::new(&synth(&pows)); + assert_eq!( + form.groups.len(), + 2, + "column 0 covers three products, the square is the other group" + ); + for _ in 0..4 { + let v: Vec = (0..COLS).map(|_| random_e(&mut s)).collect(); + let dot = |terms: &Lin| terms.iter().fold(F192::ZERO, |acc, &(c, w)| acc + w * v[c]); + let quad = form + .groups + .iter() + .fold(F192::ZERO, |acc, (l, right)| acc + v[*l] * dot(right)); + assert_eq!(quad, synth(&pows).eval(&v, true)); + assert_eq!(form.constant + dot(&form.lin) + quad, synth(&pows).eval(&v, false)); } - let base: Vec> = (0..4) - .map(|j| (0..8).map(|i| F64(13 * i + 17 * j + 1)).collect()) - .collect(); - check(&base); - let ext: Vec> = base - .iter() - .map(|c| c.iter().map(|v| F192::new(v.0, 3 * v.0, 7 * v.0)).collect()) - .collect(); - check(&ext); } - /// A third, attached "identity": the linear form `vals[1]`, whose claimed sum - /// is an evaluation of column 1 rather than zero. - fn synth_eval_attached(pows: &[F192], v: &[T], quadratic: bool) -> F192 { - synth_eval(pows, v, quadratic) + if quadratic { F192::ZERO } else { v[1].mul_e(pows[2]) } + /// `[Σ eqr·P(end), Σ eqr·Q(lo + hi)]` by evaluating the summand row by row. + fn naive_message(pows: &[F192], cols: &[Vec], eqr: &[F192], half: usize, at_one: bool) -> [F192; 2] { + (0..half).fold([F192::ZERO; 2], |acc, i| { + let row = |off: usize| -> Vec { cols.iter().map(|c| c[i + off]).collect() }; + let (lo, hi) = (row(0), row(half)); + let sum: Vec = lo.iter().zip(&hi).map(|(&a, &b)| a + b).collect(); + let end = if at_one { &hi } else { &lo }; + [ + acc[0] + eqr[i] * synth(pows).eval(end, false), + acc[1] + eqr[i] * synth(pows).eval(&sum, true), + ] + }) } - /// Tables of different heights in one batch: each table's claims are its - /// columns' true evaluations, at a prefix of the tallest table's point. + /// Every round of a table's life, at row counts on both sides of one vector: the join's + /// message and state, then each fold and joined message, against the naive formulas. #[test] - fn ragged_batch_claims_are_true_evaluations() { - let taus = [5usize, 3, 5, 0, 1]; - let cols: Vec>> = taus - .iter() - .enumerate() - .map(|(salt, &tau)| { - let a: Vec = (0..1u64 << tau).map(|i| F64(i + salt as u64)).collect(); - let b: Vec = (0..1u64 << tau).map(|i| F64(3 * i + 1 + salt as u64)).collect(); - let ab: Vec = a.iter().zip(&b).map(|(&x, &y)| x * y).collect(); - vec![a.clone(), b, ab, a] - }) - .collect(); - let airs: Vec> = taus - .iter() - .map(|&tau| Air { - tau, - n_constraints: 2, - eval: Box::new(synth_eval::), - eval_k: Box::new(synth_eval::), - }) - .collect(); - let xi = F192::new(0x9e37_79b9_7f4a_7c15, 0x1234_5678_9abc_def0, 7); - let zeta: Vec = (0..5).map(|i| F192::new(i + 3, 0x5555 * (i + 1), i + 11)).collect(); - let views: Vec> = cols.iter().map(|t| t.iter().map(|c| &c[..]).collect()).collect(); - let mut ps = ProverState::from_label(b"zc-test"); - let claims = prove(&airs, &views, xi, &zeta, &[F192::ZERO; 5], &mut ps); - let tallest = &claims[0].chi; - for ((claim, &tau), table) in claims.iter().zip(&taus).zip(&cols) { - assert_eq!(claim.chi, tallest[..tau]); - for (eval, col) in claim.evals.iter().zip(table) { - assert_eq!(*eval, primitives::multilinear::mle_eval(col, &claim.chi)); + fn round_messages_match_naive_evaluation() { + let mut s = 7; + let pows: Vec = (0..3).map(|_| random_e(&mut s)).collect(); + let form = Form::new(&synth(&pows)); + for m in [1usize, 3, 4, 5, 11] { + let kcols = random_columns(&mut s, 2 << m); + let views: Vec<&[F64]> = kcols.iter().map(|c| &c[..]).collect(); + let mut naive_cols: Vec> = kcols + .iter() + .map(|c| c.iter().map(|&x| F192::from(x)).collect()) + .collect(); + let mut eqr = eq_table(&(0..m).map(|_| random_e(&mut s)).collect::>()); + let mut state: Option = None; + let mut half = 1usize << m; + while half >= 1 { + for at_one in [false, true] { + let got = match &state { + None => message_k(&form, &views, &eqr, half, at_one), + Some(st) => message_e(&form, st, &eqr, half, at_one), + }; + assert_eq!( + got, + naive_message(&pows, &naive_cols, &eqr, half, at_one), + "m {m} half {half}" + ); + } + let r = random_e(&mut s); + shrink_eq_high(&mut eqr); + match &mut state { + None => state = Some(join(&form, &views, r, half)), + Some(st) => fold_state(st, r, half), + } + for c in &mut naive_cols { + *c = (0..half).map(|i| interp(c[i], c[i + half], r)).collect(); + } + half /= 2; } } } diff --git a/crates/leanvm_core/src/cpu/mod.rs b/crates/leanvm_core/src/cpu/mod.rs index 39ebbec1b..a60e18908 100644 --- a/crates/leanvm_core/src/cpu/mod.rs +++ b/crates/leanvm_core/src/cpu/mod.rs @@ -10,7 +10,6 @@ use std::collections::HashMap; -use crate::colval::ColVal; use crate::constraints; use crate::leaf::{self, Block, ColumnClaim, Coord}; use crate::pcs; @@ -306,126 +305,38 @@ fn table_spans() -> TableSpans { .collect() } -/// The per-table inputs to the table sumcheck (§constraints), in schema order. -/// The airs carry every committed column of their table, so a constraint indexes the -/// value array directly and each table's three bus forms can be -/// evaluated on the same values. The identities take the air's own `η`-range; the -/// three forms take the shared powers at [`xi_form_base`], folded into the forms' -/// coefficients once rather than multiplied onto every row's form value. -fn airs( +/// The per-table inputs to the table sumcheck (§constraints), in schema order: each +/// table's committed columns, its summand as one form over them (its identities +/// weighted by its own `η`-range, its three bus forms by the shared powers at +/// [`xi_form_base`]), and what that summand sums to, its bus forms' `η`-weighted +/// shares, since its identities vanish. +fn airs<'a>( taus: &[usize; tables::N_TABLES], - forms: &[Vec; 3], - form_pows: [F192; 3], -) -> Vec> { + bus: &leaf::BusProof, + cols: &'a [&'a [F64]], + spans: &[(usize, usize)], + xi: F192, +) -> Vec> { + let base = xi_form_base(); + let pows = primitives::field::powers(xi, base + 3); + let offsets = constraints::xi_offsets(tables::tables().iter().map(|t| t.n_constraints())); tables::tables() .iter() - .zip(taus) + .zip(taus.iter().zip(spans)) .enumerate() - .map(|(t, (&table, &tau))| { - // One form, not three: the batch adds the three sides' evaluations - // anyway, and summing them here is a setup cost against a dot product - // and a product list per row per node. - let bus = RowForm::new(&leaf::BusForm::sum((0..3).map(|s| forms[s][t].scaled(form_pows[s])))); - let bus_k = bus.clone(); + .map(|(t, (&table, (&tau, &(start, n))))| { + let identities = table.constraint(&pows[offsets[t]..][..table.n_constraints()]); + let bus_forms = (0..3).map(|s| bus.forms[s][t].scaled(pows[base + s])); constraints::Air { tau, - n_constraints: table.n_constraints(), - eval: Box::new(move |p, vals, quadratic| { - let air = ::lift(table.eval_constraint(p, vals, quadratic)); - ::reduce(air ^ bus.eval_unreduced(vals, quadratic)) - }), - // The same expression over K columns: the identity's K-only products - // stay 64-bit and the bus form becomes a mixed dot product. - eval_k: Box::new(move |p, vals, quadratic| { - let air = ::lift(table.eval_constraint_k(p, vals, quadratic)); - ::reduce(air ^ bus_k.eval_unreduced(vals, quadratic)) - }), + cols: &cols[start..start + n], + form: leaf::BusForm::sum(bus_forms.chain([identities])), + sigma: (0..3).fold(F192::ZERO, |acc, s| acc + pows[base + s] * bus.sigmas[s][t]), } }) .collect() } -/// A table's bus form as the batch evaluates it on every row of every round: its -/// nonzero coefficients, and its products grouped by a column they share, -/// `Σ_a v_a·(Σ_b c_ab·v_b)`, so a column that several products share multiplies -/// once rather than once per product. The arithmetic is exact, so the value is the form's. -#[derive(Clone)] -struct RowForm { - linear: Vec<(usize, F192)>, - constant: F192, - groups: Vec<(usize, Vec<(usize, F192)>)>, -} - -impl RowForm { - fn new(form: &leaf::BusForm) -> Self { - let n = form.coeffs.len(); - let linear = (0..n) - .filter(|&i| !form.coeffs[i].is_zero()) - .map(|i| (i, form.coeffs[i])) - .collect(); - let mut prods = form.prods.clone(); - let mut groups: Vec<(usize, Vec<(usize, F192)>)> = Vec::new(); - // Greedily factor out the column the most remaining products share. - while !prods.is_empty() { - let mut shared = vec![0usize; n]; - for &(a, b, _) in &prods { - shared[a] += 1; - if b != a { - shared[b] += 1; - } - } - let col = (0..n) - .max_by_key(|&c| (shared[c], std::cmp::Reverse(c))) - .expect("a column"); - let (with, rest): (Vec<_>, Vec<_>) = prods.into_iter().partition(|&(a, b, _)| a == col || b == col); - groups.push(( - col, - with.into_iter() - .map(|(a, b, c)| (if a == col { b } else { a }, c)) - .collect(), - )); - prods = rest; - } - Self { - linear, - constant: form.constant, - groups, - } - } - - /// The form at one row, unreduced, or only its degree-two terms when `quadratic`. - #[inline(always)] - fn eval_unreduced(&self, vals: &[T], quadratic: bool) -> T::Unreduced { - let mut acc = if quadratic { - T::lift(F192::ZERO) - } else { - self.linear - .iter() - .fold(T::lift(self.constant), |acc, &(i, c)| acc ^ vals[i].mul_e_unreduced(c)) - }; - for (a, terms) in &self.groups { - if let [(b, c)] = terms[..] { - acc ^= (vals[*a] * vals[b]).mul_e_unreduced(c); - continue; - } - let inner = terms - .iter() - .fold(T::lift(F192::ZERO), |acc, &(b, c)| acc ^ vals[b].mul_e_unreduced(c)); - acc ^= vals[*a].mul_e_unreduced(T::reduce(inner)); - } - acc - } -} - -/// Each table's claimed sum: its identities vanish, so what its summand comes to -/// is its three bus forms, `η`-weighted. Prover-side only, to build the waiting -/// line each round; the verifier needs just their total, which it derives. -fn sigmas(bus: &[Vec; 3], form_pows: [F192; 3]) -> Vec { - (0..tables::tables().len()) - .map(|t| (0..3).fold(F192::ZERO, |acc, s| acc + form_pows[s] * bus[s][t])) - .collect() -} - /// Where the three bus forms sit in the batch's `η`-powers: the last three, AFTER /// every table's identity range, and shared by all tables rather than one triple /// per table. That sharing is what keeps the batch tied to the bus: with a common @@ -439,13 +350,6 @@ pub fn xi_form_base() -> usize { tables::tables().iter().map(|t| t.n_constraints()).sum() } -/// The three shared form powers `η^{base}, η^{base+1}, η^{base+2}`. -fn xi_form_pows(xi: F192) -> [F192; 3] { - let base = xi_form_base(); - let pows = primitives::field::powers(xi, base + 3); - [pows[base], pows[base + 1], pows[base + 2]] -} - /// The KECCAK table's claims on its output lanes, which are the Keccak reduction's claim on /// the permutation's output: the batch claims every column of the table at one point. fn keccak_output_claim(table_claims: &[constraints::Claims]) -> keccak_gkr::OutputClaim { @@ -564,8 +468,8 @@ pub fn prove(program: &Program, public_input: [F192; 2], log_inv_rate: usize) -> // reduction's output claim (outputs), so no separate pin claims are needed. let (owners, spans) = bus_wiring(program, &w.layout); // The columns are windows into `w.q`, so both stages read them in place: the - // table sumcheck lifts each K-column into a fresh `E` copy on the round it - // joins and never writes the K-columns back. + // table sumcheck builds a table's few `E` state columns from them on the round + // it joins and never writes the K-columns back. let (bus, table_claims) = { let l = &w.layout; let cols = w.columns(); @@ -574,23 +478,10 @@ pub fn prove(program: &Program, public_input: [F192; 2], log_inv_rate: usize) -> }); let table_claims = crate::stage!("Prove constraints", || { // One sumcheck for all six tables (§constraints). - let table_cols: Vec> = spans - .iter() - .map(|&(base, n)| (0..n).map(|c| cols[base + c]).collect()) - .collect(); // The eq point is the bus GKR's ζ, not a fresh one: that is what lets the // batch settle the bus forms alongside the constraints. let xi = ps.sample(); - let form_pows = xi_form_pows(xi); - let sigma = sigmas(&bus.sigmas, form_pows); - constraints::prove( - &airs(&l.taus, &bus.forms, form_pows), - &table_cols, - xi, - &bus.point, - &sigma, - &mut ps, - ) + constraints::prove(&airs(&l.taus, &bus, &cols, &spans, xi), &bus.point, &mut ps) }); (bus, table_claims) }; diff --git a/crates/leanvm_core/src/leaf.rs b/crates/leanvm_core/src/leaf.rs index d35ff2bda..6f7efb531 100644 --- a/crates/leanvm_core/src/leaf.rs +++ b/crates/leanvm_core/src/leaf.rs @@ -7,7 +7,6 @@ //! g-powers, separators); the fingerprint challenges `α, β` are `E`-valued, so a //! leaf accumulates via the mixed `mul_base` product (2 PMULL per coordinate). -use crate::colval::ColVal; use crate::gkr; use crate::transcript::ProverState; use primitives::field::{F64, F192, F192Unreduced, Weights8, dot_base, g_pow, index_mle, mul_by_g}; @@ -497,7 +496,8 @@ pub struct BusForm { } impl BusForm { - fn new(n_cols: usize) -> Self { + /// The zero form over `n_cols` columns. + pub(crate) fn new(n_cols: usize) -> Self { Self { coeffs: vec![F192::ZERO; n_cols], prods: Vec::new(), @@ -516,12 +516,8 @@ impl BusForm { } } - /// The pointwise sum of several forms over the same columns. - /// - /// Evaluating the sum is evaluating each and adding, and the constraint batch - /// only ever wants a table's total, so its three bus sides collapse to one - /// dot product and one product list: the row loop then reads the column - /// values once for all three rather than once each. + /// The pointwise sum of several forms over the same columns. Their products are + /// listed together, a pair named in two forms twice. pub fn sum(forms: impl IntoIterator) -> Self { let mut forms = forms.into_iter(); let mut out = forms.next().expect("a table has at least one bus form"); @@ -531,31 +527,40 @@ impl BusForm { *slot += c; } out.constant += form.constant; - for (a, b, c) in form.prods { - match out.prods.iter_mut().find(|p| (p.0, p.1) == (a, b)) { - Some(p) => p.2 += c, - None => out.prods.push((a, b, c)), - } - } + out.prods.extend(form.prods); } - out.prods.retain(|p| p.2 != F192::ZERO); out } + /// The form at one point, `evals` the columns' values there; only its degree-two terms when `quadratic`. + #[cfg(test)] + pub(crate) fn eval(&self, evals: &[F192], quadratic: bool) -> F192 { + let linear = if quadratic { F192::ZERO } else { self.linear(evals) }; + self.prods + .iter() + .fold(linear, |acc, &(a, b, c)| acc + evals[a] * evals[b] * c) + } + + /// `constant + Σ coeffs[i]·evals[i]`. + fn linear(&self, evals: &[F192]) -> F192 { + self.coeffs + .iter() + .zip(evals) + .fold(self.constant, |acc, (&w, &v)| acc + w * v) + } + /// What the form sums to over the table's rows against `eq(ζ, ·)`, the target the /// zerocheck settles. The linear part factors through the columns' evaluations at /// `ζ`, which is the whole point of a form; a product coordinate does NOT, so /// `prod_sums` supplies `Σ_z eq(ζ,z)·col_a(z)·col_b(z)` for each pair it uses. fn sum_at(&self, evals: &[F192], prod_sums: &[(usize, usize, F192)]) -> F192 { - self.prods - .iter() - .fold(F192::dot(&self.coeffs, evals, self.constant), |acc, &(a, b, c)| { - let s = prod_sums - .iter() - .find(|p| (p.0, p.1) == (a, b)) - .expect("every pair was summed"); - acc + c * s.2 - }) + self.prods.iter().fold(self.linear(evals), |acc, &(a, b, c)| { + let s = prod_sums + .iter() + .find(|p| (p.0, p.1) == (a, b)) + .expect("every pair was summed"); + acc + c * s.2 + }) } } diff --git a/crates/leanvm_core/src/lib.rs b/crates/leanvm_core/src/lib.rs index ee7ddf0f7..a7e8a513c 100644 --- a/crates/leanvm_core/src/lib.rs +++ b/crates/leanvm_core/src/lib.rs @@ -20,7 +20,6 @@ //! - [`keccak`]: the `KECCAK` glue: the Keccak GKR reduction over the same commitment. //! - [`vmhash`]: VM-native hashing (the one-block hash every Merkle node uses). -pub mod colval; pub mod constraints; pub mod cpu; pub mod gkr; diff --git a/crates/leanvm_core/src/tables.rs b/crates/leanvm_core/src/tables.rs index f6012b2a0..e41c0f452 100644 --- a/crates/leanvm_core/src/tables.rs +++ b/crates/leanvm_core/src/tables.rs @@ -11,31 +11,20 @@ //! the `DEREF` store, the `JUMP` successors are each written out as the degree-2 //! bus coordinate that carries them (§sec:m3), which leaves `JUMP`'s is-nonzero //! indicator as the one identity any table still has. Every identity is `K`-valued, -//! so a relation on machine words is written out lane by lane; after the round a -//! table joins the batch its columns are `E`-valued, which is what -//! `eval_constraint` takes. +//! so a relation on machine words is written out lane by lane. A table hands its +//! identities to the table sumcheck as a degree-2 form ([`Table::constraint`]). -use crate::colval::ColVal; use crate::cpu::{Drow, Jrow, Krow, Memory, Op, Rows, Srow, Trace}; +use crate::leaf::BusForm; use crate::leaf::Coord::{self, Col, Const, GCol, Prod}; use primitives::field::{F64, F192, mul_by_g}; // ---- the identities ---------------------------------------------------------- // -// Each is written ONCE, generic over the column type: `F64` in the round a table -// joins the batch, `F192` afterwards (see [`ColVal`]). Products of two `K` -// columns stay 64-bit and an `η`-power multiplies through `mul_e`. -// // EVERY identity is `K`-valued, like the columns it reads and like the bus // coordinates (§sec:air). A relation between machine WORDS is therefore written out // lane by lane, with the tower multiplication unrolled by hand ([`TOWER_LANES`]), -// rather than assembled into one `E` equation. A table's identity slice is then a -// mixed dot product against its `η`-range: in the round a table joins the batch (the -// batch's largest) that is one 64-bit product per identity plus three PMULL for its -// `η`-power, with ONE reduction for the whole slice ([`ColVal::dot`]). One `E`-valued -// identity would instead cost a full `E×E` product against its `η`-power, plus a -// reduction of its own, and the lanes it bundles are the same lane polynomials -// either way. +// rather than assembled into one `E` equation. /// The tower product `x·y` in `E = K[y]/(y³+y+1)` as three lane sums: lane `i` is /// `Σ x_j·y_k` over `TOWER_LANES[i]`. The five partial sums of §sec:tab-mul fold @@ -53,28 +42,10 @@ const TOWER_LANES: [&[(usize, usize)]; 3] = [ /// as a coordinate ([`arith_result`]), and `JUMP`'s condition is `K`-valued, so no /// identity assembles a tower product any more. #[cfg(test)] -fn tower_lane(lane: usize, x: [T; 3], y: [T; 3]) -> T { - TOWER_LANES[lane].iter().fold(T::ZERO, |acc, &(j, k)| acc + x[j] * y[k]) -} - -/// `JUMP`'s two identities: `b = cond·w` and `cond·(b+1) = 0` (§sec:tab-jump). -/// -/// The two relations together force `b = [cond ≠ 0]`: when `cond ≠ 0` the second -/// gives `b = 1` (and the first `w = cond⁻¹`); when `cond = 0` the first gives -/// `b = 0`. The two selections need no identity: the state push carries each as its -/// own degree-2 coordinate (§sec:m3). -/// -/// The condition is `K`-valued, so both identities are single-lane. Its memory -/// flush carries literal zeros above the low limb (`memory_k`), so a word outside -/// `K` cannot balance the bus; the interpreter rejects one outright. -fn jump_identity(pows: &[F192], cols: &[T], quadratic: bool) -> F192 { - use jump::*; - let (b, b1) = if quadratic { - (T::ZERO, cols[B]) - } else { - (cols[B], cols[B] + T::ONE) - }; - T::dot(pows, &[b + cols[V_COND] * cols[W], cols[V_COND] * b1], F192::ZERO) +fn tower_lane(lane: usize, x: [F64; 3], y: [F64; 3]) -> F64 { + TOWER_LANES[lane] + .iter() + .fold(F64::ZERO, |acc, &(j, k)| acc + x[j] * y[k]) } // ---- shared bus vocabulary --------------------------------------------------- @@ -346,33 +317,23 @@ pub trait Table: Sync { fn padding(&self) -> Vec { vec![F64::ZERO; self.n_committed_columns()] } - /// How many identities [`eval_constraint`](Table::eval_constraint) folds. + /// How many identities [`constraint`](Table::constraint) folds. /// Sizes this table's slice of the batch's disjoint `xi`-range (§constraints). /// Defaults to none, which is every table but `JUMP`: a relation whose value /// rides the bus as a coordinate needs no identity to tie it (§sec:m3). fn n_constraints(&self) -> usize { 0 } - /// Evaluate the table's degree-2 constraint at one row, reading column values - /// by local index from `cols` (e.g. `cols[jump::V_COND]`) and weighting identity - /// `i` by `pows[i]`, this table's slice of the batch's `xi`-powers. The slice is - /// is exactly [`n_constraints`](Table::n_constraints) long: an identity indexed - /// past its end panics rather than silently reaching into the next table's - /// range. The table sumcheck carries every committed column of a table, in - /// local order, so `cols` is indexed directly. With `quadratic=false` it returns - /// `0` on every valid row (§sec:air); `true` selects only the degree-two terms. - /// The default is the constraint-free case; a table that declares constraints - /// and forgets to evaluate them trips the assert instead of dropping them. - fn eval_constraint(&self, pows: &[F192], _cols: &[F192], _quadratic: bool) -> F192 { - assert!(pows.is_empty(), "a table with constraints must evaluate them"); - F192::ZERO - } - /// The same identity over `K`-valued columns, for the round a table joins the - /// batch, before its columns have been folded into `E` (§sec:air). Both entry - /// points delegate to one generic definition per table, so they cannot drift. - fn eval_constraint_k(&self, pows: &[F192], _cols: &[F64], _quadratic: bool) -> F192 { - assert!(pows.is_empty(), "a table with constraints must evaluate them"); - F192::ZERO + /// The table's identities as one degree-2 form over its columns, in local order + /// (e.g. `jump::V_COND`), identity `i` weighted by `pows[i]`, this table's slice of + /// the batch's `xi`-powers. The form vanishes on every valid row (§sec:air). The + /// slice is exactly [`n_constraints`](Table::n_constraints) long: an identity + /// indexed past its end panics rather than reaching into the next table's range. + /// The default is the constraint-free case; a table that declares constraints and + /// forgets to state them trips the assert instead of dropping them. + fn constraint(&self, pows: &[F192]) -> BusForm { + assert!(pows.is_empty(), "a table with constraints must state them"); + BusForm::new(self.n_committed_columns()) } /// Declare the table's bus interactions. fn flushes(&self, f: &mut FlushBuilder); @@ -776,11 +737,23 @@ impl Table for JumpTable { fn n_constraints(&self) -> usize { 2 // the two indicator identities; the selections ride the state push } - fn eval_constraint(&self, pows: &[F192], cols: &[F192], quadratic: bool) -> F192 { - jump_identity(pows, cols, quadratic) - } - fn eval_constraint_k(&self, pows: &[F192], cols: &[F64], quadratic: bool) -> F192 { - jump_identity(pows, cols, quadratic) + /// `b = cond·w` and `cond·(b+1) = 0` (§sec:tab-jump), weighted `pows[0]` and `pows[1]`. + /// + /// The two relations together force `b = [cond ≠ 0]`: when `cond ≠ 0` the second + /// gives `b = 1` (and the first `w = cond⁻¹`); when `cond = 0` the first gives + /// `b = 0`. The two selections need no identity: the state push carries each as its + /// own degree-2 coordinate (§sec:m3). + /// + /// The condition is `K`-valued, so both identities are single-lane. Its memory + /// flush carries literal zeros above the low limb (`memory_k`), so a word outside + /// `K` cannot balance the bus; the interpreter rejects one outright. + fn constraint(&self, pows: &[F192]) -> BusForm { + use jump::*; + let mut form = BusForm::new(N); + form.coeffs[B] = pows[0]; + form.coeffs[V_COND] = pows[1]; + form.prods = vec![(V_COND, W, pows[0]), (V_COND, B, pows[1])]; + form } fn flushes(&self, f: &mut FlushBuilder) { use jump::*; @@ -1029,7 +1002,12 @@ mod tests { row[jump::V_COND] = cond; row[jump::W] = w; row[jump::B] = if cond.is_zero() { F64::ZERO } else { F64::ONE }; - assert_eq!(jump_identity(&pows, &row, false), F192::ZERO, "cond = {cond:?}"); + let row: Vec = row.into_iter().map(F192::from).collect(); + assert_eq!( + JumpTable.constraint(&pows).eval(&row, false), + F192::ZERO, + "cond = {cond:?}" + ); } } } diff --git a/crates/primitives/src/field/f192x8.rs b/crates/primitives/src/field/f192x8.rs index 904af5820..e3e882cb8 100644 --- a/crates/primitives/src/field/f192x8.rs +++ b/crates/primitives/src/field/f192x8.rs @@ -130,6 +130,15 @@ mod portable { } } + impl core::ops::Add for F64x8 { + type Output = Self; + + #[inline] + fn add(self, rhs: Self) -> Self { + Self(core::array::from_fn(|i| self.0[i] + rhs.0[i])) + } + } + impl F192x8 { /// The eight mixed products `self_i * k_i`, unreduced. #[inline] @@ -138,6 +147,14 @@ mod portable { } } + /// The embedding of the base field. + impl From for F192x8 { + #[inline] + fn from(k: F64x8) -> Self { + Self(k.0.map(F192::from)) + } + } + /// The unreduced products of eight elements, lane by lane, on a target without the wide multiplier. #[derive(Clone, Copy, Debug)] pub struct Prod8([F192Unreduced; 8]); @@ -328,6 +345,15 @@ mod avx512 { } } + impl core::ops::Add for F64x8 { + type Output = Self; + + #[inline(always)] + fn add(self, rhs: Self) -> Self { + Self(xor(self.0, rhs.0)) + } + } + impl F192x8 { /// The eight mixed products `self_i * k_i`, unreduced: six carry-less multiplications. #[inline(always)] @@ -337,6 +363,23 @@ mod avx512 { } } + /// The embedding of the base field. + impl From for F192x8 { + #[inline(always)] + fn from(k: F64x8) -> Self { + // SAFETY: the module is compiled only with avx512f enabled. + unsafe { + let zero = _mm512_setzero_si512(); + // Each word is its element's `c0`. + Self { + lo: _mm512_unpacklo_epi64(k.0, zero), + hi: _mm512_unpackhi_epi64(k.0, zero), + c2: zero, + } + } + } + } + impl Acc8 { pub const ZERO: Self = Self { even: [from_words([0; 8]); 3], @@ -649,7 +692,13 @@ mod tests { odd.to_array(), core::array::from_fn(|j| if j < 4 { x0[2 * j + 1] } else { x1[2 * (j - 4) + 1] }) ); - let ka = F64x8::from_array(&k0); + + let (ka, kb) = (F64x8::from_array(&k0), F64x8::from_array(&k1)); + assert_eq!( + F192x8::from(ka + kb).to_array(), + core::array::from_fn(|j| F192::from(k0[j] + k1[j])) + ); + assert_eq!(F192x8::from(ka).to_array(), k0.map(F192::from)); assert_eq!( xa.mul_base_unreduced(ka).reduce_lanes().to_array(), core::array::from_fn(|j| x0[j].mul_base(k0[j])) diff --git a/crates/primitives/src/field/gf2_64x3.rs b/crates/primitives/src/field/gf2_64x3.rs index 00d4394d4..2c4441429 100644 --- a/crates/primitives/src/field/gf2_64x3.rs +++ b/crates/primitives/src/field/gf2_64x3.rs @@ -265,26 +265,6 @@ pub fn mul_unreduced4(a: [F192; 4], b: [F192; 4]) -> [F192Unreduced; 4] { std::array::from_fn(|i| a[i].mul_unreduced(b[i])) } -/// Eight mixed products `t * k[i]` by one shared `E` scalar. -/// -/// On AVX-512 this is six CLMULs for all eight, against twenty-four one at a time. -#[inline(always)] -pub fn mul_base8(t: F192, k: [F64; 8]) -> [F192; 8] { - #[cfg(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx512f"))] - // SAFETY: both features are enabled at compile time, and `f192x8` fails the build without avx512bw. - return unsafe { x86_64::mul_base8(t, k) }; - #[cfg(all( - target_arch = "x86_64", - target_feature = "vpclmulqdq", - target_feature = "avx2", - not(target_feature = "avx512f") - ))] - // SAFETY: both features are enabled at compile time. - return unsafe { x86_64::mul_base8_ymm(t, k) }; - #[cfg(not(all(target_arch = "x86_64", target_feature = "vpclmulqdq", target_feature = "avx2")))] - k.map(|k| t.mul_base(k)) -} - /// Eight `E` weights, laid out once for [`dot_base`]. /// /// Pairs share a 128-bit lane, so the kernel loads them with no shuffle: @@ -951,30 +931,6 @@ pub mod x86_64 { ] } - /// Eight mixed products by one scalar. - /// - /// # Safety - /// - /// Callable only where vpclmulqdq, avx512f and avx512bw are enabled. - #[cfg(all(target_feature = "vpclmulqdq", target_feature = "avx512f"))] - #[inline] - #[target_feature(enable = "vpclmulqdq", enable = "avx512f", enable = "avx512bw")] - pub fn mul_base8(t: F192, k: [F64; 8]) -> [F192; 8] { - // `t` in every lane: [c0, c1] for both pair registers, [c2, c2] for the last. - let t01 = _mm512_broadcast_i32x4(pair(t.c0, t.c1)); - let t2 = _mm512_set1_epi64(t.c2 as i64); - let kv = crate::simd::load(&k, 0); - let [e0, e1, o0, o1, e2, o2] = mul_by_pairs(t01, t01, t2, kv); - // Gather each product's low and high halves into qword-wise vectors, then reduce. - let even = bits::<__m512i, [u64; 8]>(reduce_halves512(e0, e1)); - let odd = bits::<__m512i, [u64; 8]>(reduce_halves512(o0, o1)); - let c2 = bits::<__m512i, [u64; 8]>(reduce_halves512(e2, o2)); - std::array::from_fn(|i| { - let (j, w) = (i / 2, if i % 2 == 0 { &even } else { &odd }); - F192::new(w[2 * j], w[2 * j + 1], c2[i]) - }) - } - /// The mixed inner product over packed weights. /// /// # Safety @@ -1028,36 +984,6 @@ pub mod x86_64 { ] } - /// [`super::mul_base8`] in two YMM halves. - /// - /// # Safety - /// - /// Callable only where vpclmulqdq and avx2 are enabled. - #[cfg(all( - target_feature = "vpclmulqdq", - target_feature = "avx2", - not(target_feature = "avx512f") - ))] - #[inline] - #[target_feature(enable = "vpclmulqdq", enable = "avx2")] - pub fn mul_base8_ymm(t: F192, k: [F64; 8]) -> [F192; 8] { - let t01 = _mm256_broadcastsi128_si256(pair(t.c0, t.c1)); - let t2 = _mm256_set1_epi64x(t.c2 as i64); - let mut out = [F192::ZERO; 8]; - for (out, k) in out.as_chunks_mut::<4>().0.iter_mut().zip(k.as_chunks::<4>().0) { - let kv = bits::<[F64; 4], __m256i>(*k); - let [e0, e1, o0, o1, e2, o2] = mul_by_pairs_ymm(t01, t01, t2, kv); - let even = bits::<__m256i, [u64; 4]>(reduce_halves256(e0, e1)); - let odd = bits::<__m256i, [u64; 4]>(reduce_halves256(o0, o1)); - let c2 = bits::<__m256i, [u64; 4]>(reduce_halves256(e2, o2)); - for (i, slot) in out.iter_mut().enumerate() { - let (j, w) = (i / 2, if i % 2 == 0 { &even } else { &odd }); - *slot = F192::new(w[2 * j], w[2 * j + 1], c2[i]); - } - } - out - } - /// [`super::dot_base`] in two YMM halves. /// /// # Safety @@ -1240,13 +1166,6 @@ mod tests { fn batched_mixed_products_match_scalar() { let mut rng = Rng::new(9); for _ in 0..200 { - let t = rng.ext(); - // Edge scalars sit next to random ones, so every lane sees a zero and an all-ones word. - let mut k: [F64; 8] = std::array::from_fn(|_| F64(rng.ext().c0)); - k[1] = F64(0); - k[6] = F64(u64::MAX); - assert_eq!(mul_base8(t, k), k.map(|k| t.mul_base(k))); - let w: [F192; 16] = std::array::from_fn(|_| rng.ext()); let packed = [ Weights8::new(w[..8].try_into().unwrap()), diff --git a/crates/primitives/src/field/mod.rs b/crates/primitives/src/field/mod.rs index 4a078c452..f418a81fb 100644 --- a/crates/primitives/src/field/mod.rs +++ b/crates/primitives/src/field/mod.rs @@ -30,7 +30,7 @@ pub use f192x8::{ xor_into8, }; pub use gf2_64::F64; -pub use gf2_64x3::{F192, F192Unreduced, Weights8, dot_base, mul_base8, mul_unreduced4, mul2, mul4}; +pub use gf2_64x3::{F192, F192Unreduced, Weights8, dot_base, mul_unreduced4, mul2, mul4}; // --------------------------------------------------------------------------- // leanVM g-power helpers: domain separators / opcodes as x^k, and the g-power diff --git a/crates/primitives/src/multilinear.rs b/crates/primitives/src/multilinear.rs index aac697022..a8be9bd2a 100644 --- a/crates/primitives/src/multilinear.rs +++ b/crates/primitives/src/multilinear.rs @@ -6,7 +6,7 @@ //! `K`-valued (`F64`) while randomness is `E`-valued (`F192`), so the first //! fold of a committed table also lifts it into `E`. -use crate::field::{F64, F192, Weights8, dot_base, mul_base8, mul4}; +use crate::field::{F64, F192, Weights8, dot_base, mul4}; /// Multilinear interpolation in one variable over `E`: `lo + t·(lo+hi)`, the /// char-2 form of `(1−t)·lo + t·hi`. @@ -113,69 +113,15 @@ fn fold_low_k(table: &[F64], chi: F192) -> Vec { .collect() } -/// Bind the highest variable of a `K`-table and lift the result into `E`. -/// -/// Eight entries share one batched mixed product ([`mul_base8`]). -pub fn fold_high_k(table: &[F64], chi: F192) -> Vec { - assert!(table.len().is_multiple_of(2), "a table of whole pairs"); - let (lo, hi) = table.split_at(table.len() / 2); - // SAFETY: the fill writes every element before it reads any. - unsafe { - crate::filled_vec(lo.len(), |out| { - let (out8, out_tail) = out.as_chunks_mut::<8>(); - let ((lo8, lo_tail), (hi8, hi_tail)) = (lo.as_chunks::<8>(), hi.as_chunks::<8>()); - for ((o, l), h) in out8.iter_mut().zip(lo8).zip(hi8) { - let p = mul_base8(chi, std::array::from_fn(|i| l[i] + h[i])); - for i in 0..8 { - o[i] = F192::from(l[i]) + p[i]; - } - } - for ((o, &l), &h) in out_tail.iter_mut().zip(lo_tail).zip(hi_tail) { - *o = interp_k(l, h, chi); - } - }) - } -} - -/// Bind the highest free variable of `table` to `chi` in place: `table[i] = -/// interp(table[i], table[i + half], chi)`. Binding from the top down leaves the -/// low variables, the ones every table of a batch shares, for last. -pub fn fold_high_inplace(table: &mut Vec, chi: F192) { - debug_assert_eq!(table.len() % 2, 0); - let half = table.len() / 2; - { - // Split once rather than index twice: indexing reloads the data pointer - // and length through the container every iteration, since nothing proves - // they do not alias the elements, and pays a bounds check for it. - let (lo, hi) = (**table).split_at_mut(half); - interp_into(lo, hi, chi); - } - table.truncate(half); -} - -/// `lo[i] = interp(lo[i], hi[i], chi)`, four products per batch. -fn interp_into(lo: &mut [F192], hi: &[F192], chi: F192) { - let ((lo4, lo_tail), (hi4, hi_tail)) = (lo.as_chunks_mut::<4>(), hi.as_chunks::<4>()); - for (l, h) in lo4.iter_mut().zip(hi4) { - let p = mul4([chi; 4], std::array::from_fn(|i| l[i] + h[i])); - for i in 0..4 { - l[i] += p[i]; - } - } - for (l, h) in lo_tail.iter_mut().zip(hi_tail) { - *l = interp(*l, *h, chi); - } -} - /// Marginalize the lowest variable out of an `eq` table (in place). `eq(r_0, 0) + /// eq(r_0, 1) = 1`, so summing adjacent entries drops `r_0` with no multiplies, /// versus `2^{n-1}` to rebuild the table. pub fn shrink_eq_low(table: &mut Vec) { let half = table.len() / 2; { - // Sliced, as in `fold_high_inplace`: reading the pair and writing the - // sum through one slice drops the per-iteration reload and its bounds - // check. The write index trails the read, so the in-place walk is sound. + // Through one slice rather than the `Vec`: reading the pair and writing the + // sum drops the per-iteration reload and its bounds check. The write index + // trails the read, so the in-place walk is sound. let t: &mut [F192] = table; for i in 0..half { let (a, b) = (t[2 * i], t[2 * i + 1]); @@ -190,7 +136,7 @@ pub fn shrink_eq_low(table: &mut Vec) { pub fn shrink_eq_high(table: &mut Vec) { let half = table.len() / 2; { - // Sliced, as in `fold_high_inplace`. + // Split once, as in `shrink_eq_low`, so the loop pays no reload or bounds check. let (lo, hi) = (**table).split_at_mut(half); for (l, h) in lo.iter_mut().zip(&*hi) { *l += *h; @@ -287,18 +233,6 @@ mod tests { let point: Vec<_> = (0..n).map(|i| F192::from(F64(i % 2))).collect(); assert_eq!(mle_eval(&table, &point), fold(&point)); assert_eq!(mle_eval_par(&table, &point), fold(&point)); - // An odd length leaves a scalar tail after the batched folds. - let chi = point.first().copied().unwrap_or(F192::Y) + F192::Y; - if n > 0 { - let half = table.len() / 2; - let want: Vec<_> = (0..half).map(|i| interp_k(table[i], table[i + half], chi)).collect(); - assert_eq!(&*fold_high_k(&table, chi), &want[..]); - let lifted: Vec = table.iter().map(|&k| F192::from(k) * chi).collect(); - let mut got = lifted.clone(); - fold_high_inplace(&mut got, chi); - let want: Vec<_> = (0..half).map(|i| interp(lifted[i], lifted[i + half], chi)).collect(); - assert_eq!(got, want); - } } } } diff --git a/crates/rec_aggregation/guests/lean_ethereum.py b/crates/rec_aggregation/guests/lean_ethereum.py index cbce8fc50..b6e0bd2e6 100644 --- a/crates/rec_aggregation/guests/lean_ethereum.py +++ b/crates/rec_aggregation/guests/lean_ethereum.py @@ -2239,14 +2239,14 @@ def verify_tables(fs0, fs1, fs2, fs3, fs4, cursor, pi_0, pi_1, zeta, g_bus_mu, d claim_cplen_g[GEN ** claim_idx] = tau_g # cplen = tau_t claim_idx += 1 # The table's AIR constraint at the final point (col_evals is indexed by - # local column index; the formulas mirror tables.rs eval_constraint). Every + # local column index; the formulas mirror tables.rs Table::constraint). Every # value relation now rides the bus as a degree-2 coordinate, so only JUMP's # is-nonzero indicator is left with an identity of its own. constraint_eval = 0 if t == TABLE_JUMP: # `b = c*w` and `c*(b+1) = 0`. The condition is K-valued, its memory read # carrying literal zeros above the low limb, so both identities are - # single-lane (tables.rs jump_identity). Local columns: v_cond at 5, w at + # single-lane (tables.rs JumpTable::constraint). Local columns: v_cond at 5, w at # 12, the indicator b at 13. c = col_evals[5] b = col_evals[13]