Skip to content

Commit 02a332a

Browse files
committed
block whir: the finish packs its tables as it builds them, phase A takes them narrow (cut b2)
After b1 the 4.13x high-water is the finish itself (BIG 561: at p5's end the finish holds its tables, 24.6 GiB at eight bytes a cell, beside held_narrow and its op lists). b2 keeps those tables packed from the moment each is generated. - multilinear: NarrowColumns::pack_row_major packs a row-major trace column by column (#1013's NarrowMain::pack, the same layout), and TraceData::new_narrow holds packed columns from the start. - stark: TraceTable::pack_main_narrow / narrow_main / take_narrow_main (the packed main trace, no 64-bit copy kept); CommittedTable::from_narrow (a layout whose committed columns are the main columns in order). - trace builder: StreamSkip::pack and WindowedTraceBuilder::pack_finished_tables: each table the finish generates is packed as it is made, chunk by chunk. KECCAK, ECSM, ECDAS and an unchunked KECCAK_RND stay wide (the block cuts them after the build). - block whir: BlockOptions::pack_finished (production on, BLOCK_WHIR_PACK_FINISHED=0|1): a packed table is laid out narrow (table_of_narrow: the packed columns become the table's, its preprocessed columns checked word for word against the program's), with no transposition and no wide column made. The BLOCK REST PACKED line counts them. - phase A: a group holding a narrow table is uploaded packed and widened on the card (upload_group), the commit reads handles (a host fallback widens), and M4's packer skips what is narrow already. The memory log counts held bytes. The tables, groups and proof bytes are the same: a test proves test_keccak_multi with b1 and b2 on and off, inline and on three layout workers, and compares the partition, the statement and (deterministic grind) the bytes; a packed table whose first preprocessed column differs in one word is refused (mutation: with the comparison off, that test fails).
1 parent 6c05dfe commit 02a332a

12 files changed

Lines changed: 661 additions & 107 deletions

File tree

‎crypto/multilinear/src/constraint_argument.rs‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -286,6 +286,57 @@ impl<F: IsField + 'static, E: IsField + 'static> TraceData<F, E> {
286286
})
287287
}
288288

289+
/// [`Self::new`] over columns held narrow from the start: the packed
290+
/// columns are the committed ones, and no field element is made for them
291+
/// on the host unless a host reader asks. Refused unless `F` is Goldilocks
292+
/// (whose elements are their raw words), the rows are a power of two, and
293+
/// the public tables are the rows' height and the ones `kinds` asks for.
294+
pub fn new_narrow(
295+
packed: NarrowColumns,
296+
kinds: Vec<FactorKind>,
297+
public: Vec<Mle<E>>,
298+
) -> Result<Self, Error> {
299+
if std::any::TypeId::of::<F>()
300+
!= std::any::TypeId::of::<math::field::goldilocks::GoldilocksField>()
301+
{
302+
return Err(Error::EmptyPolynomial);
303+
}
304+
if packed.cols() == 0 {
305+
return Err(Error::EmptyPolynomial);
306+
}
307+
if !packed.rows().is_power_of_two() {
308+
return Err(Error::NotPowerOfTwo(packed.rows()));
309+
}
310+
let num_vars = packed.rows().trailing_zeros() as usize;
311+
for got in public.iter().map(Mle::num_vars) {
312+
if got != num_vars {
313+
return Err(Error::VariableCountMismatch {
314+
expected: num_vars,
315+
got,
316+
});
317+
}
318+
}
319+
let wanted = kinds.iter().filter(|k| k.source().is_none()).count();
320+
if public.len() != wanted {
321+
return Err(Error::QueryCountMismatch {
322+
expected: wanted,
323+
got: public.len(),
324+
});
325+
}
326+
Ok(Self {
327+
num_columns: packed.cols(),
328+
num_vars,
329+
columns: Held::Narrow {
330+
packed,
331+
widened: std::sync::OnceLock::new(),
332+
},
333+
public,
334+
device: std::sync::Mutex::new(None),
335+
resident: None,
336+
kinds,
337+
})
338+
}
339+
289340
/// The committed columns on the host: widened here the first time a
290341
/// narrow table is asked ([`crate::narrow::host_widens`] counts it), and
291342
/// that copy kept until [`Self::drop_widened`].

‎crypto/multilinear/src/narrow.rs‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,60 @@ impl NarrowColumns {
106106
Self::from_parts(rows, widths, data)
107107
}
108108

