Skip to content

fix: optimize Cycle size to improve tracer memory usage #852

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Aug 8, 2025
Merged

fix: optimize Cycle size to improve tracer memory usage #852

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
474c500
fix: add a stack canary to prevent stack overflows
sagar-a16z Aug 7, 2025
a08f832
fix: tracer tests
sagar-a16z Aug 8, 2025
da3b33f
fix: wasm
sagar-a16z Aug 8, 2025
cc045a5
fix: test
sagar-a16z Aug 8, 2025
75f188b
fix: optimize Cycle size to improve tracer memory usage
sagar-a16z Aug 8, 2025
3a1f896
fix fmt
sagar-a16z Aug 8, 2025
1e844fe
fix: verifier test fixtures
sagar-a16z Aug 8, 2025
ebf51c6
Merge branch 'main' into sagar/optimize_cycle_size
moodlezoup Aug 8, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions common/src/constants.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
pub const XLEN: usize = 32;
const RISCV_REGISTER_COUNT: u64 = 32;
const VIRTUAL_REGISTER_COUNT: u64 = 32; // see Section 6.1 of Jolt paper
pub const REGISTER_COUNT: u64 = RISCV_REGISTER_COUNT + VIRTUAL_REGISTER_COUNT;
const RISCV_REGISTER_COUNT: u8 = 32;
const VIRTUAL_REGISTER_COUNT: u8 = 32; // see Section 6.1 of Jolt paper
pub const REGISTER_COUNT: u8 = RISCV_REGISTER_COUNT + VIRTUAL_REGISTER_COUNT;
pub const BYTES_PER_INSTRUCTION: usize = 4;

pub const RAM_START_ADDRESS: u64 = 0x80000000;
Expand All @@ -18,7 +18,7 @@ pub const DEFAULT_MAX_INPUT_SIZE: u64 = 4096;
pub const DEFAULT_MAX_OUTPUT_SIZE: u64 = 4096;
pub const DEFAULT_MAX_TRACE_LENGTH: u64 = 1 << 24;

pub const fn virtual_register_index(index: u64) -> u64 {
pub const fn virtual_register_index(index: u8) -> u8 {
index + VIRTUAL_REGISTER_COUNT
}

Expand Down
7 changes: 4 additions & 3 deletions jolt-core/src/zkvm/bytecode/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ pub struct BytecodePreprocessing {
/// See Section 6.1 of the Jolt paper, "Reflecting the program counter". The virtual address
/// is the one used to keep track of the next (potentially virtual) instruction to execute.
/// Key: (ELF address, virtual sequence index or 0)
pub virtual_address_map: BTreeMap<(usize, usize), usize>,
pub virtual_address_map: BTreeMap<(usize, u16), usize>,
pub d: usize,
}

Expand All @@ -54,8 +54,9 @@ impl BytecodePreprocessing {
(instr.address, instr.virtual_sequence_remaining.unwrap_or(0)),
virtual_address
),
None
);
None,
"Virtual address map already contains entry for address: {:#X}, virtual sequence: {:?}. map size: {}",
instr.address, instr.virtual_sequence_remaining, virtual_address_map.len());
virtual_address += 1;
}

