diff --git a/AGENTS.md b/AGENTS.md index 1f7b4789b..daf03b5e8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -106,7 +106,7 @@ The same verification algorithm is written out three times. Any change to snark The crate walks the protocol once for two entry points. `verify_execution` checks a proof and returns its `Summary` (the bytecode point, flock's claims, where flock's scalars end, and every query's own Merkle path). `replay_unchecked` returns the same `Summary` and checks nothing: every check sits behind `Transcript::checked`, so a new check goes there or the replay pays for it, and nothing else may depend on the mode. **`aggregate` does not verify its children.** It assumes them checked and only replays each (in parallel, one single-threaded replay per worker) to read out the witness `gen_verify` hands the guest, computing the bytecode claim's value itself. A child the replay cannot walk is refused as `InvalidChild`: announced sizes out of range, a stream or openings of the wrong length or shape, and so most tampering, since the pruned openings must sit at the query positions the transcript draws (a child whose statement does not match its proof typically fails as `InvalidChild(Snark(MerkleOpeningMissing))`). So is a child whose fresh claim on flock's matrices is false. Any other invalid child is caught by the guest, which verifies every child in the VM: its run fails and `aggregate` returns `AggregationError::ProofError`. A caller that has not checked a child calls `EthereumProof::verify` on it first. -Understand the third before changing the verifier. `guests/lean_ethereum.py` is zkDSL, not runnable Python. `lean_compiler` lowers it to the six-opcode, write-once-memory VM, so the prover proves every verifier step. Its size and instruction mix are what the recursion benchmark reports first. It verifies raw signatures of both schemes: a node's coverage table has one contiguous region per scheme and one for DA roots, so the one range check a write already needs also keeps a signature off the other scheme's declared claims, and the statement's signer lists say which scheme verified which claim. Both lists are `(pubkey, message)` pairs, a claim's message riding its own four-cell slot, hashed and rebuilt by the same helpers (`claim_list_digest`, `child_claim_list_digest`); the signer-set digest is one compression over the two list digests, each plain BLAKE2s of its claims: a run-time length rides the byte counter because the counter is a memory operand, split as `doc/leanvm` §sec:prog-byte-counter describes. Both lists count claims rather than distinct signers, a pubkey claiming once per message it signed. Both schemes' addresses are built per signature in-circuit: leanXMSS's from the leaf index the signature carries (its bits are hinted and pinned Boolean, one decomposition serving both layers, the statement recording no leaf index), leanSPHINCS's from the index its message digest picks. Both share one WOTS+C walk (`walk_group`, the digits dispatched three at a time) and one Merkle step (`is_child`, the two children hinted in order). A leanSPHINCS value is one cell, as a leanXMSS one is, so a hash's output feeds the next where it lands. What the leanSPHINCS verifier pays for is dispatch, and it is written around it: every digest field that selects something (a codeword, a tree's leaf, a subtree's WOTS key, the one-time digits three at a time) is hinted in the exponent, range checked, and dispatched to an arm specialized on it, so a chain's steps, a Merkle pair's order and every address are the arm's constants; the arm returns the field as the digest holds it, and the sum of those must be the digest, which is what binds a hinted dispatch. Only the index and the digest's unused bits are read as bits. A forest WOTS key is a function specialized per codeword (`sp_forest_key`); the other arms are inline, which is where the bytecode goes. `DBG_PROF=1` (with `DBG_PROF_DUMP` and `DBG_DISASM`) gives the per-function and per-pc cycle profile this was tuned with. Two consequences: +Understand the third before changing the verifier. `guests/lean_ethereum.py` is zkDSL, not runnable Python. `lean_compiler` lowers it to the six-opcode, write-once-memory VM, so the prover proves every verifier step. Its size and instruction mix are what the recursion benchmark reports first. It verifies raw signatures of both schemes: a node's coverage table has one contiguous region per scheme and one for DA roots, so the one range check a write already needs also keeps a signature off the other scheme's declared claims, and the statement's signer lists say which scheme verified which claim. Both lists are `(pubkey, message)` pairs, a claim's message riding its own four-cell slot, hashed and rebuilt by the same helpers (`claim_list_digest`, `child_claim_list_digest`); the signer-set digest is one compression over the two list digests, each plain BLAKE2s of its claims: a run-time length rides the byte counter because the counter is a memory operand, split as `doc/leanvm` §sec:prog-byte-counter describes. Both lists count claims rather than distinct signers, a pubkey claiming once per message it signed. Both schemes' addresses are built per signature in-circuit: leanXMSS's from the leaf index the signature carries (its bits are hinted and pinned Boolean, one decomposition serving both layers, the statement recording no leaf index), leanSPHINCS's from the index its message digest picks. Both share one WOTS+C walk (`walk_group`, the digits dispatched three at a time) and one Merkle step (`is_child`, the two children hinted in order). A leanSPHINCS value is one cell, as a leanXMSS one is, so a hash's output feeds the next where it lands. What the leanSPHINCS verifier pays for is dispatch, and it is written around it: every digest field that selects something (a codeword, a tree's leaf, a subtree's WOTS key, the one-time digits three at a time) is hinted in the exponent, range checked, and dispatched to an arm specialized on it, so a chain's steps, a Merkle pair's order and every address are the arm's constants; the arm returns the field as the digest holds it, and the sum of those must be the digest, which is what binds a hinted dispatch. Only the index is read as bits; the digest's unused bits and the encoding counter are hinted bytes in the exponent, each range checked and read off a table of the 256 byte values (`sp_from_bytes`). A forest WOTS key is a function specialized per codeword (`sp_forest_key`); the other arms are inline, which is where the bytecode goes. `DBG_PROF=1` (with `DBG_PROF_DUMP` and `DBG_DISASM`) gives the per-function and per-pc cycle profile this was tuned with. Two consequences: - The guest is **self-referential**: it verifies proofs of itself, so `unified_guest` compiles it to a fixed point on its own log size. The digest needs no fixed point, riding the statement instead of the code, which is also what lets one bytecode serve any inner size and PCS rate. - It does not verify *quite* everything in-circuit. Three claims on fixed polynomials (stacked bytecode, flock's A0/B0) are deferred. Each node batches its children's carried claims with the fresh ones its verifications raise, `2n` per polynomial down to one; only the root's are discharged natively, by `EthereumProof::verify` (explained in `doc/leanvm/`). diff --git a/README.md b/README.md index 1ae54a900..83e603921 100644 --- a/README.md +++ b/README.md @@ -77,12 +77,12 @@ cargo leanvm aggregate --xmss 1290 --log-inv-rate 1 --repeat 3 ``` aggregation, 1,290 leanXMSS signatures - cycles (VM steps) : 1,336,043 = 2^20.35 - details : SET 2^18.207 (22.7%) BLAKE2S 2^17.992 (19.5%) MUL 2^17.943 (18.9%) XOR2^17.866 (17.9%) DEREF 2^17.437 (13.3%) JUMP 2^16.672 (7.8%) MEMORY 2^21.125 BYTECODE 2^17.961 TOTAL_COMMITTED 2^26.672 - proof size : 318.4 KiB - proving time : 0.947 s ± 0.8% peak memory 15.46 GiB - per signature : 1,362.468 signatures/s - verifying : 5.387 ms + cycles (VM steps) : 1,333,892 = 2^20.347 + details : SET 2^18.21 (22.7%) BLAKE2S 2^17.992 (19.5%) MUL 2^17.943 (18.9%) XOR 2^17.851 (17.7%) DEREF 2^17.437 (13.3%) JUMP 2^16.672 (7.8%) MEMORY 2^21.124 BYTECODE 2^18.13 TOTAL_COMMITTED 2^26.676 + proof size : 320.9 KiB + proving time : 0.912 s ± 3.0% peak memory 15.44 GiB + per signature : 1,414.874 signatures/s + verifying : 5.64 ms ``` ### leanSPHINCS aggregation @@ -95,12 +95,12 @@ cargo leanvm aggregate --sphincs 850 --log-inv-rate 1 --repeat 3 ``` aggregation, 850 leanSPHINCS signatures - cycles (VM steps) : 1,870,199 = 2^20.835 - details : SET 2^18.992 (27.9%) XOR 2^18.696 (22.7%) MUL 2^18.208 (16.2%) BLAKE2S2^17.998 (14.0%) DEREF 2^17.879 (12.9%) JUMP 2^16.855 (6.3%) MEMORY 2^21.403 BYTECODE 2^17.961 TOTAL_COMMITTED 2^26.775 - proof size : 325.8 KiB - proving time : 1 s ± 3.9% peak memory 16.102 GiB - per signature : 850.082 signatures/s - verifying : 5.113 ms + cycles (VM steps) : 1,698,709 = 2^20.696 + details : SET 2^18.897 (28.7%) XOR 2^18.491 (21.7%) DEREF 2^18.062 (16.1%) BLAKE2S 2^17.998 (15.4%) MUL 2^17.621 (11.9%)JUMP 2^16.678 (6.2%) MEMORY 2^21.323 BYTECODE 2^18.13 TOTAL_COMMITTED 2^26.778 + proof size : 325.2 KiB + proving time : 0.965 s ± 2.1% peak memory 16.017 GiB + per signature : 880.823 signatures/s + verifying : 5.787 ms ``` ### data availability @@ -111,12 +111,12 @@ cargo leanvm aggregate --blobs 16 --log-inv-rate 1 --repeat 3 ``` aggregation, 16 blobs - cycles (VM steps) : 2,989,506 = 2^21.511 - details : MUL 2^19.899 (32.7%) XOR 2^19.809 (30.7%) DEREF 2^19.138 (19.3%) JUMP 2^18.299 (10.8%) SET 2^16.787 (3.8%) BLAKE2S 2^16.295 (2.7%) MEMORY 2^21.695 BYTECODE 2^17.737 TOTAL_COMMITTED 2^26.695 - proof size : 324.0 KiB - proving time : 0.986 s ± 11.1% peak memory 12.534 GiB - blob throughput : 16.235 blobs/s, 2.029 MiB/s - verifying : 6.573 ms + cycles (VM steps) : 2,993,056 = 2^21.513 + details : MUL 2^19.899 (32.7%) XOR 2^19.809 (30.7%) DEREF 2^19.138 (19.3%) JUMP 2^18.299 (10.8%) SET 2^16.842 (3.9%) BLAKE2S 2^16.295 (2.7%) MEMORY 2^21.697 BYTECODE 2^18.13 TOTAL_COMMITTED 2^26.699 + proof size : 321.0 KiB + proving time : 0.873 s ± 2.9% peak memory 12.249 GiB + blob throughput : 18.322 blobs/s, 2.29 MiB/s + verifying : 10.283 ms ``` ### recursion @@ -127,11 +127,11 @@ cargo leanvm recursion --n 2 --xmss-per-leaf 900 --log-inv-rate 2 --repeat 3 ``` recursion 2→1, over leaves of 900 leanXMSS signatures - cycles (VM steps) : 599,460 = 2^19.193 - details : MUL 2^17.877 (40.2%) DEREF 2^17.091 (23.3%) XOR 2^16.807 (19.1%) SET 2^15.956 (10.6%) JUMP 2^14.494 (3.8%) BLAKE2S 2^14.126 (3.0%) MEMORY 2^19.559 BYTECODE 2^17.901 TOTAL_COMMITTED 2^24.533 - proof size : 201.5 KiB - proving time : 0.291 s ± 3.4% peak memory 14.749 GiB - verifying : 6.271 ms + cycles (VM steps) : 457,176 = 2^18.802 + details : MUL 2^17.441 (38.9%) XOR 2^16.697 (23.2%) DEREF 2^16.575 (21.4%) SET 2^15.183 (8.1%) JUMP 2^14.309 (4.4%) BLAKE2S 2^14.126 (3.9%) MEMORY 2^19.314 BYTECODE 2^18.13 TOTAL_COMMITTED 2^24.428 + proof size : 199.1 KiB + proving time : 0.276 s ± 1.5% peak memory 14.694 GiB + verifying : 6.945 ms ``` ### hashing @@ -143,17 +143,17 @@ BENCH_REPEAT=3 BENCH_COOLDOWN=2 FLOCK_N_LOG=18 cargo bench -p flock --bench hash ``` Flock BLAKE2s batch proving, 262,144 compressions (2^18 slots) setup (preprocessing, excluded) : 0.0 ms - witness-gen : 34.8 ms ± 26.8% 6.1% - commit : 102.2 ms ± 1.3% 17.8% - zerocheck : 244.9 ms ± 7.5% 42.6% - lincheck : 20.2 ms ± 3.4% 3.5% - pcs opening : 172.3 ms ± 1.3% 30.0% + witness-gen : 26.8 ms ± 54.6% 4.9% + commit : 98.5 ms ± 3.3% 18.1% + zerocheck : 234.4 ms ± 6.4% 43.1% + lincheck : 19.0 ms ± 3.2% 3.5% + pcs opening : 164.7 ms ± 1.4% 30.3% other : 0.0 ms 0.0% ------------------------------------------ - prove TOTAL (witness excluded) : 539.7 ms ± 3.3% 93.9% - verify : 2.0 ms - throughput : 485,765 compressions/s ± 3.3% - (~2416.7 leanXMSS/s equivalent at 201 compressions/signature) + prove TOTAL (witness excluded) : 516.6 ms ± 2.5% 95.1% + verify (scalar, one thread) : 2.3 ms + throughput : 507,431 compressions/s ± 2.5% + (~2524.5 leanXMSS/s equivalent at 201 compressions/signature) ``` ### Fibonacci @@ -164,11 +164,11 @@ cargo leanvm fibonacci --n 2000000 --log-inv-rate 1 --repeat 3 ``` Fibonacci (in the exponent, i.e. modulo 2^64 - 1), N = 2,000,000 - cycles (VM steps) : 2,127,880 - details : MUL 2^20.944 (98.9%) SET 2^13.288 (0.5%) DEREF 2^12.967 (0.4%) JUMP 2^10.968 (0.1%) XOR 2^10.966 (0.1%) MEMORY 2^20.96 BYTECODE 2^11.352 TOTAL_COMMITTED 2^25.26 - proof size : 286.0 KiB - proving : 0.344 s ± 4.6% 6,181,087 cycles/s peak memory 5.199 GiB - verifying : 2.218 ms + cycles (VM steps) : 2,040,016 + details : MUL 2^20.944 (98.9%) SET 2^13.288 (0.5%) DEREF 2^12.967 (0.4%) JUMP 2^10.968 (0.1%) XOR 2^10.966 (0.1%) MEMORY2^20.959 BYTECODE 2^10.017 TOTAL_COMMITTED 2^25.259 + proof size : 284.9 KiB + proving : 0.319 s ± 6.4% 6,402,981 cycles/s peak memory 3.388 GiB + verifying : 3.709 ms ``` ## SNARK machinery diff --git a/crates/leanvm_core/src/cpu/execute.rs b/crates/leanvm_core/src/cpu/execute.rs index 011a744ae..1f89ea789 100644 --- a/crates/leanvm_core/src/cpu/execute.rs +++ b/crates/leanvm_core/src/cpu/execute.rs @@ -11,7 +11,7 @@ use primitives::{ }; pub struct Execution { - pub mem: Vec, // data memory after the run, write-once (size cells, power of two) + pub mem: Memory, // data memory after the run, write-once (size cells, power of two) pub cycles: usize, // number of instructions the run executed (trace length) pub mem_used: usize, // cells actually touched, before the power-of-two pad of `mem` /// Cells an instruction read before anything wrote them, so the value it read was @@ -28,6 +28,8 @@ pub struct Execution { /// and an arithmetic back-solve writes its operand before reading it. pub unconstrained_reads: Vec, pub(crate) trace: Trace, // rows + final access-count columns, emitted in the same walk + /// `g^0, g^1, ..` as far as the run indexed them: every pc, frame base and cell it reached. + pub(crate) gpow: Vec, } impl Execution { @@ -37,6 +39,74 @@ impl Execution { } } +/// `log2` of the cells a [`Memory`] page holds. +const PAGE_LOG: usize = 16; +const PAGE: usize = 1 << PAGE_LOG; +// The padded image is a whole number of pages. +const _: () = assert!(PAGE_LOG <= MIN_LOG_MEM); + +/// A cell's write state. A cell not `Written` holds ZERO; a `DEREF` between two +/// unwritten cells links them, and a write to either reaches the other at once. +#[derive(Clone, Copy, PartialEq, Eq)] +#[repr(u8)] +enum State { + Unwritten = 0, + Linked, + Written, +} + +/// One cell: its word, its access count as an exponent (the column holds `g^count`), +/// and its write state, side by side so one access touches one line. All zero is an +/// unwritten cell holding ZERO, never accessed. +#[derive(Clone, Copy)] +#[repr(C)] +pub(crate) struct Cell { + pub(crate) word: F192, + pub(crate) count: u32, + state: State, +} + +impl Cell { + pub(crate) const UNWRITTEN: Self = Self { + word: F192::ZERO, + count: 0, + state: State::Unwritten, + }; +} + +/// The write-once data memory: `len` cells, a power of two, in pages, so that the +/// image grows without moving what it holds. Cells past the last page are unwritten +/// and never accessed. +pub struct Memory { + pages: Vec>, + len: usize, +} + +impl Memory { + /// Cells in the image, its power-of-two pad included. + pub fn len(&self) -> usize { + self.len + } + + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + /// The word in cell `i`. + pub fn word(&self, i: usize) -> F192 { + self.page(i >> PAGE_LOG).map_or(F192::ZERO, |p| p[i & (PAGE - 1)].word) + } + + /// Page `p`'s cells, `None` past the last one. + pub(crate) fn page(&self, p: usize) -> Option<&[Cell]> { + self.pages.get(p).map(|p| &p[..]) + } + + pub(crate) fn page_log() -> usize { + PAGE_LOG + } +} + /// Why a run has no execution the interpreter can find: the program, on this public /// input and advice, asks for something the machine cannot do. #[derive(Clone, Debug, PartialEq, Eq)] @@ -84,6 +154,8 @@ pub enum Fault { /// The longest run the interpreter executes. const MAX_STEPS: usize = 100_000_000; +// A row touches at most nine cells, so every access count fits the `u32` exponent. +const _: () = assert!(9 * MAX_STEPS < u32::MAX as usize); impl ExecError { fn new(program: &Program, pc: u32, fault: Fault) -> Self { @@ -155,32 +227,37 @@ fn in_k(what: &'static str, v: F192) -> Result { as_addr(v).ok_or(Fault::NotInK { what, value: v }) } -fn pop_witness<'a>( - witness: &'a HashMap>>, - positions: &mut HashMap<&'a str, usize>, +/// A named witness stream and the next entry `hint_witness` pops from it. Resolved +/// once per run, so a pop indexes it rather than hashing its name. +struct Stream<'a> { name: &'a str, - len: u32, -) -> Result<&'a [F192], Fault> { - let entries = witness - .get(name) - .ok_or_else(|| Fault::Witness(format!("no witness stream `{name}` (Program::set_witness)")))?; - let position = positions.entry(name).or_default(); - let entry = entries.get(*position).ok_or_else(|| { - Fault::Witness(format!( - "witness stream `{name}` exhausted (needs entry {}, has {})", - *position + 1, - entries.len() - )) - })?; - if entry.len() != len as usize { - return Err(Fault::Witness(format!( - "witness `{name}` entry {} holds {} values, the destination {len}", - *position, - entry.len() - ))); + entries: Option<&'a [Vec]>, + next: usize, +} + +impl<'a> Stream<'a> { + fn pop(&mut self, len: u32) -> Result<&'a [F192], Fault> { + let name = self.name; + let entries = self + .entries + .ok_or_else(|| Fault::Witness(format!("no witness stream `{name}` (Program::set_witness)")))?; + let entry = entries.get(self.next).ok_or_else(|| { + Fault::Witness(format!( + "witness stream `{name}` exhausted (needs entry {}, has {})", + self.next + 1, + entries.len() + )) + })?; + if entry.len() != len as usize { + return Err(Fault::Witness(format!( + "witness `{name}` entry {} holds {} values, the destination {len}", + self.next, + entry.len() + ))); + } + self.next += 1; + Ok(entry) } - *position += 1; - Ok(entry) } impl Program { @@ -201,73 +278,84 @@ impl Program { .max(1 << MIN_LOG_MEM); let mut g = GPow::new(seed); - // Dense write-once data memory (read path stays a vector for speed), the - // per-cell access count (g^{count}, default g^0 = 1), and each cell's state. + // The write-once data memory, each cell's access count and write state with its word. let n0 = self.main_frame.max(2) as usize; let mut m = Mem { - cells: vec![F192::ZERO; n0], - state: vec![State::Unwritten; n0], - count: vec![F64::ONE; n0], + pages: Vec::new(), + used: 0, + counts: 0, links: HashMap::new(), unwritten_reads: Vec::new(), program: self, dbg_pc: 0, dbg_hint: None, }; + m.ensure(n0 - 1); // Seed the public input into m[0], m[1] (addresses g^0, g^1, §sec:e2e-pi). m.put(0, public_input[0])?; m.put(1, public_input[1])?; - // Per-pc bytecode execution count (g^{count}). - let mut bytecode_count: Vec = vec![F64::ONE; self.prog.len()]; + // Per-pc bytecode execution count, as an exponent of `g`. + let mut bytecode_count: Vec = vec![0; self.prog.len()]; let mut next_free = self.main_frame; let (mut pc, mut fp) = (0u32, 0u32); let mut steps = 0usize; + // Witness streams, resolved by name once so a pop indexes one. + let mut streams: Vec = Vec::new(); + let mut stream_of: HashMap<&str, u32> = HashMap::new(); // Per-pc hint index. `self.hints` is keyed by pc, so probing it each step // costs a hash of the counter for what is almost always a miss; the // program has one bytecode slot per pc, so flatten the map into a dense - // table: 0 = no hints, otherwise the 1-based index into `hint_lists`. + // table: 0 = no hints, otherwise the 1-based index into `hint_lists`, whose + // entries pair a pc's hints with where their streams start in `hint_streams`. let mut hint_at = vec![0u32; self.prog.len()]; - let mut hint_lists: Vec<&[RHint]> = Vec::with_capacity(self.hints.len()); + let mut hint_lists: Vec<(&[RHint], usize)> = Vec::with_capacity(self.hints.len()); + let mut hint_streams: Vec = Vec::new(); for (&hpc, hs) in &self.hints { - hint_lists.push(hs.as_slice()); + hint_lists.push((hs.as_slice(), hint_streams.len())); hint_at[hpc as usize] = hint_lists.len() as u32; + for h in hs { + let (RHint::WitnessStack { name, .. } | RHint::WitnessHeap { name, .. }) = h else { + hint_streams.push(u32::MAX); + continue; + }; + let s = *stream_of.entry(name).or_insert_with(|| { + streams.push(Stream { + name, + entries: self.witness.get(name).map(Vec::as_slice), + next: 0, + }); + streams.len() as u32 - 1 + }); + hint_streams.push(s); + } } // `DBG_PROF=1`: per-pc step counts, printed as a per-function cycle // profile after the run (needs `fn_ranges`, i.e. a compiled program). let mut prof: Option> = std::env::var("DBG_PROF").is_ok().then(|| vec![0u64; self.prog.len()]); - // Per-stream cursor into the named witness data (`hint_witness` pops - // sequentially). - let mut witness_positions: HashMap<&str, usize> = HashMap::new(); // Baby-step table for `hint_decompose_bits_exponent`, built on first use. let mut dlog_cache: Option<(GPow, F64)> = None; // Per-opcode trace rows, accumulated during the walk and assembled into the // `Trace` once the run finishes (alongside the final count columns). - let mut xor: Vec = Vec::new(); - let mut mul: Vec = Vec::new(); - let mut set: Vec = Vec::new(); - let mut deref: Vec = Vec::new(); - let mut jump: Vec = Vec::new(); - let mut blake2s: Vec = Vec::new(); - - // A cell that is not `Written` holds ZERO. A `DEREF` between two unwritten cells - // links them, and a write to either reaches the other at once. - #[derive(Clone, Copy, PartialEq, Eq)] - enum State { - Unwritten, - Linked, - Written, - } + let mut xor: Rows = Rows::new(); + let mut mul: Rows = Rows::new(); + let mut set: Rows = Rows::new(); + let mut deref: Rows = Rows::new(); + let mut jump: Rows = Rows::new(); + let mut blake2s: Rows = Rows::new(); - // The three dense per-cell vectors, kept in lockstep. Every method on the hot + // The memory's pages, grown as the run reaches them. Every method on the hot // path is `#[inline(always)]`: they sit in the interpreter's opcode loop. struct Mem<'a> { - cells: Vec, - state: Vec, - count: Vec, + pages: Vec>, + /// One past the highest cell accessed or reserved: every cell below it is + /// in a page. + used: usize, + /// The highest access count any cell reached. + counts: u32, /// The cells a `DEREF` linked each cell to. links: HashMap>, /// Cells an instruction read before anything wrote them. @@ -283,39 +371,57 @@ impl Program { dbg_hint: Option<&'static str>, } impl Mem<'_> { - // Grow the dense vectors so `idx` is in range. All accessed cells - // satisfy cell < next_free after their frame's allocation, so this - // only ever extends. + // Cover cell `idx`. All accessed cells satisfy cell < next_free after + // their frame's allocation, so this only ever extends. #[inline(always)] fn ensure(&mut self, idx: usize) { - if idx >= self.cells.len() { - let n = idx + 1; - self.cells.resize(n, F192::ZERO); - self.state.resize(n, State::Unwritten); - self.count.resize(n, F64::ONE); + if idx >= self.used { + self.grow(idx); } } + #[cold] + #[inline(never)] + fn grow(&mut self, idx: usize) { + self.used = idx + 1; + while self.pages.len() <= idx >> PAGE_LOG { + // SAFETY: all zero is a valid `Cell`, the unwritten one. + self.pages + .push(unsafe { Box::<[Cell]>::new_zeroed_slice(PAGE).assume_init() }); + } + } + /// Cell `c`, which `ensure` has covered. + #[inline(always)] + fn at(&mut self, c: u32) -> &mut Cell { + let c = c as usize; + &mut self.pages[c >> PAGE_LOG][c & (PAGE - 1)] + } + #[inline(always)] + fn covered(&self, c: u32) -> Option<&Cell> { + let c = c as usize; + (c < self.used).then(|| &self.pages[c >> PAGE_LOG][c & (PAGE - 1)]) + } #[inline(always)] fn written(&self, cell: u32) -> bool { - self.state.get(cell as usize) == Some(&State::Written) + self.covered(cell).is_some_and(|c| c.state == State::Written) } // A hint's read, which writes nothing: what a hint computes is advice the // instructions go on to check. #[inline(always)] fn get(&self, cell: u32) -> F192 { - self.cells.get(cell as usize).copied().unwrap_or(F192::ZERO) + self.covered(cell).map_or(F192::ZERO, |c| c.word) } // An instruction's read. An unwritten cell reads as ZERO, which is then written // there: the row is filled from the final image, which has to agree. #[inline(always)] fn read(&mut self, cell: u32) -> F192 { - let c = cell as usize; - self.ensure(c); - if self.state[c] != State::Written { - self.state[c] = State::Written; + self.ensure(cell as usize); + let c = self.at(cell); + let word = c.word; + if c.state != State::Written { + c.state = State::Written; self.unwritten_reads.push(cell); } - self.cells[c] + word } // Write-once store, carried to every cell linked to this one. #[inline(always)] @@ -328,14 +434,14 @@ impl Program { // `put` on one cell; whether it was linked. #[inline(always)] fn store(&mut self, cell: u32, v: F192) -> Result { - let c = cell as usize; - self.ensure(c); - let s = self.state[c]; - if s == State::Written && self.cells[c] != v { + self.ensure(cell as usize); + let c = self.at(cell); + let s = c.state; + if s == State::Written && c.word != v { return Err(self.conflict(cell, v)); } - self.cells[c] = v; - self.state[c] = State::Written; + c.word = v; + c.state = State::Written; Ok(s == State::Linked) } #[cold] @@ -353,8 +459,9 @@ impl Program { fn link(&mut self, a: u32, b: u32) { for (x, y) in [(a, b), (b, a)] { self.ensure(x as usize); - if self.state[x as usize] == State::Unwritten { - self.state[x as usize] = State::Linked; + let c = self.at(x); + if c.state == State::Unwritten { + c.state = State::Linked; } self.links.entry(x).or_default().push(y); } @@ -363,21 +470,21 @@ impl Program { fn conflict(&self, cell: u32, v: F192) -> ExecError { let fault = Fault::Conflict { cell, - had: self.cells[cell as usize], + had: self.get(cell), new: v, hint: self.dbg_hint, }; ExecError::new(self.program, self.dbg_pc, fault) } - // Read the running access count and advance it by ×g (the free increment). - // ×g is ×x, i.e. `mul_by_g`, a shift+fold rather than a PMULL; this runs on every - // memory access (several million per run), so the cheap form matters. + // Read the running access count and advance it: the column's `×g`, on the + // exponent. A run makes fewer than `2^32` accesses in all, so it cannot wrap. #[inline(always)] - fn bump_access_count(&mut self, cell: u32) -> F64 { + fn bump_access_count(&mut self, cell: u32) -> u32 { self.ensure(cell as usize); - let cell_idx = cell as usize; - let count = self.count[cell_idx]; - self.count[cell_idx] = mul_by_g(count); + let c = self.at(cell); + let count = c.count; + c.count = count + 1; + self.counts = self.counts.max(count + 1); count } } @@ -449,8 +556,8 @@ impl Program { // Apply the hints scheduled before this instruction. if hint_at[pc as usize] != 0 { - let hs = hint_lists[hint_at[pc as usize] as usize - 1]; - for h in hs { + let (hs, first) = hint_lists[hint_at[pc as usize] as usize - 1]; + for (h, &stream) in hs.iter().zip(&hint_streams[first..]) { m.dbg_hint = Some(match h { RHint::FrameAddress { .. } => "FrameAddress", RHint::AllocFrames { .. } => "AllocFrames", @@ -524,7 +631,7 @@ impl Program { RHint::Print { label, cell } => { let c = fp + cell; if m.written(c) { - let v = m.cells[c as usize]; + let v = m.get(c); // Small integers and small g-powers overlap (8 = x^3 // = g^3): show every reading that applies. Only a // K-valued word (extension limbs 0) can be a g-power. @@ -550,17 +657,15 @@ impl Program { eprintln!("[print] {label} = "); } } - RHint::WitnessStack { name, base, len } => { - let values = - pop_witness(&self.witness, &mut witness_positions, name, *len).map_err(fail)?; + RHint::WitnessStack { base, len, .. } => { + let values = streams[stream as usize].pop(*len).map_err(fail)?; for (k, &value) in values.iter().enumerate() { m.put(fp + base + k as u32, value)?; } } - RHint::WitnessHeap { name, ptr, lo, len } => { + RHint::WitnessHeap { ptr, lo, len, .. } => { let b = heap_base(&m, &mut g, fp + ptr, "hint_witness heap pointer").map_err(fail)?; - let values = - pop_witness(&self.witness, &mut witness_positions, name, *len).map_err(fail)?; + let values = streams[stream as usize].pop(*len).map_err(fail)?; for (k, &value) in values.iter().enumerate() { m.put(b + lo + k as u32, value)?; } @@ -639,7 +744,7 @@ impl Program { let bytecode_read = { let v = bytecode_count[pc as usize]; - bytecode_count[pc as usize] = mul_by_g(v); + bytecode_count[pc as usize] = v + 1; v }; @@ -721,9 +826,9 @@ impl Program { // unwritten for good: only the validity of `a2` matters there. DerefMode::Cell => { if m.written(a2) { - m.put(a3, m.cells[a2 as usize])?; + m.put(a3, m.get(a2))?; } else if m.written(a3) { - m.put(a2, m.cells[a3 as usize])?; + m.put(a2, m.get(a3))?; } else { m.link(a2, a3); } @@ -894,20 +999,21 @@ impl Program { } } - let mut trace = Trace { + let counts = bytecode_count.iter().copied().max().unwrap_or(0).max(m.counts) as usize + 1; + let trace = Trace { xor, mul, set, deref, jump, blake2s, - mem_count: m.count, bytecode_count, + counts, }; // Pad memory to a power of two (the boundary tables read a dense image), at // least 2^MIN_LOG_MEM cells (doc §Memory), doubled until the stacked witness // reaches `min_log_committed`. - let mem_used = m.cells.len(); + let mem_used = m.used; let mut cells = mem_used.next_power_of_two().max(1 << MIN_LOG_MEM); let taus = super::layout::table_logs(trace.row_counts()); let log_bytecode = crate::log2_strict_usize(self.prog.len()); @@ -916,14 +1022,16 @@ impl Program { cells *= 2; } assert!(cells <= 1 << MAX_LOG_MEM, "data memory exceeds 2^{MAX_LOG_MEM} cells"); - m.cells.resize(cells, F192::ZERO); - trace.mem_count.resize(cells, F64::ONE); Ok(Execution { - mem: m.cells, + mem: Memory { + pages: m.pages, + len: cells, + }, cycles: steps, mem_used, unconstrained_reads: m.unwritten_reads, trace, + gpow: g.into_powers(), }) } } diff --git a/crates/leanvm_core/src/cpu/hints.rs b/crates/leanvm_core/src/cpu/hints.rs index e7dfaea99..a60f51722 100644 --- a/crates/leanvm_core/src/cpu/hints.rs +++ b/crates/leanvm_core/src/cpu/hints.rs @@ -70,6 +70,11 @@ impl GPow { self.pow.len() } + /// The forward table, `g^0` to `g^{covered - 1}`. + pub(crate) fn into_powers(self) -> Vec { + self.pow + } + /// Extend the forward table to cover index `upto`. pub fn grow_to(&mut self, upto: usize) { assert!(upto < MAX_CELLS as usize, "address space overflow (program too large)"); diff --git a/crates/leanvm_core/src/cpu/layout.rs b/crates/leanvm_core/src/cpu/layout.rs index 79cf56d2a..ff7927eac 100644 --- a/crates/leanvm_core/src/cpu/layout.rs +++ b/crates/leanvm_core/src/cpu/layout.rs @@ -233,71 +233,41 @@ pub fn bytecode_columns(prog: &[Op]) -> [Vec; 8] { let gpow = primitives::field::g_powers((max_op + 1).max(2)); let g_at = |i: u32| gpow[i as usize]; // operand g-power - let opcode = |op: &Op| match op { - Op::Xor { .. } => OP_XOR, - Op::Mul { .. } => OP_MUL, - Op::Set { .. } => OP_SET, - Op::Deref { .. } => OP_DEREF, - Op::Jump { .. } => OP_JUMP, - Op::Blake2s { .. } => OP_BLAKE2S, - }; - let operands = |op: &Op| -> (F64, F64, F64) { + // The opcode, then seven operand slots: offsets as g-powers, the `SET` immediate's + // limbs (c2 in the fpc slot), the `DEREF` store-mode flags, and `BLAKE2s`'s four + // input words, chaining-value base, output base and metadata cell. Unused slots are zero. + let row = |op: &Op| -> [F64; 8] { + let z = F64::ZERO; match *op { - Op::Xor { a, b, c } | Op::Mul { a, b, c } => (g_at(a), g_at(b), g_at(c)), - // The immediate's first two K-limbs ride operand slots o2/o3; c2 - // rides the fpc slot below. - Op::Set { o, k } => (g_at(o), F64(k.c0), F64(k.c1)), - Op::Deref { o1, o2, o3, .. } => (g_at(o1), g_at(o2), g_at(o3)), - Op::Jump { oc, od, of } => (g_at(oc), g_at(od), g_at(of)), - // BLAKE2s's first three input-word offsets; the last two ride the - // fpc/ffp bytecode slots below. - Op::Blake2s { ins, .. } => (g_at(ins[0]), g_at(ins[1]), g_at(ins[2])), + Op::Xor { a, b, c } => [OP_XOR, g_at(a), g_at(b), g_at(c), z, z, z, z], + Op::Mul { a, b, c } => [OP_MUL, g_at(a), g_at(b), g_at(c), z, z, z, z], + Op::Set { o, k } => [OP_SET, g_at(o), F64(k.c0), F64(k.c1), F64(k.c2), z, z, z], + Op::Deref { o1, o2, o3, mode } => [OP_DEREF, g_at(o1), g_at(o2), g_at(o3), mode.f_pc(), mode.f_fp(), z, z], + Op::Jump { oc, od, of } => [OP_JUMP, g_at(oc), g_at(od), g_at(of), z, z, z, z], + Op::Blake2s { ins, cv, out, md } => [ + OP_BLAKE2S, + g_at(ins[0]), + g_at(ins[1]), + g_at(ins[2]), + g_at(ins[3]), + g_at(cv), + g_at(out), + g_at(md), + ], } }; - // The 4th/5th bytecode operand slots: the two DEREF store-mode flags, or - // BLAKE2s's remaining input word / chaining-value base (0 elsewhere). - let fpc = |op: &Op| match op { - Op::Deref { mode, .. } => mode.f_pc(), - Op::Blake2s { ins, .. } => g_at(ins[3]), - Op::Set { k, .. } => F64(k.c2), - _ => F64::ZERO, - }; - let ffp = |op: &Op| match op { - Op::Deref { mode, .. } => mode.f_fp(), - Op::Blake2s { cv, .. } => g_at(*cv), - _ => F64::ZERO, - }; - // The 6th/7th bytecode operand slots: BLAKE2s's output base and its metadata - // cell (0 elsewhere). - let extra0 = |op: &Op| match op { - Op::Blake2s { out, .. } => g_at(*out), - _ => F64::ZERO, - }; - let extra1 = |op: &Op| match op { - Op::Blake2s { md, .. } => g_at(*md), - _ => F64::ZERO, - }; - // The program is PUBLIC (not committed): eight public columns over the - // program cube, embedded in the bytecode seed/finalize blocks below. - let column = |f: &(dyn Fn(&Op) -> F64 + Sync)| parallel::map_collect(prog.len(), |i| f(&prog[i])); - let prog_op: Vec = column(&opcode); - let prog_o1: Vec = column(&|o| operands(o).0); - let prog_o2: Vec = column(&|o| operands(o).1); - let prog_o3: Vec = column(&|o| operands(o).2); - let prog_fpc: Vec = column(&fpc); - let prog_ffp: Vec = column(&ffp); - let prog_extra0: Vec = column(&extra0); - let prog_extra1: Vec = column(&extra1); - [ - prog_op, - prog_o1, - prog_o2, - prog_o3, - prog_fpc, - prog_ffp, - prog_extra0, - prog_extra1, - ] + // The program is PUBLIC (not committed): eight public columns over the program + // cube, embedded in the bytecode seed/finalize blocks below, written in one pass. + // SAFETY: the pass below writes every slot of every column. + let mut columns: [Vec; 8] = std::array::from_fn(|_| unsafe { primitives::uninit_vec(prog.len()) }); + let dst = columns.each_mut().map(|c| parallel::SendPtr(c.as_mut_ptr())); + parallel::for_each(prog.len(), |i| { + for (column, value) in dst.iter().zip(row(&prog[i])) { + // SAFETY: instruction `i` writes slot `i` of each column, in bounds and once. + unsafe { column.add(i).write(value) }; + } + }); + columns } /// The stacked bytecode polynomial: the eight columns at their bus tuple @@ -443,10 +413,14 @@ impl Program { let log_mem = crate::log2_strict_usize(cells); let sch = schema(); - // Precompute g^0..g^{span-1} once so every address/pc/operand fill is an - // O(1) lookup instead of an O(log) power. - let span = cells.max(bytecode_size); - let gpow = primitives::field::g_powers(span); + // g^0..g^{span-1}, so every pc/frame/operand/count fill is an O(1) lookup instead + // of an O(log) power: a frame base or an operand is below the cells the run used. + // The run's own table covers them unless a count outgrew it. + let span = exec.mem_used.max(bytecode_size).max(tr.counts); + let gpow = match exec.gpow.get(..span) { + Some(gpow) => std::borrow::Cow::Borrowed(gpow), + None => std::borrow::Cow::Owned(primitives::field::g_powers(span)), + }; // The public layout (flush/count blocks, placements, boundary, taus) is a pure // function of the program + announced sizes + public input, with no committed @@ -458,7 +432,7 @@ impl Program { row_counts.iter().all(|&r| r <= 1 << MAX_LOG_ROWS), "a table exceeds 2^{MAX_LOG_ROWS} rows" ); - let pi = [exec.mem[0], exec.mem[1]]; + let pi = [exec.mem.word(0), exec.mem.word(1)]; let l = layout(&self.prog, log_mem, row_counts, pi); // The stacked witness is written exactly ONCE: allocate it, carve one window @@ -504,11 +478,41 @@ impl Program { // be written: the stack is uninitialized, so one left out would be read // as indeterminate bytes rather than caught by a length mismatch. const _: () = assert!(N_SHARED == 6, "a new shared column needs a fill here"); - parallel::fill(windows[MEM_LO], |i| F64(exec.mem[i].c0)); - parallel::fill(windows[MEM_HI], |i| F64(exec.mem[i].c1)); - parallel::fill(windows[MEM_TOP], |i| F64(exec.mem[i].c2)); - parallel::fill(windows[MFCNT], |i| tr.mem_count[i]); // counts ended at g^{A[i]} - parallel::fill(windows[BFCNT], |i| tr.bytecode_count[i]); // … at g^{A[pc]} + // One pass over the image fills all four of its columns, streamed out a + // cache line of each at a time: nothing reads them before the commitment. + const LINE: usize = 8; + const BLOCK: usize = 1 << 10; + let page_log = Memory::page_log(); + let image = [MEM_LO, MEM_HI, MEM_TOP, MFCNT].map(|c| { + assert_eq!(windows[c].len(), cells, "a memory column spans the image"); + parallel::SendPtr(windows[c].as_mut_ptr()) + }); + parallel::for_each_chunk(cells / BLOCK, |first, end| { + let stream = primitives::stream::Stream::new(); + for b in first..end { + let at = b * BLOCK; + let page = exec.mem.page(at >> page_log); + for line in (at..at + BLOCK).step_by(LINE) { + let mut buf = [[F64::ZERO; LINE]; 4]; + for j in 0..LINE { + let i = line + j; + let cell = page.map_or(&Cell::UNWRITTEN, |p| &p[i & ((1 << page_log) - 1)]); + // counts ended at g^{A[i]} + let values = [cell.word.c0, cell.word.c1, cell.word.c2].map(F64); + buf[0][j] = values[0]; + buf[1][j] = values[1]; + buf[2][j] = values[2]; + buf[3][j] = gpow[cell.count as usize]; + } + for (column, values) in image.iter().zip(&buf) { + // SAFETY: block `b` owns cells `b·BLOCK..(b+1)·BLOCK` of every + // column, in bounds as `cells` is a whole number of blocks. + unsafe { stream.copy_uninit(column.add(line), values) }; + } + } + } + }); + parallel::fill(windows[BFCNT], |i| gpow[tr.bytecode_count[i] as usize]); // … at g^{A[pc]} }); // flock's packed BLAKE2s witness q_flock, ALWAYS committed in this same stack: // built from the executed BLAKE2s rows in order (row j = flock instance j), @@ -519,17 +523,17 @@ impl Program { // words are the nine cells they read, in the finished (write-once) // memory image. let blocks: Vec<_> = parallel::map_collect(tr.blake2s.len(), |i| { - let r = &tr.blake2s[i]; + let r = tr.blake2s.get(i); let a = tables::blake2s_addresses(&self.prog, r); let chunk = |c0: u32, c1: u32| { - let (w0, w1) = (exec.mem[c0 as usize], exec.mem[c1 as usize]); + let (w0, w1) = (exec.mem.word(c0 as usize), exec.mem.word(c1 as usize)); [F64(w0.c0), F64(w0.c1), F64(w1.c0), F64(w1.c1)] }; crate::hash_flock::compression( chunk(a[0], a[1]), chunk(a[2], a[3]), chunk(a[4], a[4] + 1), - exec.mem[a[6] as usize], + exec.mem.word(a[6] as usize), ) }); crate::hash_flock::build_qflock_prepared(&blocks, windows[QFLOCK]) diff --git a/crates/leanvm_core/src/cpu/mod.rs b/crates/leanvm_core/src/cpu/mod.rs index 8c2b4bba6..164a9b7e2 100644 --- a/crates/leanvm_core/src/cpu/mod.rs +++ b/crates/leanvm_core/src/cpu/mod.rs @@ -27,10 +27,11 @@ pub mod hints; mod isa; pub mod layout; mod trace; -pub use execute::{ExecError, Execution, Fault}; +pub(crate) use execute::Cell; +pub use execute::{ExecError, Execution, Fault, Memory}; pub use isa::{DerefMode, Op}; pub use layout::*; -pub(crate) use trace::{Brow, Drow, Jrow, Srow, Trace, Xrow}; +pub(crate) use trace::{Brow, Drow, Jrow, Rows, Srow, Trace, Xrow}; /// Witness-gen `BLAKE2s` compression: the four message cells' eight /// words are laid out little-endian into 64 bytes, combined with the supplied @@ -343,7 +344,7 @@ fn airs( // 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 = leaf::BusForm::sum((0..3).map(|s| forms[s][t].scaled(form_pows[s]))); + let bus = RowForm::new(&leaf::BusForm::sum((0..3).map(|s| forms[s][t].scaled(form_pows[s])))); let bus_k = bus.clone(); constraints::Air { tau, @@ -363,6 +364,78 @@ fn airs( .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. @@ -807,8 +880,8 @@ mod tests { // The output cells hold the compression of the two inputs under the // pi-supplied chaining value (two 128-bit chunks). let d = blake2s_compress(a, b, cv_lanes(pi[0], pi[1]), md()); - assert_eq!(exec.mem[6], cell(d[0], d[1])); - assert_eq!(exec.mem[7], cell(d[2], d[3])); + assert_eq!(exec.mem.word(6), cell(d[0], d[1])); + assert_eq!(exec.mem.word(7), cell(d[2], d[3])); let (proof, _) = prove(&program, pi, crate::pcs::TEST_LOG_INV_RATE).unwrap(); verify(&program, &pi, &proof).expect("a hand-assembled program verifies"); @@ -859,8 +932,8 @@ mod tests { let exec = program.execute(pi).unwrap(); let d = blake2s_compress(h, h, cv_lanes(pi[0], pi[1]), md()); - assert_eq!(exec.mem[4], cell(d[0], d[1])); - assert_eq!(exec.mem[5], cell(d[2], d[3])); + assert_eq!(exec.mem.word(4), cell(d[0], d[1])); + assert_eq!(exec.mem.word(5), cell(d[2], d[3])); } /// A 192-bit-word MUL: the E-product of two full machine words. Full-limb @@ -879,6 +952,6 @@ mod tests { let program = Program::from_bytecode(prog, 5); let pi = [w(1), w(2)]; let exec = program.execute(pi).unwrap(); - assert_eq!(exec.mem[4], x * y, "MUL computes the E product"); + assert_eq!(exec.mem.word(4), x * y, "MUL computes the E product"); } } diff --git a/crates/leanvm_core/src/cpu/trace.rs b/crates/leanvm_core/src/cpu/trace.rs index d83ca703f..60f98a431 100644 --- a/crates/leanvm_core/src/cpu/trace.rs +++ b/crates/leanvm_core/src/cpu/trace.rs @@ -8,39 +8,40 @@ //! still holds what each accessed cell held at the time it was accessed. The //! rows are the interpreter's largest write stream, so what they do not carry //! they do not pay for, twice: once writing them and once reading them back. - -use primitives::field::F64; +//! +//! An access count is held as its exponent: the column holds `g^k`, the row `k`, +//! and the fill looks the power up. /// `XOR` or `MUL` row: the three cells are `fp·g^{a,b,c}`. pub(crate) struct Xrow { pub(crate) pc: u32, pub(crate) fp: u32, // frame base: address = fp + offset, operand = g^offset - pub(crate) ra: F64, - pub(crate) rb: F64, - pub(crate) rc: F64, - pub(crate) bytecode_read: F64, + pub(crate) ra: u32, + pub(crate) rb: u32, + pub(crate) rc: u32, + pub(crate) bytecode_read: u32, } pub(crate) struct Srow { pub(crate) pc: u32, pub(crate) fp: u32, - pub(crate) r: F64, - pub(crate) bytecode_read: F64, + pub(crate) r: u32, + pub(crate) bytecode_read: u32, } pub(crate) struct Drow { pub(crate) pc: u32, pub(crate) fp: u32, - pub(crate) r1: F64, - pub(crate) r2: F64, - pub(crate) r3: F64, - pub(crate) bytecode_read: F64, + pub(crate) r1: u32, + pub(crate) r2: u32, + pub(crate) r3: u32, + pub(crate) bytecode_read: u32, } pub(crate) struct Jrow { pub(crate) pc: u32, pub(crate) fp: u32, - pub(crate) rc: F64, - pub(crate) rd: F64, - pub(crate) rf: F64, - pub(crate) bytecode_read: F64, + pub(crate) rc: u32, + pub(crate) rd: u32, + pub(crate) rf: u32, + pub(crate) bytecode_read: u32, } /// `BLAKE2s` row: the nine per-cell memory access counts of the four @@ -51,23 +52,66 @@ pub(crate) struct Jrow { pub(crate) struct Brow { pub(crate) pc: u32, pub(crate) fp: u32, - pub(crate) ra: [F64; 2], // per-cell counts for the two a input cells - pub(crate) rb: [F64; 2], // … the two b input cells - pub(crate) rcv: [F64; 2], // … the two cv input cells - pub(crate) rc: [F64; 2], // … the two c output cells - pub(crate) rmd: F64, // … and the metadata cell - pub(crate) bytecode_read: F64, + pub(crate) ra: [u32; 2], // per-cell counts for the two a input cells + pub(crate) rb: [u32; 2], // … the two b input cells + pub(crate) rcv: [u32; 2], // … the two cv input cells + pub(crate) rc: [u32; 2], // … the two c output cells + pub(crate) rmd: u32, // … and the metadata cell + pub(crate) bytecode_read: u32, +} + +/// Rows in fixed-size chunks, so the trace grows without ever moving what it holds: +/// a `Vec` would copy every row written so far each time it doubled. +pub(crate) struct Rows { + chunks: Vec>, + len: usize, +} + +const CHUNK_LOG: usize = 14; +const CHUNK: usize = 1 << CHUNK_LOG; + +impl Rows { + pub(crate) fn new() -> Self { + Self { + chunks: Vec::new(), + len: 0, + } + } + + #[inline(always)] + pub(crate) fn push(&mut self, row: T) { + if self.len & (CHUNK - 1) == 0 { + self.chunks.push(Vec::with_capacity(CHUNK)); + } + // SAFETY: a chunk was pushed whenever the last one filled, so there is one. + unsafe { self.chunks.last_mut().unwrap_unchecked() }.push(row); + self.len += 1; + } + + pub(crate) fn len(&self) -> usize { + self.len + } + + #[inline(always)] + pub(crate) fn get(&self, i: usize) -> &T { + &self.chunks[i >> CHUNK_LOG][i & (CHUNK - 1)] + } + + pub(crate) fn iter(&self) -> impl DoubleEndedIterator { + self.chunks.iter().flatten() + } } pub(crate) struct Trace { - pub(crate) xor: Vec, - pub(crate) mul: Vec, - pub(crate) set: Vec, - pub(crate) deref: Vec, - pub(crate) jump: Vec, - pub(crate) blake2s: Vec, - pub(crate) mem_count: Vec, // per-cell running access count g^{count}; final = g^{A[i]} - pub(crate) bytecode_count: Vec, // per-pc running execution count g^{count}; final = g^{A[pc]} + pub(crate) xor: Rows, + pub(crate) mul: Rows, + pub(crate) set: Rows, + pub(crate) deref: Rows, + pub(crate) jump: Rows, + pub(crate) blake2s: Rows, + pub(crate) bytecode_count: Vec, // per-pc execution count; the column holds g^{A[pc]} + /// One past the highest access count, of a cell or a pc: the `g`-powers a fill looks up. + pub(crate) counts: usize, } impl Trace { diff --git a/crates/leanvm_core/src/gkr.rs b/crates/leanvm_core/src/gkr.rs index 5783d40a6..81ca15835 100644 --- a/crates/leanvm_core/src/gkr.rs +++ b/crates/leanvm_core/src/gkr.rs @@ -5,11 +5,16 @@ //! sumcheck has degree four. An odd-depth tree starts with one binary layer. //! Leaves and every layer are `E`-valued (the bus fingerprints mix `K`-columns //! into `E` upstream, [`crate::leaf`]). +//! +//! On AVX-512 the large rounds run eight pairs of rows at a time in transposed form +//! (`vertical`), and skip the runs of rows known to be the identity. + +use std::ops::Range; use crate::PAR_THRESHOLD; use crate::transcript::ProverState; use primitives::field::{F192, F192Unreduced, mul_unreduced4, mul2, mul4}; -use primitives::multilinear::{eq_table, interp, shrink_eq_low}; +use primitives::multilinear::{eq_table, interp}; use primitives::stream::Stream; /// Rows per parallel window: enough tasks to keep every thread fed, but not so few @@ -19,41 +24,106 @@ fn window_rows(total: usize) -> usize { total.div_ceil(tasks).clamp(64, 1 << 10) } +fn add4(a: [F192; 4], b: [F192; 4]) -> [F192; 4] { + std::array::from_fn(|i| a[i] + b[i]) +} + +/// `Σ_t f(t)` over `0..tasks`, in parallel when asked. +fn sum_tasks(tasks: usize, parallel: bool, f: impl Fn(usize) -> [F192; 4] + Sync) -> [F192; 4] { + if parallel { + parallel::map_reduce(tasks, || [F192::ZERO; 4], f, add4) + } else { + (0..tasks).map(f).fold([F192::ZERO; 4], add4) + } +} + +/// `out[i] = ∏ current[4i..4i+4]`, for `out.len() == current.len() / 4`. +pub(crate) fn products_into(current: &[F192], out: &mut [F192]) { + debug_assert_eq!(out.len(), current.len() / 4); + #[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] + let done = vertical::products_into(current, out); + #[cfg(not(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq")))] + let done = 0; + for (row, slot) in out.iter_mut().enumerate().skip(done) { + let [left, right] = mul2( + [current[4 * row], current[4 * row + 2]], + [current[4 * row + 1], current[4 * row + 3]], + ); + *slot = left * right; + } +} + +/// The level above `current`: the product of each four-tuple, an incomplete last one +/// padded with the identity. +fn next_level(current: &[F192]) -> Vec { + /// Rows per task, staged before they are published. + const ROWS: usize = 1 << 9; + let full_rows = current.len() / 4; + // SAFETY: the products below write `next[..full_rows]`, and the tail its last slot. + let mut next: Vec = unsafe { primitives::uninit_vec(current.len().div_ceil(4)) }; + if full_rows >= PAR_THRESHOLD { + // The next level and the GKR read the output long after this pass evicts it. + parallel::chunks_mut(&mut next[..full_rows], ROWS, |task, destination| { + let mut stage = [F192::ZERO; ROWS]; + let stage = &mut stage[..destination.len()]; + let start = 4 * ROWS * task; + products_into(¤t[start..start + 4 * stage.len()], stage); + Stream::new().copy(destination, stage); + }); + } else { + products_into(¤t[..4 * full_rows], &mut next[..full_rows]); + } + let tail = ¤t[4 * full_rows..]; + if !tail.is_empty() { + let child = |index| tail.get(index).copied().unwrap_or(F192::ONE); + let [left, right] = mul2([child(0), child(2)], [child(1), child(3)]); + next[full_rows] = left * right; + } + next +} + +/// A tree's leaves, with the two levels above them, which the leaf builder forms while +/// the leaves are in cache, and its runs of identity leaves. +pub struct Leaves { + pub values: Vec, + /// `fours[i] = ∏ values[4i..4i+4]`, and `sixteens` likewise over `fours`, an + /// incomplete last four-tuple padded with the identity. + pub fours: Vec, + pub sixteens: Vec, + /// Runs of identity leaves among `values`, sorted, which the GKR rounds above them + /// need neither read nor write. + pub ones: Vec>, +} + +impl Leaves { + pub fn new(values: Vec) -> Self { + let fours = next_level(&values); + let sixteens = next_level(&fours); + Self { + values, + fours, + sixteens, + ones: Vec::new(), + } + } +} + /// Build only the levels consumed by radix four: `0,2,4,…`, plus a final /// binary root when the logical depth is odd. -fn build_layers(leaves: Vec, mu: usize) -> Vec> { - assert!(!leaves.is_empty()); - assert!(leaves.len() <= 1usize << mu); +fn build_layers(leaves: Leaves, mu: usize) -> Vec> { // At mu = 22 the leaf level alone is hundreds of megabytes, and every level // dies with the proof. let mut layers: Vec> = (0..=mu).map(|_| Vec::new()).collect(); - layers[0] = leaves; + layers[0] = leaves.values; let mut level = 0; - while level + 2 <= mu { - let current = &layers[level]; - let full_rows = current.len() / 4; - let product = |row: usize| { - let [left, right] = mul2( - [current[4 * row], current[4 * row + 2]], - [current[4 * row + 1], current[4 * row + 3]], - ); - left * right - }; - let mut next: Vec = if current.len() == 1 { - Vec::from_iter([current[0]]) - } else if current.len() == 2 { - Vec::from_iter([current[0] * current[1]]) - } else if full_rows >= PAR_THRESHOLD { - primitives::par_collect(full_rows, product) - } else { - (0..full_rows).map(product).collect() - }; - if !current.len().is_multiple_of(4) && current.len() > 2 { - let row = full_rows; - let child = |index| current.get(4 * row + index).copied().unwrap_or(F192::ONE); - let [left, right] = mul2([child(0), child(2)], [child(1), child(3)]); - next.push(left * right); + for (given, below) in [(2, leaves.fours), (4, leaves.sixteens)] { + if given <= mu { + layers[given] = below; + level = given; } + } + while level + 2 <= mu { + let next = next_level(&layers[level]); level += 2; layers[level] = next; } @@ -67,6 +137,62 @@ fn build_layers(leaves: Vec, mu: usize) -> Vec> { layers } +/// `eq(vars, ·)` as the tensor product of a table over its low `lo_bits` variables +/// and one over the rest. A sum against it weights each term by the low factor and +/// each block of `2^lo_bits` terms by the high one, so the full table is never built, +/// nor marginalized round after round. +struct SplitEq { + lo: Vec, + hi: Vec, + lo_bits: usize, + /// `lo` eight entries at a time, for the vertical kernels. + #[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] + lo8: Vec, +} + +impl SplitEq { + fn new(vars: &[F192]) -> Self { + let lo_bits = vars.len().div_ceil(2).max(vars.len().min(3)); + let (low, high) = vars.split_at(lo_bits); + let lo = eq_table(low); + Self { + #[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] + lo8: lo + .as_chunks::<8>() + .0 + .iter() + .map(primitives::field::x8::F192x8::load) + .collect(), + lo, + hi: eq_table(high), + lo_bits, + } + } + + /// `Σ_{p ∈ range} eq(p)·f(p)`, where `term(p, w)` is `f(p)` already weighted by `w`, + /// the low factor of `eq(p)`. + #[inline] + fn sum(&self, range: Range, term: impl Fn(usize, F192) -> [F192Unreduced; 4]) -> [F192; 4] { + let mask = (1usize << self.lo_bits) - 1; + let mut total = [F192::ZERO; 4]; + let mut start = range.start; + while start < range.end { + let block = start >> self.lo_bits; + let end = range.end.min((block + 1) << self.lo_bits); + let block_sum = (start..end).fold([F192Unreduced::ZERO; 4], |mut acc, p| { + let t = term(p, self.lo[p & mask]); + for i in 0..4 { + acc[i] ^= t[i]; + } + acc + }); + total = add4(total, mul4([self.hi[block]; 4], block_sum.map(F192Unreduced::reduce))); + start = end; + } + total + } +} + #[inline(always)] fn quartic_summand(lines: [[F192; 2]; 4], equality: F192) -> [F192Unreduced; 4] { let [left0, left2, right0, right2] = mul4( @@ -88,15 +214,43 @@ fn quartic_summand(lines: [[F192; 2]; 4], equality: F192) -> [F192Unreduced; 4] mul_unreduced4([equality; 4], [c0 + at_one, c2, c3, c4]) } +/// The rows of a level `shift` binary levels up whose children are all in `ones`. +#[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] +fn coarsen(ones: &[Range], shift: usize) -> Vec> { + ones.iter() + .map(|r| r.start.div_ceil(1 << shift)..r.end >> shift) + .filter(|r| !r.is_empty()) + .collect() +} + +/// Whether `ones` covers `rows`. +#[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] +fn covered(ones: &[Range], rows: Range) -> bool { + let i = ones.partition_point(|r| r.start <= rows.start); + i > 0 && ones[i - 1].end >= rows.end +} + +/// Below this logical height a layer runs the scalar kernels. +#[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] +const VERTICAL_ROWS: usize = 1 << 10; + /// Two binary product levels contracted into one degree-four layer. struct QuaternaryLayerState { /// Four child tables interleaved in their original order. This lets the /// prover consume a product-tree level without first transposing it. values: Vec, - next: Vec, /// Logical row count after identity padding. `values` stores an arbitrary /// prefix; every omitted row is the constant four-tuple one. logical_rows: usize, + /// Once folded by the vertical kernels, the rows live here instead of in `values`. + #[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] + groups: Vec, + #[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] + spare: Vec, + /// Runs of rows known to be the identity, sorted: the vertical kernels skip + /// them, and leave their groups unwritten. + #[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] + ones: Vec>, } impl QuaternaryLayerState { @@ -106,63 +260,93 @@ impl QuaternaryLayerState { values.resize(4 * values.len().max(1).div_ceil(4), F192::ONE); debug_assert_eq!(values.len() % 4, 0); debug_assert!(values.len() <= 4 * width); - let rows = (values.len() / 4).div_ceil(2); Self { values, - // SAFETY: the first `fold` writes every slot of `next[..4 * rows]` before - // any read (its windows cover the full pairs, its tail block the odd row), - // and neither `round_message` nor `children` reads `next`. - next: unsafe { primitives::uninit_vec(4 * rows) }, logical_rows: width, + #[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] + groups: Vec::new(), + #[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] + spare: Vec::new(), + #[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] + ones: Vec::new(), } } /// `(q(0)+q(1), [X²]q, [X³]q, [X⁴]q)`. - fn round_message(&self, equality: &[F192]) -> [F192; 4] { + fn round_message(&self, eq: &SplitEq) -> [F192; 4] { + #[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] + if self.logical_rows >= VERTICAL_ROWS { + return vertical::message(&self.values, &self.ones, eq); + } + self.round_message_scalar(eq) + } + + fn fold(&mut self, challenge: F192) { + #[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] + self.leave_groups(); + self.fold_scalar(challenge); + } + + /// The vertical kernels' view of the rows, and the buffer the fold writes. + #[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] + fn split(&mut self) -> (vertical::Input<'_>, &mut Vec) { + let data = if self.groups.is_empty() { + vertical::Data::Rows(&self.values) + } else { + vertical::Data::Groups(&self.groups) + }; + (vertical::Input { data, ones: &self.ones }, &mut self.spare) + } + + /// Move the rows back into `values` for the scalar kernels. + #[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] + fn leave_groups(&mut self) { + if !self.groups.is_empty() { + self.values = vertical::to_rows(&self.groups, &self.ones); + self.groups = Vec::new(); + self.spare = Vec::new(); + } + } + + fn round_message_scalar(&self, eq: &SplitEq) -> [F192; 4] { let stored_rows = self.values.len() / 4; let full_pairs = stored_rows / 2; - let summand = |row: usize| -> [F192Unreduced; 4] { - let (lo, hi) = (8 * row, 8 * row + 4); + let summand = |pair: usize, weight: F192| -> [F192Unreduced; 4] { + let (lo, hi) = (8 * pair, 8 * pair + 4); let lines = [0, 1, 2, 3].map(|child| { let at_zero = self.values[lo + child]; [at_zero, at_zero + self.values[hi + child]] }); - quartic_summand(lines, equality[row]) - }; - let xor = |mut left: [F192Unreduced; 4], right: [F192Unreduced; 4]| { - for coefficient in 0..4 { - left[coefficient] ^= right[coefficient]; - } - left + quartic_summand(lines, weight) }; let rows = window_rows(full_pairs); - let window = |index: usize| -> [F192Unreduced; 4] { + let window = |index: usize| { let base = index * rows; - (base..(base + rows).min(full_pairs)).fold([F192Unreduced::ZERO; 4], |sum, row| xor(sum, summand(row))) - }; - let windows = full_pairs.div_ceil(rows); - let mut message = if full_pairs >= PAR_THRESHOLD { - parallel::map_reduce(windows, || [F192Unreduced::ZERO; 4], window, xor) - } else { - (0..windows).map(window).fold([F192Unreduced::ZERO; 4], xor) + eq.sum(base..(base + rows).min(full_pairs), summand) }; + let mut message = sum_tasks(full_pairs.div_ceil(rows), full_pairs >= PAR_THRESHOLD, window); if !stored_rows.is_multiple_of(2) { - let lo = 8 * full_pairs; - let lines = [0, 1, 2, 3].map(|child| { - let at_zero = self.values[lo + child]; - [at_zero, at_zero + F192::ONE] + let tail = eq.sum(full_pairs..full_pairs + 1, |pair, weight| { + let lo = 8 * pair; + let lines = [0, 1, 2, 3].map(|child| { + let at_zero = self.values[lo + child]; + [at_zero, at_zero + F192::ONE] + }); + quartic_summand(lines, weight) }); - message = xor(message, quartic_summand(lines, equality[full_pairs])); + message = add4(message, tail); } - message.map(F192Unreduced::reduce) + message } - fn fold(&mut self, challenge: F192) { + fn fold_scalar(&mut self, challenge: F192) { let stored_rows = self.values.len() / 4; let full_rows = stored_rows / 2; let rows = stored_rows.div_ceil(2); - self.next.truncate(4 * rows); - let (values, next) = (&self.values, &mut self.next); + // SAFETY: the windows below write every slot of `next[..4 * full_rows]` and the + // tail block the odd row, before anything reads it. + let mut next: Vec = unsafe { primitives::uninit_vec(4 * rows) }; + let values = &self.values; let fold_row = |row: usize| -> [F192; 4] { // One slice, not eight indexes: the bounds checks and the // index-by-24 multiplies fall out. @@ -211,19 +395,20 @@ impl QuaternaryLayerState { [challenge; 4], ); for child in 0..4 { - self.next[4 * full_rows + child] = self.values[lo + child] + folds[child]; + next[4 * full_rows + child] = self.values[lo + child] + folds[child]; } } - std::mem::swap(&mut self.values, &mut self.next); + self.values = next; self.logical_rows /= 2; } - fn fold_and_message(&mut self, challenge: F192, equality: &[F192]) -> [F192; 4] { + fn fold_and_message_scalar(&mut self, challenge: F192, eq: &SplitEq) -> [F192; 4] { let stored_rows = self.values.len() / 4; let rows = stored_rows.div_ceil(2); - self.next.truncate(4 * rows); + // SAFETY: the tasks below write every slot of `next[..4 * rows]` before anything reads it. + let mut next: Vec = unsafe { primitives::uninit_vec(4 * rows) }; let values = &self.values; - let dst = parallel::SendPtr(self.next.as_mut_ptr()); + let dst = parallel::SendPtr(next.as_mut_ptr()); const PAIRS: usize = 16; let pairs = rows.div_ceil(2); let task = |index: usize| { @@ -241,8 +426,7 @@ impl QuaternaryLayerState { stage[offset + i] = left[i] + product[i]; } } - let mut message = [F192Unreduced::ZERO; 4]; - for pair in first..end { + let message = eq.sum(first..end, |pair, weight| { let lo = 8 * (pair - first); let left = &stage[lo..lo + 4]; let right = if 2 * pair + 1 < rows { @@ -251,11 +435,8 @@ impl QuaternaryLayerState { &[F192::ONE; 4] }; let lines = std::array::from_fn(|i| [left[i], left[i] + right[i]]); - let terms = quartic_summand(lines, equality[pair]); - for i in 0..4 { - message[i] ^= terms[i]; - } - } + quartic_summand(lines, weight) + }); // The next round reads the destination; this round reads only the local stage. let stream = Stream::new(); let len = 4 * (end_row - 2 * first); @@ -263,21 +444,10 @@ impl QuaternaryLayerState { unsafe { stream.copy(dst.slice(8 * first, len), &stage[..len]) }; message }; - let xor = |mut a: [F192Unreduced; 4], b: [F192Unreduced; 4]| { - for i in 0..4 { - a[i] ^= b[i]; - } - a - }; - let tasks = pairs.div_ceil(PAIRS); - let message = if rows >= PAR_THRESHOLD { - parallel::map_reduce(tasks, || [F192Unreduced::ZERO; 4], task, xor) - } else { - (0..tasks).map(task).fold([F192Unreduced::ZERO; 4], xor) - }; - std::mem::swap(&mut self.values, &mut self.next); + let message = sum_tasks(pairs.div_ceil(PAIRS), rows >= PAR_THRESHOLD, task); + self.values = next; self.logical_rows /= 2; - message.map(F192Unreduced::reduce) + message } fn children(&self) -> [F192; 4] { @@ -287,6 +457,271 @@ impl QuaternaryLayerState { } } +/// Fold every tree by `challenge`, returning the next round's messages. +fn fold_and_message(trees: &mut [QuaternaryLayerState; 3], challenge: F192, eq: &SplitEq) -> [[F192; 4]; 3] { + #[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] + { + if trees[0].logical_rows >= VERTICAL_ROWS { + let [a, b, c] = trees.each_mut().map(QuaternaryLayerState::split); + let messages = vertical::fold([a.0, b.0, c.0], challenge, eq, [a.1, b.1, c.1]); + for tree in trees.iter_mut() { + std::mem::swap(&mut tree.groups, &mut tree.spare); + tree.values = Vec::new(); + tree.logical_rows /= 2; + tree.ones = coarsen(&tree.ones, 1); + } + return messages; + } + for tree in trees.iter_mut() { + tree.leave_groups(); + } + } + trees.each_mut().map(|tree| tree.fold_and_message_scalar(challenge, eq)) +} + +/// The AVX-512 kernels: eight pairs of rows at a time, transposed so that every +/// product is lane-wise ([`primitives::field::x8`]). After its first fold a layer +/// keeps its rows transposed, in [`Group`]s, until it is small enough for the scalar +/// kernels. +#[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] +mod vertical { + use super::{SplitEq, add4, coarsen, covered, sum_tasks}; + use primitives::field::F192; + use primitives::field::x8::{F192x8, Sum8}; + use primitives::stream::Stream; + use std::ops::Range; + + /// Eight pairs of rows: `g[4h + c]` lane `p` is child `c` of row `2p + h`. + pub(super) type Group = [F192x8; 8]; + + /// Groups per parallel task. + const TASK: usize = 32; + + /// Group `g` of a row table, the identity past its end. + #[inline(always)] + fn load(values: &[F192], g: usize) -> Group { + match values.get(64 * g..64 * g + 64) { + Some(run) => F192x8::load_runs(run.try_into().unwrap()), + None => { + let mut padded = [F192::ONE; 64]; + let tail = &values[(64 * g).min(values.len())..]; + padded[..tail.len()].copy_from_slice(tail); + F192x8::load_runs(&padded) + } + } + } + + /// Add eight pairs' `eq`-weighted summands. The weight scales child 0's line, and + /// with it every product, so it costs two products rather than one per coefficient. + #[inline(always)] + fn add_message(g: &Group, eq: F192x8, acc: &mut [Sum8; 4]) { + let a0 = eq * g[0]; + let d0 = eq * (g[0] + g[4]); + let left0 = a0 * g[1]; + let left2 = d0 * (g[1] + g[5]); + let left1 = (a0 + d0) * g[5] + left0 + left2; + let right0 = g[2] * g[3]; + let right2 = (g[2] + g[6]) * (g[3] + g[7]); + let right1 = g[6] * g[7] + right0 + right2; + let c0 = left0.mul_sum(right0); + let c4 = left2.mul_sum(right2); + let middle = left1.mul_sum(right1); + let at_one = (left0 + left1 + left2).mul_sum(right0 + right1 + right2); + let cross_even = (left0 + left2).mul_sum(right0 + right2); + let cross_high = (left1 + left2).mul_sum(right1 + right2); + acc[0] ^= c0 ^ at_one; + acc[1] ^= cross_even ^ c0 ^ c4 ^ middle; + acc[2] ^= cross_high ^ middle ^ c4; + acc[3] ^= c4; + } + + /// A task's message: group sums under one high `eq` factor, flushed when it changes. + struct Message<'a> { + eq: &'a SplitEq, + block: usize, + sums: [Sum8; 4], + total: [F192; 4], + } + + impl<'a> Message<'a> { + fn new(eq: &'a SplitEq, first: usize) -> Self { + Self { + eq, + block: first >> (eq.lo_bits - 3), + sums: [Sum8::ZERO; 4], + total: [F192::ZERO; 4], + } + } + + #[inline(always)] + fn add(&mut self, g: usize, group: &Group) { + let shift = self.eq.lo_bits - 3; + if g >> shift != self.block { + self.flush(); + self.block = g >> shift; + } + add_message(group, self.eq.lo8[g & ((1 << shift) - 1)], &mut self.sums); + } + + fn flush(&mut self) { + let sums = std::mem::replace(&mut self.sums, [Sum8::ZERO; 4]).map(Sum8::reduce); + let scaled = primitives::field::mul4([self.eq.hi[self.block]; 4], sums); + self.total = add4(self.total, scaled); + } + + fn finish(mut self) -> [F192; 4] { + self.flush(); + self.total + } + } + + /// The round message of a row table; identity pairs add nothing. + pub(super) fn message(values: &[F192], ones: &[Range], eq: &SplitEq) -> [F192; 4] { + let groups = values.len().div_ceil(64); + sum_tasks(groups.div_ceil(TASK), groups > TASK, |task| { + let range = task * TASK..((task + 1) * TASK).min(groups); + let mut message = Message::new(eq, range.start); + for g in range { + if !covered(ones, 16 * g..16 * g + 16) { + message.add(g, &load(values, g)); + } + } + message.finish() + }) + } + + /// Fold sixteen pairs, `a`'s then `b`'s, into the eight pairs their sixteen rows make. + #[inline(always)] + fn fold_pairs(a: &Group, b: &Group, r: F192x8) -> Group { + let fold = |g: &Group, c: usize| g[c] + r * (g[c] + g[4 + c]); + let mut out = [F192x8::splat(F192::ZERO); 8]; + for c in 0..4 { + (out[c], out[4 + c]) = F192x8::deinterleave(fold(a, c), fold(b, c)); + } + out + } + + pub(super) enum Data<'a> { + Rows(&'a [F192]), + Groups(&'a [Group]), + } + + /// A layer's rows as a fold reads them, with the runs known to be the identity. + pub(super) struct Input<'a> { + pub(super) data: Data<'a>, + pub(super) ones: &'a [Range], + } + + impl Input<'_> { + fn len(&self) -> usize { + match self.data { + Data::Rows(values) => values.len().div_ceil(64), + Data::Groups(groups) => groups.len(), + } + } + + /// Group `g`, the identity past the end or where `ones` says so. + #[inline(always)] + fn get(&self, g: usize) -> Group { + let ones = [F192x8::splat(F192::ONE); 8]; + if covered(self.ones, 16 * g..16 * g + 16) { + return ones; + } + match self.data { + Data::Rows(values) => load(values, g), + Data::Groups(groups) => groups.get(g).copied().unwrap_or(ones), + } + } + } + + /// Fold every tree's rows by `challenge` into its `out`, in one dispatch for the + /// three, returning the next round's messages. + pub(super) fn fold( + inputs: [Input; 3], + challenge: F192, + eq: &SplitEq, + outs: [&mut Vec; 3], + ) -> [[F192; 4]; 3] { + /// Groups below which a dispatch costs more than it saves. + const PARALLEL: usize = 8 * TASK; + let n_out = inputs.each_ref().map(|input| input.len().div_ceil(2)); + let mut outs = outs; + let dst: [parallel::SendPtr; 3] = std::array::from_fn(|t| { + let out = &mut *outs[t]; + if out.capacity() < n_out[t] { + *out = Vec::new(); + // SAFETY: the tasks below write every group before anything reads it. + *out = unsafe { primitives::uninit_vec(n_out[t]) }; + } else { + // SAFETY: as above; `Group` is plain data. + unsafe { out.set_len(n_out[t]) }; + } + parallel::SendPtr(out.as_mut_ptr()) + }); + let tasks: Vec<(usize, usize)> = (0..3) + .flat_map(|t| (0..n_out[t]).step_by(TASK).map(move |g| (t, g))) + .collect(); + let r = F192x8::splat(challenge); + let ones = inputs.each_ref().map(|input| coarsen(input.ones, 1)); + let task = |i: usize| { + let (t, first) = tasks[i]; + let stream = Stream::new(); + let mut message = Message::new(eq, first); + for g in first..(first + TASK).min(n_out[t]) { + // Identity rows fold to identity rows, which add nothing, and their + // group is left unwritten. + if covered(&ones[t], 16 * g..16 * g + 16) { + continue; + } + let pairs = fold_pairs(&inputs[t].get(2 * g), &inputs[t].get(2 * g + 1), r); + message.add(g, &pairs); + // The next round reads it, long after this round's pass evicted it. + // SAFETY: tasks own disjoint groups of the outputs. + unsafe { stream.copy(dst[t].slice(g, 1), std::slice::from_ref(&pairs)) }; + } + let mut messages = [[F192::ZERO; 4]; 3]; + messages[t] = message.finish(); + messages + }; + let add = |a: [[F192; 4]; 3], b: [[F192; 4]; 3]| std::array::from_fn(|t| add4(a[t], b[t])); + if n_out.iter().sum::() > PARALLEL { + parallel::map_reduce(tasks.len(), || [[F192::ZERO; 4]; 3], task, add) + } else { + (0..tasks.len()).map(task).fold([[F192::ZERO; 4]; 3], add) + } + } + + pub(super) fn to_rows(groups: &[Group], ones: &[Range]) -> Vec { + let mut values = vec![F192::ZERO; 64 * groups.len()]; + for (g, group) in groups.iter().enumerate() { + if covered(ones, 16 * g..16 * g + 16) { + values[64 * g..64 * g + 64].fill(F192::ONE); + continue; + } + for (j, lanes) in group.iter().enumerate() { + for (p, e) in lanes.to_array().into_iter().enumerate() { + values[64 * g + 8 * p + j] = e; + } + } + } + values + } + + /// `out[i] = ∏ current[4i..4i+4]` for the leading multiple of sixteen rows; returns how many. + pub(super) fn products_into(current: &[F192], out: &mut [F192]) -> usize { + let product = |a: F192x8, b: F192x8, c: F192x8, d: F192x8| (a * b) * (c * d); + let (chunks, _) = out.as_chunks_mut::<16>(); + for (slot, run) in chunks.iter_mut().zip(current.as_chunks::<64>().0) { + let g = F192x8::load_runs(run); + let (lo, hi) = F192x8::interleave(product(g[0], g[1], g[2], g[3]), product(g[4], g[5], g[6], g[7])); + let (first, second) = slot.split_at_mut(8); + lo.store(first.try_into().unwrap()); + hi.store(second.try_into().unwrap()); + } + 16 * chunks.len() + } +} + /// The result of a batched grand-product proof: the three roots and leaf /// evaluations, all reduced to one shared point. Under [`RootShape::FirstTwoShared`], /// `roots[0] == roots[1]` by construction rather than by a check. @@ -309,12 +744,17 @@ pub enum RootShape { } /// Prove three identity-padded grand products as one RLC-batched radix-four GKR. -pub fn prove_product_triple(leaves: [Vec; 3], ps: &mut ProverState, shape: RootShape) -> ProductTriple { - let mu = crate::log2_ceil_usize(leaves[0].len()); +pub fn prove_product_triple(leaves: [Leaves; 3], ps: &mut ProverState, shape: RootShape) -> ProductTriple { + let mu = crate::log2_ceil_usize(leaves[0].values.len()); assert!( - leaves.iter().all(|lane| !lane.is_empty() && lane.len() <= 1 << mu), + leaves + .iter() + .all(|lane| !lane.values.is_empty() && lane.values.len() <= 1 << mu), "batched trees must be nonempty prefixes of the first tree's logical tree" ); + // In leaves; a layer's rows over level `l` are runs of `2^(l + 2)` leaves. + #[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] + let identities = leaves.each_ref().map(|lane| lane.ones.clone()); let mut layers = leaves.map(|lane| build_layers(lane, mu)); let roots = [layers[0][mu][0], layers[1][mu][0], layers[2][mu][0]]; match shape { @@ -359,16 +799,18 @@ pub fn prove_product_triple(leaves: [Vec; 3], ps: &mut ProverState, shape: } let width = 1usize << round_count; + #[allow(unused_mut)] let mut trees = [0, 1, 2].map(|tree| QuaternaryLayerState::new(std::mem::take(&mut layers[tree][layer - 2]), width)); - let mut equality = if round_count > 0 { - eq_table(&point[1..]) - } else { - Vec::new() - }; + #[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] + for (tree, ones) in trees.iter_mut().zip(&identities) { + tree.ones = coarsen(ones, layer); + } let mut round_point = Vec::with_capacity(round_count); + // Round `j` weights its pairs by `eq(point[1 + j..], ·)`. let mut messages = if round_count > 0 { - trees.each_ref().map(|tree| tree.round_message(&equality)) + let eq = SplitEq::new(&point[1..]); + trees.each_ref().map(|tree| tree.round_message(&eq)) } else { [[F192::ZERO; 4]; 3] }; @@ -383,13 +825,12 @@ pub fn prove_product_triple(leaves: [Vec; 3], ps: &mut ProverState, shape: ps.add_scalars(&coeffs); let challenge = ps.sample(); round_point.push(challenge); - shrink_eq_low(&mut equality); if round + 1 < round_count { - messages = trees.each_mut().map(|tree| tree.fold_and_message(challenge, &equality)); + let eq = SplitEq::new(&point[2 + round..]); + messages = fold_and_message(&mut trees, challenge, &eq); } else { - // The last shrink exhausts `equality`, so the final round has no - // table to weight a message by and needs the fold alone. Both - // kernels stay for that reason; `fold` is not dead. + // The final round has no next message to weight, and needs the + // fold alone. Both kernels stay for that reason; `fold` is not dead. for tree in &mut trees { tree.fold(challenge); } @@ -434,17 +875,18 @@ mod tests { folded[0] } + fn element(i: usize, salt: usize) -> F192 { + F192::new((17 * i + salt + 1) as u64, (i * i + 3) as u64, (5 * i + 7) as u64) + } + #[test] fn quartic_round_message_matches_direct_evaluation() { - for width in [2, 4, 8, 16] { - let below: Vec = (0..4 * width) - .map(|i| F192::new((17 * i + width + 1) as u64, (i * i + 3) as u64, (5 * i + 7) as u64)) - .collect(); + for width in [2, 4, 8, 16, 1 << 11] { + let below: Vec = (0..4 * width).map(|i| element(i, width)).collect(); let state = QuaternaryLayerState::new(below, width); - let equality: Vec = (0..width / 2) - .map(|i| F192::new((31 * i + 5) as u64, (7 * i + 1) as u64, (11 * i + 9) as u64)) - .collect(); - let [difference, c2, c3, c4] = state.round_message(&equality); + let point: Vec = (0..width.ilog2() - 1).map(|i| element(31 * i as usize, 5)).collect(); + let equality = eq_table(&point); + let [difference, c2, c3, c4] = state.round_message(&SplitEq::new(&point)); let direct = |point: F192| { (0..width / 2).fold(F192::ZERO, |sum, row| { let values = [0, 1, 2, 3] @@ -463,37 +905,54 @@ mod tests { } } + /// The fused and the vertical kernels against the scalar fold and message, every round. #[test] fn fused_fold_matches_separate_fold_and_message() { - for width in [4usize, 16, 1 << 14] { - for len in [1, 4, 5, 7, 8, 9, 31, 32, 33, 4 * width - 5, 4 * width - 1, 4 * width] { - if len > 4 * width { - continue; - } - let values: Vec = (0..len) - .map(|i| F192::new((17 * i + 1) as u64, (i * i + 3) as u64, (5 * i + 7) as u64)) - .collect(); - let mut reference = QuaternaryLayerState::new(values.to_vec(), width); - let mut fused = QuaternaryLayerState::new(values, width); + for width in [4usize, 16, 1 << 11, 1 << 14] { + let lens = [1, 4, 5, 7, 8, 9, 31, 32, 33, 4 * width - 5, 4 * width - 1, 4 * width]; + for three in lens.windows(3).filter(|l| l.iter().all(|&n| n <= 4 * width)) { + let tree = |n: usize| QuaternaryLayerState::new((0..n).map(|i| element(i, n)).collect(), width); + let mut reference = [three[0], three[1], three[2]].map(tree); + let mut fused = [three[0], three[1], three[2]].map(tree); let point: Vec = (0..width.ilog2() - 1) .map(|i| F192::new(31 + u64::from(i), 7, 11)) .collect(); - let mut equality = eq_table(&point); - while reference.logical_rows > 2 { - let challenge = F192::new(reference.logical_rows as u64, 13, 19); - reference.fold(challenge); - shrink_eq_low(&mut equality); - let message = fused.fold_and_message(challenge, &equality); - assert_eq!(message, reference.round_message(&equality), "width={width}, len={len}"); - assert_eq!(&*fused.values, &*reference.values, "width={width}, len={len}"); + let eq = SplitEq::new(&point); + for (f, r) in fused.iter().zip(&reference) { + assert_eq!(f.round_message(&eq), r.round_message_scalar(&eq)); + } + let mut round = 0; + while reference[0].logical_rows > 2 { + round += 1; + let challenge = F192::new(reference[0].logical_rows as u64, 13, 19); + let eq = SplitEq::new(&point[round..]); + let messages = fold_and_message(&mut fused, challenge, &eq); + for (r, message) in reference.iter_mut().zip(messages) { + r.fold_scalar(challenge); + assert_eq!(message, r.round_message_scalar(&eq), "width={width}, lens={three:?}"); + } + } + for (f, r) in fused.iter_mut().zip(&mut reference) { + r.fold(F192::Y); + f.fold(F192::Y); + assert_eq!(f.children(), r.children()); } - reference.fold(F192::Y); - fused.fold(F192::Y); - assert_eq!(fused.children(), reference.children()); } } } + #[test] + fn level_products_match_scalar() { + for len in [4, 5, 63, 64, 65, 4 << 12, (4 << 12) + 3, (4 << 13) - 1] { + let current: Vec = (0..len).map(|i| element(i, 3)).collect(); + let want: Vec = current + .chunks(4) + .map(|c| c.iter().fold(F192::ONE, |acc, &v| acc * v)) + .collect(); + assert_eq!(next_level(¤t), want, "len={len}"); + } + } + #[test] fn radix_four_claims_are_true_at_even_and_odd_depths() { for mu in 0..=10 { @@ -506,7 +965,11 @@ mod tests { .each_ref() .map(|lane| lane.iter().copied().fold(F192::ONE, |product, value| product * value)); let mut ps = ProverState::from_label(b"radix-four-gkr-test"); - let proved = prove_product_triple(leaves.each_ref().map(|l| l.to_vec()), &mut ps, RootShape::Distinct); + let proved = prove_product_triple( + leaves.each_ref().map(|l| Leaves::new(l.to_vec())), + &mut ps, + RootShape::Distinct, + ); assert_eq!(proved.roots, expected_roots); for lane in 0..3 { assert_eq!(proved.values[lane], mle_eval_e(&leaves[lane], &proved.point)); @@ -530,7 +993,7 @@ mod tests { }); let mut sparse_ps = ProverState::from_label(b"sparse-radix-four-gkr-test"); let proved = prove_product_triple( - leaves.each_ref().map(|l| l.to_vec()), + leaves.each_ref().map(|l| Leaves::new(l.to_vec())), &mut sparse_ps, RootShape::Distinct, ); @@ -546,8 +1009,11 @@ mod tests { } let proof = sparse_ps.into_proof(); let mut dense_ps = ProverState::from_label(b"sparse-radix-four-gkr-test"); - let dense_proved = - prove_product_triple(dense.each_ref().map(|l| l.to_vec()), &mut dense_ps, RootShape::Distinct); + let dense_proved = prove_product_triple( + dense.each_ref().map(|l| Leaves::new(l.to_vec())), + &mut dense_ps, + RootShape::Distinct, + ); assert_eq!(dense_proved.roots, proved.roots); assert_eq!(dense_proved.point, proved.point); assert_eq!(dense_proved.values, proved.values); diff --git a/crates/leanvm_core/src/leaf.rs b/crates/leanvm_core/src/leaf.rs index 18cc7232d..563ff6a2a 100644 --- a/crates/leanvm_core/src/leaf.rs +++ b/crates/leanvm_core/src/leaf.rs @@ -7,13 +7,13 @@ //! 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::PAR_THRESHOLD; use crate::colval::ColVal; use crate::gkr; use crate::transcript::ProverState; -use primitives::field::{F64, F192, F192Unreduced, g_pow, index_mle}; -use primitives::multilinear::{eq_eval, eq_table, mle_eval}; +use primitives::field::{F64, F192, F192Unreduced, Weights8, dot_base, g_pow, index_mle, mul_by_g}; +use primitives::multilinear::{eq_eval, eq_table, mle_eval, mle_eval_par}; use std::collections::HashMap; +use std::ops::Range; use std::sync::Arc; /// One tuple coordinate as a function of the block's row `z`. @@ -211,85 +211,244 @@ fn push_terms<'a>(c: &'a Coord, w: F192, terms: &mut Vec>, constant: &m } } -/// Build one side's leaf vector: block `b` row `z` holds `β − Σ_i w_i c_i(z)` for -/// the fingerprint weights `w = eq(α⃗, ·)`, followed implicitly by the identity `1` +/// One block's leaf formula: `const_part + Σ terms(z)` on its real rows, the identity +/// on its padding rows. +struct LeafBlock<'a> { + side: usize, + offset: usize, + kappa: usize, + rows: usize, + const_part: F192, + terms: Vec>, +} + +impl LeafBlock<'_> { + /// The leaf of real row `z`, whose index column value `g^z` is `index`. + #[inline] + fn leaf(&self, z: usize, cols: &[&[F64]], index: F64) -> F192 { + // The α-weighted coordinate sum defers its reductions: each mixed + // product contributes its three raw limb products (3 PMULL, no + // reduction tail), one combined reduction per row at the end, + // bit-identical to summing reduced `mul_base` terms. + let mut acc = F192Unreduced::ZERO; + for t in &self.terms { + acc ^= match t { + Term::Col(i, c) => c.mul_base_unreduced(cols[*i][z]), + Term::Prod(i, j, c) => c.mul_base_unreduced(cols[*i][z] * cols[*j][z]), + Term::Index(c) => c.mul_base_unreduced(index), + Term::Public(vals, c) => c.mul_base_unreduced(vals[z]), + }; + } + self.const_part + acc.reduce() + } + + /// The leaves of rows `z0..z0 + out.len()`. + fn fill(&self, z0: usize, out: &mut [F192], cols: &[&[F64]]) { + let real = self.rows.saturating_sub(z0).min(out.len()); + let (real_out, padding) = out.split_at_mut(real); + // A padding row's leaf is the identity, so it leaves the product alone. + padding.fill(F192::ONE); + #[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] + let done = self.fill8(z0, real_out, cols); + #[cfg(not(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq")))] + let done = 0; + let mut index = g_pow(z0 + done); + for (i, slot) in real_out.iter_mut().enumerate().skip(done) { + *slot = self.leaf(z0 + i, cols, index); + index = mul_by_g(index); + } + } + + /// [`fill`](Self::fill) eight rows at a time, every term one lane-wise mixed + /// product ([`primitives::field::x8`]); returns how many rows it wrote. + #[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] + fn fill8(&self, z0: usize, out: &mut [F192], cols: &[&[F64]]) -> usize { + use primitives::field::x8::{F64x8, F192x8, Mixed8}; + let eight = |source: &[F64], z: usize| F64x8::load(source[z..z + 8].try_into().unwrap()); + // `g^z` for the eight rows, stepped by `g^8`. + let mut index = F64x8::new(std::array::from_fn(|i| g_pow(z0 + i))); + let step = F64x8::splat(g_pow(8)); + let coeffs: Vec = self + .terms + .iter() + .map(|t| match t { + Term::Col(_, c) | Term::Prod(_, _, c) | Term::Index(c) | Term::Public(_, c) => F192x8::splat(*c), + }) + .collect(); + let constant = F192x8::splat(self.const_part); + let (chunks, _) = out.as_chunks_mut::<8>(); + for (k, slot) in chunks.iter_mut().enumerate() { + let z = z0 + 8 * k; + let mut acc = Mixed8::ZERO; + for (t, &c) in self.terms.iter().zip(&coeffs) { + match t { + Term::Col(i, _) => acc.add(c, eight(cols[*i], z)), + Term::Prod(i, j, _) => acc.add(c, eight(cols[*i], z) * eight(cols[*j], z)), + Term::Index(_) => acc.add(c, index), + Term::Public(vals, _) => acc.add(c, eight(vals, z)), + } + } + index = index * step; + (acc.reduce() + constant).store(slot); + } + 8 * chunks.len() + } +} + +/// `dst[i] = ∏ src[4i..4i+4]`, an incomplete last four-tuple padded with the identity. +fn padded_products(src: &[F192], dst: &mut [F192]) { + let full = src.len() / 4; + gkr::products_into(&src[..4 * full], &mut dst[..full]); + if full < dst.len() { + dst[full] = src[4 * full..].iter().fold(F192::ONE, |acc, &x| acc * x); + } +} + +/// Build the three sides' leaves: block `b` row `z` holds `β − Σ_i w_i c_i(z)` for the +/// side's fingerprint weights `w = eq(α⃗, ·)`, followed implicitly by the identity `1` /// up to `2^μ`. The row-invariant weights and constant coordinates are folded once -/// per block into `const_part`. `gpow` supplies `g^z` for [`Coord::Index`] and is -/// the caller's, shared by the three sides. -pub fn build_leaves( - blocks: &[Block], - lay: &Layout, - cols: &[&[F64]], - w: &[F192], - beta: F192, - gpow: &[F64], -) -> Vec { - let explicit = blocks - .iter() - .enumerate() - .map(|(b, block)| lay.offsets[b] + (1usize << block.kappa)) - .max() - .unwrap_or(1); - debug_assert!(explicit <= 1usize << lay.mu); - // `stack_offsets` packs the power-of-two blocks contiguously from zero, so the - // blocks tile `0..explicit` and every slot below is written by one of them: the - // identity fill would be overwritten in full, and this is the largest buffer in - // the proof. The `covered` test is what licenses skipping it, so a layout that - // ever left a hole falls back to filling rather than reading uninitialized rows. - // Capacity is rounded to whole four-tuples because `gkr::QuaternaryLayerState` - // pads this level to that before reading it, and growing it here would copy it. - let covered: usize = blocks.iter().map(|blk| 1usize << blk.kappa).sum(); - let capacity = explicit.next_multiple_of(4); - let mut leaves = if covered == explicit { - // SAFETY: the per-block fills below cover `0..explicit` exactly, and each - // joins before this function returns. - let mut values = unsafe { primitives::uninit_vec(capacity) }; - values.truncate(explicit); - values - } else { - let mut values = Vec::with_capacity(capacity); - values.resize(explicit, F192::ONE); - values - }; - for (b, blk) in blocks.iter().enumerate() { - let mut const_part = beta; - let mut terms: Vec = Vec::with_capacity(blk.coords.len()); - for (i, c) in blk.coords.iter().enumerate() { - push_terms(c, w[i], &mut terms, &mut const_part); +/// per block into its `const_part`. +/// +/// The blocks of one height and row count (a table's push, pull and count blocks, +/// which read the same columns) are built together, a chunk of rows at a time, so a +/// column is read from memory once however many blocks flush it. Each chunk's +/// four- and sixteen-tuple products, the GKR's levels two and four, are formed while +/// it is in cache. +pub fn build_leaves(sides: [(&[Block], &Layout, &[F192], F192); 3], cols: &[&[F64]]) -> [gkr::Leaves; 3] { + /// Rows staged at a time: one block's chunk stays in L1 while its products are formed. + const CHUNK: usize = 1 << 11; + /// Leaves per task, whatever the group's block count, so that tasks balance. + const TASK_LEAVES: usize = 1 << 15; + /// Blocks below this height are left to a serial pass, being too small to own + /// whole sixteen-tuples. + const CHUNKED: usize = 4; + + let mut blocks: Vec = Vec::new(); + let mut lens = [1usize; 3]; + for (side, &(side_blocks, lay, w, beta)) in sides.iter().enumerate() { + // `stack_offsets` packs the power-of-two blocks contiguously from zero, + // largest first, so they tile `0..len`, and the blocks of height `CHUNKED` or + // more precede the smaller ones. + if !side_blocks.is_empty() { + lens[side] = side_blocks.iter().map(|blk| 1usize << blk.kappa).sum(); } - let row = |z: usize| -> F192 { - // The α-weighted coordinate sum defers its reductions: each mixed - // product contributes its three raw limb products (3 PMULL, no - // reduction tail), one combined reduction per row at the end, - // bit-identical to summing reduced `mul_base` terms. - let mut acc = F192Unreduced::ZERO; - for t in &terms { - acc ^= match t { - Term::Col(i, c) => c.mul_base_unreduced(cols[*i][z]), - Term::Prod(i, j, c) => c.mul_base_unreduced(cols[*i][z] * cols[*j][z]), - Term::Index(c) => c.mul_base_unreduced(gpow[z]), - Term::Public(vals, c) => c.mul_base_unreduced(vals[z]), - }; + for (b, blk) in side_blocks.iter().enumerate() { + assert!( + lay.offsets[b] + (1 << blk.kappa) <= lens[side], + "the blocks tile the leaves" + ); + let mut const_part = beta; + let mut terms: Vec = Vec::with_capacity(blk.coords.len()); + for (i, c) in blk.coords.iter().enumerate() { + push_terms(c, w[i], &mut terms, &mut const_part); } - const_part + acc.reduce() - }; - let off = lay.offsets[b]; - let (real, padding) = leaves[off..off + (1usize << blk.kappa)].split_at_mut(blk.rows); - if real.len() >= PAR_THRESHOLD { - parallel::fill(real, row); - } else { - for (z, slot) in real.iter_mut().enumerate() { - *slot = row(z); + blocks.push(LeafBlock { + side, + offset: lay.offsets[b], + kappa: blk.kappa, + rows: blk.rows, + const_part, + terms, + }); + } + debug_assert!(lens[side] <= 1usize << lay.mu); + } + + // Capacity is rounded to whole four-tuples because `gkr::QuaternaryLayerState` + // pads this level to that before reading it, and growing it here would copy it. + let mut values: [Vec; 3] = lens.map(|n| { + // SAFETY: the blocks write every leaf of `0..n` (one identity leaf for a side + // without blocks, below) before anything reads it. + let mut v = unsafe { primitives::uninit_vec(n.next_multiple_of(4)) }; + v.truncate(n); + v + }); + // SAFETY: the chunks write the products of the tuples inside blocks of height + // `CHUNKED` or more, and the serial pass below the rest. + let mut fours: [Vec; 3] = lens.map(|n| unsafe { primitives::uninit_vec(n.div_ceil(4)) }); + let mut sixteens: [Vec; 3] = lens.map(|n| unsafe { primitives::uninit_vec(n.div_ceil(16)) }); + + let mut shapes: Vec<(usize, usize)> = blocks + .iter() + .filter(|b| b.kappa >= CHUNKED) + .map(|b| (b.kappa, b.rows)) + .collect(); + shapes.sort_unstable(); + shapes.dedup(); + let groups: Vec> = shapes + .iter() + .map(|&shape| blocks.iter().filter(|b| (b.kappa, b.rows) == shape).collect()) + .collect(); + let tasks: Vec<(usize, Range)> = shapes + .iter() + .zip(&groups) + .enumerate() + .flat_map(|(g, (&(kappa, _), group))| { + let height = 1usize << kappa; + let rows = (TASK_LEAVES / group.len()).next_power_of_two().min(height); + (0..height).step_by(rows).map(move |z| (g, z..z + rows)) + }) + .collect(); + let ptrs = [&mut values, &mut fours, &mut sixteens] + .map(|level| level.each_mut().map(|v| parallel::SendPtr(v.as_mut_ptr()))); + parallel::for_each(tasks.len(), |task| { + let (group, ref rows) = tasks[task]; + let mut stage = [F192::ZERO; CHUNK]; + let mut stage4 = [F192::ZERO; CHUNK / 4]; + let mut stage16 = [F192::ZERO; CHUNK / 16]; + // The GKR reads all three long after this pass has evicted them. + let stream = primitives::stream::Stream::new(); + for z0 in rows.clone().step_by(CHUNK) { + let len = CHUNK.min(rows.len()); + let (leaves, fours, sixteens) = (&mut stage[..len], &mut stage4[..len / 4], &mut stage16[..len / 16]); + for blk in &groups[group] { + blk.fill(z0, leaves, cols); + gkr::products_into(leaves, fours); + gkr::products_into(fours, sixteens); + let at = blk.offset + z0; + // SAFETY: a block's chunks are disjoint ranges of its side's leaves, blocks + // are disjoint, and a chunk owns the tuples it starts, `at` being a multiple + // of sixteen. + unsafe { + stream.copy(ptrs[0][blk.side].slice(at, len), leaves); + stream.copy(ptrs[1][blk.side].slice(at / 4, len / 4), fours); + stream.copy(ptrs[2][blk.side].slice(at / 16, len / 16), sixteens); + } } } - // A padding row's leaf is the identity, so it leaves the product alone. - if padding.len() >= PAR_THRESHOLD { - parallel::fill(padding, |_| F192::ONE); - } else { - padding.fill(F192::ONE); + }); + + for (side, v) in values.iter_mut().enumerate() { + let mut tail = lens[side]; + for blk in blocks.iter().filter(|b| b.side == side && b.kappa < CHUNKED) { + blk.fill(0, &mut v[blk.offset..blk.offset + (1 << blk.kappa)], cols); + tail = tail.min(blk.offset); } + if sides[side].0.is_empty() { + v[0] = F192::ONE; + } + // The tuples among the smallest blocks, the last ones padded with the identity. + padded_products(&v[tail..], &mut fours[side][tail / 4..]); + padded_products(&fours[side][tail / 4..], &mut sixteens[side][tail / 16..]); } - leaves + let [v0, v1, v2] = values; + let [f0, f1, f2] = fours; + let [s0, s1, s2] = sixteens; + [(0, v0, f0, s0), (1, v1, f1, s1), (2, v2, f2, s2)].map(|(side, values, fours, sixteens)| { + // The padding rows of the chunked blocks, by offset. + let mut ones: Vec<_> = blocks + .iter() + .filter(|blk| blk.side == side && blk.kappa >= CHUNKED && blk.rows < 1 << blk.kappa) + .map(|blk| blk.offset + blk.rows..blk.offset + (1 << blk.kappa)) + .collect(); + ones.sort_unstable_by_key(|r| r.start); + gkr::Leaves { + values, + fours, + sixteens, + ones, + } + }) } /// One table's bus contribution on one side, as a form over that table's committed @@ -322,8 +481,7 @@ impl BusForm { /// The same form scaled by `w`. Every coefficient is `E`-valued already, so /// folding the side's `η`-power in here costs three multiplies once per table - /// instead of one per [`eval_unreduced`](Self::eval_unreduced), which the - /// zerocheck calls per row per round. + /// instead of one per row and round of the zerocheck. pub fn scaled(&self, w: F192) -> Self { Self { coeffs: self.coeffs.iter().map(|&c| c * w).collect(), @@ -358,22 +516,6 @@ impl BusForm { out } - /// The form at one point, unreduced: `evals` are the columns' values there. - /// This is what the zerocheck evaluates, per row while a table is unfolded and - /// at the sumcheck point after, so a table's several forms share one reduction - /// rather than paying one per term. - /// `quadratic` selects only degree-two terms, for a sumcheck round coefficient. - pub fn eval_unreduced(&self, evals: &[T], quadratic: bool) -> T::Unreduced { - self.prods.iter().fold( - if quadratic { - T::lift(F192::ZERO) - } else { - T::dot_unreduced(&self.coeffs, evals) ^ T::lift(self.constant) - }, - |acc, &(a, b, c)| acc ^ (evals[a] * evals[b]).mul_e_unreduced(c), - ) - } - /// 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 @@ -559,10 +701,12 @@ fn decompose_prove( } } } - let vals: Vec = parallel::map_collect(jobs.len(), |i| { - let (col, kappa) = jobs[i]; - mle_eval(cols[col], &zeta[..kappa]) - }); + // One evaluation at a time, each across the whole pool: a few long columns, one + // per worker, would leave the rest of it idle. + let vals: Vec = jobs + .iter() + .map(|&(col, kappa)| mle_eval_par(cols[col], &zeta[..kappa])) + .collect(); // Pass 2: replay in the original order; duplicates reuse the recorded claim. let mut fresh_iter = jobs.iter().zip(vals.iter()); @@ -678,36 +822,24 @@ pub fn prove_balance( let w = fingerprint_weights(&alphas); let count_w = fingerprint_weights(&[F192::ZERO; N_TUPLE_BITS]); let beta = ps.sample(); - // The `g^z` table backing `Coord::Index`, built once for the three sides: push - // and pull would otherwise build the same table twice and count, which carries - // no `Index` coordinate, would build one it never reads. - let index_k = [push, pull, count] - .into_iter() - .flatten() - .filter(|b| b.coords.iter().any(|c| matches!(c, Coord::Index))) - .map(|b| b.kappa) - .max(); - let gpow = index_k.map_or_else(Vec::new, |k| primitives::field::g_powers(1usize << k)); - // Three independent leaf vectors, built one after another: each `build_leaves` - // already fans its own blocks out across the whole pool, so nesting a - // three-way outer split on top would only add a barrier. - let [push_leaves, pull_leaves, count_leaves] = crate::stage!("Bus leaves", || { - [ - build_leaves(push, &push_lay, cols, &w, beta, &gpow), - build_leaves(pull, &pull_lay, cols, &w, beta, &gpow), - build_leaves(count, &count_lay, cols, &count_w, F192::ZERO, &gpow), - ] - }); - // Leaf construction keeps the all-one padding implicit; decomposition uses the full logical depth. + // The leaves keep the all-one padding implicit; decomposition uses the full logical depth. count_lay.mu = push_lay.mu; + let leaves = crate::stage!("Bus leaves", || { + build_leaves( + sides( + [push, pull, count], + [&push_lay, &pull_lay, &count_lay], + &w, + &count_w, + beta, + ), + cols, + ) + }); // All three trees run as ONE RLC-batched GKR (equal μ: push/pull match // block-for-block, count is padded), so every claim lands on ONE point ζ. let bus_gkr = crate::stage!("Bus GKR", || { - gkr::prove_product_triple( - [push_leaves, pull_leaves, count_leaves], - ps, - gkr::RootShape::FirstTwoShared, - ) + gkr::prove_product_triple(leaves, ps, gkr::RootShape::FirstTwoShared) }); // Framework blocks keep their per-column claims (deduped: push/pull share ζ); @@ -746,7 +878,12 @@ pub fn prove_balance( ); } }); - let (table_evals, prod_sums) = tables_and_prods_at(cols, tables, &forms, &bus_gkr.point); + let (table_evals, prod_sums) = crate::stage!("Bus tables", || tables_and_prods_at( + cols, + tables, + &forms, + &bus_gkr.point + )); let sigmas: [Vec; 3] = std::array::from_fn(|s| { let sigmas: Vec = forms[s] .iter() @@ -772,17 +909,69 @@ pub fn prove_balance( } } -/// Every table's committed columns at `ζ[..τ_t]`, and, for every column pair its -/// forms multiply, `Σ_z eq(ζ[..τ], z)·col_a(z)·col_b(z)`. +/// `Σ_z eq(point, z)·c(z)` for every column `c`, then `Σ_z eq(point, z)·a(z)·b(z)` for +/// every pair `(a, b)`, all in one parallel pass over the rows. /// -/// One eq table per table, streamed once past every column and every pair at -/// the same time. Evaluated apart these are `n_cols + n_pairs` fold ladders -/// over the same table at the same point, and a ladder lifts every `K` word it -/// reads into an `E` it writes and reads again, where a dot against the weights -/// moves the column's own eight bytes. Pairs are deduped across the three sides -/// and the several blocks that carry the same address, so an address costs one -/// pass however often it is flushed. `tables[t] = (base, n_cols)` in the global -/// schema; a pair names LOCAL column indices. +/// `eq` factors into a table over the low variables, packed for [`dot_base`] and +/// held in L1, and one over the rest that weights each block's reduced inner sums, +/// so every column moves only its own eight bytes a row and no `E` table of its +/// length is ever written. +fn eq_sums(point: &[F192], columns: &[&[F64]], pairs: &[(&[F64], &[F64])]) -> Vec { + const LOW: usize = 10; + if point.len() < 3 { + let eq = eq_table(point); + let dot = |f: &dyn Fn(usize) -> F64| { + eq.iter() + .enumerate() + .fold(F192::ZERO, |acc, (z, &w)| acc + w.mul_base(f(z))) + }; + return columns + .iter() + .map(|c| dot(&|z| c[z])) + .chain(pairs.iter().map(|(a, b)| dot(&|z| a[z] * b[z]))) + .collect(); + } + let low_vars = point.len().min(LOW); + let low: Vec = eq_table(&point[..low_vars]) + .as_chunks::<8>() + .0 + .iter() + .map(Weights8::new) + .collect(); + let high = eq_table(&point[low_vars..]); + let width = 1usize << low_vars; + parallel::fold_reduce( + high.len(), + || vec![F192::ZERO; columns.len() + pairs.len()], + |acc, row| { + let span = row * width..(row + 1) * width; + let (singles, products) = acc.split_at_mut(columns.len()); + for (slot, c) in singles.iter_mut().zip(columns) { + *slot += high[row] * dot_base(&low, &c[span.clone()]).reduce(); + } + let mut product = [F64::ZERO; 1 << LOW]; + for (slot, (a, b)) in products.iter_mut().zip(pairs) { + for ((p, &x), &y) in product.iter_mut().zip(&a[span.clone()]).zip(&b[span.clone()]) { + *p = x * y; + } + *slot += high[row] * dot_base(&low, &product[..width]).reduce(); + } + }, + |mut left, right| { + for (slot, part) in left.iter_mut().zip(right) { + *slot += part; + } + left + }, + ) +} + +/// Every table's committed columns at `ζ[..τ_t]`, and, for every column pair its +/// forms multiply, `Σ_z eq(ζ[..τ], z)·col_a(z)·col_b(z)`, in one pass per table +/// ([`eq_sums`]). Pairs are deduped across the three sides and the several blocks +/// that carry the same address, so an address costs one product however often it is +/// flushed. `tables[t] = (base, n_cols)` in the global schema; a pair names LOCAL +/// column indices. #[allow(clippy::type_complexity)] fn tables_and_prods_at( cols: &[&[F64]], @@ -790,10 +979,6 @@ fn tables_and_prods_at( forms: &[Vec; 3], zeta: &[F192], ) -> (Vec>, Vec>) { - /// Rows per task: the eq slice a task reads stays in L2 while every column - /// and pair accumulator sweeps past it. - const ROWS: usize = 1 << 12; - tables .iter() .enumerate() @@ -805,48 +990,133 @@ fn tables_and_prods_at( .collect(); pairs.sort_unstable(); pairs.dedup(); - - let eq = eq_table(&zeta[..tau]); - let n_acc = n_cols + pairs.len(); - let sums = parallel::fold_reduce( - (1usize << tau).div_ceil(ROWS), - || vec![F192Unreduced::ZERO; n_acc], - |acc, chunk| { - let lo = chunk * ROWS; - let weights = &eq[lo..(lo + ROWS).min(1 << tau)]; - let span = |c: usize| &cols[base + c][lo..lo + weights.len()]; - for (c, slot) in acc[..n_cols].iter_mut().enumerate() { - for (&w, &v) in weights.iter().zip(span(c)) { - *slot ^= w.mul_base_unreduced(v); - } - } - for (&(a, b), slot) in pairs.iter().zip(&mut acc[n_cols..]) { - for ((&w, &x), &y) in weights.iter().zip(span(a)).zip(span(b)) { - *slot ^= w.mul_base_unreduced(x * y); - } - } - }, - |mut left, right| { - for (slot, part) in left.iter_mut().zip(right) { - *slot ^= part; - } - left - }, - ); - let evals = sums[..n_cols].iter().map(|s| s.reduce()).collect(); + let columns = &cols[base..base + n_cols]; + let products: Vec<(&[F64], &[F64])> = pairs.iter().map(|&(a, b)| (columns[a], columns[b])).collect(); + let mut sums = eq_sums(&zeta[..tau], columns, &products); let prods = pairs .iter() .zip(&sums[n_cols..]) - .map(|(&(a, b), s)| (a, b, s.reduce())) + .map(|(&(a, b), &s)| (a, b, s)) .collect(); - (evals, prods) + sums.truncate(n_cols); + (sums, prods) }) .unzip() } #[cfg(test)] mod tests { - use super::soundness_bits; + use super::*; + use crate::gkr::{Leaves, RootShape, prove_product_triple}; + + /// A coordinate's value at row `z`, from its definition. + fn coord_value(c: &Coord, cols: &[&[F64]], z: usize) -> F64 { + match c { + Coord::Const(v) => *v, + Coord::Col(i) => cols[*i][z], + Coord::GCol(i, k) => g_pow(*k as usize) * cols[*i][z], + Coord::Prod(i, j, k) => g_pow(*k as usize) * cols[*i][z] * cols[*j][z], + Coord::Index => g_pow(z), + Coord::Public(vals) => vals[z], + Coord::Sum(cs) => cs.iter().fold(F64::ZERO, |acc, c| acc + coord_value(c, cols, z)), + } + } + + /// The leaves by their definition, every block at its offset. + fn naive_leaves(blocks: &[Block], lay: &Layout, cols: &[&[F64]], w: &[F192], beta: F192) -> Vec { + let len = blocks.iter().map(|b| 1 << b.kappa).sum::().max(1); + let mut leaves = vec![F192::ONE; len]; + for (b, blk) in blocks.iter().enumerate() { + for z in 0..blk.rows { + leaves[lay.offsets[b] + z] = blk + .coords + .iter() + .zip(w) + .fold(beta, |acc, (c, &wi)| acc + wi.mul_base(coord_value(c, cols, z))); + } + } + leaves + } + + /// The builder's leaves and products are their definitions, and its identity runs + /// leave the proof as it is. + #[test] + fn built_leaves_match_the_definition() { + let mut seed = 0x9E37_79B9_7F4A_7C15u64; + let mut next = move || { + seed ^= seed << 13; + seed ^= seed >> 7; + seed ^= seed << 17; + seed + }; + let columns: Vec> = (0..6).map(|_| (0..1 << 12).map(|_| F64(next())).collect()).collect(); + let cols: Vec<&[F64]> = columns.iter().map(Vec::as_slice).collect(); + let public = Arc::new((0..1 << 11).map(|_| F64(next())).collect::>()); + let block = |kappa: usize, rows: usize, coords: Vec| Block { + kappa, + rows, + coords, + padding: Vec::new(), + }; + use Coord::*; + let wide = |count: Coord| { + vec![ + Const(F64(5)), + Index, + count, + Prod(0, 2, 1), + Sum(vec![Col(4), GCol(5, 3)]), + ] + }; + let narrow = |count: Coord| vec![Const(F64(7)), Col(3), count, Public(public.clone())]; + let push = vec![ + block(12, 3001, wide(GCol(1, 1))), + block(11, 1 << 11, narrow(Col(2))), + block(2, 3, vec![Const(F64(1)), Col(4)]), + block(0, 1, vec![Const(F64(2))]), + ]; + let pull = vec![ + block(12, 3001, wide(Col(1))), + block(11, 1 << 11, narrow(GCol(2, 1))), + block(1, 2, vec![Const(F64(3)), Col(0)]), + ]; + let count = vec![block(12, 3001, vec![Col(1)]), block(11, 1 << 11, vec![Col(2)])]; + let lays = [&push, &pull, &count].map(|b| layout(b)); + let w = fingerprint_weights(&[F192::new(3, 1, 4), F192::new(1, 5, 9), F192::new(2, 6, 5), F192::Y]); + let count_w = fingerprint_weights(&[F192::ZERO; N_TUPLE_BITS]); + let beta = F192::new(27, 18, 28); + let sides = sides( + [&push, &pull, &count], + [&lays[0], &lays[1], &lays[2]], + &w, + &count_w, + beta, + ); + let want: [Vec; 3] = sides.map(|(blocks, lay, w, beta)| naive_leaves(blocks, lay, &cols, w, beta)); + + let built = build_leaves(sides, &cols); + for (got, want) in built.iter().zip(&want) { + let reference = Leaves::new(want.clone()); + assert_eq!( + (&got.values, &got.fours, &got.sixteens), + (&reference.values, &reference.fours, &reference.sixteens) + ); + assert!( + got.ones + .iter() + .flat_map(Clone::clone) + .all(|i| got.values[i] == F192::ONE) + ); + } + assert!(built.iter().all(|leaves| !leaves.ones.is_empty())); + let mut proofs = [ProverState::from_label(b"leaves"), ProverState::from_label(b"leaves")]; + let [ps0, ps1] = &mut proofs; + let planned = prove_product_triple(built, ps0, RootShape::Distinct); + let reference = prove_product_triple(want.map(Leaves::new), ps1, RootShape::Distinct); + assert_eq!((planned.point, planned.values), (reference.point, reference.values)); + let [a, b] = proofs.map(|ps| ps.into_proof().stream); + assert_eq!(a, b); + } /// The bound is `(N_TUPLE_BITS + 1)·2^mu` plus the GKR terms: only the bus /// DEPTH costs bits now, the multilinear fingerprint having fixed each factor's diff --git a/crates/leanvm_core/src/tables.rs b/crates/leanvm_core/src/tables.rs index a46eafd6f..0b1824530 100644 --- a/crates/leanvm_core/src/tables.rs +++ b/crates/leanvm_core/src/tables.rs @@ -16,7 +16,7 @@ //! `eval_constraint` takes. use crate::colval::ColVal; -use crate::cpu::{Brow, Drow, Jrow, Op, Srow, Trace}; +use crate::cpu::{Brow, Drow, Jrow, Memory, Op, Rows, Srow, Trace}; use crate::leaf::Coord::{self, Col, Const, GCol, Prod}; use primitives::field::{F64, F192, mul_by_g}; @@ -192,7 +192,7 @@ impl FlushBuilder { /// image (for read values), and `g^0..` for O(1) address/operand lookups. pub struct FillCtx<'a> { pub(crate) trace: &'a Trace, - pub(crate) mem: &'a [F192], + pub(crate) mem: &'a Memory, pub(crate) gpow: &'a [F64], pub(crate) prog: &'a [Op], /// This table's height `2^tau`, the length of every window in `out`. Rows past the @@ -212,7 +212,7 @@ pub type ColumnOut<'a> = &'a mut [F64]; impl<'a> FillCtx<'a> { pub(crate) fn new( trace: &'a Trace, - mem: &'a [F192], + mem: &'a Memory, gpow: &'a [F64], prog: &'a [Op], height: usize, @@ -246,7 +246,7 @@ impl<'a> FillCtx<'a> { /// Write local column `at`: `f` over the trace rows, then its padding value to the /// end of the window. - fn col(&self, out: &mut [ColumnOut], rows: &[R], at: usize, f: impl Fn(&R) -> F64 + Sync) { + fn col(&self, out: &mut [ColumnOut], rows: &Rows, at: usize, f: impl Fn(&R) -> F64 + Sync) { self.cols(out, rows, at, |r| [f(r)]); } @@ -257,11 +257,11 @@ impl<'a> FillCtx<'a> { fn cols( &self, out: &mut [ColumnOut], - rows: &[R], + rows: &Rows, at: usize, f: impl Fn(&R) -> [F64; N] + Sync, ) { - self.cols_at(out, rows.len(), at, |i| f(&rows[i])); + self.cols_at(out, rows.len(), at, |i| f(rows.get(i))); } /// [`Self::cols`] over row indices, for values held in a side buffer rather @@ -282,20 +282,35 @@ impl<'a> FillCtx<'a> { }); assert!(n_rows <= n, "a table's rows must fit its cube"); let padding: [F64; N] = std::array::from_fn(|k| self.padding[at + k]); - parallel::for_each(n, |i| { - let v = if i < n_rows { f(i) } else { padding }; - for (k, p) in dst.iter().enumerate() { - // SAFETY: distinct `i` write disjoint in-bounds slots of each of the - // `N` windows, each exactly once, and the dispatch blocks until - // every write is finished. - unsafe { p.add(i).write(v[k]) }; + // A cache line of rows is gathered column by column, then streamed out: the + // stack is read by a later pass, so its lines need not be fetched. + const LINE: usize = 8; + parallel::for_each_chunk(n.div_ceil(LINE), |first, end| { + let stream = primitives::stream::Stream::new(); + for b in first..end { + let start = b * LINE; + let len = LINE.min(n - start); + let mut buf = [[F64::ZERO; LINE]; N]; + for j in 0..len { + let i = start + j; + let v = if i < n_rows { f(i) } else { padding }; + for k in 0..N { + buf[k][j] = v[k]; + } + } + for (p, column) in dst.iter().zip(&buf) { + // SAFETY: block `b` writes rows `start..start + len` of each of the + // `N` windows, in bounds and disjoint from every other block's, and + // the dispatch blocks until every write is finished. + unsafe { stream.copy_uninit(p.add(start), &column[..len]) }; + } } }); } /// The three `K`-lanes of the 192-bit word in memory cell `addr`. fn limbs(&self, addr: u32) -> [F64; 3] { - let w = self.mem[addr as usize]; + let w = self.mem.word(addr as usize); [F64(w.c0), F64(w.c1), F64(w.c2)] } } @@ -535,8 +550,8 @@ impl Table for Arith { vb[2], ] }); - ctx.cols(out, rows, RA, |r| [r.ra, r.rb, r.rc]); - ctx.col(out, rows, RBC, |r| r.bytecode_read); + ctx.cols(out, rows, RA, |r| [r.ra, r.rb, r.rc].map(|k| ctx.g_at(k))); + ctx.col(out, rows, RBC, |r| ctx.g_at(r.bytecode_read)); } } @@ -594,8 +609,8 @@ impl Table for SetTable { let k = imm(r).1; [F64(k.c0), F64(k.c1), F64(k.c2)] }); - ctx.col(out, rows, R, |r| r.r); - ctx.col(out, rows, RBC, |r| r.bytecode_read); + ctx.col(out, rows, R, |r| ctx.g_at(r.r)); + ctx.col(out, rows, RBC, |r| ctx.g_at(r.bytecode_read)); } } @@ -674,7 +689,7 @@ impl Table for DerefTable { ctx.col(out, rows, FP, |r| ctx.g_at(r.fp)); debug_assert!( rows.iter().all(|r| { - let p = ctx.mem[(r.fp + ins(r).0) as usize]; + let p = ctx.mem.word((r.fp + ins(r).0) as usize); p.c1 == 0 && p.c2 == 0 }), "deref pointer must be K-valued" @@ -690,14 +705,14 @@ impl Table for DerefTable { ctx.g_at(o3), mode.f_pc(), mode.f_fp(), - F64(ctx.mem[(r.fp + o1) as usize].c0), + F64(ctx.mem.word((r.fp + o1) as usize).c0), v3[0], v3[1], v3[2], ] }); - ctx.cols(out, rows, R1, |r| [r.r1, r.r2, r.r3]); - ctx.col(out, rows, RBC, |r| r.bytecode_read); + ctx.cols(out, rows, R1, |r| [r.r1, r.r2, r.r3].map(|k| ctx.g_at(k))); + ctx.col(out, rows, RBC, |r| ctx.g_at(r.bytecode_read)); } } @@ -776,7 +791,7 @@ impl Table for JumpTable { Op::Jump { oc, od, of } => (oc, od, of), op => unreachable!("a JUMP row's pc {} holds {op:?}", r.pc), }; - let cell = |r: &Jrow, o: u32| ctx.mem[(r.fp + o) as usize]; + let cell = |r: &Jrow, o: u32| ctx.mem.word((r.fp + o) as usize); let cond = |r: &Jrow| cell(r, ins(r).0); ctx.col(out, rows, PC, |r| ctx.g_at(r.pc)); ctx.col(out, rows, FP, |r| ctx.g_at(r.fp)); @@ -814,8 +829,8 @@ impl Table for JumpTable { } let mut inv = acc.inv(); let mut w = vec![F192::ZERO; rows.len()]; - for (i, r) in rows.iter().enumerate().rev() { - let c = cond(r); + for i in (0..rows.len()).rev() { + let c = cond(rows.get(i)); if !c.is_zero() { w[i] = inv * prefix[i]; inv *= c; @@ -824,8 +839,8 @@ impl Table for JumpTable { (w, b) }; ctx.cols_at(out, rows.len(), W, |i| [F64(w[i].c0), b[i]]); - ctx.cols(out, rows, RC, |r| [r.rc, r.rd, r.rf]); - ctx.col(out, rows, RBC, |r| r.bytecode_read); + ctx.cols(out, rows, RC, |r| [r.rc, r.rd, r.rf].map(|k| ctx.g_at(k))); + ctx.col(out, rows, RBC, |r| ctx.g_at(r.bytecode_read)); } } @@ -947,7 +962,7 @@ impl Table for Blake2sTable { // the four message cells, then the cv pair, the output pair and the // metadata. A cell's two lanes are one read, so each group of four takes two. let word_pair = |c0: u32, c1: u32| { - let (w0, w1) = (ctx.mem[c0 as usize], ctx.mem[c1 as usize]); + let (w0, w1) = (ctx.mem.word(c0 as usize), ctx.mem.word(c1 as usize)); [F64(w0.c0), F64(w0.c1), F64(w1.c0), F64(w1.c1)] }; ctx.cols(out, rows, V_M0, |r| { @@ -967,15 +982,16 @@ impl Table for Blake2sTable { word_pair(a[4], a[4] + 1) }); ctx.cols(out, rows, MD0, |r| { - let md = ctx.mem[ad(r)[6] as usize]; + let md = ctx.mem.word(ad(r)[6] as usize); [F64(md.c0), F64(md.c1)] }); ctx.cols(out, rows, R_M0, |r| { [ r.ra[0], r.ra[1], r.rb[0], r.rb[1], r.rcv[0], r.rcv[1], r.rc[0], r.rc[1], r.rmd, ] + .map(|k| ctx.g_at(k)) }); - ctx.col(out, rows, RBC, |r| r.bytecode_read); + ctx.col(out, rows, RBC, |r| ctx.g_at(r.bytecode_read)); } } diff --git a/crates/pcs/src/ntt.rs b/crates/pcs/src/ntt.rs index 9f8ff7749..5d6a25cc9 100644 --- a/crates/pcs/src/ntt.rs +++ b/crates/pcs/src/ntt.rs @@ -10,6 +10,7 @@ use primitives::field::F8; pub mod additive_ntt_f64; pub mod inv_table; pub use additive_ntt_f64::{AdditiveNttF64, RowSink, transpose_lane_major}; +pub(crate) use additive_ntt_f64::{prefetch, transposed_butterfly_lanes}; pub use inv_table::InvNttTableByteSingleGf8; /// Twiddle recurrence used to build the next subspace layer's evaluation points: diff --git a/crates/pcs/src/ntt/additive_ntt_f64.rs b/crates/pcs/src/ntt/additive_ntt_f64.rs index ec907e964..b35b6f170 100644 --- a/crates/pcs/src/ntt/additive_ntt_f64.rs +++ b/crates/pcs/src/ntt/additive_ntt_f64.rs @@ -258,8 +258,14 @@ impl AdditiveNttF64 { // Where the deep pass starts. // // every layer fits one L3 sub-block -> deep pass only, one sweep - // otherwise -> the deep pass keeps one L2 sub-block of layers - let deep = if log_d - start <= fit3 { fit3 } else { fit2 }; + // otherwise -> as few gathered passes as an L3 deep pass needs, + // then the smallest deep pass, at least L2-sized, they leave + let rest = log_d - start; + let deep = if rest <= fit3 { + fit3 + } else { + fit2.max(rest.saturating_sub((rest - fit3).div_ceil(fit2) * fit2)) + }; let deep_start = log_d.saturating_sub(deep).max(par_log).max(start); // A buffer this large is evicted before the next pass reads it back. @@ -379,35 +385,44 @@ impl AdditiveNttF64 { // Without one, a task is one (block, residue) pair. let n_tasks = if msg.is_some() { step } else { step << layer }; let base = parallel::SendPtr(data.as_mut_ptr()); + // With a message, every block of a residue starts from the same rows: they are gathered + // once, and each block but the last transforms a copy. + let n_blocks = if msg.is_some() { 1usize << layer } else { 1 }; + let len = rows * num_ntts; 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 L2-resident scratch of 2^g rows (two with copies) serves every group of the task. + with_scratch(if n_blocks > 1 { 2 * len } else { len }, |scratch| { + let (gathered, work) = scratch.split_at_mut(len); // 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; - 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() { + // Row `i` of residue `r`, as a word offset inside its block. + let row = |r: usize, i: usize| (r + (i << log_step)) * num_ntts; + let block_off = |block: usize| (block << (log_d - layer)) * num_ntts; + // Gather: the scattered rows become one contiguous 2^g-row buffer. + let gather = |dst: &mut [F64], block: usize, r: usize| { + let src_row = |i: usize| match msg { + Some(m) => m.0.wrapping_add(row(r, i)).cast_const(), + None => base.0.wrapping_add(block_off(block) + row(r, i)).cast_const(), + }; + for (i, dst) in dst.chunks_exact_mut(num_ntts).enumerate() { + // No prefetcher follows rows this far apart. + if i + PREFETCH_ROWS < rows { + prefetch(src_row(i + PREFETCH_ROWS), num_ntts); + } // 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 { - match msg { - 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); + // - A residue reads all its message rows before it writes any of block 0. + dst.copy_from_slice(unsafe { std::slice::from_raw_parts(src_row(i), num_ntts) }); } - // 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); - // Scatter: every row returns to its place. - for (i, src) in scratch.chunks_exact(num_ntts).enumerate() { + }; + // Transform, then scatter every row back to its place. + let finish = |buf: &mut [F64], block: usize, r: usize| { + // A (layer + g)-layer domain whose sub-block index is the global block. + self.run_layers(buf, layer + g, num_ntts, layer, layer + g, layer, block); + for (i, src) in buf.chunks_exact(num_ntts).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(block) + row(r, i), num_ntts) }; match &stream { Some(s) => s.copy(dst, src), None => dst.copy_from_slice(src), @@ -416,13 +431,18 @@ impl AdditiveNttF64 { }; for t in lo..hi { if msg.is_some() { - // Block 0 may be the message itself, so it is transformed last. - for block in (0..1usize << layer).rev() { - group(block, t); + gather(gathered, 0, t); + // Block 0 may be the message itself, so it is written last. + for block in (1..n_blocks).rev() { + work.copy_from_slice(gathered); + finish(work, block, t); } + finish(gathered, 0, t); } else { // Task index = block * step + residue. - group(t >> log_step, t & (step - 1)); + let (block, r) = (t >> log_step, t & (step - 1)); + gather(gathered, block, r); + finish(gathered, block, r); } } }); @@ -666,6 +686,25 @@ pub fn transpose_lane_major(out: &mut [F64], msg: &[F64], n_lanes: usize, log_ro }); } +/// Rows a gather fetches ahead of the one it copies. +const PREFETCH_ROWS: usize = 8; + +/// Hint the cache lines of `len` words at `p` into L1. A hint only: it never faults. +#[inline(always)] +pub(crate) fn prefetch(p: *const F64, len: usize) { + #[cfg(target_arch = "x86_64")] + { + use core::arch::x86_64::{_MM_HINT_T0, _mm_prefetch}; + let (start, end) = (p as usize & !63, p as usize + 8 * len); + for line in (start..end).step_by(64) { + // SAFETY: a prefetch never faults, whatever the address. + unsafe { _mm_prefetch::<_MM_HINT_T0>(line as *const i8) }; + } + } + #[cfg(not(target_arch = "x86_64"))] + let _ = (p, len); +} + /// Words one L2-resident unit of work may span: a gathered row group or a deep sub-block. /// /// # Why this value @@ -804,6 +843,20 @@ fn butterfly_interleaved_block(block: &mut [F64], twiddle: F64, block_size_half: } } +/// The transposed butterfly over every lane of a (top row, bottom row) pair: +/// `s = u + v; new_u = s; new_v = v + s*t`. +pub(crate) fn transposed_butterfly_lanes(top: &mut [F64], bot: &mut [F64], twiddle: F64) { + assert_eq!(top.len(), bot.len()); + // The forward butterfly of (v, s) leaves v + s*t in v and s + (v + s*t) in s, which one more XOR undoes. + for (u, v) in top.iter_mut().zip(bot.iter()) { + *u += *v; + } + butterfly_lanes(bot, top, twiddle); + for (u, v) in top.iter_mut().zip(bot.iter()) { + *u += *v; + } +} + /// 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. /// @@ -817,18 +870,22 @@ 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 tail = top.len() % 8; + // SAFETY: the target features are enabled at compile time; each full + // iteration reads and writes exactly eight elements from both rows, and + // the masked tail only the `tail` elements left. 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); } - } - for lane in 8 * vectors..top.len() { - let v = bot[lane]; - let new_u = top[lane] + v * twiddle; - top[lane] = new_u; - bot[lane] = v + new_u; + if tail != 0 { + butterfly_lanes_avx512_masked( + top.as_mut_ptr().add(8 * vectors), + bot.as_mut_ptr().add(8 * vectors), + twiddle.0, + (1u8 << tail) - 1, + ); + } } } #[cfg(all( @@ -839,18 +896,22 @@ fn butterfly_lanes(top: &mut [F64], bot: &mut [F64], twiddle: F64) { ))] { let vectors = top.len() / 4; - // SAFETY: the target features are enabled at compile time and each - // iteration reads and writes exactly four elements from both rows. + let tail = top.len() % 4; + // SAFETY: the target features are enabled at compile time; each full + // iteration reads and writes exactly four elements from both rows, and + // the masked tail only the `tail` elements left. unsafe { for i in 0..vectors { butterfly_lanes_avx2(top.as_mut_ptr().add(4 * i), bot.as_mut_ptr().add(4 * i), twiddle.0); } - } - for lane in 4 * vectors..top.len() { - let v = bot[lane]; - let new_u = top[lane] + v * twiddle; - top[lane] = new_u; - bot[lane] = v + new_u; + if tail != 0 { + butterfly_lanes_avx2_masked( + top.as_mut_ptr().add(4 * vectors), + bot.as_mut_ptr().add(4 * vectors), + twiddle.0, + tail, + ); + } } } #[cfg(all(target_arch = "aarch64", target_feature = "aes"))] @@ -907,35 +968,76 @@ fn butterfly_lanes(top: &mut [F64], bot: &mut [F64], twiddle: F64) { #[target_feature(enable = "vpclmulqdq", enable = "avx2")] unsafe fn butterfly_lanes_avx2(top: *mut F64, bot: *mut F64, twiddle: u64) { use core::arch::x86_64::*; + // SAFETY: the caller supplies valid four-element rows and the function's + // target features cover every intrinsic below. + unsafe { + let (new_u, new_v) = butterfly_ymm(_mm256_loadu_si256(top.cast()), _mm256_loadu_si256(bot.cast()), twiddle); + _mm256_storeu_si256(top.cast(), new_u); + _mm256_storeu_si256(bot.cast(), new_v); + } +} + +/// [`butterfly_lanes_avx2`] over the first `n < 4` lanes only. +/// +/// # Safety +/// Requires VPCLMULQDQ + AVX2; `top` and `bot` must each address `n` readable +/// and writable F64 values. +#[cfg(all( + target_arch = "x86_64", + target_feature = "vpclmulqdq", + target_feature = "avx2", + not(target_feature = "avx512f") +))] +#[inline] +#[target_feature(enable = "vpclmulqdq", enable = "avx2")] +unsafe fn butterfly_lanes_avx2_masked(top: *mut F64, bot: *mut F64, twiddle: u64, n: usize) { + use core::arch::x86_64::*; + // SAFETY: masked-off lanes are neither read nor written, and the caller + // supplies the `n` lanes that are. + unsafe { + let mask = _mm256_cmpgt_epi64(_mm256_set1_epi64x(n as i64), _mm256_setr_epi64x(0, 1, 2, 3)); + let u = _mm256_maskload_epi64(top.cast(), mask); + let v = _mm256_maskload_epi64(bot.cast(), mask); + let (new_u, new_v) = butterfly_ymm(u, v, twiddle); + _mm256_maskstore_epi64(top.cast(), mask, new_u); + _mm256_maskstore_epi64(bot.cast(), mask, new_v); + } +} + +/// Four butterflies `u' = u + v * t`, `v' = v + u'` in registers. +#[cfg(all( + target_arch = "x86_64", + target_feature = "vpclmulqdq", + target_feature = "avx2", + not(target_feature = "avx512f") +))] +#[inline] +#[target_feature(enable = "vpclmulqdq", enable = "avx2")] +fn butterfly_ymm( + u: core::arch::x86_64::__m256i, + v: core::arch::x86_64::__m256i, + twiddle: u64, +) -> (core::arch::x86_64::__m256i, core::arch::x86_64::__m256i) { + use core::arch::x86_64::*; #[inline] #[target_feature(enable = "vpclmulqdq", enable = "avx2")] - unsafe fn reduce(p: __m256i, r: __m256i) -> __m256i { + fn reduce(p: __m256i, r: __m256i) -> __m256i { let t = _mm256_clmulepi64_epi128::<0x01>(p, r); let u = _mm256_clmulepi64_epi128::<0x01>(t, r); _mm256_xor_si256(_mm256_xor_si256(p, t), u) } - // SAFETY: the caller supplies valid four-element rows and the function's - // target features cover every intrinsic below. - unsafe { - let u = _mm256_loadu_si256(top.cast()); - let v = _mm256_loadu_si256(bot.cast()); - let tw = _mm256_set1_epi64x(twiddle as i64); - let r = _mm256_set1_epi64x(0x1b); - - let even = reduce(_mm256_clmulepi64_epi128::<0x00>(v, tw), r); - let odd = reduce(_mm256_clmulepi64_epi128::<0x11>(v, tw), r); - // The odd products land in each lane's low half; swap them up, then take - // the odd qwords (32-bit elements 2, 3, 6, 7) from them. - let odd = _mm256_shuffle_epi32::<0x4e>(odd); - let product = _mm256_blend_epi32::<0b1100_1100>(even, odd); - - let new_u = _mm256_xor_si256(u, product); - let new_v = _mm256_xor_si256(v, new_u); - _mm256_storeu_si256(top.cast(), new_u); - _mm256_storeu_si256(bot.cast(), new_v); - } + let tw = _mm256_set1_epi64x(twiddle as i64); + let r = _mm256_set1_epi64x(0x1b); + let even = reduce(_mm256_clmulepi64_epi128::<0x00>(v, tw), r); + let odd = reduce(_mm256_clmulepi64_epi128::<0x11>(v, tw), r); + // The odd products land in each lane's low half; swap them up, then take + // the odd qwords (32-bit elements 2, 3, 6, 7) from them. + let odd = _mm256_shuffle_epi32::<0x4e>(odd); + let product = _mm256_blend_epi32::<0b1100_1100>(even, odd); + let new_u = _mm256_xor_si256(u, product); + (new_u, _mm256_xor_si256(v, new_u)) } /// Eight F64 butterflies as four independent NEON lane-pair reductions. @@ -1014,6 +1116,48 @@ 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. /// +/// # 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] +#[target_feature(enable = "vpclmulqdq", enable = "avx512f", enable = "avx2")] +unsafe fn butterfly_lanes_avx512(top: *mut F64, bot: *mut F64, twiddle: u64) { + use core::arch::x86_64::*; + // SAFETY: + // - The caller supplies two valid eight-word rows. + // - This function's target features cover every intrinsic below. + unsafe { + let (new_u, new_v) = butterfly_zmm(_mm512_loadu_si512(top.cast()), _mm512_loadu_si512(bot.cast()), twiddle); + _mm512_storeu_si512(top.cast(), new_u); + _mm512_storeu_si512(bot.cast(), new_v); + } +} + +/// [`butterfly_lanes_avx512`] over the lanes set in `mask` only. +/// +/// # Safety +/// +/// - Requires VPCLMULQDQ and AVX-512F. +/// - Each pointer must address the masked lanes, readable and writable; the others are never touched. +#[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_masked(top: *mut F64, bot: *mut F64, twiddle: u64, mask: u8) { + use core::arch::x86_64::*; + // SAFETY: masked loads and stores touch only the lanes the caller vouches for. + unsafe { + let u = _mm512_maskz_loadu_epi64(mask, top.cast()); + let v = _mm512_maskz_loadu_epi64(mask, bot.cast()); + let (new_u, new_v) = butterfly_zmm(u, v, twiddle); + _mm512_mask_storeu_epi64(top.cast(), mask, new_u); + _mm512_mask_storeu_epi64(bot.cast(), mask, new_v); + } +} + +/// Eight butterflies `u' = u + v * t`, `v' = v + u'` in registers, one per 64-bit lane. +/// /// # Algorithm /// /// ```text @@ -1030,52 +1174,39 @@ unsafe fn butterfly_lanes_neon_8(top: *mut F64, bot: *mut F64, twiddle: u64) { /// - 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. #[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) { +fn butterfly_zmm( + u: core::arch::x86_64::__m512i, + v: core::arch::x86_64::__m512i, + twiddle: u64, +) -> (core::arch::x86_64::__m512i, core::arch::x86_64::__m512i) { use core::arch::x86_64::*; - // SAFETY: - // - The caller supplies two valid eight-word rows. - // - This function's target features cover every intrinsic below. - 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. - let spill = _mm512_ternarylogic_epi64::( - _mm512_srli_epi64::<63>(hi), - _mm512_srli_epi64::<61>(hi), - _mm512_srli_epi64::<60>(hi), - ); - // 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 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. + let spill = _mm512_ternarylogic_epi64::( + _mm512_srli_epi64::<63>(hi), + _mm512_srli_epi64::<61>(hi), + _mm512_srli_epi64::<60>(hi), + ); + // 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)); + // 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))); + (new_u, _mm512_xor_si512(v, new_u)) } /// Two F64 butterflies with a shared twiddle, NEON-resident end to end. diff --git a/crates/pcs/src/whir.rs b/crates/pcs/src/whir.rs index 820647f4d..9b0d5b06b 100644 --- a/crates/pcs/src/whir.rs +++ b/crates/pcs/src/whir.rs @@ -1155,12 +1155,11 @@ pub(crate) fn recursive_prover_with_prepared_basis( ps.hint_merkle(PrunedMerklePaths::prune(l0_tree, block_len_0, &queries_0, l0_row)); drop(span); - // Induce basis_0 from the L0 opens. L0 dominates the induce phase, where - // the sparse-prefix transposed-NTT path wins; the dispatcher auto-selects - // it (deeper levels stay dense), mirroring the original. + // Induce basis_0 from the L0 opens. The dispatcher picks the sparse-prefix + // transposed NTT wherever a level opens more queries than its blowup. let sks_vks_n1 = eval_sk_at_vks(n1); let span = tracing::info_span!("Induce", level = 0).entered(); - let (basis_0_induced, enforced_sum_0) = induce_sumcheck_poly_auto_base( + let (basis_0_induced, enforced_sum_0) = induce_sumcheck_poly_auto( n1, log_inv_rate_0, &sks_vks_n1, @@ -1281,8 +1280,15 @@ pub(crate) fn recursive_prover_with_prepared_basis( let sks_vks_i = eval_sk_at_vks(n_next); let span = tracing::info_span!("Induce", level = i + 1).entered(); - let (basis_i_induced, enforced_sum_i) = - induce_sumcheck_poly(n_next, &sks_vks_i, &opened_rows_i, &level_rs, &queries_i, &weights_i); + let (basis_i_induced, enforced_sum_i) = induce_sumcheck_poly_auto( + n_next, + config.log_inv_rates[i + 1], + &sks_vks_i, + &opened_rows_i, + &level_rs, + &queries_i, + &weights_i, + ); drop(span); let span = tracing::info_span!("Introduce", level = i + 1).entered(); @@ -1452,8 +1458,7 @@ mod tests { let sks_vks = eval_sk_at_vks(log_msg_cols); let dense = induce_sumcheck_poly(log_msg_cols, &sks_vks, &rows, &v_challenges, &qs, &weights); - let via_ntt = - induce_sumcheck_poly_via_ntt_base(log_msg_cols, log_inv_rate, &rows, &v_challenges, &qs, &weights); + let via_ntt = induce_sumcheck_poly_via_ntt(log_msg_cols, log_inv_rate, &rows, &v_challenges, &qs, &weights); assert_eq!(dense.1, via_ntt.1, "enforced_sum mismatch"); assert_eq!(&*dense.0, &*via_ntt.0, "basis_poly mismatch"); } diff --git a/crates/pcs/src/whir_induce.rs b/crates/pcs/src/whir_induce.rs index 75887f98e..cdf5c1878 100644 --- a/crates/pcs/src/whir_induce.rs +++ b/crates/pcs/src/whir_induce.rs @@ -368,19 +368,21 @@ fn transpose_layers_ext_windowed( } } -/// Sparse-prefix variant of [`transpose_forward_ntt_ext`]: the input has only -/// `positions.len()` nonzeros and the first `k` transpose steps (forward -/// layers `log_d-1 .. log_d-k`, pairing distances `1 .. 2^(k-1)`) mix only -/// WITHIN `2^k`-aligned windows. We process just the windows that contain a -/// nonzero (a dense `2^k` transpose each, disjoint so window-parallel), -/// densify, then run the remaining steps as full dense sweeps. Output is -/// identical to `transpose_forward_ntt_ext` on the scattered input. Mirror -/// of `whir::transpose_forward_ntt_sparse`. +/// Sparse-prefix variant of [`transpose_forward_ntt_ext`], keeping only the first +/// `keep` outputs: the input has only `positions.len()` nonzeros, and the first +/// `k` transpose steps (forward layers `log_d-1 .. log_d-k`, pairing distances +/// `1 .. 2^(k-1)`) mix only WITHIN `2^k`-aligned windows. We process just the +/// windows that contain a nonzero (a dense `2^k` transpose each, disjoint so +/// window-parallel), then the remaining steps in two sweeps of the buffer (see +/// [`transpose_dense_rest`]). Output is identical to the first `keep` entries of +/// `transpose_forward_ntt_ext` on the scattered input. Mirror of +/// `whir::transpose_forward_ntt_sparse`. fn transpose_forward_ntt_sparse_ext( ntt: &AdditiveNttF64, positions: &[usize], values: &[F192], log_d: usize, + keep: usize, ) -> Vec { let _span = tracing::info_span!( "NTT", @@ -389,8 +391,20 @@ fn transpose_forward_ntt_sparse_ext( nonzero = positions.len() ) .entered(); + transpose_sparse_windowed(ntt, positions, values, log_d, keep, TRANSPOSE_WINDOW_LOG) +} + +fn transpose_sparse_windowed( + ntt: &AdditiveNttF64, + positions: &[usize], + values: &[F192], + log_d: usize, + keep: usize, + window_log: usize, +) -> Vec { use std::collections::HashMap; let n = 1usize << log_d; + assert!(keep.is_power_of_two() && keep <= n); // No prefix for small domains: just scatter + full dense transpose. let k = if log_d >= 12 { 8usize.min(log_d) } else { 0 }; @@ -403,6 +417,7 @@ fn transpose_forward_ntt_sparse_ext( if log_d > 0 { transpose_forward_ntt_ext(ntt, &mut data, log_d); } + data.truncate(keep); return data; } @@ -438,29 +453,122 @@ fn transpose_forward_ntt_sparse_ext( } } }); + win_vec.sort_unstable_by_key(|(w, _)| *w); + transpose_dense_rest(ntt, &win_vec, k, log_d, keep, window_log) +} - // Densify (active windows only; the rest stay zero, which is the correct - // post-step-(k-1) state for an all-zero window). - // SAFETY: zero is a valid F192, and the inactive windows must read as zero. - let mut data = vec![F192::ZERO; n]; - for (w, buf) in &win_vec { - data[(w << k)..((w + 1) << k)].copy_from_slice(buf); - } +/// Rows, as a log, of the cache-resident windows of [`transpose_dense_rest`]. +const TRANSPOSE_WINDOW_LOG: usize = 15; + +/// Consecutive residues one gathered group of [`transpose_dense_rest`] takes: eight E values are three whole cache lines. +const TRANSPOSE_RESIDUES: usize = 8; - // Remaining steps s = k..log_d-1 = forward layers (log_d-1-k) .. 0, dense. - transpose_layers_ext(ntt, &mut data, log_d, (0..(log_d - k)).rev()); +/// The dense steps after the sparse prefix, in two sweeps of the buffer. +/// +/// ```text +/// window pass: one task per 2^w-row window, built from the sparse windows inside it, +/// every layer whose blocks fit the window run there +/// gathered pass: one task per group of rows r + i * 2^w, for 8 consecutive residues r, +/// the remaining top layers run in a scratch, layer 0 only for the outputs kept +/// ``` +fn transpose_dense_rest( + ntt: &AdditiveNttF64, + windows: &[(usize, Vec)], + k: usize, + log_d: usize, + keep: usize, + window_log: usize, +) -> Vec { + use crate::ntt::{prefetch, transposed_butterfly_lanes}; + let n = 1usize << log_d; + let w_log = window_log.min(log_d).max(k); + let win = 1usize << w_log; + // The top layers left after the windows, run as one gathered transform of 2^top rows. + let top = log_d - w_log; + let as_words = |e: &mut [F192]| { + // SAFETY: an F192 is three F64 words with no padding. + unsafe { std::slice::from_raw_parts_mut(e.as_mut_ptr().cast::(), 3 * e.len()) } + }; + + // SAFETY: the window pass writes every slot before the gathered pass reads it. + let mut data = unsafe { primitives::uninit_vec::(n) }; + parallel::chunks_mut(&mut data, win, |c, window| { + window.fill(F192::ZERO); + let base = c << w_log; + let first = windows.partition_point(|(w, _)| (*w << k) < base); + for (w, buf) in windows[first..].iter().take_while(|(w, _)| (*w << k) < base + win) { + window[(*w << k) - base..][..buf.len()].copy_from_slice(buf); + } + // Layers whose blocks fit the window, block sizes 2^(k+1) .. 2^w. + for layer in (top..log_d - k).rev() { + let block_size = 1usize << (log_d - layer); + for (b, block) in window.chunks_mut(block_size).enumerate() { + let (lo, hi) = block.split_at_mut(block_size / 2); + transposed_butterfly_lanes( + as_words(lo), + as_words(hi), + ntt.twiddle(layer, (base >> (log_d - layer)) + b), + ); + } + } + }); + if top == 0 { + data.truncate(keep); + return data; + } + assert!(keep < n, "layer 0 keeps only its top half"); + // Rows kept from each group: those below `keep`. + + let rows = 1usize << top; + let kept_rows = (keep >> w_log).max(1); + let group = TRANSPOSE_RESIDUES.min(win); + let base = parallel::SendPtr(data.as_mut_ptr()); + parallel::for_each_chunk(win / group, |lo, hi| { + let mut scratch = vec![F192::ZERO; rows * group]; + for g in lo..hi { + let r = g * group; + let at = |i: usize| (i << w_log) + r; + for (i, dst) in scratch.chunks_exact_mut(group).enumerate() { + if i + 8 < rows { + prefetch(base.0.wrapping_add(at(i + 8)).cast::(), 3 * group); + } + // SAFETY: the window pass wrote every row, and this group alone owns these residues. + dst.copy_from_slice(unsafe { base.slice(at(i), group) }); + } + // Layer `layer` pairs rows 2^(top-1-layer) apart, inside blocks of 2^(top-layer) rows. + for layer in (0..top).rev() { + let block_rows = 1usize << (top - layer); + for (b, block) in scratch.chunks_mut(block_rows * group).enumerate() { + let (lo, hi) = block.split_at_mut(block_rows / 2 * group); + if kept_rows <= block_rows / 2 { + // Every later layer reads only the first `kept_rows` of each block, all in its top + // half, whose output is the sum. + for (a, &b) in lo[..kept_rows * group].iter_mut().zip(hi.iter()) { + *a += b; + } + } else { + transposed_butterfly_lanes(as_words(lo), as_words(hi), ntt.twiddle(layer, b)); + } + } + } + for (i, src) in scratch.chunks_exact(group).take(kept_rows).enumerate() { + // SAFETY: as above. + unsafe { base.slice(at(i), group) }.copy_from_slice(src); + } + } + }); + data.truncate(keep); data } /// `F^T`-based fast path for [`induce_sumcheck_poly`]: scatter per-query /// E-weights into the codeword domain, apply `F^T` with K-twiddles, keep the /// low `2^log_msg_cols` outputs. Byte-identical output to the dense path -/// (pinned by `induce_via_ntt_matches_dense`). Mirror of -/// `whir::induce_sumcheck_poly_via_ntt` with the L0 mixed row dot. -pub(crate) fn induce_sumcheck_poly_via_ntt_base( +/// (pinned by `induce_via_ntt_matches_dense`). +pub(crate) fn induce_sumcheck_poly_via_ntt( log_msg_cols: usize, log_inv_rate: usize, - opened_rows: &[Vec], + opened_rows: &[Vec], v_challenges: &[F192], queries: &[usize], weights: &[F192], @@ -470,13 +578,14 @@ pub(crate) fn induce_sumcheck_poly_via_ntt_base( let block_len = 1usize << log_block; let n_queries = queries.len(); assert_eq!(opened_rows.len(), n_queries); + assert!(queries.iter().all(|&q| q < block_len), "queries lie in the codeword"); let eq = build_eq_table_ext(v_challenges); debug_assert_eq!(weights.len(), n_queries); let mut enforced_sum = F192::ZERO; for i in 0..n_queries { - enforced_sum += F64::dot(&opened_rows[i], &eq) * weights[i]; + enforced_sum += T::dot(&opened_rows[i], &eq) * weights[i]; } let mut coeffs = if log_block == 0 { @@ -488,40 +597,61 @@ pub(crate) fn induce_sumcheck_poly_via_ntt_base( c } else { let ntt = AdditiveNttF64::standard(log_block); - transpose_forward_ntt_sparse_ext(&ntt, queries, weights, log_block) + transpose_forward_ntt_sparse_ext(&ntt, queries, weights, log_block, n) }; coeffs.truncate(n); (coeffs, enforced_sum) } -/// Query density threshold for choosing the sparse NTT over dense expansion. -const NTT_QUERIES_PER_BLOWUP: usize = 8; - -/// Cost-based dispatch: the NTT overtakes the dense expansion once a level -/// opens more than [`NTT_QUERIES_PER_BLOWUP`] queries per unit of blowup. +/// Cost-based dispatch: the dense expansion costs a pass over the `2^log_msg_cols` +/// outputs per query, the NTT a few passes over the `2^log_inv_rate` times larger +/// codeword domain, so the NTT wins once a level opens more queries than its blowup. #[inline] pub(crate) fn induce_use_ntt_heuristic(log_msg_cols: usize, log_inv_rate: usize, n_queries: usize) -> bool { - log_msg_cols >= 12 && n_queries > NTT_QUERIES_PER_BLOWUP * (1usize << log_inv_rate) + log_msg_cols >= 12 && n_queries > 1usize << log_inv_rate } /// Dispatch between the dense [`induce_sumcheck_poly`] and the sparse -/// [`induce_sumcheck_poly_via_ntt_base`] for L0 (base-field rows). Mirror of -/// `whir::induce_sumcheck_poly_auto`: in the recursive PCS this fires -/// only at the top level (large message domain, many queries); deeper levels -/// stay dense. Both paths produce identical output, so a mis-dispatch only -/// costs time. -pub(crate) fn induce_sumcheck_poly_auto_base( +/// [`induce_sumcheck_poly_via_ntt`]. Both paths produce identical output, so a +/// mis-dispatch only costs time. +pub(crate) fn induce_sumcheck_poly_auto( log_msg_cols: usize, log_inv_rate: usize, sks_vks: &[F64], - opened_rows: &[Vec], + opened_rows: &[Vec], v_challenges: &[F192], queries: &[usize], weights: &[F192], ) -> (Vec, F192) { if induce_use_ntt_heuristic(log_msg_cols, log_inv_rate, queries.len()) { - induce_sumcheck_poly_via_ntt_base(log_msg_cols, log_inv_rate, opened_rows, v_challenges, queries, weights) + induce_sumcheck_poly_via_ntt(log_msg_cols, log_inv_rate, opened_rows, v_challenges, queries, weights) } else { induce_sumcheck_poly(log_msg_cols, sks_vks, opened_rows, v_challenges, queries, weights) } } + +#[cfg(test)] +mod tests { + use super::*; + use primitives::test_rng::Rng; + + /// The sparse transpose against the dense one, with windows small enough to leave a gathered pass. + #[test] + fn sparse_transpose_matches_dense() { + let mut rng = Rng::new(0x7A); + for (log_d, window_log, keep_log) in [(13usize, 9usize, 12usize), (14, 11, 11), (12, 12, 11), (16, 8, 15)] { + let ntt = AdditiveNttF64::standard(log_d); + let n = 1usize << log_d; + let positions: Vec = (0..50).map(|_| rng.next_u64() as usize % n).collect(); + let values = rng.ext_vec(positions.len()); + let mut want = vec![F192::ZERO; n]; + for (&p, &v) in positions.iter().zip(&values) { + want[p] += v; + } + transpose_forward_ntt_ext(&ntt, &mut want, log_d); + want.truncate(1 << keep_log); + let got = transpose_sparse_windowed(&ntt, &positions, &values, log_d, 1 << keep_log, window_log); + assert_eq!(got, want, "log_d={log_d}, window_log={window_log}"); + } + } +} diff --git a/crates/primitives/src/field/gf2_64x3.rs b/crates/primitives/src/field/gf2_64x3.rs index 1a108010c..c05de9301 100644 --- a/crates/primitives/src/field/gf2_64x3.rs +++ b/crates/primitives/src/field/gf2_64x3.rs @@ -352,7 +352,7 @@ impl F192Unreduced { /// Build from three 128-bit coefficients. #[inline] - fn from_wide(coeffs: [u128; 3]) -> Self { + pub(crate) fn from_wide(coeffs: [u128; 3]) -> Self { Self { coeffs: coeffs.map(|c| [c as u64, (c >> 64) as u64]), } @@ -656,7 +656,7 @@ pub mod x86_64 { /// d2 = p0 ^ p1 ^ p02 /// ``` #[inline(always)] - fn fold(p0: V, p1: V, p2: V, p01: V, p02: V, p12: V, xor: impl Fn(V, V) -> V) -> [V; 3] { + pub(crate) fn fold(p0: V, p1: V, p2: V, p01: V, p02: V, p12: V, xor: impl Fn(V, V) -> V) -> [V; 3] { // Shared by d0 and d1. let q = xor(p0, p12); [xor(xor(q, p1), p2), xor(q, p01), xor(xor(p0, p1), p02)] @@ -737,7 +737,7 @@ pub mod x86_64 { #[cfg(all(target_feature = "vpclmulqdq", target_feature = "avx512f"))] #[inline] #[target_feature(enable = "avx512f")] - unsafe fn reduce_lanes512(lo: __m512i, hi: __m512i) -> __m512i { + pub(crate) unsafe fn reduce_lanes512(lo: __m512i, hi: __m512i) -> __m512i { // The bits of `hi * 0x1B` shifted past x^63. let spill = _mm512_xor_si512( _mm512_xor_si512(_mm512_srli_epi64::<63>(hi), _mm512_srli_epi64::<61>(hi)), diff --git a/crates/primitives/src/field/mod.rs b/crates/primitives/src/field/mod.rs index b3088641e..ad211f29b 100644 --- a/crates/primitives/src/field/mod.rs +++ b/crates/primitives/src/field/mod.rs @@ -17,6 +17,8 @@ pub mod gf2_64; pub mod gf2_64x3; pub mod gf2_8; pub mod phi8_tower; +#[cfg(all(target_arch = "x86_64", target_feature = "avx512f", target_feature = "vpclmulqdq"))] +pub mod x8; pub use gf2_8::F8; pub use gf2_64::F64; diff --git a/crates/primitives/src/field/x8.rs b/crates/primitives/src/field/x8.rs new file mode 100644 index 000000000..a828be23c --- /dev/null +++ b/crates/primitives/src/field/x8.rs @@ -0,0 +1,434 @@ +//! Eight `F192` lanes held coefficient-major in three AVX-512 registers: register `k` +//! holds coefficient `k` of every lane, so a product is lane-wise and needs no packing. +//! +//! A product is Karatsuba's six CLMULs on the even lanes and six on the odd ones, +//! each CLMUL selecting one qword per 128-bit lane. The even and odd halves come back +//! as 128-bit partials, which interleave into the lanes again for the reduction. A sum +//! that is itself summed across lanes ([`Sum8`]) skips the interleave and the reduction. + +use core::arch::x86_64::*; +use core::ops::{Add, AddAssign, BitXor, BitXorAssign, Mul}; + +use super::gf2_64::{F64, reduce}; +use super::gf2_64x3::F192; +use super::gf2_64x3::x86_64::{fold, reduce_lanes512}; + +/// Eight elements; register `k` is coefficient `k` of lanes `0..8`. +#[derive(Clone, Copy, Debug)] +#[repr(C, align(64))] +pub struct F192x8([__m512i; 3]); + +/// The sum over lanes of unreduced products, as four 128-bit partials per coefficient. +#[derive(Clone, Copy, Debug)] +pub struct Sum8([__m512i; 3]); + +/// Eight `K` elements in one register. +#[derive(Clone, Copy, Debug)] +pub struct F64x8(__m512i); + +/// Per-lane unreduced sums of mixed products `c · k`, `c ∈ E` and `k ∈ K`. +#[derive(Clone, Copy, Debug)] +pub struct Mixed8 { + even: [__m512i; 3], + odd: [__m512i; 3], +} + +#[inline(always)] +fn xor(a: __m512i, b: __m512i) -> __m512i { + // SAFETY: the module's cfg enables avx512f. + unsafe { _mm512_xor_si512(a, b) } +} + +#[inline(always)] +fn vector(v: [i64; 8]) -> __m512i { + // SAFETY: the module's cfg enables avx512f; the array is eight qwords. + unsafe { _mm512_loadu_si512(v.as_ptr().cast()) } +} + +#[inline(always)] +fn load_u64x8(src: *const u64) -> __m512i { + // SAFETY: callers pass a pointer to eight readable qwords. + unsafe { _mm512_loadu_si512(src.cast()) } +} + +/// The lane-wise base reduction of the even-lane and odd-lane 128-bit partials. +#[inline(always)] +fn reduce_halves(even: __m512i, odd: __m512i) -> __m512i { + // SAFETY: the module's cfg enables avx512f. + unsafe { reduce_lanes512(_mm512_unpacklo_epi64(even, odd), _mm512_unpackhi_epi64(even, odd)) } +} + +/// Qword `i` of the result is `a[idx[i]]` for `idx[i] < 8`, else `b[idx[i] - 8]`. +#[inline(always)] +fn select(a: __m512i, idx: [i64; 8], b: __m512i) -> __m512i { + // SAFETY: the module's cfg enables avx512f. + unsafe { _mm512_permutex2var_epi64(a, vector(idx), b) } +} + +/// Qword `i` of the result is `c[idx[i]]` where `mask` has bit `i`, else `src[i]`. +#[inline(always)] +fn select_masked(src: __m512i, mask: u8, idx: [i64; 8], c: __m512i) -> __m512i { + // SAFETY: the module's cfg enables avx512f. + unsafe { _mm512_mask_permutexvar_epi64(src, mask, vector(idx), c) } +} + +/// The permutes of an eight-element block between its 24 interleaved qwords and the +/// three coefficient registers. Qword `q` of the block is coefficient `q % 3` of +/// element `q / 3`. +const fn aos_to_planes(k: usize) -> ([i64; 8], u8, [i64; 8]) { + let (mut pair, mut mask, mut third) = ([0i64; 8], 0u8, [0i64; 8]); + let mut e = 0; + while e < 8 { + let q = 3 * e + k; + if q < 16 { + pair[e] = q as i64; + } else { + mask |= 1 << e; + third[e] = (q - 16) as i64; + } + e += 1; + } + (pair, mask, third) +} + +const fn planes_to_aos(o: usize) -> ([i64; 8], u8, [i64; 8]) { + let (mut pair, mut mask, mut third) = ([0i64; 8], 0u8, [0i64; 8]); + let mut i = 0; + while i < 8 { + let q = 8 * o + i; + let (e, k) = (q / 3, q % 3); + if k < 2 { + pair[i] = (e + 8 * k) as i64; + } else { + mask |= 1 << i; + third[i] = e as i64; + } + i += 1; + } + (pair, mask, third) +} + +const LOAD: [([i64; 8], u8, [i64; 8]); 3] = [aos_to_planes(0), aos_to_planes(1), aos_to_planes(2)]; +const STORE: [([i64; 8], u8, [i64; 8]); 3] = [planes_to_aos(0), planes_to_aos(1), planes_to_aos(2)]; + +/// Transpose an 8x8 matrix of qwords: row `r` of the input is column `r` of the output. +#[inline(always)] +fn transpose(r: [__m512i; 8]) -> [__m512i; 8] { + // SAFETY: the module's cfg enables avx512f. + unsafe { + 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) + } + }); + // `u[4h + j]`: rows `4h..4h+4` at qwords `j` and `j + 4`. + let u: [__m512i; 8] = std::array::from_fn(|i| { + let (h, s, o) = (i / 4, (i / 2) % 2, i % 2); + let (a, b) = (t[4 * h + o], t[4 * h + 2 + o]); + if s == 0 { + _mm512_shuffle_i64x2::<0x88>(a, b) + } else { + _mm512_shuffle_i64x2::<0xDD>(a, b) + } + }); + std::array::from_fn(|c| { + let (a, b) = (u[c % 4], u[4 + c % 4]); + if c < 4 { + _mm512_shuffle_i64x2::<0x88>(a, b) + } else { + _mm512_shuffle_i64x2::<0xDD>(a, b) + } + }) + } +} + +impl F192x8 { + #[inline(always)] + pub fn splat(e: F192) -> Self { + // SAFETY: the module's cfg enables avx512f. + unsafe { + Self([ + _mm512_set1_epi64(e.c0 as i64), + _mm512_set1_epi64(e.c1 as i64), + _mm512_set1_epi64(e.c2 as i64), + ]) + } + } + + /// Eight consecutive elements. + #[inline(always)] + pub fn load(src: &[F192; 8]) -> Self { + let q = src.as_ptr().cast::(); + // SAFETY: the block is 24 qwords. + let v = unsafe { [load_u64x8(q), load_u64x8(q.add(8)), load_u64x8(q.add(16))] }; + Self(std::array::from_fn(|k| { + let (pair, mask, third) = LOAD[k]; + select_masked(select(v[0], pair, v[1]), mask, third, v[2]) + })) + } + + /// The lanes as eight consecutive elements. + #[inline(always)] + pub fn store(self, dst: &mut [F192; 8]) { + let q = dst.as_mut_ptr().cast::(); + for (o, &(pair, mask, third)) in STORE.iter().enumerate() { + let v = select_masked(select(self.0[0], pair, self.0[1]), mask, third, self.0[2]); + // SAFETY: the block is 24 writable qwords. + unsafe { _mm512_storeu_si512(q.add(8 * o).cast(), v) }; + } + } + + pub fn to_array(self) -> [F192; 8] { + let mut out = [F192::ZERO; 8]; + self.store(&mut out); + out + } + + /// Eight runs of eight consecutive elements, transposed: `out[j]` lane `p` is `src[8p + j]`. + #[inline(always)] + pub fn load_runs(src: &[F192; 64]) -> [Self; 8] { + let q = src.as_ptr().cast::(); + // Block `b` of run `p` is its qwords `8b..8b+8`; transposing a block across runs + // gives every run's qword `8b + i` in one register. + // SAFETY: the 64 elements are 192 readable qwords. + let blocks: [[__m512i; 8]; 3] = + std::array::from_fn(|b| transpose(std::array::from_fn(|p| load_u64x8(unsafe { q.add(24 * p + 8 * b) })))); + std::array::from_fn(|j| Self(std::array::from_fn(|k| blocks[(3 * j + k) / 8][(3 * j + k) % 8]))) + } + + /// The even lanes of `a` then of `b`, and the odd lanes of `a` then of `b`. + #[inline(always)] + pub fn deinterleave(a: Self, b: Self) -> (Self, Self) { + const EVEN: [i64; 8] = [0, 2, 4, 6, 8, 10, 12, 14]; + const ODD: [i64; 8] = [1, 3, 5, 7, 9, 11, 13, 15]; + ( + Self(std::array::from_fn(|k| select(a.0[k], EVEN, b.0[k]))), + Self(std::array::from_fn(|k| select(a.0[k], ODD, b.0[k]))), + ) + } + + /// The inverse of [`deinterleave`](Self::deinterleave): lanes `a0 b0 a1 b1 …` over two registers. + #[inline(always)] + pub fn interleave(a: Self, b: Self) -> (Self, Self) { + const LO: [i64; 8] = [0, 8, 1, 9, 2, 10, 3, 11]; + const HI: [i64; 8] = [4, 12, 5, 13, 6, 14, 7, 15]; + ( + Self(std::array::from_fn(|k| select(a.0[k], LO, b.0[k]))), + Self(std::array::from_fn(|k| select(a.0[k], HI, b.0[k]))), + ) + } + + /// The even-lane and odd-lane unreduced products, y-folded. + #[inline(always)] + fn products(self, b: Self) -> ([__m512i; 3], [__m512i; 3]) { + let [a0, a1, a2] = self.0; + let [b0, b1, b2] = b.0; + let xa = [a0, a1, a2, xor(a0, a1), xor(a0, a2), xor(a1, a2)]; + let xb = [b0, b1, b2, xor(b0, b1), xor(b0, b2), xor(b1, b2)]; + // SAFETY: the module's cfg enables avx512f and vpclmulqdq. + let (even, odd) = unsafe { + ( + std::array::from_fn::<_, 6, _>(|i| _mm512_clmulepi64_epi128::<0x00>(xa[i], xb[i])), + std::array::from_fn::<_, 6, _>(|i| _mm512_clmulepi64_epi128::<0x11>(xa[i], xb[i])), + ) + }; + let [p0, p1, p2, p01, p02, p12] = even; + let even = fold(p0, p1, p2, p01, p02, p12, xor); + let [p0, p1, p2, p01, p02, p12] = odd; + (even, fold(p0, p1, p2, p01, p02, p12, xor)) + } + + /// The lane-wise products, summed over the lanes and left unreduced. + #[inline(always)] + pub fn mul_sum(self, b: Self) -> Sum8 { + let (even, odd) = self.products(b); + Sum8(std::array::from_fn(|k| xor(even[k], odd[k]))) + } +} + +impl Add for F192x8 { + type Output = Self; + #[inline(always)] + fn add(self, rhs: Self) -> Self { + Self(std::array::from_fn(|k| xor(self.0[k], rhs.0[k]))) + } +} + +impl AddAssign for F192x8 { + #[inline(always)] + fn add_assign(&mut self, rhs: Self) { + *self = *self + rhs; + } +} + +impl Mul for F192x8 { + type Output = Self; + #[inline(always)] + fn mul(self, rhs: Self) -> Self { + let (even, odd) = self.products(rhs); + Self(std::array::from_fn(|k| reduce_halves(even[k], odd[k]))) + } +} + +impl Sum8 { + // SAFETY: an all-zero bit pattern is a valid `__m512i`. + pub const ZERO: Self = Self(unsafe { std::mem::zeroed() }); + + /// The lane sum, reduced. + #[inline] + pub fn reduce(self) -> F192 { + let [c0, c1, c2] = self.0.map(|s| { + // SAFETY: the module's cfg enables avx512f. + let words: [u64; 8] = unsafe { std::mem::transmute(s) }; + let (lo, hi) = ( + words[0] ^ words[2] ^ words[4] ^ words[6], + words[1] ^ words[3] ^ words[5] ^ words[7], + ); + reduce(u128::from(hi) << 64 | u128::from(lo)) + }); + F192::new(c0, c1, c2) + } +} + +impl BitXor for Sum8 { + type Output = Self; + #[inline(always)] + fn bitxor(self, rhs: Self) -> Self { + Self(std::array::from_fn(|k| xor(self.0[k], rhs.0[k]))) + } +} + +impl BitXorAssign for Sum8 { + #[inline(always)] + fn bitxor_assign(&mut self, rhs: Self) { + *self = *self ^ rhs; + } +} + +impl F64x8 { + #[inline(always)] + pub fn new(k: [F64; 8]) -> Self { + Self::load(&k) + } + + #[inline(always)] + pub fn load(k: &[F64; 8]) -> Self { + Self(load_u64x8(k.as_ptr().cast())) + } + + #[inline(always)] + pub fn splat(k: F64) -> Self { + // SAFETY: the module's cfg enables avx512f. + Self(unsafe { _mm512_set1_epi64(k.0 as i64) }) + } + + pub fn to_array(self) -> [F64; 8] { + // SAFETY: both are 64 bytes. + unsafe { std::mem::transmute(self.0) } + } +} + +impl Mul for F64x8 { + type Output = Self; + #[inline(always)] + fn mul(self, rhs: Self) -> Self { + // SAFETY: the module's cfg enables avx512f and vpclmulqdq. + Self(unsafe { + reduce_halves( + _mm512_clmulepi64_epi128::<0x00>(self.0, rhs.0), + _mm512_clmulepi64_epi128::<0x11>(self.0, rhs.0), + ) + }) + } +} + +impl Default for Mixed8 { + fn default() -> Self { + Self::ZERO + } +} + +impl Mixed8 { + // SAFETY: an all-zero bit pattern is a valid `__m512i`. + pub const ZERO: Self = Self { + even: unsafe { std::mem::zeroed() }, + odd: unsafe { std::mem::zeroed() }, + }; + + /// Add `c · k[i]` to lane `i`; `c` is a [`F192x8::splat`]. + #[inline(always)] + pub fn add(&mut self, c: F192x8, k: F64x8) { + for i in 0..3 { + // SAFETY: the module's cfg enables avx512f and vpclmulqdq. Immediate bit 4 + // picks the scalar's odd qword; the splat is the same in both. + unsafe { + self.even[i] = xor(self.even[i], _mm512_clmulepi64_epi128::<0x00>(c.0[i], k.0)); + self.odd[i] = xor(self.odd[i], _mm512_clmulepi64_epi128::<0x10>(c.0[i], k.0)); + } + } + } + + #[inline(always)] + pub fn reduce(self) -> F192x8 { + F192x8(std::array::from_fn(|i| reduce_halves(self.even[i], self.odd[i]))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_rng::Rng; + + fn elements(rng: &mut Rng) -> [F192; 8] { + let mut e: [F192; 8] = std::array::from_fn(|_| rng.ext()); + e[3] = F192::new(u64::MAX, u64::MAX, u64::MAX); + e[5] = F192::ZERO; + e + } + + #[test] + fn lanes_match_scalar() { + let mut rng = Rng::new(17); + for _ in 0..200 { + let (a, b) = (elements(&mut rng), elements(&mut rng)); + let (va, vb) = (F192x8::load(&a), F192x8::load(&b)); + assert_eq!(va.to_array(), a); + assert_eq!((va * vb).to_array(), std::array::from_fn(|i| a[i] * b[i])); + assert_eq!((va + vb).to_array(), std::array::from_fn(|i| a[i] + b[i])); + let dot = (0..8).fold(F192::ZERO, |acc, i| acc + a[i] * b[i]); + assert_eq!(va.mul_sum(vb).reduce(), dot); + assert_eq!((va.mul_sum(vb) ^ vb.mul_sum(va)).reduce(), F192::ZERO); + assert_eq!(F192x8::splat(a[0]).to_array(), [a[0]; 8]); + + let (even, odd) = F192x8::deinterleave(va, vb); + let both: Vec = a.iter().chain(&b).copied().collect(); + assert_eq!(even.to_array(), std::array::from_fn(|i| both[2 * i])); + assert_eq!(odd.to_array(), std::array::from_fn(|i| both[2 * i + 1])); + let (lo, hi) = F192x8::interleave(even, odd); + assert_eq!((lo.to_array(), hi.to_array()), (a, b)); + + let c = rng.ext(); + let k: [F64; 8] = std::array::from_fn(|_| F64(rng.ext().c0)); + let l: [F64; 8] = std::array::from_fn(|_| F64(rng.ext().c1)); + let mut acc = Mixed8::ZERO; + acc.add(F192x8::splat(c), F64x8::new(k)); + acc.add(F192x8::splat(a[1]), F64x8::new(k) * F64x8::load(&l)); + let want = std::array::from_fn(|i| c.mul_base(k[i]) + a[1].mul_base(k[i] * l[i])); + assert_eq!(acc.reduce().to_array(), want); + assert_eq!((F64x8::new(k) * F64x8::splat(l[0])).to_array(), k.map(|k| k * l[0])); + } + } + + #[test] + fn runs_transpose() { + let mut rng = Rng::new(5); + let src: [F192; 64] = std::array::from_fn(|_| rng.ext()); + let runs = F192x8::load_runs(&src); + for (j, run) in runs.iter().enumerate() { + assert_eq!(run.to_array(), std::array::from_fn(|p| src[8 * p + j])); + } + } +} diff --git a/crates/rec_aggregation/guests/lean_ethereum.py b/crates/rec_aggregation/guests/lean_ethereum.py index 175eaa8df..93b1153de 100644 --- a/crates/rec_aggregation/guests/lean_ethereum.py +++ b/crates/rec_aggregation/guests/lean_ethereum.py @@ -34,6 +34,8 @@ # Six challenges compose the F2-linear map that batches the 192 transposed # ring-switch coordinates. RING_MAP_SHIFTS = [32, 16, 8, 4, 2, 1] +RS_CHUNK = 8 # ring-switch rows per loop iteration +RS_COUNT_BITS = 6 # a ring-switch coordinate count is below 2^RS_COUNT_BITS # Exponent bit-widths: an announced 32-bit count decomposes into COUNT_BITS bits, # its top bit constrained to zero so the native strict 32-bit bound holds; any # structural size (sums of 2^kappa, packing offsets) fits SIZE_BITS bits; and a @@ -78,7 +80,9 @@ LAYER_COUNT = 5 LAYER_ROW = 6 LAYER_POS = 7 -LAYER_SLOTS = 8 +LAYER_HALF = 8 # g^(k // 2) at radix-four layer k +LAYER_PARITY = 9 # k % 2, as a field element +LAYER_SLOTS = 10 # One OOD sample of a WHIR level. OOD_BETA = 0 OOD_Y = 1 @@ -225,6 +229,9 @@ LIG_LOG_MSG_COLS_CAP = LIG_LOG_MSG_COLS_CAP_PLACEHOLDER YR_LOG_CAP = YR_LOG_CAP_PLACEHOLDER MAX_STACK_LOG = LIG_MIN_LOG_SIZE + LIG_N_LOG_SIZES - 1 +SEL_NIBBLES = (MAX_STACK_LOG + 3) // 4 # four selector bits share one table lookup +SEL_EQ_STRIDE = 30 # a nibble's eq tables over its top 1..4 coords +EQ_CHAIN_ARMS = SIZE_BITS + SLOT_STRIDE_LOG + 1 # eq prefix chain lengths LIG_N_LEVELS = LIG_N_LEVELS_PLACEHOLDER LIG_YR_LEVEL = LIG_YR_LEVEL_PLACEHOLDER LIG_YR_LOG_LEN = LIG_YR_LOG_LEN_PLACEHOLDER @@ -401,10 +408,17 @@ SP_SUBNODE_MUL = 2 ** 10 * HI_MUL SP_FNODE_LEVEL_MUL = 2 ** 3 * HI_MUL SP_FNODE_MUL = 2 ** 6 * HI_MUL +# Values bounded a byte at a time: each byte hinted in the exponent, range checked, +# and read as the bits it spells off a table of the SP_BYTE_VALUES bytes. +SP_BYTE_VALUES = 256 # The encoding counter, LE_32 in the low four bytes of its lane: bounded by -# decomposing exactly that many bits, so the guest accepts no preimage the native -# verifier cannot parse. -SP_COUNTER_BITS = 32 +# rebuilding it from exactly that many bytes, so the guest accepts no preimage the +# native verifier cannot parse. +SP_COUNTER_BYTES = 4 +# The bits past the second cell's fields, likewise: a byte at a time, the last one +# holding the rest. +SP_UNUSED_BYTES = (SP_UNUSED_BITS + 7) // 8 +SP_UNUSED_TOP = 2 ** (SP_UNUSED_BITS - 8 * (SP_UNUSED_BYTES - 1)) # --------------------------------------------------------------- node capacities # MAX_KEYS caps the coverage table's slots, both schemes' declared claims and their @@ -587,18 +601,18 @@ def exponent_tables(): return g_logs_pow2, g_squares -def g_power_of_word(value, g_squares, nbits: Const): - # g^value for a concrete integer `value` < 2^nbits: advice-decompose its bits, - # tie them back to the word, and assemble Π g^(bit_j·2^j). - bits = HeapBuf(GEN ** nbits) - hint_decompose_bits(bits, value, nbits) +def g_power_of_word(value): + # g^value for a concrete integer `value` < 2^LOG_WORD_BITS: advice-decompose its + # bits, tie them back to the word, and assemble Π g^(bit_j·2^j). + bits = StackBuf(LOG_WORD_BITS) + hint_decompose_bits(bits, value, LOG_WORD_BITS) word = 0 g_value = GEN ** 0 - for j in unroll(0, nbits): - bit = bits[GEN ** j] - assert bit * bit == bit + for j in unroll(0, LOG_WORD_BITS): + bit = bits[j] + bits[j] = bit * bit word += bit * (2 ** j) - g_value *= (1 + bit * (g_squares[GEN ** j] + 1)) + g_value *= 1 + bit * (1 + GEN ** (2 ** j)) assert word == value return g_value @@ -698,16 +712,16 @@ def decode_query_bits(squeezed_word, positions_out, bit_ptrs_out, depth: Const): return -def grind_check(state_0, state_1, nonce, nbits_g): +def grind_check(state_0, state_1, nonce, nbits: Const): # WHIR fold/query grinding: digest = H(H(state, POW_BASE), (nonce, POW_NONCE)), - # whose low nbits (nbits_g = g^nbits) must be zero. The PoW window of + # whose low nbits must be zero. The PoW window of # transcript::pow_bits_ok is `digest.0 & ((1 << bits) - 1)` with nbits < 64, so # it lives entirely in the digest's FIRST 64-bit lane: only that lane is # advice-decomposed and verified, not all FIELD_BITS of the cell. The caller # absorbs the full field nonce afterwards. The honest prover searches the # deterministic u64 subset while verification permits the full field: each # candidate still costs one hash and succeeds with probability 2^-bits. - if nbits_g == GEN ** 0: + if nbits == 0: assert nonce == 0 # native canonical zero-work nonce st = [state_0, state_1] base = StackBuf(2) @@ -718,19 +732,16 @@ def grind_check(state_0, state_1, nonce, nbits_g): hint_f192_limbs(lanes, out[0]) high = (out[0] + lanes[0]) * Y_INV assert_in_k(lanes[0], high) - # Frame cells for the unrolled pass (no DEREF per bit), named by `addr` for the - # zero-check walk, whose bound is runtime and so must index a pointer. + # The lane is rebuilt from its bits at and above nbits only, which IS the claim + # that its low nbits coordinates are zero; the hinted low bits go unread. lane_bits = StackBuf(BASE_FIELD_BITS) hint_decompose_bits(lane_bits, lanes[0], BASE_FIELD_BITS) acc = 0 - for i in unroll(0, BASE_FIELD_BITS): + for i in unroll(nbits, BASE_FIELD_BITS): b = lane_bits[i] lane_bits[i] = b * b # booleanity, as a write-once pin acc += b * COORD_BASIS[i] - assert acc == lanes[0] # the bits ARE the lane's coordinates, so the pins below bind it - lane_ptr = addr(lane_bits) - for xb in mul_range(1, nbits_g): - assert lane_ptr[xb] == 0 + assert acc == lanes[0] return @@ -797,33 +808,76 @@ def lag64(z, out, node_base: Const): return -def eq_prefix_chain(chain, seed, a, b, count_g): - # Prefix products of eq(a_k, b_k) = 1 + a_k + b_k from `seed`, so a reader picks - # the partial product up at its own certified length. Entry t is written from - # inputs with index < t only, so a garbage tail past a buffer's written extent - # cannot corrupt any shorter prefix. - chain[GEN ** 0] = seed - for xk in mul_range(1, count_g): - chain[xk * GEN] = chain[xk] * (1 + a[xk] + b[xk]) +def eq_prefix_chain(chain, seed, a, b1, count_g): + # Prefix products of eq(a_k, b_k) = a_k + b1_k, b1 = 1 + b, from `seed`, so a + # reader picks the partial product up at its own certified length. Entry t is + # written from inputs with index < t only, so a garbage tail past a buffer's + # written extent cannot corrupt any shorter prefix. + assert log(count_g) < EQ_CHAIN_ARMS + done = match(log(count_g), range(0, EQ_CHAIN_ARMS), lambda n: eq_prefix_run(chain, seed, a, b1, n)) # a match needs a target return -def rs_eq_run(chain, z_vals, point, count_g): +def eq_prefix_run(chain, seed, a, b1, n: Const): + chain[GEN ** 0] = seed + v = seed + for k in unroll(0, n): + v = v * (a[GEN ** k] + b1[GEN ** k]) + chain[GEN ** (k + 1)] = v + return 0 + + +def plus_one_run(out, src, n: Const): + for k in unroll(0, n): + out[GEN ** k] = 1 + src[GEN ** k] + return 0 + + +def rs_eq_run(chain, z_vals, point1, count_g): # One run of the telescoped ring-switch product E = sum_k c_k * prod_j # (z_j^(2^k) + 1 + ris_j): coordinate x multiplies row k by (z^(2^k) + 1 + # point_x), z evolving by squaring per row. The runtime coordinates walk # OUTSIDE and the fixed Frobenius powers inside, so nothing stores a z-power - # table. - for xk in mul_range(1, count_g): - zv = z_vals[xk] - one_plus = 1 + point[xk] - row = chain * xk ** BASE_FIELD_BITS + # table. Coordinates go two an iteration, so a row is read and written once + # per pair; count = 2 * half + odd is pinned from advice bits, and an odd + # count's last coordinate runs after the loop. Returns the final row. + bits = StackBuf(RS_COUNT_BITS) + hint_decompose_bits_exponent(bits, count_g, RS_COUNT_BITS) + odd = bits[0] + bits[0] = odd * odd + half_g = GEN ** 0 + for k in unroll(1, RS_COUNT_BITS): + b = bits[k] + bits[k] = b * b + half_g *= 1 + b * (1 + GEN ** (2 ** (k - 1))) + odd_g = 1 + odd * (1 + GEN) + assert half_g * half_g * odd_g == count_g + for xp in mul_range(1, half_g): + pair = xp * xp + zs = z_vals * pair + ps = point1 * pair + zv_0 = zs[GEN ** 0] + zv_1 = zs[GEN ** 1] + one_plus_0 = ps[GEN ** 0] + one_plus_1 = ps[GEN ** 1] + row = chain * xp ** BASE_FIELD_BITS + nxt = row * GEN ** BASE_FIELD_BITS + for k in unroll(0, BASE_FIELD_BITS): + nxt[GEN ** k] = row[GEN ** k] * (zv_0 + one_plus_0) * (zv_1 + one_plus_1) + if k != BASE_FIELD_BITS - 1: + zv_0 *= zv_0 + zv_1 *= zv_1 + if odd == 1: + last = half_g * half_g + zv = z_vals[last] + one_plus = point1[last] + row = chain * half_g ** BASE_FIELD_BITS nxt = row * GEN ** BASE_FIELD_BITS for k in unroll(0, BASE_FIELD_BITS): nxt[GEN ** k] = row[GEN ** k] * (zv + one_plus) if k != BASE_FIELD_BITS - 1: zv *= zv - return + return chain * (half_g * odd_g) ** BASE_FIELD_BITS def fold_final_msg(msg, point, log_len: Const): @@ -862,18 +916,29 @@ def sumcheck_round4(state_0, state_1, msg_cursor, claim): return fs[0], fs[1], msg_cursor, c0 + y * (c1 + y * (c2 + y * c3)), y -def sumcheck_round5(state_0, state_1, msg_cursor, claim, prev_challenge): - # One GKR round. The prover sends every coefficient but c0, which the round's - # pulled-out eq factor leaves fixed: `c0 + prev_challenge * (c1 + ... + c4) == - # claim`. - fs = [state_0, state_1] - fs, c1, msg_cursor = fs_next(fs, msg_cursor) - fs, c2, msg_cursor = fs_next(fs, msg_cursor) - fs, c3, msg_cursor = fs_next(fs, msg_cursor) - fs, c4, msg_cursor = fs_next(fs, msg_cursor) - fs, y = squeeze(fs) - c0 = claim + prev_challenge * (c1 + c2 + c3 + c4) - return fs[0], fs[1], msg_cursor, c0 + y * (c1 + y * (c2 + y * (c3 + y * c4))), y +def gkr_rounds(rd, pt, n: Const): + # n GKR rounds, from the round record at `rd` (fields GKR_ROUNDS_CAP apart) to the + # one n further. The prover sends every coefficient but c0, which the round's + # pulled-out eq factor leaves fixed: c0 + r (c1 + ... + c4) == claim, r being the + # previous point's coordinate pt[q]; the challenge goes to pt[MU_CAP + 2 + q]. + fs = [rd[GEN ** (ROUND_FS0 * GKR_ROUNDS_CAP)], rd[GEN ** (ROUND_FS1 * GKR_ROUNDS_CAP)]] + cur = rd[GEN ** (ROUND_CURSOR * GKR_ROUNDS_CAP)] + claim = rd[GEN ** (ROUND_CLAIM * GKR_ROUNDS_CAP)] + for q in unroll(0, n): + fs, c1, cur = fs_next(fs, cur) + fs, c2, cur = fs_next(fs, cur) + fs, c3, cur = fs_next(fs, cur) + fs, c4, cur = fs_next(fs, cur) + fs, rk = squeeze(fs) + c0 = claim + pt[GEN ** q] * (c1 + c2 + c3 + c4) + pt[GEN ** (MU_CAP + 2 + q)] = rk + claim = c0 + rk * (c1 + rk * (c2 + rk * (c3 + rk * c4))) + rd_next = rd * GEN ** n + rd_next[GEN ** (ROUND_FS0 * GKR_ROUNDS_CAP)] = fs[0] + rd_next[GEN ** (ROUND_FS1 * GKR_ROUNDS_CAP)] = fs[1] + rd_next[GEN ** (ROUND_CURSOR * GKR_ROUNDS_CAP)] = cur + rd_next[GEN ** (ROUND_CLAIM * GKR_ROUNDS_CAP)] = claim + return def batch_sumcheck(fs0, fs1, msgs, running, point, n_rounds: Const): @@ -896,47 +961,71 @@ def batch_sumcheck(fs0, fs1, msgs, running, point, n_rounds: Const): @inline -def verify_merkle_path(leaf_0, leaf_1, direction_bits, depth: Const): +def verify_merkle_path(leaf, direction_bits, depth: Const): # Hinted child pairs are hashed in order; the boolean query bit selects the # child that must equal the running node, binding each link of the path. path = StackBuf(LIG_PATH_CAP) hint_witness(path[0:4 * depth], "merkle_children") + # The selected pair starts at path_ptr * g^(2*bit) = path_ptr + bit * path_step. path_ptr = addr(path) - node_0 = leaf_0 - node_1 = leaf_1 + path_step = path_ptr * (1 + GEN ** 2) + node = leaf for level in unroll(0, depth): - dir_bit = direction_bits[GEN ** level] - selected = path_ptr * GEN ** (4 * level) * (1 + dir_bit * (1 + GEN ** 2)) - selected[1] = node_0 - selected[GEN] = node_1 - parent = StackBuf(2) - blake2s(path[4 * level:4 * level + 2], path[4 * level + 2:4 * level + 4], parent) - node_0 = parent[0] - node_1 = parent[1] - return node_0, node_1 - - -@inline -def hash_cap_node(cap, index): - children = cap * index ** 4 - parent = StackBuf(2) - blake2s([children[1], children[GEN]], [children[GEN ** 2], children[GEN ** 3]], parent) - cap[index * index] = parent[0] - cap[GEN * index * index] = parent[1] - return + selected = path_ptr + direction_bits[GEN ** level] * path_step + selected[GEN ** (4 * level)] = node[0] + selected[GEN ** (4 * level + 1)] = node[1] + node = StackBuf(2) + blake2s(path[4 * level:4 * level + 2], path[4 * level + 2:4 * level + 4], node) + return node def verify_merkle_cap(cap, flags, depth: Const): + # Unrolled over the cap's nodes, so every address is an immediate. Every node's + # hinted digest (node k at heap cells 2k, 2k+1) is read once into a frame run, + # at 2(k - 1), so hashing an internal node writes its digest onto (and so + # asserts) the hinted value. An active (flagged) node is hashed; an inactive one + # pins its children's flags to zero, so every active node's parent is active + # too, up to the root, which is always hashed. The queries flag the parents of + # the cap leaves they reach. The run is sized per depth, StackBuf sizes being + # parse-time. if depth != 0: - hash_cap_node(cap, GEN) - flags[GEN] = 1 - for parent in mul_range(GEN, GEN ** (2 ** (depth - 1))): - for side in unroll(0, 2): - child = parent * parent * GEN ** side - if flags[child] != 0: - flags[parent] = 1 # every active node forces its parent to be hashed - hash_cap_node(cap, child) - return cap[GEN ** 2], cap[GEN ** 3] + n_leaves = 2 ** depth + if const(depth == 1): + runs = StackBuf(4 * 2 - 2) + elif const(depth == 2): + runs = StackBuf(4 * 4 - 2) + elif const(depth == 3): + runs = StackBuf(4 * 8 - 2) + elif const(depth == 4): + runs = StackBuf(4 * 16 - 2) + elif const(depth == 5): + runs = StackBuf(4 * 32 - 2) + elif const(depth == 6): + runs = StackBuf(4 * 64 - 2) + elif const(depth == 7): + runs = StackBuf(4 * 128 - 2) + else: + runs = StackBuf(4 * 256 - 2) + for node in unroll(1, 2 * n_leaves): + runs[2 * node - 2] = cap[GEN ** (2 * node)] + runs[2 * node - 1] = cap[GEN ** (2 * node + 1)] + for node in unroll(1, n_leaves): + left = 4 * node - 2 + if const(node == 1): + blake2s(runs[left:left + 2], runs[left + 2:left + 4], runs[0:2]) + else: + if flags[GEN ** node] != 0: + blake2s(runs[left:left + 2], runs[left + 2:left + 4], runs[2 * node - 2:2 * node]) + else: + if const(2 * node // n_leaves == 0): # node has flagged children + flags[GEN ** (2 * node)] = 0 + flags[GEN ** (2 * node + 1)] = 0 + root_0 = runs[0] + root_1 = runs[1] + else: + root_0 = cap[GEN ** 2] + root_1 = cap[GEN ** 3] + return root_0, root_1 # ============================== the stacked WHIR opening ============================ @@ -951,19 +1040,13 @@ def opening_fold(fs0, fs1, cursor, c0, c1, c2): return fs[0], fs[1], cursor, claim, c0, c2, r -def opening_ood_point(fs0, fs1, point, n_g): - # Share the squeeze loop across point dimensions. - states = HeapBuf((n_g * GEN) ** 2) - states[1] = fs0 - states[GEN] = fs1 - for x in mul_range(1, n_g): - state = states * x * x - r, f0, f1 = squeeze_step(state[1], state[GEN]) - point[x] = r - state[GEN ** 2] = f0 - state[GEN ** 3] = f1 - last = states * n_g * n_g - return last[1], last[GEN] +def opening_ood_point(fs0, fs1, point, n: Const): + # One specialization per point dimension, the state riding frame cells. + fs = [fs0, fs1] + for x in unroll(0, n): + fs, r = squeeze(fs) + point[GEN ** x] = r + return fs[0], fs[1] def opening_final_message(fs0, fs1, cursor, out, n: Const): @@ -975,73 +1058,136 @@ def opening_final_message(fs0, fs1, cursor, out, n: Const): def opening_row_weights(point, out, folds: Const, reverse: Const): - for i in unroll(0, 2 ** folds): + # The row's eq weights as their two tensor factors: weight i is + # out[i % 2^lo] * out[2^lo + i // 2^lo], the low `lo` coordinates first. At + # level 0 lane i is interleaving index 2^folds - 1 - i, which reverses both. + if const(folds // 3 == 0): + lo = folds + else: + lo = 3 + hi = folds - lo + for i in unroll(0, 2 ** lo): if reverse == 1: - slot = 2 ** folds - 1 - i + slot = 2 ** lo - 1 - i else: slot = i - out[GEN ** i] = eq_weight(point, folds, slot, 0) + out[GEN ** i] = eq_weight(point, lo, slot, 0) + for i in unroll(0, 2 ** hi): + if reverse == 1: + slot = 2 ** hi - 1 - i + else: + slot = i + out[GEN ** (2 ** lo + i)] = eq_weight(point * GEN ** lo, hi, slot, 0) return -def opening_queries(cap, flags, query_weights, query_bit_ptrs, row_eq_weights, n_queries_g, base: Const, interleave: Const, blocks: Const, depth: Const, cap_depth: Const): - # Specialize by row and path shape so opening configurations share query code. - query_sum_chain = HeapBuf(n_queries_g * GEN) - query_sum_chain[GEN ** 0] = 0 - for xe in mul_range(1, n_queries_g): - if base == 1: - row_len = interleave +@inline +def opening_query(cap, flags, row_eq_weights, direction_bits, base: Const, interleave: Const, blocks: Const, depth: Const, cap_depth: Const): + # One query: hash the hinted row up its Merkle path into the shared cap, and + # return its dot product with the row's eq weights. + if base == 1: + row_len = interleave + else: + row_len = 3 * interleave + row = StackBuf(LIG_ROW_CAP) + hint_witness(row[0:row_len], "merkle_leaf_rows") + row_dot = 0 + packed_row = StackBuf(LIG_PACKED_ROW_CAP) + if base == 1: + # Packing proves each hinted lane is in K before hashing or folding it. + for jb in unroll(0, interleave // 2): + packed_row[jb] = pack64x2(row[2 * jb], row[2 * jb + 1]) + else: + # Pack the checked tower limbs into the leaf's contiguous byte image. + for jb in unroll(0, 3 * interleave // 2): + packed_row[jb] = pack64x2(row[2 * jb], row[2 * jb + 1]) + # The weights' tensor factors (opening_row_weights): a run of lo_len row + # entries shares one hi factor. + if const(interleave // 8 == 0): + lo_len = interleave + else: + lo_len = 8 + for h in unroll(0, interleave // lo_len): + part = 0 + for l in unroll(0, lo_len): + jw = lo_len * h + l + if base == 1: + row_word = row[jw] + elif 3 * jw % 2 == 0: + # limbs (3w, 3w+1) are a pack; add Y^2 * limb(3w+2). + row_word = packed_row[3 * jw // 2] + Y_TOWER * Y_TOWER * row[3 * jw + 2] + else: + # limbs (3w+1, 3w+2) are a pack; shift it by Y and add limb(3w). + row_word = row[3 * jw] + Y_TOWER * packed_row[(3 * jw + 1) // 2] + part += row_word * row_eq_weights[GEN ** l] + row_dot += row_eq_weights[GEN ** (lo_len + h)] * part + # Hash the packed row as full BLAKE2s blocks. + leaf_hash_state = StackBuf(2) + blake2s(packed_row[0:2], packed_row[2:4], leaf_hash_state, counter=64, final=1 // blocks) + for jb in unroll(1, blocks): + leaf_digest = StackBuf(2) + blake2s(packed_row[4 * jb:4 * jb + 2], packed_row[4 * jb + 2:4 * jb + 4], leaf_digest, cv=leaf_hash_state, counter=64 * (jb + 1), final=(jb + 1) // blocks) + leaf_hash_state = leaf_digest + path_depth = depth - cap_depth + node = verify_merkle_path(leaf_hash_state, direction_bits, path_depth) + if cap_depth != 0: + # The reached cap leaf is node 2*parent + bit at cells 2*node, 2*node + 1, + # where parent = 2^(cap_depth-1) + q: q in the exponent, the constant + # parts folded into the address immediates. + if const(cap_depth == 1): + flags[GEN] = 1 + leaf = cap * (1 + direction_bits[GEN ** path_depth] * (1 + GEN ** 2)) else: - row_len = 3 * interleave - row = StackBuf(LIG_ROW_CAP) - hint_witness(row[0:row_len], "merkle_leaf_rows") - row_dot = 0 - packed_row = StackBuf(LIG_PACKED_ROW_CAP) + q = 1 + direction_bits[GEN ** (path_depth + 1)] * (1 + GEN) + for bit in unroll(1, cap_depth - 1): + q *= 1 + direction_bits[GEN ** (path_depth + 1 + bit)] * (1 + GEN ** (2 ** bit)) + flags_q = flags * q + flags_q[GEN ** (2 ** (cap_depth - 1))] = 1 # the cap check propagates this obligation to the root + q2 = q * q + leaf = cap * q2 * q2 * (1 + direction_bits[GEN ** path_depth] * (1 + GEN ** 2)) + leaf[GEN ** (2 ** (cap_depth + 1))] = node[0] + leaf[GEN ** (2 ** (cap_depth + 1) + 1)] = node[1] + else: + cap[GEN ** 2] = node[0] + cap[GEN ** 3] = node[1] + return row_dot + + +def opening_queries(cap, flags, lam, lam_n, query_bit_ptrs, row_eq_weights, n_iters_g, odd: Const, base: Const, interleave: Const, blocks: Const, depth: Const, cap_depth: Const): + # Specialize by row and path shape so opening configurations share query code. + # Base-field rows (level 0) run two queries an iteration, sharing the loop + # overhead and the weight reads; an odd count runs its first query alone, + # ahead of the pairs. Extension rows run one query an iteration. + # Query k is weighted lam^k: the chain runs h <- (h + dot_k) / lam, ending at + # sum_k lam^(k - n) dot_k, and one product by lam_n = lam^n restores it. + mu = 1 / lam + chain = HeapBuf(n_iters_g * GEN) + if odd == 1: + chain[GEN ** 0] = opening_query(cap, flags, row_eq_weights, query_bit_ptrs[GEN ** 0], base, interleave, blocks, depth, cap_depth) * mu + else: + chain[GEN ** 0] = 0 + for xi in mul_range(1, n_iters_g): if base == 1: - # Packing proves each hinted lane is in K before hashing or folding it. - for jb in unroll(0, interleave // 4): - e0 = row[4 * jb] - e1 = row[4 * jb + 1] - e2 = row[4 * jb + 2] - e3 = row[4 * jb + 3] - packed_row[2 * jb] = pack64x2(e0, e1) - packed_row[2 * jb + 1] = pack64x2(e2, e3) - row_dot += e0 * row_eq_weights[GEN ** (4 * jb)] + e1 * row_eq_weights[GEN ** (4 * jb + 1)] + e2 * row_eq_weights[GEN ** (4 * jb + 2)] + e3 * row_eq_weights[GEN ** (4 * jb + 3)] + bit_ptrs = query_bit_ptrs * xi * xi + d0 = opening_query(cap, flags, row_eq_weights, bit_ptrs[GEN ** odd], base, interleave, blocks, depth, cap_depth) + d1 = opening_query(cap, flags, row_eq_weights, bit_ptrs[GEN ** (odd + 1)], base, interleave, blocks, depth, cap_depth) + chain[xi * GEN] = ((chain[xi] + d0) * mu + d1) * mu else: - # Pack the checked tower limbs into the leaf's contiguous byte image. - for jb in unroll(0, 3 * interleave // 4): - packed_row[2 * jb] = pack64x2(row[4 * jb], row[4 * jb + 1]) - packed_row[2 * jb + 1] = pack64x2(row[4 * jb + 2], row[4 * jb + 3]) - for jw in unroll(0, interleave): - if 3 * jw % 2 == 0: - # limbs (3w, 3w+1) are a pack; add Y^2 * limb(3w+2). - row_word = packed_row[3 * jw // 2] + Y_TOWER * Y_TOWER * row[3 * jw + 2] - else: - # limbs (3w+1, 3w+2) are a pack; shift it by Y and add limb(3w). - row_word = row[3 * jw] + Y_TOWER * packed_row[(3 * jw + 1) // 2] - row_dot += row_word * row_eq_weights[GEN ** jw] - # Hash the packed row as full BLAKE2s blocks. - leaf_hash_state = StackBuf(2) - blake2s(packed_row[0:2], packed_row[2:4], leaf_hash_state, counter=64, final=1 // blocks) - for jb in unroll(1, blocks): - leaf_digest = StackBuf(2) - blake2s(packed_row[4 * jb:4 * jb + 2], packed_row[4 * jb + 2:4 * jb + 4], leaf_digest, cv=leaf_hash_state, counter=64 * (jb + 1), final=(jb + 1) // blocks) - leaf_hash_state = leaf_digest - query_sum_chain[xe * GEN] = query_sum_chain[xe] + query_weights[xe] * row_dot - direction_bits = query_bit_ptrs[xe] - path_depth = depth - cap_depth - node_0, node_1 = verify_merkle_path(leaf_hash_state[0], leaf_hash_state[1], direction_bits, path_depth) - if cap_depth != 0: - parent = GEN ** (2 ** (cap_depth - 1)) - for bit in unroll(0, cap_depth - 1): - parent *= 1 + direction_bits[GEN ** (path_depth + 1 + bit)] * (1 + GEN ** (2 ** bit)) - flags[parent] = 1 # the cap check propagates this obligation to the root - cap_index = parent * parent * (1 + direction_bits[GEN ** path_depth] * (1 + GEN)) - else: - cap_index = GEN - cap[cap_index * cap_index] = node_0 - cap[GEN * cap_index * cap_index] = node_1 - return query_sum_chain[n_queries_g] + d = opening_query(cap, flags, row_eq_weights, query_bit_ptrs[xi], base, interleave, blocks, depth, cap_depth) + chain[xi * GEN] = (chain[xi] + d) * mu + return chain[n_iters_g] * lam_n + + +@inline +def residual_query(basis_chain, row_d, vanish: Const, n_cols: Const): + # prod_t (d_t + chain_t) over one query's message-column coordinates, chain_0 + # being the query position. + prefix_eq = row_d[GEN ** 0] + basis_chain + for t in unroll(1, n_cols): + # subspace-vanishing recurrence for the novel-basis point + basis_chain *= basis_chain + LIG_VANISH_VALS[vanish + t - 1] + prefix_eq *= row_d[GEN ** t] + basis_chain + return prefix_eq def open_stacked(m_idx: Const, fs0, fs1, target, commit_root_0, commit_root_1, cursor): @@ -1124,7 +1270,7 @@ def open_stacked(m_idx: Const, fs0, fs1, target, commit_root_0, commit_root_1, c # level's message-column dimension. for os in unroll(0, LIG_OOD_SAMPLES[ml + 1]): oz = ood_z * GEN ** (((lvl + 1) * LIG_MAX_OOD_SAMPLES + os) * LIG_LOG_MSG_COLS_CAP) - f0, f1 = opening_ood_point(fs[0], fs[1], oz, GEN ** LIG_LOG_MSG_COLS[ml]) + f0, f1 = opening_ood_point(fs[0], fs[1], oz, LIG_LOG_MSG_COLS[ml]) fs = [f0, f1] sample = ood * GEN ** ((lvl + 1) * ood_stride + os * OOD_SLOTS) fs, ood_y, msg_cursor = fs_next(fs, msg_cursor) @@ -1136,7 +1282,7 @@ def open_stacked(m_idx: Const, fs0, fs1, target, commit_root_0, commit_root_1, c q_nonce = msg_cursor[GEN ** 0] # raw transport word: bound by the DS_POW_NONCE absorb below msg_cursor = msg_cursor * GEN if LIG_QUERY_GRIND_BITS[ml] != 0: - grind_check(fs[0], fs[1], q_nonce, GEN ** LIG_QUERY_GRIND_BITS[ml]) + grind_check(fs[0], fs[1], q_nonce, LIG_QUERY_GRIND_BITS[ml]) else: assert q_nonce == 0 fs = absorb_nonce(fs, q_nonce) @@ -1183,7 +1329,7 @@ def open_stacked(m_idx: Const, fs0, fs1, target, commit_root_0, commit_root_1, c level_roots[GEN ** (2 * lvl)] = root_0 level_roots[GEN ** (2 * lvl + 1)] = root_1 - level_query_sum = opening_queries(cap, flags, query_weights * GEN ** (lvl * max_q), query_bit_ptrs * GEN ** pos_off, row_eq_weights, GEN ** n_queries, 1 // (lvl + 1), interleave, LIG_LEAF_BLOCKS[ml], depth, cap_depth) + level_query_sum = opening_queries(cap, flags, lam, lam_pow, query_bit_ptrs * GEN ** pos_off, row_eq_weights, GEN ** (n_queries // (1 + 1 // (lvl + 1))), n_queries % 2 * (1 // (lvl + 1)), 1 // (lvl + 1), interleave, LIG_LEAF_BLOCKS[ml], depth, cap_depth) # Every level, including the last, ties its commitment in through an intro # message. The level's claims then enter the running one with powers of @@ -1239,15 +1385,21 @@ def open_stacked(m_idx: Const, fs0, fs1, target, commit_root_0, commit_root_1, c # entries above m without reading uninitialized cells. lane_folds = LIG_FOLDS[m_idx * LIG_MAX_LEVELS] fold_head = n_folds - lane_folds + # point1 = 1 + point, the form every eq factor reads. point = HeapBuf(SIZE_BITS + SLOT_STRIDE_LOG) + point1 = HeapBuf(SIZE_BITS + SLOT_STRIDE_LOG) for j in unroll(0, fold_head): point[GEN ** j] = fold_challenges[GEN ** (lane_folds + j)] + point1[GEN ** j] = 1 + fold_challenges[GEN ** (lane_folds + j)] for j in unroll(0, yr_log): point[GEN ** (fold_head + j)] = tail_challenges[GEN ** j] + point1[GEN ** (fold_head + j)] = 1 + tail_challenges[GEN ** j] for j in unroll(0, lane_folds): point[GEN ** (fold_head + yr_log + j)] = fold_challenges[GEN ** j] + point1[GEN ** (fold_head + yr_log + j)] = 1 + fold_challenges[GEN ** j] for j in unroll(n_folds + yr_log, SIZE_BITS + SLOT_STRIDE_LOG): point[GEN ** j] = 0 + point1[GEN ** j] = 1 # ---- per-level induced bases at the single terminal point ---- # Every query of a level runs the SAME product shape over its message-column @@ -1256,22 +1408,29 @@ def open_stacked(m_idx: Const, fs0, fs1, target, commit_root_0, commit_root_1, c # 1 + c_t * (1 + chain_t * inv_t) == (1 + c_t) + (c_t * inv_t) * chain_t # a coordinate's two coefficients depend on the challenge and the baked # vanishing inverse alone, so they hoist out of the query loop (one row a level, - # fold coords then tail coords) and each query is one multiply-add a - # coordinate. - basis_a = HeapBuf(GEN ** (n_levels * LIG_LOG_MSG_COLS_CAP)) - basis_b = HeapBuf(GEN ** (n_levels * LIG_LOG_MSG_COLS_CAP)) + # fold coords then tail coords). Factoring b_t = c_t * inv_t out of every + # coordinate leaves (a_t / b_t) + chain_t per query, so a query is one add and + # one multiply a coordinate, and the level's product of the b_t scales its sum + # once. A zero challenge makes b_t zero and the division unsatisfiable, which + # rejects rather than accepts. + basis_d = HeapBuf(GEN ** (n_levels * LIG_LOG_MSG_COLS_CAP)) + level_scale = HeapBuf(GEN ** n_levels) for lvl in unroll(0, n_levels): ml = m_idx * LIG_MAX_LEVELS + lvl prefix_len = LIG_RESIDUAL_PREFIX_LEN[ml] vanish = m_idx * LIG_MAX_VANISH_LEN + LIG_VANISH_OFF[ml] + scale = level_betas[GEN ** lvl] for t in unroll(0, prefix_len): fold_c = fold_challenges[GEN ** (LIG_RESIDUAL_FOLD_OFF[ml] + t)] - basis_a[GEN ** (lvl * LIG_LOG_MSG_COLS_CAP + t)] = 1 + fold_c - basis_b[GEN ** (lvl * LIG_LOG_MSG_COLS_CAP + t)] = fold_c * LIG_VANISH_INVS[vanish + t] + coef_b = fold_c * LIG_VANISH_INVS[vanish + t] + basis_d[GEN ** (lvl * LIG_LOG_MSG_COLS_CAP + t)] = (1 + fold_c) / coef_b + scale *= coef_b for j in unroll(0, yr_log): tail_c = tail_challenges[GEN ** j] - basis_a[GEN ** (lvl * LIG_LOG_MSG_COLS_CAP + prefix_len + j)] = 1 + tail_c - basis_b[GEN ** (lvl * LIG_LOG_MSG_COLS_CAP + prefix_len + j)] = tail_c * LIG_VANISH_INVS[vanish + prefix_len + j] + coef_b = tail_c * LIG_VANISH_INVS[vanish + prefix_len + j] + basis_d[GEN ** (lvl * LIG_LOG_MSG_COLS_CAP + prefix_len + j)] = (1 + tail_c) / coef_b + scale *= coef_b + level_scale[GEN ** lvl] = scale inner_chain = HeapBuf(GEN ** (n_levels + 1)) inner_chain[GEN ** 0] = 0 for lvl in unroll(0, n_levels): @@ -1280,20 +1439,26 @@ def open_stacked(m_idx: Const, fs0, fs1, target, commit_root_0, commit_root_1, c basis_row = lvl * LIG_LOG_MSG_COLS_CAP residual_chain = HeapBuf(GEN ** (max_q + 1)) residual_chain[GEN ** 0] = 0 - for xr in mul_range(1, GEN ** LIG_QUERIES[ml]): + # Two queries an iteration share the row's reads and constants. + for xr in mul_range(1, GEN ** (LIG_QUERIES[ml] // 2)): ml = m_idx * LIG_MAX_LEVELS + lvl # rebound: the body captures by value vanish = m_idx * LIG_MAX_VANISH_LEN + LIG_VANISH_OFF[ml] basis_row = lvl * LIG_LOG_MSG_COLS_CAP max_q = LIG_MAX_QUERIES[m_idx] - basis_chain = query_positions[GEN ** LIG_POSITIONS_OFF[ml] * xr] - prefix_eq = basis_a[GEN ** basis_row] + basis_b[GEN ** basis_row] * basis_chain - for t in unroll(1, LIG_LOG_MSG_COLS[ml]): - # subspace-vanishing recurrence for the novel-basis point - basis_chain *= (basis_chain + LIG_VANISH_VALS[vanish + t - 1]) - prefix_eq *= basis_a[GEN ** (basis_row + t)] + basis_b[GEN ** (basis_row + t)] * basis_chain - residual_chain[xr * GEN] = residual_chain[xr] + query_weights[GEN ** (lvl * max_q)* xr] * prefix_eq - # accumulate beta_lvl * (per-level residual sum) into the grand residual - inner_chain[GEN ** (lvl + 1)] = inner_chain[GEN ** lvl] + level_betas[GEN ** lvl] * residual_chain[GEN ** LIG_QUERIES[ml]] + pair = xr * xr + positions = query_positions * GEN ** LIG_POSITIONS_OFF[ml] * pair + weights = query_weights * GEN ** (lvl * max_q) * pair + row_d = basis_d * GEN ** basis_row + eq_0 = residual_query(positions[GEN ** 0], row_d, vanish, LIG_LOG_MSG_COLS[ml]) + eq_1 = residual_query(positions[GEN ** 1], row_d, vanish, LIG_LOG_MSG_COLS[ml]) + residual_chain[xr * GEN] = residual_chain[xr] + weights[GEN ** 0] * eq_0 + weights[GEN ** 1] * eq_1 + level_sum = residual_chain[GEN ** (LIG_QUERIES[ml] // 2)] + if const(LIG_QUERIES[ml] % 2 == 1): + last = LIG_QUERIES[ml] - 1 + eq_last = residual_query(query_positions[GEN ** (LIG_POSITIONS_OFF[ml] + last)], basis_d * GEN ** basis_row, vanish, LIG_LOG_MSG_COLS[ml]) + level_sum = level_sum + query_weights[GEN ** (lvl * max_q + last)] * eq_last + # accumulate beta_lvl * prod_t b_t * (per-level residual sum) into the grand residual + inner_chain[GEN ** (lvl + 1)] = inner_chain[GEN ** lvl] + level_scale[GEN ** lvl] * level_sum # Explicit OOD eq bases at the same terminal point. ood_inner = 0 @@ -1309,7 +1474,7 @@ def open_stacked(m_idx: Const, fs0, fs1, target, commit_root_0, commit_root_1, c for t in unroll(0, yr_log): scalar *= (1 + oz[GEN ** (z_folded + t)] + tail_challenges[GEN ** t]) ood_inner += scalar - return sumcheck_target, point, inner_chain[GEN ** n_levels] + ood_inner, yr_at_tail + return sumcheck_target, point, point1, inner_chain[GEN ** n_levels] + ood_inner, yr_at_tail # ============================== inner-proof verification ============================ @@ -1342,6 +1507,8 @@ def verify_bus_gkr(fs0, fs1, cursor, g_bus_mu, zeta): layers[GEN ** LAYER_COUNT] = root_count layers[GEN ** LAYER_ROW] = gkr_pts layers[GEN ** LAYER_POS] = GEN ** 0 + layers[GEN ** LAYER_HALF] = GEN ** 0 + layers[GEN ** LAYER_PARITY] = 0 # Contract two binary product levels at a time. pair_bounds[g^d] = g^(d//2) is # the radix-four layer count, and shift is g exactly when the depth is odd, in @@ -1379,33 +1546,60 @@ def verify_bus_gkr(fs0, fs1, cursor, g_bus_mu, zeta): nxt[GEN ** (LAYER_PUSH + side)] = evals[2 * side] + c0 * (evals[2 * side] + evals[2 * side + 1]) nxt[GEN ** LAYER_ROW] = nextrow nxt[GEN ** LAYER_POS] = GEN + nxt[GEN ** LAYER_HALF] = GEN ** 0 + nxt[GEN ** LAYER_PARITY] = 0 for x_pair in mul_range(1, pair_bounds[g_bus_mu]): - x_layer = x_pair * x_pair * depth_shift[g_bus_mu] + sq = x_pair * x_pair + x_layer = sq * depth_shift[g_bus_mu] layer = layers * x_layer ** LAYER_SLOTS layer_fs = [layer[GEN ** LAYER_FS0], layer[GEN ** LAYER_FS1]] layer_fs, lam = squeeze(layer_fs) point_row = layer[GEN ** LAYER_ROW] round_pos = layer[GEN ** LAYER_POS] nextrow = point_row * GEN ** MU_CAP - head = rounds * round_pos ** ROUND_SLOTS - head[GEN ** ROUND_FS0] = layer_fs[0] - head[GEN ** ROUND_FS1] = layer_fs[1] - head[GEN ** ROUND_CURSOR] = layer[GEN ** LAYER_CURSOR] - head[GEN ** ROUND_CLAIM] = layer[GEN ** LAYER_PUSH] + lam * (layer[GEN ** LAYER_PULL] + lam * layer[GEN ** LAYER_COUNT]) - for x_round in mul_range(1, x_layer): - rd = rounds * (round_pos * x_round) ** ROUND_SLOTS - nfs0, nfs1, ncur, nclaim, rk = sumcheck_round5(rd[GEN ** ROUND_FS0], rd[GEN ** ROUND_FS1], rd[GEN ** ROUND_CURSOR], rd[GEN ** ROUND_CLAIM], point_row[x_round]) - nextrow[x_round * GEN ** 2] = rk - rd_next = rd * GEN ** ROUND_SLOTS - rd_next[GEN ** ROUND_FS0] = nfs0 - rd_next[GEN ** ROUND_FS1] = nfs1 - rd_next[GEN ** ROUND_CURSOR] = ncur - rd_next[GEN ** ROUND_CLAIM] = nclaim - tail = rounds * (round_pos * x_layer) ** ROUND_SLOTS - tail_fs = [tail[GEN ** ROUND_FS0], tail[GEN ** ROUND_FS1]] - tcur = tail[GEN ** ROUND_CURSOR] - tclaim = tail[GEN ** ROUND_CLAIM] + # Round i's record keeps field k at rounds * g^(k GKR_ROUNDS_CAP + i), so a + # round's record is one product away from the counter. + head = rounds * round_pos + head[GEN ** (ROUND_FS0 * GKR_ROUNDS_CAP)] = layer_fs[0] + head[GEN ** (ROUND_FS1 * GKR_ROUNDS_CAP)] = layer_fs[1] + head[GEN ** (ROUND_CURSOR * GKR_ROUNDS_CAP)] = layer[GEN ** LAYER_CURSOR] + head[GEN ** (ROUND_CLAIM * GKR_ROUNDS_CAP)] = layer[GEN ** LAYER_PUSH] + lam * (layer[GEN ** LAYER_PULL] + lam * layer[GEN ** LAYER_COUNT]) + # 2k rounds, or 2k + 1 when the depth is odd: four a loop step, the state + # between them staying in the frame, then two when k is odd, then the odd one. + half = layer[GEN ** LAYER_HALF] + parity = layer[GEN ** LAYER_PARITY] + for x4 in mul_range(1, half): + step = x4 ** 4 + rd = head * step + pt = point_row * step + round_fs = [rd[GEN ** (ROUND_FS0 * GKR_ROUNDS_CAP)], rd[GEN ** (ROUND_FS1 * GKR_ROUNDS_CAP)]] + cur = rd[GEN ** (ROUND_CURSOR * GKR_ROUNDS_CAP)] + claim = rd[GEN ** (ROUND_CLAIM * GKR_ROUNDS_CAP)] + for q in unroll(0, 4): + # As in gkr_rounds. + round_fs, c1, cur = fs_next(round_fs, cur) + round_fs, c2, cur = fs_next(round_fs, cur) + round_fs, c3, cur = fs_next(round_fs, cur) + round_fs, c4, cur = fs_next(round_fs, cur) + round_fs, rk = squeeze(round_fs) + c0 = claim + pt[GEN ** q] * (c1 + c2 + c3 + c4) + pt[GEN ** (MU_CAP + 2 + q)] = rk # nextrow[round * g^2] + claim = c0 + rk * (c1 + rk * (c2 + rk * (c3 + rk * c4))) + rd_next = rd * GEN ** 4 + rd_next[GEN ** (ROUND_FS0 * GKR_ROUNDS_CAP)] = round_fs[0] + rd_next[GEN ** (ROUND_FS1 * GKR_ROUNDS_CAP)] = round_fs[1] + rd_next[GEN ** (ROUND_CURSOR * GKR_ROUNDS_CAP)] = cur + rd_next[GEN ** (ROUND_CLAIM * GKR_ROUNDS_CAP)] = claim + if parity != 0: + quad = half ** 4 + gkr_rounds(head * quad, point_row * quad, 2) + if x_layer != sq: + gkr_rounds(head * sq, point_row * sq, 1) + tail = head * x_layer + tail_fs = [tail[GEN ** (ROUND_FS0 * GKR_ROUNDS_CAP)], tail[GEN ** (ROUND_FS1 * GKR_ROUNDS_CAP)]] + tcur = tail[GEN ** (ROUND_CURSOR * GKR_ROUNDS_CAP)] + tclaim = tail[GEN ** (ROUND_CLAIM * GKR_ROUNDS_CAP)] evals = StackBuf(4 * N_GKR_SIDES) # the four children of each side, in side order for i in unroll(0, 4 * N_GKR_SIDES): tail_fs, ev, tcur = fs_next(tail_fs, tcur) @@ -1429,6 +1623,8 @@ def verify_bus_gkr(fs0, fs1, cursor, g_bus_mu, zeta): nxt[GEN ** (LAYER_PUSH + side)] = lo + c1 * (lo + hi) nxt[GEN ** LAYER_ROW] = nextrow nxt[GEN ** LAYER_POS] = round_pos * x_layer * GEN + nxt[GEN ** LAYER_HALF] = half * (1 + parity * (1 + GEN)) + nxt[GEN ** LAYER_PARITY] = 1 + parity last = layers * g_bus_mu ** LAYER_SLOTS fs = [last[GEN ** LAYER_FS0], last[GEN ** LAYER_FS1]] cursor = last[GEN ** LAYER_CURSOR] @@ -1472,6 +1668,36 @@ def padding_weight(rows, zeta, tau: Const): return (1 + top) * pad, lead, rest +def index_mle(zeta, n: Const): + # The Index coordinate's MLE over the block's n = κ low coords: + # prod_t (1 + zeta_t * (1 + g^(2^t))). + v = GEN ** 0 + for t in unroll(0, n): + if t == 0: + v = 1 + zeta[GEN ** t] * INDEX_MLE_FACTORS[t] + else: + v *= 1 + zeta[GEN ** t] * INDEX_MLE_FACTORS[t] + return v + + +def block_selector(sel_bits, zeta1_hi, idxc_hi, n: Const): + # A bus block's eq over its n = mu - κ selector bits, the offset's bits above κ, + # each boolean-pinned, and g^offset rebuilt from those bits alone (weights + # g^(2^(κ+k)) = idxc_hi[k] + 1), which the caller pins to the certified offset. + eq = GEN ** 0 + rebuilt = GEN ** 0 + for k in unroll(0, n): + sbit = sel_bits[GEN ** k] + sel_bits[GEN ** k] = sbit * sbit # booleanity as a write-once pin + if k == 0: + eq = sbit + zeta1_hi[GEN ** k] # eq over GF(2) is 1 + b + z + rebuilt = 1 + sbit * idxc_hi[GEN ** k] + else: + eq *= sbit + zeta1_hi[GEN ** k] + rebuilt *= 1 + sbit * idxc_hi[GEN ** k] + return eq, rebuilt + + def verify_tables(fs0, fs1, cursor, pi_0, pi_1, zeta, g_bus_mu, dims_g, table_rows, block_kappa, g_squares, fp_w, beta, claim_pool, claim_cplen_g, chi, claim_push, claim_pull, claim_count): # Settle the bus against the six tables, in four steps: certify each side's # leaf-cube tiling, decompose the three GKR leaf values over it, run the ONE @@ -1551,6 +1777,8 @@ def verify_tables(fs0, fs1, cursor, pi_0, pi_1, zeta, g_bus_mu, dims_g, table_ro idxc_tab = HeapBuf(SIZE_BITS) # INDEX_MLE_FACTORS[t] = 1 + g^(2^t) for t in unroll(0, SIZE_BITS): idxc_tab[GEN ** t] = INDEX_MLE_FACTORS[t] + zeta1 = HeapBuf(g_bus_mu) # 1 + zeta, so an eq factor 1 + b + z is one XOR + zeta1_done = match(log(g_bus_mu), range(0, SIZE_BITS), lambda n: plus_one_run(zeta1, zeta, n)) # a match needs a target bus_table_total = StackBuf(N_GKR_SIDES) # per side, what its tables' blocks owe block_eq_hi = StackBuf(N_BLOCKS) # every block's eq_hi, reused below block_index_mle = HeapBuf(N_BLOCKS) # per push block with an Index coord @@ -1571,23 +1799,15 @@ def verify_tables(fs0, fs1, cursor, pi_0, pi_1, zeta, g_bus_mu, dims_g, table_ro # g^offset from those high bits alone (weights g^(2^(κ+k))) and # asserting it equals block_off_g pins the bits AND the κ-alignment # in one shot; the low κ bit cells are written but never read. - sel_len_g = g_bus_mu / kappa_g # g^(mu - κ) - assert log(sel_len_g) < SIZE_BITS - zeta_hi = zeta * kappa_g offset_bits = HeapBuf(GEN ** SIZE_BITS) hint_decompose_bits_exponent(offset_bits, block_off_g[GEN ** b], SIZE_BITS) + sel_len_g = g_bus_mu / kappa_g # g^(mu - κ) + assert log(sel_len_g) < SIZE_BITS sel_bits = offset_bits * kappa_g - eq_chain = HeapBuf(MU_CAP + 2) - goff_chain = HeapBuf(MU_CAP + 2) # rebuild g^offset from the high bits - eq_chain[GEN ** 0] = 1 - goff_chain[GEN ** 0] = 1 - for xk in mul_range(1, sel_len_g): - sbit = sel_bits[xk] - sel_bits[xk] = sbit * sbit # booleanity as a write-once pin - eq_chain[xk * GEN] = eq_chain[xk] * (1 + sbit + zeta_hi[xk]) # eq over GF(2) is 1 + b + z - goff_chain[xk * GEN] = goff_chain[xk] * (1 + sbit * (g_squares[kappa_g * xk] + 1)) - eq_hi = eq_chain[sel_len_g] - assert goff_chain[sel_len_g] == block_off_g[GEN ** b] # bits == offset >> κ, κ-aligned + zeta1_hi = zeta1 * kappa_g + idxc_hi = idxc_tab * kappa_g + eq_hi, goff = match(log(sel_len_g), range(0, SIZE_BITS), lambda n: block_selector(sel_bits, zeta1_hi, idxc_hi, n)) + assert goff == block_off_g[GEN ** b] # bits == offset >> κ, κ-aligned selector_sum += eq_hi block_eq_hi[b] = eq_hi # A TABLE's block streams no value here: the table sumcheck settles its @@ -1620,12 +1840,7 @@ def verify_tables(fs0, fs1, cursor, pi_0, pi_1, zeta, g_bus_mu, dims_g, table_ro if s == PULL_SIDE: coord_val = block_index_mle[GEN ** (b - SIDE_BLOCK_START[PULL_SIDE])] else: - # Index-coord MLE: prod_t (1 + zeta_t * (1 + g^(2^t))) - idx_chain = HeapBuf(MU_CAP + 2) - idx_chain[GEN ** 0] = 1 - for xt in mul_range(1, kappa_g): - idx_chain[xt * GEN] = idx_chain[xt] * (1 + zeta[xt] * idxc_tab[xt]) - coord_val = idx_chain[kappa_g] + coord_val = match(log(kappa_g), range(0, SIZE_BITS), lambda n: index_mle(zeta, n)) if s == PUSH_SIDE: block_index_mle[GEN ** b] = coord_val if COORD_TYPE[ci] == COORD_KIND_PUBLIC: @@ -1740,9 +1955,7 @@ def verify_tables(fs0, fs1, cursor, pi_0, pi_1, zeta, g_bus_mu, dims_g, table_ro claim = zc_last[GEN ** ROUND_CLAIM] # peq[g^tau] = eq(zeta[..tau], chi[..tau]), as a prefix chain. zc_peq = HeapBuf(g_zc_n * GEN) - zc_peq[GEN ** 0] = 1 - for xi in mul_range(1, g_zc_n): - zc_peq[xi * GEN] = zc_peq[xi] * (1 + zeta[xi] + chi[xi]) + eq_prefix_chain(zc_peq, 1, chi, zeta1, g_zc_n) # Per table: every committed column's evaluation (pooled), its AIR constraint # at its own reduced point chi[..tau_t], weighted into the batch's final claim. air_acc = 0 @@ -1963,6 +2176,16 @@ def verify_flock(fs0, fs1, cursor, tau_blake2s_g, zerocheck_chis, lincheck_rs, z return fs[0], fs[1], cursor, zerocheck_z, lincheck_alpha, matrix_eval +@inline +def ring_map(lin_eval, map_challenges): + for stage in unroll(0, len(RING_MAP_SHIFTS)): + frobenius = lin_eval + for k in unroll(0, RING_MAP_SHIFTS[stage]): + frobenius *= frobenius + lin_eval += map_challenges[GEN ** stage] * frobenius + return lin_eval + + def certify_placement(kappa_base, g_squares): # Certify the native order: descending kappa, then ascending column index. # Accumulate g^offset, with each column advancing it by g^(2^kappa). @@ -1999,46 +2222,83 @@ def certify_placement(kappa_base, g_squares): return g_total, col_off_g, col_kappa_g -def column_selector(offset, point, kappa: Const): +def column_selector(offset, sel_eq, sel_g, kappa: Const): # Rebuild the offset using only bits above the column's kappa, certifying its # alignment. The same bits select the column at the complete opening point. - # Both offset and point are zero above m, so extending to MAX_STACK_LOG adds - # only factors eq(0, 0) = 1. - bits = StackBuf(MAX_STACK_LOG) - hint_decompose_bits_exponent(bits, offset, MAX_STACK_LOG) + # Both offset and point are zero above m, so extending to 4 * SEL_NIBBLES adds + # only factors eq(0, 0) = 1. The bits go a nibble at a time, top bit first: + # their g-power index u (bit 4j+3-t weighing 2^t) reads the nibble's eq over + # its coordinates and its g^offset share off the tables of selector_tables. + bits = StackBuf(4 * SEL_NIBBLES) + hint_decompose_bits_exponent(bits, offset, 4 * SEL_NIBBLES) rebuilt = GEN ** 0 selector = GEN ** 0 - for k in unroll(kappa, MAX_STACK_LOG): - bit = bits[k] - bits[k] = bit * bit - rebuilt *= 1 + bit * (1 + GEN ** (2 ** k)) - selector *= 1 + bit + point[GEN ** k] + for j in unroll(kappa // 4, SEL_NIBBLES): + width = 4 + if const(j == kappa // 4): + width = 4 * j + 4 - kappa + idx = GEN ** 0 + for t in unroll(0, width): + bit = bits[4 * j + 3 - t] + bits[4 * j + 3 - t] = bit * bit + idx *= 1 + bit * (1 + GEN ** (2 ** t)) + selector *= sel_eq[GEN ** (SEL_EQ_STRIDE * j + 2 ** width - 2) * idx] + rebuilt *= sel_g[GEN ** (16 * j) * idx] assert rebuilt == offset return selector -def check_opening_terminal(zeta, chi, rm, g_bus_mu, g_zc_n, g_log_mem, tau_blake2s_g, claim_cplen_g, lam_pool, col_offsets, col_kappas, z_vals, c_table, point, inner_total, yr_at_tail, sumcheck_target): +def selector_tables(point, point1): + # Per nibble j of the selector bits, the eq tables over its top 1..4 + # coordinates (coordinate 4j+3-t is index bit t; the top w form the level-w + # table at 2^w - 2), and the g-powers g^(sum_t u_t 2^(4j+3-t)) of its indices. + sel_eq = HeapBuf(SEL_NIBBLES * SEL_EQ_STRIDE) + sel_g = HeapBuf(SEL_NIBBLES * 16) + for j in unroll(0, SEL_NIBBLES): + levels = StackBuf(SEL_EQ_STRIDE) + levels[0] = point1[GEN ** (4 * j + 3)] + levels[1] = point[GEN ** (4 * j + 3)] + for t in unroll(1, 4): + r = point[GEN ** (4 * j + 3 - t)] + r1 = point1[GEN ** (4 * j + 3 - t)] + for i in unroll(0, 2 ** t): + levels[2 ** (t + 1) - 2 + i] = levels[2 ** t - 2 + i] * r1 + levels[2 ** (t + 1) - 2 + 2 ** t + i] = levels[2 ** t - 2 + i] * r + for i in unroll(0, SEL_EQ_STRIDE): + sel_eq[GEN ** (SEL_EQ_STRIDE * j + i)] = levels[i] + for u in unroll(0, 16): + g_share = GEN ** 0 + for t in unroll(0, 4): + if (u // 2 ** t) % 2 == 1: + g_share *= GEN ** (2 ** (4 * j + 3 - t)) + sel_g[GEN ** (16 * j + u)] = g_share + return sel_eq, sel_g + + +def check_opening_terminal(zeta, chi, rm, g_bus_mu, g_zc_n, g_log_mem, tau_blake2s_g, claim_cplen_g, lam_pool, col_offsets, col_kappas, z_vals, c_table, point, point1, inner_total, yr_at_tail, sumcheck_target): # Evaluate each transparent weight at the complete point in witness order. # A claim is its low point followed by the certified column's selector bits; # q_flock slots prepend their fixed slot bits to the low point. + # point1 = 1 + point holds each coordinate's eq factor against a zero bit. zeta_eq_chain = HeapBuf(SIZE_BITS + 1) - eq_prefix_chain(zeta_eq_chain, 1, zeta, point, g_bus_mu) + eq_prefix_chain(zeta_eq_chain, 1, zeta, point1, g_bus_mu) chi_eq_chain = HeapBuf(SIZE_BITS + 1) - eq_prefix_chain(chi_eq_chain, 1, chi, point, g_zc_n) + eq_prefix_chain(chi_eq_chain, 1, chi, point1, g_zc_n) chi_slot_eq_chain = HeapBuf(SIZE_BITS + 1) - eq_prefix_chain(chi_slot_eq_chain, 1, chi, point * GEN ** SLOT_STRIDE_LOG, tau_blake2s_g) + eq_prefix_chain(chi_slot_eq_chain, 1, chi, point1 * GEN ** SLOT_STRIDE_LOG, tau_blake2s_g) pi_chain = HeapBuf(SIZE_BITS + 1) pi_chain[GEN ** 0] = 1 - pi_chain[GEN ** 1] = 1 + rm + point[GEN ** 0] + pi_chain[GEN ** 1] = rm + point1[GEN ** 0] for xk in mul_range(GEN, g_log_mem): - pi_chain[xk * GEN] = pi_chain[xk] * (1 + point[xk]) + pi_chain[xk * GEN] = pi_chain[xk] * point1[xk] pi_eq = pi_chain[g_log_mem] + sel_eq, sel_g = selector_tables(point, point1) selectors = StackBuf(N_COMMITTED_COLS) for c in unroll(0, N_COMMITTED_COLS): offset = col_offsets[GEN ** c] kappa_g = col_kappas[GEN ** c] - selectors[c] = match(log(kappa_g), range(0, N_COLUMN_LOGS), lambda kappa: column_selector(offset, point, kappa)) + selectors[c] = match(log(kappa_g), range(0, N_COLUMN_LOGS), lambda kappa: column_selector(offset, sel_eq, sel_g, kappa)) inner_sum = inner_total for j in unroll(0, N_CLAIMS): @@ -2053,9 +2313,18 @@ def check_opening_terminal(zeta, chi, rm, g_bus_mu, g_zc_n, g_log_mem, tau_blake if CLAIM_POINT_BUF[j] == POINT_BUF_RHO: low_eq = chi_eq_chain[cplen_g] if CLAIM_POINT_BUF[j] == POINT_BUF_QFLOCK_RHO: + # Top bit first, so claims sharing high slot bits share partial products. slot_eq = GEN ** 0 - for k in unroll(0, SLOT_STRIDE_LOG): - slot_eq *= 1 + CLAIM_QFLOCK_SLOT_BITS[SLOT_STRIDE_LOG * j + k] + point[GEN ** k] + for kk in unroll(0, SLOT_STRIDE_LOG): + k = SLOT_STRIDE_LOG - 1 - kk + if CLAIM_QFLOCK_SLOT_BITS[SLOT_STRIDE_LOG * j + k] == 1: + slot_factor = point[GEN ** k] + else: + slot_factor = point1[GEN ** k] + if kk == 0: + slot_eq = slot_factor + else: + slot_eq *= slot_factor low_eq = slot_eq * chi_slot_eq_chain[cplen_g] nlow = cplen_g * GEN ** SLOT_STRIDE_LOG assert nlow == col_kappas[GEN ** CLAIM_COMMITTED_COL[j]] @@ -2066,8 +2335,7 @@ def check_opening_terminal(zeta, chi, rm, g_bus_mu, g_zc_n, g_log_mem, tau_blake prod_chains = HeapBuf((qflockv_g * GEN) ** BASE_FIELD_BITS) for k in unroll(0, BASE_FIELD_BITS): prod_chains[GEN ** k] = 1 - rs_eq_run(prod_chains, z_vals, point, qflockv_g) - prod_final = prod_chains * qflockv_g ** BASE_FIELD_BITS + prod_final = rs_eq_run(prod_chains, z_vals, point1, qflockv_g) rs_weight = 0 for k in unroll(0, BASE_FIELD_BITS): rs_weight += c_table[GEN ** k] * prod_final[GEN ** k] @@ -2107,10 +2375,10 @@ def verify_sub(pi_0, pi_1, seed_0, seed_1, g_logs_pow2, g_squares, defer_out): fs, x, cursor = fs_next(fs, cursor) table_rows[GEN ** t] = x fs, log_inv_rate, cursor = fs_next(fs, cursor) - rate_sel = g_power_of_word(log_inv_rate, g_squares, LOG_WORD_BITS) / GEN + rate_sel = g_power_of_word(log_inv_rate) / GEN assert log(rate_sel) < LIG_N_RATES dims_g = HeapBuf(N_TABLES + 1) # [g^log_mem, g^tau_0 .. g^tau_{N_TABLES-1}] - g_log_mem = g_power_of_word(log_mem_word, g_squares, LOG_WORD_BITS) + g_log_mem = g_power_of_word(log_mem_word) assert log(g_log_mem) < COUNT_BITS assert log(g_log_mem / GEN ** MIN_LOG_MEM) < COUNT_BITS # native MIN_LOG_MEM <= log_mem dims_g[GEN ** 0] = g_log_mem @@ -2198,23 +2466,22 @@ def verify_sub(pi_0, pi_1, seed_0, seed_1, g_logs_pow2, g_squares, defer_out): for k in unroll(0, shift): coefficient *= coefficient c_table[GEN ** (slot * 2 * shift + shift)] = map_challenge * coefficient - # Evaluate the claim and combine its 64 packing rows: the running x-power and - # the running sum ride one two-slot chain. - rs_chain = HeapBuf(((2 ** K_SKIP) + 1) * PAIR_SLOTS) - rs_chain[GEN ** 0] = GEN ** 0 # x^i + # Evaluate the claim and combine its 64 packing rows, sum_i x^i L(slice_i), + # RS_CHUNK rows an iteration in Horner form: the running x-power and the + # running sum ride one two-slot chain. + rs_chain = HeapBuf(((2 ** K_SKIP) // RS_CHUNK + 1) * PAIR_SLOTS) + rs_chain[GEN ** 0] = GEN ** 0 # x^(RS_CHUNK j) rs_chain[GEN ** 1] = 0 # the running sum - for x_round in mul_range(1, GEN ** (2 ** K_SKIP)): - lin_eval = z_partial[x_round] - for stage in unroll(0, len(RING_MAP_SHIFTS)): - frobenius = lin_eval - for k in unroll(0, RING_MAP_SHIFTS[stage]): - frobenius *= frobenius - lin_eval += map_challenges[GEN ** stage] * frobenius + for x_round in mul_range(1, GEN ** ((2 ** K_SKIP) // RS_CHUNK)): + rows = z_partial * x_round ** RS_CHUNK + chunk = ring_map(rows[GEN ** (RS_CHUNK - 1)], map_challenges) + for r in unroll(1, RS_CHUNK): + chunk = chunk * 2 + ring_map(rows[GEN ** (RS_CHUNK - 1 - r)], map_challenges) row = rs_chain * x_round ** PAIR_SLOTS x_pow = row[GEN ** 0] - row[GEN ** PAIR_SLOTS] = x_pow * 2 - row[GEN ** (PAIR_SLOTS + 1)] = row[GEN ** 1] + x_pow * lin_eval - rs_end = rs_chain * (GEN ** (2 ** K_SKIP)) ** PAIR_SLOTS + row[GEN ** PAIR_SLOTS] = x_pow * const(2 ** RS_CHUNK) + row[GEN ** (PAIR_SLOTS + 1)] = row[GEN ** 1] + x_pow * chunk + rs_end = rs_chain * (GEN ** ((2 ** K_SKIP) // RS_CHUNK)) ** PAIR_SLOTS transposed_claim = rs_end[GEN ** 1] # Suffix point for the transparent weight. for t in unroll(0, LINCHECK_ROUNDS): @@ -2250,13 +2517,13 @@ def verify_sub(pi_0, pi_1, seed_0, seed_1, g_logs_pow2, g_squares, defer_out): # dispatch independently for every inner proof in a mixed-rate batch. config_sel = size_sel * rate_sel ** LIG_N_LOG_SIZES assert log(config_sel) < LIG_N_CANDIDATES - sumcheck_target, point, inner_total, yr_at_tail = match(log(config_sel), range(0, LIG_N_CANDIDATES), lambda m_idx: open_stacked(m_idx, fs[0], fs[1], target, commit_root_0, commit_root_1, cursor)) + sumcheck_target, point, point1, inner_total, yr_at_tail = match(log(config_sel), range(0, LIG_N_CANDIDATES), lambda m_idx: open_stacked(m_idx, fs[0], fs[1], target, commit_root_0, commit_root_1, cursor)) # `stream` is a fixed-capacity witness transport. The shape fixes the exact # consumed prefix, whose every word is transcript-bound; the unused suffix # is outside the recursively verified proof and intentionally unconstrained. # ---- generalized eval_b terminal (runtime claim shapes) ---- - check_opening_terminal(zeta, chi, rm, g_bus_mu, g_zc_n, g_log_mem, tau_blake2s_g, claim_cplen_g, lam_pool, col_offsets, col_kappas, z_vals, c_table, point, inner_total, yr_at_tail, sumcheck_target) + check_opening_terminal(zeta, chi, rm, g_bus_mu, g_zc_n, g_log_mem, tau_blake2s_g, claim_cplen_g, lam_pool, col_offsets, col_kappas, z_vals, c_table, point, point1, inner_total, yr_at_tail, sumcheck_target) # ---- export this sub-proof's deferred-claim data to the caller (FRESH_*) ---- for k in unroll(0, BYTECODE_LOG): @@ -2442,9 +2709,9 @@ def verify_lean_xmss(signer): digit_product, encoded = match(log(group), range(0, DIGIT_GROUPS), lambda k: walk_group(steps, pp, step_md[0], tips, DIGIT_GROUP * p, DIGIT_GROUP, k, 1, powers, digit_product, encoded)) group = hint_witness("digits") assert log(group) < CHAIN_LENGTH - digit_product, encoded = match(log(group), range(0, CHAIN_LENGTH), lambda k: walk_group(steps, pp, step_md[0], tips, LX_V - 1, 1, k, 1, powers, digit_product, encoded)) + # The digits land where the digest already is, which is the assert. + digit_product, digest[0] = match(log(group), range(0, CHAIN_LENGTH), lambda k: walk_group(steps, pp, step_md[0], tips, LX_V - 1, 1, k, 1, powers, digit_product, encoded)) assert digit_product == GEN ** LX_TARGET_SUM - assert encoded == digest[0] # The one-time leaf, then the path to the layer's root. leaf = StackBuf(WORDS_PER_BLOCK) @@ -2483,46 +2750,68 @@ def sp_forest_key(at, pp, t: Const): @inline -def sp_fold(node, at, pp, md, out, slot: Const, index: Const, height: Const, tw: Const, level_mul: Const, node_mul: Const): +def sp_fold(node, at, pp, md, out, slot: Const, index: Const, height: Const, tw: Const, level_mul: Const, node_mul: Const, own_mul: Const): # Fold leaf `index` of a tree of the forest up its hinted path, into the # tree's two top nodes, out[slot] and out[slot + 2], each beside a cell nothing # reads. No hash joins them: the path's last element is the top node the fold # does not reach. The index is the arm's constant, so the order of each pair - # and every node's address are the program's. + # and every node's address are the program's. `at` is the leaf's own address, + # own_mul its field's unit: a node's level and index take that field's place. cur = node for level in unroll(0, height - 2): sibling = hint_witness("sp_forest_nodes") parent = StackBuf(WORDS_PER_BLOCK) if const((index // 2 ** level) % 2 == 0): - blake2s([pp, at + const(tw + (level + 1) * level_mul + (index // 2 ** (level + 1)) * node_mul)], [cur, sibling], parent, md=md) + blake2s([pp, at + (const(index * own_mul) + const(tw + (level + 1) * level_mul + (index // 2 ** (level + 1)) * node_mul))], [cur, sibling], parent, md=md) else: - blake2s([pp, at + const(tw + (level + 1) * level_mul + (index // 2 ** (level + 1)) * node_mul)], [sibling, cur], parent, md=md) + blake2s([pp, at + (const(index * own_mul) + const(tw + (level + 1) * level_mul + (index // 2 ** (level + 1)) * node_mul))], [sibling, cur], parent, md=md) cur = parent[0] sibling = hint_witness("sp_forest_nodes") if const((index // 2 ** (height - 2)) % 2 == 0): - blake2s([pp, at + const(tw + (height - 1) * level_mul + (index // 2 ** (height - 1)) * node_mul)], [cur, sibling], out[slot + 2 * (index // 2 ** (height - 1)):slot + 2 * (index // 2 ** (height - 1)) + 2], md=md) + blake2s([pp, at + (const(index * own_mul) + const(tw + (height - 1) * level_mul + (index // 2 ** (height - 1)) * node_mul))], [cur, sibling], out[slot + 2 * (index // 2 ** (height - 1)):slot + 2 * (index // 2 ** (height - 1)) + 2], md=md) else: - blake2s([pp, at + const(tw + (height - 1) * level_mul + (index // 2 ** (height - 1)) * node_mul)], [sibling, cur], out[slot + 2 * (index // 2 ** (height - 1)):slot + 2 * (index // 2 ** (height - 1)) + 2], md=md) + blake2s([pp, at + (const(index * own_mul) + const(tw + (height - 1) * level_mul + (index // 2 ** (height - 1)) * node_mul))], [sibling, cur], out[slot + 2 * (index // 2 ** (height - 1)):slot + 2 * (index // 2 ** (height - 1)) + 2], md=md) hint_witness(out[slot + 2 - 2 * (index // 2 ** (height - 1)):slot + 3 - 2 * (index // 2 ** (height - 1))], "sp_forest_nodes") return @inline -def sp_subtree_arm(node, sub_at, pp, md, subtops, c: Const, j: Const, a: Const): - # The path of WOTS key `a` in its subtree. Returns `a` as its addresses hold it - # and as the digest does, which is what binds the dispatch. - sp_fold(node, sub_at, pp, md, subtops, 4 * j, a, SP_SUBTREE_HEIGHT, SP_TW_FSUBNODE, SP_SUBNODE_LEVEL_MUL, SP_SUBNODE_MUL) - return const(a * SP_KEY_MUL), const(a) * COORD_BASIS[SP_KEY_OFF + (2 * c + j) * SP_SUBTREE_HEIGHT] +def sp_subtree_arm(node, key_at, pp, md, subtops, c: Const, j: Const, a: Const): + # The path of WOTS key `a` in subtree j, from the key's address. Returns the + # subtree and `a` as its addresses hold them, and `a` as the digest does, which + # is what binds the dispatch. + sp_fold(node, key_at, pp, md, subtops, 4 * j, a, SP_SUBTREE_HEIGHT, SP_TW_FSUBNODE, SP_SUBNODE_LEVEL_MUL, SP_SUBNODE_MUL, SP_KEY_MUL) + return const(j * SP_SUBTREE_MUL + a * SP_KEY_MUL), const(a) * COORD_BASIS[SP_KEY_OFF + (2 * c + j) * SP_SUBTREE_HEIGHT] @inline -def sp_tree_arm(node, tree_at, pp, md, tops, c: Const, s: Const): - # The path of leaf `s` in tree c, likewise. - sp_fold(node, tree_at, pp, md, tops, 4 * c, s, SP_TREE_HEIGHT, SP_TW_FNODE, SP_FNODE_LEVEL_MUL, SP_FNODE_MUL) - return const(s * SP_LEAF_MUL), const(s) * COORD_BASIS[SP_LEAF_OFF + c * SP_TREE_HEIGHT] +def sp_tree_arm(node, leaf_at, pp, md, tops, c: Const, s: Const): + # The path of leaf `s` in tree c, from the leaf's address, likewise. + sp_fold(node, leaf_at, pp, md, tops, 4 * c, s, SP_TREE_HEIGHT, SP_TW_FNODE, SP_FNODE_LEVEL_MUL, SP_FNODE_MUL, SP_LEAF_MUL) + return const(c * HI_MUL + s * SP_LEAF_MUL), const(s) * COORD_BASIS[SP_LEAF_OFF + c * SP_TREE_HEIGHT] + +@inline +def sp_from_bytes(byte_bits, marks, n: Const, top: Const): + # The integer whose n bytes, low first, are hinted in the exponent at marks[k]: + # each range checked (the top one below `top`, SP_BYTE_VALUES or SP_UNUSED_TOP) + # and read off `byte_bits` as the bits it spells, so the integer is below + # top * 2^(8 (n - 1)). + for t in unroll(0, n): + k = n - 1 - t + if const(t == 0): + if const(top == SP_BYTE_VALUES): + assert log(marks[k]) < SP_BYTE_VALUES + else: + assert log(marks[k]) < SP_UNUSED_TOP + value = byte_bits[marks[k]] + else: + assert log(marks[k]) < SP_BYTE_VALUES + value = byte_bits[marks[k]] + value * GEN ** 8 + return value -def verify_sig_lean_sphincs(signer): + +def verify_sig_lean_sphincs(signer, byte_bits): # `signer` is one 4-cell entry of the leanSPHINCS coverage table: the pubkey's root and # public parameter, then the message THAT signer signed, the two bound together # by the signer-set digest, as for leanXMSS. @@ -2545,84 +2834,71 @@ def verify_sig_lean_sphincs(signer): # The digest's first cell is the sixteen codewords and its second is bit # fields. A field a dispatch takes (a codeword, a tree's leaf, a subtree's # WOTS key) is hinted in the exponent and bound by what its arm returns, the - # field as the digest holds it. The index and the unused bits are advice bits, - # pinned Boolean. `bound` gathers what the second cell's fields weigh. + # field as the digest holds it. The index is advice bits, pinned Boolean, and + # the unused bits hinted bytes (sp_from_bytes). `bound` gathers what the second + # cell's fields weigh. bits = StackBuf(SP_H) hint_decompose_bits(bits, digest[1], SP_H) for i in unroll(0, SP_H): bits[i] = bits[i] * bits[i] - high = StackBuf(SP_UNUSED_BITS) - hint_decompose_bits(high, digest[1], SP_UNUSED_BITS, SP_USED_BITS) - for i in unroll(0, SP_UNUSED_BITS): - high[i] = high[i] * high[i] # nodes[l] is idx >> l: the index, then the node above it on each level. nodes = StackBuf(SP_H) shifts(bits, 0, SP_H, nodes) idx = nodes[0] - unused = StackBuf(SP_UNUSED_BITS) - shifts(high, 0, SP_UNUSED_BITS, unused) - bound = idx + unused[0] * COORD_BASIS[SP_USED_BITS] + unused_bytes = StackBuf(SP_UNUSED_BYTES) + hint_witness(unused_bytes, "sp_unused") + unused = sp_from_bytes(byte_bits, unused_bytes, SP_UNUSED_BYTES, SP_UNUSED_TOP) + bound = idx + unused * COORD_BASIS[SP_USED_BITS] # ---- the few-time signature: two opened WOTS keys per tree of the forest ---- # A tree's top nodes land in tops[4c] and tops[4c + 2]. tops = StackBuf(4 * SP_TREES) words = StackBuf(2 * SP_TREES) for c in unroll(0, SP_TREES): - if const(c == 0): - tree_at = idx - else: - tree_at = idx + const(c * HI_MUL) - # The tree's leaf: g^s, and s where an address holds it, which the - # dispatch on the first writes again. + # The tree's leaf: g^s, and the tree and s where an address holds them, + # which the dispatch on the first writes again. leaf_mark = StackBuf(2) hint_witness(leaf_mark, "sp_leaves") assert log(leaf_mark[0]) < SP_LEAVES - leaf_at = tree_at + leaf_mark[1] + leaf_at = idx + leaf_mark[1] subtops = StackBuf(8) for j in unroll(0, 2): - if const(j == 0): - sub_at = leaf_at - else: - sub_at = leaf_at + SP_SUBTREE_MUL + # Likewise the subtree and its WOTS key. key_mark = StackBuf(2) hint_witness(key_mark, "sp_keys") assert log(key_mark[0]) < SP_KEYS word = hint_witness("sp_words") assert log(word) < SP_WORDS - key_at = sub_at + key_mark[1] + key_at = leaf_at + key_mark[1] node, words[2 * c + j] = match(log(word), range(0, SP_WORDS), lambda t: sp_forest_key(key_at, pp, t)) - key_mark[1], key_weight = match(log(key_mark[0]), range(0, SP_KEYS), lambda a: sp_subtree_arm(node, sub_at, pp, node_md[0], subtops, c, j, a)) + key_mark[1], key_weight = match(log(key_mark[0]), range(0, SP_KEYS), lambda a: sp_subtree_arm(node, key_at, pp, node_md[0], subtops, c, j, a)) bound = bound + key_weight st = StackBuf(WORDS_PER_BLOCK) blake2s([pp, leaf_at + SP_TW_FTREE_LEAF], [subtops[0], subtops[2]], st, counter=64, final=0) leaf = StackBuf(WORDS_PER_BLOCK) blake2s([subtops[4], subtops[6]], [0, 0], leaf, cv=st, counter=96, final=1) - leaf_mark[1], leaf_weight = match(log(leaf_mark[0]), range(0, SP_LEAVES), lambda s: sp_tree_arm(leaf[0], tree_at, pp, node_md[0], tops, c, s)) - bound = bound + leaf_weight + leaf_mark[1], leaf_weight = match(log(leaf_mark[0]), range(0, SP_LEAVES), lambda s: sp_tree_arm(leaf[0], leaf_at, pp, node_md[0], tops, c, s)) + if const(c + 1 == SP_TREES): + digest[1] = bound + leaf_weight # stored where the digest already is: the assert + else: + bound = bound + leaf_weight # The codewords, a byte each: eight to a lane, low byte first. lane_0 = words[7] lane_1 = words[15] for k in unroll(1, 8): lane_0 = words[7 - k] + lane_0 * GEN ** 8 lane_1 = words[15 - k] + lane_1 * GEN ** 8 - assert lane_0 + lane_1 * Y_TOWER == digest[0] - assert bound == digest[1] + digest[0] = lane_0 + lane_1 * Y_TOWER # The few-time key: P, A and the 2 * SP_TREES top nodes. forest_key = StackBuf(WORDS_PER_BLOCK) hash_values(pp, idx + SP_TW_FKEY, tops, 2 * SP_TREES, forest_key) # ---- the one-time signature of the few-time key ---- - # D = Th(P, A, key | LE_32(c)), 52 bytes. The counter is bounded by decomposing - # exactly its bits, so the guest accepts no preimage the native verifier cannot - # parse. - ctr = hint_witness("sp_counter") - ctr_bits = StackBuf(SP_COUNTER_BITS) - hint_decompose_bits(ctr_bits, ctr, SP_COUNTER_BITS) - ctr_bound = 0 - for i in unroll(0, SP_COUNTER_BITS): - ctr_bits[i] = ctr_bits[i] * ctr_bits[i] - ctr_bound = ctr_bits[SP_COUNTER_BITS - 1 - i] + GEN * ctr_bound - assert ctr_bound == ctr + # D = Th(P, A, key | LE_32(c)), 52 bytes. The counter is rebuilt from its four + # bytes, so the guest accepts no preimage the native verifier cannot parse. + ctr_bytes = StackBuf(SP_COUNTER_BYTES) + hint_witness(ctr_bytes, "sp_counter") + ctr = sp_from_bytes(byte_bits, ctr_bytes, SP_COUNTER_BYTES, SP_BYTE_VALUES) encoding = StackBuf(WORDS_PER_BLOCK) blake2s([pp, idx + SP_TW_ENC], [forest_key[0], ctr], encoding, counter=52, final=1) @@ -2643,9 +2919,9 @@ def verify_sig_lean_sphincs(signer): # The chain the groups leave over. group = hint_witness("digits") assert log(group) < CHAIN_LENGTH - digit_product, encoded = match(log(group), range(0, CHAIN_LENGTH), lambda k: walk_group(idx, pp, step_md[0], tips, SP_V - 1, 1, k, 0, powers, digit_product, encoded)) + # The digits land where D already is, which is the assert. + digit_product, encoding[0] = match(log(group), range(0, CHAIN_LENGTH), lambda k: walk_group(idx, pp, step_md[0], tips, SP_V - 1, 1, k, 0, powers, digit_product, encoded)) assert digit_product == GEN ** SP_TARGET_SUM - assert encoded == encoding[0] # The one-time leaf, then the path from it to the public key's root. leaf = StackBuf(WORDS_PER_BLOCK) hash_values(pp, idx + TW_LEAF, tips, SP_V, leaf) @@ -3229,11 +3505,16 @@ def main(): assert log(idx) < log(lx_slots_g) cover[idx] = xi verify_lean_xmss(lx_table * (idx ** 4)) - for xj in mul_range(1, n_raw_s_g): - off_hint = hint_witness("sp_raw_index") - assert log(off_hint) < log(lean_sphincs_slots_g) - cover[lx_slots_g * off_hint] = n_raw_x_g * xj - verify_sig_lean_sphincs(lean_sphincs_table * (off_hint ** 4)) + if n_raw_s_g != 1: + # byte_bits[g^v] = v, read at a hinted byte's exponent (sp_from_bytes). + byte_bits = HeapBuf(SP_BYTE_VALUES) + for v in unroll(0, SP_BYTE_VALUES): + byte_bits[GEN ** v] = v + for xj in mul_range(1, n_raw_s_g): + off_hint = hint_witness("sp_raw_index") + assert log(off_hint) < log(lean_sphincs_slots_g) + cover[lx_slots_g * off_hint] = n_raw_x_g * xj + verify_sig_lean_sphincs(lean_sphincs_table * (off_hint ** 4), byte_bits) if n_direct_da_g == GEN: root_0, root_1, vector_0, vector_1 = da_verify(g_squares) diff --git a/crates/rec_aggregation/src/aggregation.rs b/crates/rec_aggregation/src/aggregation.rs index 27c0ed7ab..6bb661535 100644 --- a/crates/rec_aggregation/src/aggregation.rs +++ b/crates/rec_aggregation/src/aggregation.rs @@ -160,6 +160,14 @@ fn count(n: usize) -> F192 { F192::new(g_pow(n).0, 0, 0) } +/// The bytes of a `bits`-bit value, low first, each in the exponent: how the guest +/// bounds a value it rebuilds (`sp_from_bytes`). +fn exponent_bytes(value: u64, bits: usize) -> Vec { + (0..bits.div_ceil(8)) + .map(|k| count((value >> (8 * k) & 0xff) as usize)) + .collect() +} + /// A field element as the decimal `u128` literal the zkDSL parser accepts. fn dsl_u128(value: F192) -> u128 { assert_eq!(value.c2, 0, "u128 DSL literal cannot encode the top F192 limb"); @@ -1675,6 +1683,17 @@ const STREAM_CAP: usize = 8192; /// Named hint entries for a single sub-proof, ordered within each stream. type SubHints = Vec<(&'static str, Vec)>; +/// An edit to a node's hints: how a test hands the guest a witness the honest host +/// would not. +#[cfg(test)] +type Forgery = Box; + +#[cfg(test)] +thread_local! { + /// The forgery the next node this thread aggregates gets. + static FORGE: std::cell::RefCell> = const { std::cell::RefCell::new(None) }; +} + /// One `hint_witness` stream: a name and its entries, in the order the guest /// pops them. #[derive(Default)] @@ -1688,6 +1707,13 @@ impl Hints { } } + /// Entry `i` of stream `name`, for a test to forge. + #[cfg(test)] + fn entry_mut(&mut self, name: &str, i: usize) -> &mut Vec { + let (_, entries) = self.0.iter_mut().find(|(n, _)| n == name).expect("a hinted stream"); + &mut entries[i] + } + fn install(self, guest: &mut Program) { for (name, entries) in self.0 { guest.set_witness(name, entries); @@ -1891,6 +1917,17 @@ fn push_lean_sphincs_hints( hints.push("sp_forest_values", vec![pack_16_bytes(value)]); } let (idx, marks) = lean_sphincs::message_digest(pp, &sig.randomizer, message); + // The digest's bits past its fields, which only the guest's rebuild of it reads. + let digest = lean_sphincs::hasher(pp, &lean_sphincs::tweak(lean_sphincs::TWEAK_MSG, 0, 0)) + .update(message) + .update(&sig.randomizer) + .finalize(); + let unused_bits = 8 * digest.len() - lean_sphincs::DIGEST_BITS; + let unused = (0..unused_bits).fold(0u64, |value, bit| { + let position = lean_sphincs::DIGEST_BITS + bit; + value | u64::from(digest[position / 8] >> (position % 8) & 1) << bit + }); + hints.push("sp_unused", exponent_bytes(unused, unused_bits)); for word in marks.iter().flat_map(|mark| mark.words) { hints.push("sp_words", vec![count(word)]); } @@ -1899,10 +1936,15 @@ fn push_lean_sphincs_hints( let address = lean_sphincs::address(&lean_sphincs::tweak(0, field, 0)); vec![count(index), pack_16_bytes(&address)] }; - for (tree, mark_of) in sig.forest.iter().zip(&marks) { - hints.push("sp_leaves", mark(mark_of.leaf, lean_sphincs::leaf_field(mark_of.leaf))); - for (subtree, &key) in tree.subtrees.iter().zip(&mark_of.keys) { - hints.push("sp_keys", mark(key, lean_sphincs::key_field(key))); + // A leaf's address field carries its tree, a key's its subtree, so the guest + // reaches either address in one sum off the index. + for (c, (tree, mark_of)) in sig.forest.iter().zip(&marks).enumerate() { + hints.push( + "sp_leaves", + mark(mark_of.leaf, c as u32 | lean_sphincs::leaf_field(mark_of.leaf)), + ); + for (j, (subtree, &key)) in tree.subtrees.iter().zip(&mark_of.keys).enumerate() { + hints.push("sp_keys", mark(key, (128 * j as u32) | lean_sphincs::key_field(key))); for node in &subtree.path { hints.push("sp_forest_nodes", vec![pack_16_bytes(node)]); } @@ -1914,7 +1956,10 @@ fn push_lean_sphincs_hints( let forest_key = lean_sphincs::forest_recover(pp, idx as u32, &marks, &sig.forest); let codeword = lean_sphincs::encode(pp, idx as u32, &forest_key, sig.counter) .ok_or(AggregationError::MalformedRawSignature)?; - hints.push("sp_counter", vec![F192::new(u64::from(sig.counter), 0, 0)]); + hints.push( + "sp_counter", + exponent_bytes(u64::from(sig.counter), 8 * lean_sphincs::COUNTER_LEN), + ); for entry in digit_hints(&codeword) { hints.push("digits", entry); } @@ -2221,6 +2266,10 @@ pub fn aggregate_with_stats( // `2^MU_MIN`. A run smaller than that (a leaf of a few dozen signatures) grows // its SET table until it clears the floor. program.min_log_committed = MU_MIN; + #[cfg(test)] + if let Some(forge) = FORGE.with(|forge| forge.borrow_mut().take()) { + forge(&mut hints); + } hints.install(&mut program); let (proof, stats) = prove(&program, public_input, log_inv_rate).map_err(AggregationError::ProofError)?; Ok(( @@ -3119,6 +3168,32 @@ mod tests { assert_eq!(aggregate.lean_sphincs_signers.len(), 1); } + /// The leanSPHINCS digest's fields are bound to it in-circuit: a node whose hints + /// give a tree's leaf the address field of another tree, or the digest's unused + /// bits off by one, fails the guest's run. + #[test] + fn forged_lean_sphincs_digest_fields_are_rejected() { + leanvm_core::init_prover(); + let signers = get_lean_sphincs_signers(1); + let forgeries: [Forgery; 2] = [ + Box::new(|hints: &mut Hints| { + let other_tree = pack_16_bytes(&lean_sphincs::address(&lean_sphincs::tweak(0, 1, 0))); + hints.entry_mut("sp_leaves", 0)[1] += other_tree; + }), + Box::new(|hints: &mut Hints| { + hints.entry_mut("sp_unused", 0)[0] *= count(1); + }), + ]; + for forgery in forgeries { + FORGE.with(|forge| *forge.borrow_mut() = Some(forgery)); + let result = aggregate(&[], vec![], signers.clone(), &[], None, LOG_INV_RATE); + assert!( + matches!(result, Err(AggregationError::ProofError(_))), + "a forged digest field is rejected" + ); + } + } + #[test] fn aggregate_one_signer() { leanvm_core::init_prover(); @@ -3222,11 +3297,17 @@ mod tests { let source = format!( r#"{helpers} def main(): - point = HeapBuf(MAX_STACK_LOG) + point = HeapBuf(4 * SEL_NIBBLES) hint_witness(point[0:MAX_STACK_LOG], "point") + for k in unroll(MAX_STACK_LOG, 4 * SEL_NIBBLES): + point[GEN ** k] = 0 + point1 = HeapBuf(4 * SEL_NIBBLES) + for k in unroll(0, 4 * SEL_NIBBLES): + point1[GEN ** k] = 1 + point[GEN ** k] + sel_eq, sel_g = selector_tables(point, point1) offset = hint_witness("offset") kappa_g = hint_witness("kappa") - weight = match(log(kappa_g), range(0, N_COLUMN_LOGS), lambda kappa: column_selector(offset, point, kappa)) + weight = match(log(kappa_g), range(0, N_COLUMN_LOGS), lambda kappa: column_selector(offset, sel_eq, sel_g, kappa)) public = GEN ** 0 assert public[1] == weight return @@ -3285,10 +3366,10 @@ def main(): direction = addr(bits) leaf = StackBuf(2) hint_witness(leaf[0:2], "leaf") - a, b = verify_merkle_path(leaf[0], leaf[1], direction, 3) + root = verify_merkle_path(leaf, direction, 3) public = GEN ** 0 - assert public[1] == a - assert public[GEN] == b + assert public[1] == root[0] + assert public[GEN] == root[1] return "# );