109+
/// [`Self::pack`] of a row-major trace: the `cols`-wide rows of words
110+
/// `row_major`, packed column by column (the stark crate's
111+
/// `NarrowMain::pack`, a block of rows at a time so the block stays in
112+
/// cache across the columns). On the caller's thread. `None` when `cols`
113+
/// is zero or does not divide the words.
114+
pub fn pack_row_major(row_major: &[u64], cols: usize) -> Option<Self> {
115+
const BLOCK_ROWS: usize = 1 << 10;
116+
if cols == 0 || !row_major.len().is_multiple_of(cols) {
117+
return None;
118+
}
119+
let rows = row_major.len() / cols;
120+
let mut max = vec![0u64; cols];
121+
for row in row_major.chunks_exact(cols) {
122+
for (m, &v) in max.iter_mut().zip(row) {
123+
*m = (*m).max(v);
124+
}
125+
}
126+
let widths: Vec<u8> = max.into_iter().map(width_of).collect();
127+
let total: usize = widths.iter().map(|&w| rows * w as usize).sum();
128+
let mut data = vec![0u8; total];
129+
let mut columns: Vec<&mut [u8]> = Vec::with_capacity(cols);
130+
let mut rest = data.as_mut_slice();
131+
for &w in &widths {
132+
let (column, tail) = rest.split_at_mut(rows * w as usize);
133+
columns.push(column);
134+
rest = tail;
135+
}
136+
for first in (0..rows).step_by(BLOCK_ROWS) {
137+
let last = (first + BLOCK_ROWS).min(rows);
138+
let block = &row_major[first * cols..last * cols];
139+
for (c, (column, &w)) in columns.iter_mut().zip(&widths).enumerate() {
140+
let words = block.iter().skip(c).step_by(cols);
141+
let w = w as usize;
142+
let out = &mut column[first * w..last * w];
143+
match w {
144+
1 => out.iter_mut().zip(words).for_each(|(o, &v)| *o = v as u8),
145+
2 => out
146+
.chunks_exact_mut(2)
147+
.zip(words)
148+
.for_each(|(o, &v)| o.copy_from_slice(&(v as u16).to_le_bytes())),
149+
4 => out
150+
.chunks_exact_mut(4)
151+
.zip(words)
152+
.for_each(|(o, &v)| o.copy_from_slice(&(v as u32).to_le_bytes())),
153+
_ => out
154+
.chunks_exact_mut(8)
155+
.zip(words)
156+
.for_each(|(o, &v)| o.copy_from_slice(&v.to_le_bytes())),
157+
}
158+
}
159+
}
160+
Self::from_parts(rows, widths, data)
161+
}
162+
109163
/// [`Self::pack`] of field columns. `None` unless `F` is Goldilocks, whose
110164
/// elements are their raw words, or the columns are not one height.
111165
pub fn pack_columns<F: IsField + 'static>(columns: &[Mle<F>]) -> Option<Self> {
@@ -374,6 +428,50 @@ mod tests {
374428
assert_ne!(narrow.column(0), column);
375429
}
376430

431+
/// A row-major trace packs to what its columns pack to, at every width
432+
/// boundary and across the pack's row blocks; a width that does not divide
433+
/// the words packs nothing.
434+
#[test]
435+
fn a_row_major_trace_packs_as_its_columns() {
436+
let edges = [
437+
0u64,
438+
0xff,
439+
0x100,
440+
0xffff,
441+
0x1_0000,
442+
0xffff_ffff,
443+
0x1_0000_0000,
444+
u64::MAX,
445+
];
446+
for rows in [1usize, 7, 1024, 1025, 3000] {
447+
let columns: Vec<Vec<u64>> = edges
448+
.iter()
449+
.map(|&edge| {
450+
(0..rows)
451+
.map(|r| {
452+
if r == rows / 2 {
453+
edge
454+
} else {
455+
(r as u64).wrapping_mul(0x9e37_79b9_7f4a_7c15) % (edge / 2 + 1)
456+
}
457+
})
458+
.collect()
459+
})
460+
.collect();
461+
let row_major: Vec<u64> = (0..rows)
462+
.flat_map(|r| columns.iter().map(move |c| c[r]))
463+
.collect();
464+
let raw: Vec<&[u64]> = columns.iter().map(Vec::as_slice).collect();
465+
assert_eq!(
466+
NarrowColumns::pack_row_major(&row_major, edges.len()),
467+
NarrowColumns::pack(&raw),
468+
"{rows} rows"
469+
);
470+
}
471+
assert!(NarrowColumns::pack_row_major(&[1, 2, 3], 2).is_none());
472+
assert!(NarrowColumns::pack_row_major(&[1, 2], 0).is_none());
473+
}
474+
377475
/// Only Goldilocks columns widen to field elements: their elements are
378476
/// their raw words.
379477
#[test]

0 commit comments

Comments
 (0)