Expand Down
4 changes: 2 additions & 2 deletions jolt-core/src/zkvm/bytecode/read_raf_checking.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -409,7 +409,7 @@ impl<F: JoltField> ReadRafSumcheck<F> {
.chain(once(instr.operands.rd))
.chain(once(instr.operands.rs1))
.chain(once(instr.operands.rs2))
.map(|r| eq_r_register[r])
.map(|r| eq_r_register[r as usize])
.zip(gamma_powers)
.map(|(claim, gamma)| claim * gamma)
.sum::<F>()
Expand Down Expand Up @@ -464,7 +464,7 @@ impl<F: JoltField> ReadRafSumcheck<F> {

let mut linear_combination: F = F::zero();

linear_combination += eq_r_register[instr.operands.rd];
linear_combination += eq_r_register[instr.operands.rd as usize];
linear_combination += gamma_powers[1].mul_u64(unexpanded_pc as u64);
if flags[CircuitFlags::IsNoop] {
linear_combination += gamma_powers[2];
Expand Down
11 changes: 1 addition & 10 deletions jolt-core/src/zkvm/r1cs/inputs.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -227,10 +227,7 @@ impl JoltR1CSInputs {
coeffs.into()
}
JoltR1CSInputs::Rd => {
let coeffs: Vec<u8> = trace
.par_iter()
.map(|cycle| cycle.rd_write().0 as u8)
.collect();
let coeffs: Vec<u8> = trace.par_iter().map(|cycle| cycle.rd_write().0).collect();
coeffs.into()
}
JoltR1CSInputs::Imm => {
Expand Down Expand Up @@ -266,9 +263,6 @@ impl JoltR1CSInputs {
tracer::instruction::RAMAccess::Read(read) => read.value,
tracer::instruction::RAMAccess::Write(write) => write.pre_value,
tracer::instruction::RAMAccess::NoOp => 0,
tracer::instruction::RAMAccess::Atomic(_) => {
unimplemented!("Atomic instructions are mapped to virtual sequences")
}
})
.collect();
coeffs.into()
Expand All @@ -280,9 +274,6 @@ impl JoltR1CSInputs {
tracer::instruction::RAMAccess::Read(read) => read.value,
tracer::instruction::RAMAccess::Write(write) => write.post_value,
tracer::instruction::RAMAccess::NoOp => 0,
tracer::instruction::RAMAccess::Atomic(_) => {
unimplemented!("Atomic instructions are mapped to virtual sequences")
}
})
.collect();
coeffs.into()
Expand Down
84 changes: 42 additions & 42 deletions jolt-core/src/zkvm/registers/read_write_checking.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,11 +55,11 @@ struct DataBuffers<F: JoltField> {
}

struct ReadWriteCheckingProverState<F: JoltField> {
addresses: Vec<(usize, usize, usize)>,
addresses: Vec<(u8, u8, u8)>,
chunk_size: usize,
val_checkpoints: Vec<F>,
data_buffers: Vec<DataBuffers<F>>,
I: Vec<Vec<(usize, usize, F, F)>>,
I: Vec<Vec<(usize, u8, F, F)>>,
A: Vec<F>,
gruens_eq_r_prime: GruenSplitEqPolynomial<F>,
inc_cycle: MultilinearPolynomial<F>,
Expand Down Expand Up @@ -92,7 +92,7 @@ impl<F: JoltField> ReadWriteCheckingProverState<F> {
let mut delta = [0; K];
for cycle in trace_chunk.iter() {
let (k, pre_value, post_value) = cycle.rd_write();
delta[k] += post_value as i128 - pre_value as i128;
delta[k as usize] += post_value as i128 - pre_value as i128;
}
delta
})
Expand Down Expand Up @@ -148,7 +148,7 @@ impl<F: JoltField> ReadWriteCheckingProverState<F> {
let _guard = span.enter();

// Data structure described in Equation (72)
let I: Vec<Vec<(usize, usize, F, F)>> = trace
let I: Vec<Vec<(usize, u8, F, F)>> = trace
.par_chunks(chunk_size)
.enumerate()
.map(|(chunk_index, trace_chunk)| {
Expand Down Expand Up @@ -351,43 +351,43 @@ impl<F: JoltField> RegistersReadWriteChecking<F> {

let k = addresses[j].0;
unsafe {
dirty_indices.insert_unchecked(k);
dirty_indices.insert_unchecked(k as usize);
}
rs1_ra[0][k] += A[j_bound];
rs1_ra[0][k as usize] += A[j_bound];

let k = addresses[j].1;
unsafe {
dirty_indices.insert_unchecked(k);
dirty_indices.insert_unchecked(k as usize);
}
rs2_ra[0][k] += A[j_bound];
rs2_ra[0][k as usize] += A[j_bound];

let k = addresses[j].2;
unsafe {
dirty_indices.insert_unchecked(k);
dirty_indices.insert_unchecked(k as usize);
}
rd_wa[0][k] += A[j_bound];
rd_wa[0][k as usize] += A[j_bound];
}

for j in (j_prime + 1) << round..(j_prime + 2) << round {
let j_bound = j % (1 << round);

let k = addresses[j].0;
unsafe {
dirty_indices.insert_unchecked(k);
dirty_indices.insert_unchecked(k as usize);
}
rs1_ra[1][k] += A[j_bound];
rs1_ra[1][k as usize] += A[j_bound];

let k = addresses[j].1;
unsafe {
dirty_indices.insert_unchecked(k);
dirty_indices.insert_unchecked(k as usize);
}
rs2_ra[1][k] += A[j_bound];
rs2_ra[1][k as usize] += A[j_bound];

let k = addresses[j].2;
unsafe {
dirty_indices.insert_unchecked(k);
dirty_indices.insert_unchecked(k as usize);
}
rd_wa[1][k] += A[j_bound];
rd_wa[1][k as usize] += A[j_bound];
}

for k in dirty_indices.ones() {
Expand All @@ -399,8 +399,8 @@ impl<F: JoltField> RegistersReadWriteChecking<F> {
loop {
let (row, col, inc_lt, inc) = inc_iter.next().unwrap();
debug_assert_eq!(*row, j_prime);
val_j_r[0][*col] += *inc_lt;
val_j_0[*col] += *inc;
val_j_r[0][*col as usize] += *inc_lt;
val_j_0[*col as usize] += *inc;
if inc_iter.peek().unwrap().0 != j_prime {
break;
}
Expand All @@ -413,8 +413,8 @@ impl<F: JoltField> RegistersReadWriteChecking<F> {
for inc in inc_iter {
let (row, col, inc_lt, inc) = *inc;
debug_assert_eq!(row, j_prime + 1);
val_j_r[1][col] += inc_lt;
val_j_0[col] += inc;
val_j_r[1][col as usize] += inc_lt;
val_j_0[col as usize] += inc;
}

let eq_r_prime_eval = gruens_eq_r_prime.E_out_current()[j_prime / 2];
Expand Down Expand Up @@ -525,43 +525,43 @@ impl<F: JoltField> RegistersReadWriteChecking<F> {

let k = addresses[j].0;
unsafe {
dirty_indices.insert_unchecked(k);
dirty_indices.insert_unchecked(k as usize);
}
rs1_ra[0][k] += A[j_bound];
rs1_ra[0][k as usize] += A[j_bound];

let k = addresses[j].1;
unsafe {
dirty_indices.insert_unchecked(k);
dirty_indices.insert_unchecked(k as usize);
}
rs2_ra[0][k] += A[j_bound];
rs2_ra[0][k as usize] += A[j_bound];

let k = addresses[j].2;
unsafe {
dirty_indices.insert_unchecked(k);
dirty_indices.insert_unchecked(k as usize);
}
rd_wa[0][k] += A[j_bound];
rd_wa[0][k as usize] += A[j_bound];
}

for j in (j_prime + 1) << round..(j_prime + 2) << round {
let j_bound = j % (1 << round);

let k = addresses[j].0;
unsafe {
dirty_indices.insert_unchecked(k);
dirty_indices.insert_unchecked(k as usize);
}
rs1_ra[1][k] += A[j_bound];
rs1_ra[1][k as usize] += A[j_bound];

let k = addresses[j].1;
unsafe {
dirty_indices.insert_unchecked(k);
dirty_indices.insert_unchecked(k as usize);
}
rs2_ra[1][k] += A[j_bound];
rs2_ra[1][k as usize] += A[j_bound];

let k = addresses[j].2;
unsafe {
dirty_indices.insert_unchecked(k);
dirty_indices.insert_unchecked(k as usize);
}
rd_wa[1][k] += A[j_bound];
rd_wa[1][k as usize] += A[j_bound];
}

for k in dirty_indices.ones() {
Expand All @@ -573,8 +573,8 @@ impl<F: JoltField> RegistersReadWriteChecking<F> {
loop {
let (row, col, inc_lt, inc) = inc_iter.next().unwrap();
debug_assert_eq!(*row, j_prime);
val_j_r[0][*col] += *inc_lt;
val_j_0[*col] += *inc;
val_j_r[0][*col as usize] += *inc_lt;
val_j_0[*col as usize] += *inc;
if inc_iter.peek().unwrap().0 != j_prime {
break;
}
Expand All @@ -587,8 +587,8 @@ impl<F: JoltField> RegistersReadWriteChecking<F> {
for inc in inc_iter {
let (row, col, inc_lt, inc) = *inc;
debug_assert_eq!(row, j_prime + 1);
val_j_r[1][col] += inc_lt;
val_j_0[col] += inc;
val_j_r[1][col as usize] += inc_lt;
val_j_0[col as usize] += inc;
}

let x_in = (j_prime / 2) & x_bitmask;
Expand Down Expand Up @@ -872,7 +872,7 @@ impl<F: JoltField> RegistersReadWriteChecking<F> {

for i in 0..I_chunk.len() {
let (j_prime, k, inc_lt, inc) = I_chunk[i];
if let Some(bound_index) = bound_indices[k] {
if let Some(bound_index) = bound_indices[k as usize] {
if I_chunk[bound_index].0 == j_prime / 2 {
// Neighbor was already processed
debug_assert!(j_prime % 2 == 1);
Expand All @@ -890,7 +890,7 @@ impl<F: JoltField> RegistersReadWriteChecking<F> {
};

I_chunk[next_bound_index] = (j_prime / 2, k, bound_value, inc);
bound_indices[k] = Some(next_bound_index);
bound_indices[k as usize] = Some(next_bound_index);
next_bound_index += 1;
}
I_chunk.truncate(next_bound_index);
Expand Down Expand Up @@ -934,7 +934,7 @@ impl<F: JoltField> RegistersReadWriteChecking<F> {
.iter()
.enumerate()
{
ra_chunk[*k] += A[j_bound];
ra_chunk[*k as usize] += A[j_bound];
}
});
*rs1_ra = Some(MultilinearPolynomial::from(rs1_ra_evals));
Expand All @@ -956,7 +956,7 @@ impl<F: JoltField> RegistersReadWriteChecking<F> {
.iter()
.enumerate()
{
ra_chunk[*k] += A[j_bound];
ra_chunk[*k as usize] += A[j_bound];
}
});
*rs2_ra = Some(MultilinearPolynomial::from(rs2_ra_evals));
Expand All @@ -978,7 +978,7 @@ impl<F: JoltField> RegistersReadWriteChecking<F> {
.iter()
.enumerate()
{
wa_chunk[*k] += A[j_bound];
wa_chunk[*k as usize] += A[j_bound];
}
});
*rd_wa = Some(MultilinearPolynomial::from(rd_wa_evals));
Expand All @@ -997,7 +997,7 @@ impl<F: JoltField> RegistersReadWriteChecking<F> {
.for_each(|(chunk_index, (val_chunk, I_chunk))| {
for (j, k, inc_lt, _inc) in I_chunk.iter_mut() {
debug_assert_eq!(*j, chunk_index);
val_chunk[*k] += *inc_lt;
val_chunk[*k as usize] += *inc_lt;
}
});
*val = Some(MultilinearPolynomial::from(val_evals));
Expand Down
2 changes: 1 addition & 1 deletion jolt-core/src/zkvm/registers/val_evaluation.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ impl<F: JoltField> ValEvaluationSumcheck<F> {
.par_iter()
.map(|cycle| {
let instr = cycle.instruction().normalize();
eq_r_address[instr.operands.rd]
eq_r_address[instr.operands.rd as usize]
})
.collect();
let wa = MultilinearPolynomial::from(wa);
Expand Down
4 changes: 2 additions & 2 deletions jolt-core/src/zkvm/witness.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -223,7 +223,7 @@ impl CommittedPolynomial {
.map(|cycle| {
let flag = cycle.instruction().circuit_flags()
[CircuitFlags::WriteLookupOutputToRD as usize];
(cycle.rd_write().0 as u8) * (flag as u8)
(cycle.rd_write().0) * (flag as u8)
})
.collect();
coeffs.into()
Expand All @@ -233,7 +233,7 @@ impl CommittedPolynomial {
.par_iter()
.map(|cycle| {
let flag = cycle.instruction().circuit_flags()[CircuitFlags::Jump as usize];
(cycle.rd_write().0 as u8) * (flag as u8)
(cycle.rd_write().0) * (flag as u8)
})
.collect();
coeffs.into()
Expand Down
Binary file modified jolt-verifier/tests/fixtures/fib_io_device.bin
View file
Open in desktop
Binary file not shown.
Binary file modified jolt-verifier/tests/fixtures/fib_proof.bin
View file
Open in desktop
Binary file not shown.
Binary file modified jolt-verifier/tests/fixtures/jolt_verifier_preprocessing.dat
View file
Open in desktop
Binary file not shown.
5 changes: 3 additions & 2 deletions tracer/src/instruction/add.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,9 @@ declare_riscv_instr!(

impl ADD {
fn exec(&self, cpu: &mut Cpu, _: &mut <ADD as RISCVInstruction>::RAMAccess) {
cpu.x[self.operands.rd] =
cpu.sign_extend(cpu.x[self.operands.rs1].wrapping_add(cpu.x[self.operands.rs2]));
cpu.x[self.operands.rd as usize] = cpu.sign_extend(
cpu.x[self.operands.rs1 as usize].wrapping_add(cpu.x[self.operands.rs2 as usize]),
);
}
}

Expand Down
5 changes: 3 additions & 2 deletions tracer/src/instruction/addi.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,9 @@ declare_riscv_instr!(

impl ADDI {
fn exec(&self, cpu: &mut Cpu, _: &mut <ADDI as RISCVInstruction>::RAMAccess) {
cpu.x[self.operands.rd] = cpu
.sign_extend(cpu.x[self.operands.rs1].wrapping_add(normalize_imm(self.operands.imm)));
cpu.x[self.operands.rd as usize] = cpu.sign_extend(
cpu.x[self.operands.rs1 as usize].wrapping_add(normalize_imm(self.operands.imm)),
);
}
}

Expand Down
5 changes: 3 additions & 2 deletions tracer/src/instruction/addiw.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,9 @@ impl ADDIW {
// ignored and the result is the low 32 bits of the result sign-extended to 64 bits. Note,
// ADDIW rd, rs1, 0 writes the sign extension of the lower 32 bits of register rs1 into
// register rd (assembler pseudoinstruction SEXT.W).
cpu.x[self.operands.rd] =
cpu.x[self.operands.rs1].wrapping_add(normalize_imm(self.operands.imm)) as i32 as i64;
cpu.x[self.operands.rd as usize] = cpu.x[self.operands.rs1 as usize]
.wrapping_add(normalize_imm(self.operands.imm))
as i32 as i64;
}
}

Expand Down
Loading