diff --git a/crates/compiler/src/lib.rs b/crates/compiler/src/lib.rs index 97ebca3c..38b07c95 100644 --- a/crates/compiler/src/lib.rs +++ b/crates/compiler/src/lib.rs @@ -1,4 +1,4 @@ -use vm::{Bytecode, F, Label, PUBLIC_INPUT_START, ZERO_VEC_PTR, execute_bytecode}; +use vm::{Bytecode, F, Label, PUBLIC_INPUT_START, ZERO_VEC_PTR}; use crate::{ a_simplify_lang::simplify_program, b_compile_intermediate::compile_to_intermediate_bytecode, @@ -37,5 +37,5 @@ pub fn compile_program(program: &str) -> Bytecode { pub fn compile_and_run(program: &str, public_input: &[F], private_input: &[F]) { let bytecode = compile_program(program); - let _ = execute_bytecode(&bytecode, public_input, private_input); + let _ = bytecode.execute(public_input, private_input); } diff --git a/crates/vm/src/bytecode/mod.rs b/crates/vm/src/bytecode/mod.rs index 52f3585e..3bc2638b 100644 --- a/crates/vm/src/bytecode/mod.rs +++ b/crates/vm/src/bytecode/mod.rs @@ -4,6 +4,13 @@ use std::{ fmt::{Display, Formatter}, }; +use utils::{build_poseidon16, build_poseidon24, pretty_integer}; + +use crate::{ + DIMENSION, ExecutionResult, F, MAX_MEMORY_SIZE, Memory, PUBLIC_INPUT_START, RunnerError, + build_public_memory, +}; + pub mod operand; pub use operand::*; pub mod hint; @@ -23,6 +30,151 @@ pub struct Bytecode { pub ending_pc: usize, } +impl Bytecode { + #[must_use] + pub fn execute(&self, public_input: &[F], private_input: &[F]) -> ExecutionResult { + let mut std_out = String::new(); + let first = self + .execute_internal( + public_input, + private_input, + MAX_MEMORY_SIZE / 2, + false, + &mut std_out, + ) + .unwrap_or_else(|err| { + if !std_out.is_empty() { + print!("{std_out}"); + } + panic!("Error during bytecode execution: {err}"); + }); + + self.execute_internal( + public_input, + private_input, + first.no_vec_runtime_memory, + true, + &mut String::new(), + ) + .unwrap() + } + + pub(crate) fn execute_internal( + &self, + public_input: &[F], + private_input: &[F], + no_vec_runtime_memory: usize, + final_execution: bool, + std_out: &mut String, + ) -> Result { + // TODO avoid rebuilding each time + let (p16, p24) = (build_poseidon16(), build_poseidon24()); + + // set public memory + let mut memory = Memory::new(build_public_memory(public_input)); + + let public_memory_size = (PUBLIC_INPUT_START + public_input.len()).next_power_of_two(); + let mut fp = public_memory_size; + + private_input + .iter() + .copied() + .enumerate() + .try_for_each(|(i, v)| memory.set(fp + i, v))?; + + fp = (fp + private_input.len()).next_multiple_of(DIMENSION); + + let initial_ap = fp + self.starting_frame_memory; + let initial_ap_vec = + (initial_ap + no_vec_runtime_memory).next_multiple_of(DIMENSION) / DIMENSION; + + let (mut pc, mut ap, mut ap_vec) = (0, initial_ap, initial_ap_vec); + + let mut poseidon16_calls = 0; + let mut poseidon24_calls = 0; + let mut dot_product_ext_ext_calls = 0; + let mut dot_product_base_ext_calls = 0; + let mut cpu_cycles = 0; + + let mut last_checkpoint_cpu_cycles = 0; + let mut checkpoint_ap = initial_ap; + let mut checkpoint_ap_vec = ap_vec; + + let mut pcs = Vec::new(); + let mut fps = Vec::new(); + + while pc != self.ending_pc { + if pc >= self.instructions.len() { + return Err(RunnerError::PCOutOfBounds); + } + + pcs.push(pc); + fps.push(fp); + + cpu_cycles += 1; + + // hints (no PC change) + for hint in self.hints.get(&pc).unwrap_or(&vec![]) { + hint.execute( + &mut memory, + fp, + &mut ap, + &mut ap_vec, + std_out, + cpu_cycles, + &mut last_checkpoint_cpu_cycles, + &mut checkpoint_ap, + &mut checkpoint_ap_vec, + )?; + } + + // instruction + self.instructions[pc].execute( + &mut memory, + &mut fp, + &mut pc, + &p16, + &p24, + &mut poseidon16_calls, + &mut poseidon24_calls, + &mut dot_product_ext_ext_calls, + &mut dot_product_base_ext_calls, + )?; + } + + debug_assert_eq!(pc, self.ending_pc); + pcs.push(pc); + fps.push(fp); + + if final_execution { + if !std_out.is_empty() { + print!("{std_out}"); + } + print_report( + self, + public_input.len(), + private_input.len(), + cpu_cycles, + &memory, + initial_ap_vec, + ap_vec, + poseidon16_calls, + poseidon24_calls, + dot_product_ext_ext_calls, + dot_product_base_ext_calls, + ); + } + + Ok(ExecutionResult { + no_vec_runtime_memory: ap - initial_ap, + public_memory_size, + memory, + pcs, + fps, + }) + } +} + impl Display for Bytecode { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { for (pc, instruction) in self.instructions.iter().enumerate() { @@ -36,3 +188,62 @@ impl Display for Bytecode { Ok(()) } } + +fn print_report( + bc: &Bytecode, + public_in_len: usize, + private_in_len: usize, + cycles: usize, + memory: &Memory, + initial_ap_vec: usize, + ap_vec: usize, + p16_calls: usize, + p24_calls: usize, + dp_ext_ext_calls: usize, + dp_base_ext_calls: usize, +) { + println!("\nBytecode size: {}", pretty_integer(bc.instructions.len())); + println!("Public input size: {}", pretty_integer(public_in_len)); + println!("Private input size: {}", pretty_integer(private_in_len)); + println!("Executed {} instructions", pretty_integer(cycles)); + + let runtime_mem = memory.0.len() - (PUBLIC_INPUT_START + public_in_len); + println!( + "Runtime memory: {} ({:.2}% vec)", + pretty_integer(runtime_mem), + (DIMENSION * (ap_vec - initial_ap_vec)) as f64 / runtime_mem as f64 * 100.0 + ); + + let total_poseidon = p16_calls + p24_calls; + if total_poseidon > 0 { + println!( + "Poseidon2_16 calls: {}, Poseidon2_24 calls: {} (1 poseidon per {} instructions)", + pretty_integer(p16_calls), + pretty_integer(p24_calls), + cycles / total_poseidon + ); + } + if dp_ext_ext_calls > 0 { + println!( + "DotProductExtExt calls: {}", + pretty_integer(dp_ext_ext_calls) + ); + } + if dp_base_ext_calls > 0 { + println!( + "DotProductBaseExt calls: {}", + pretty_integer(dp_base_ext_calls) + ); + } + + let used_cells = memory + .0 + .iter() + .skip(PUBLIC_INPUT_START + public_in_len) + .filter(|&&x| x.is_some()) + .count(); + println!( + "Memory usage: {:.1}%", + used_cells as f64 / runtime_mem as f64 * 100.0 + ); +} diff --git a/crates/vm/src/memory.rs b/crates/vm/src/memory.rs index 6897f2e7..e2dcdfaf 100644 --- a/crates/vm/src/memory.rs +++ b/crates/vm/src/memory.rs @@ -1,7 +1,12 @@ -use p3_field::{BasedVectorSpace, ExtensionField}; +use p3_field::{BasedVectorSpace, ExtensionField, PrimeCharacteristicRing}; +use p3_symmetric::Permutation; use rayon::prelude::*; +use utils::{build_poseidon16, build_poseidon24}; -use crate::{DIMENSION, EF, F, RunnerError}; +use crate::{ + DIMENSION, EF, F, POSEIDON_16_NULL_HASH_PTR, POSEIDON_24_NULL_HASH_PTR, PUBLIC_INPUT_START, + RunnerError, ZERO_VEC_PTR, +}; pub(crate) const MAX_MEMORY_SIZE: usize = 1 << 23; @@ -86,3 +91,19 @@ impl Memory { .try_for_each(|(i, &v)| self.set(index * DIMENSION + i, v)) } } + +#[must_use] +pub fn build_public_memory(public_input: &[F]) -> Vec { + // padded to a power of two + let public_memory_len = (PUBLIC_INPUT_START + public_input.len()).next_power_of_two(); + let mut public_memory = F::zero_vec(public_memory_len); + public_memory[PUBLIC_INPUT_START..][..public_input.len()].copy_from_slice(public_input); + for pm in public_memory.iter_mut().take((ZERO_VEC_PTR + 2) * 8) { + *pm = F::ZERO; // zero vector + } + public_memory[POSEIDON_16_NULL_HASH_PTR * 8..(POSEIDON_16_NULL_HASH_PTR + 2) * 8] + .copy_from_slice(&build_poseidon16().permute([F::ZERO; 16])); + public_memory[POSEIDON_24_NULL_HASH_PTR * 8..(POSEIDON_24_NULL_HASH_PTR + 1) * 8] + .copy_from_slice(&build_poseidon24().permute([F::ZERO; 24])[16..]); + public_memory +} diff --git a/crates/vm/src/runner.rs b/crates/vm/src/runner.rs index b5019f9a..e8db0ba3 100644 --- a/crates/vm/src/runner.rs +++ b/crates/vm/src/runner.rs @@ -1,45 +1,4 @@ -use p3_field::PrimeCharacteristicRing; -use p3_symmetric::Permutation; -use utils::{build_poseidon16, build_poseidon24, pretty_integer}; - -use crate::{ - DIMENSION, F, MAX_MEMORY_SIZE, Memory, POSEIDON_16_NULL_HASH_PTR, POSEIDON_24_NULL_HASH_PTR, - PUBLIC_INPUT_START, RunnerError, ZERO_VEC_PTR, bytecode::Bytecode, -}; - -#[must_use] -pub fn execute_bytecode( - bytecode: &Bytecode, - public_input: &[F], - private_input: &[F], -) -> ExecutionResult { - let mut std_out = String::new(); - let first_exec = match execute_bytecode_helper( - bytecode, - public_input, - private_input, - MAX_MEMORY_SIZE / 2, - false, - &mut std_out, - ) { - Ok(first_exec) => first_exec, - Err(err) => { - if !std_out.is_empty() { - print!("{std_out}"); - } - panic!("Error during bytecode execution: {err}"); - } - }; - execute_bytecode_helper( - bytecode, - public_input, - private_input, - first_exec.no_vec_runtime_memory, - true, - &mut String::new(), - ) - .unwrap() -} +use crate::Memory; #[derive(Debug)] pub struct ExecutionResult { @@ -49,165 +8,3 @@ pub struct ExecutionResult { pub pcs: Vec, pub fps: Vec, } - -#[must_use] -pub fn build_public_memory(public_input: &[F]) -> Vec { - // padded to a power of two - let public_memory_len = (PUBLIC_INPUT_START + public_input.len()).next_power_of_two(); - let mut public_memory = F::zero_vec(public_memory_len); - public_memory[PUBLIC_INPUT_START..][..public_input.len()].copy_from_slice(public_input); - for pm in public_memory.iter_mut().take((ZERO_VEC_PTR + 2) * 8) { - *pm = F::ZERO; // zero vector - } - public_memory[POSEIDON_16_NULL_HASH_PTR * 8..(POSEIDON_16_NULL_HASH_PTR + 2) * 8] - .copy_from_slice(&build_poseidon16().permute([F::ZERO; 16])); - public_memory[POSEIDON_24_NULL_HASH_PTR * 8..(POSEIDON_24_NULL_HASH_PTR + 1) * 8] - .copy_from_slice(&build_poseidon24().permute([F::ZERO; 24])[16..]); - public_memory -} - -fn execute_bytecode_helper( - bytecode: &Bytecode, - public_input: &[F], - private_input: &[F], - no_vec_runtime_memory: usize, - final_execution: bool, - std_out: &mut String, -) -> Result { - let poseidon_16 = build_poseidon16(); // TODO avoid rebuilding each time - let poseidon_24 = build_poseidon24(); - - // set public memory - let mut memory = Memory::new(build_public_memory(public_input)); - - let public_memory_size = (PUBLIC_INPUT_START + public_input.len()).next_power_of_two(); - let mut fp = public_memory_size; - - for (i, value) in private_input.iter().enumerate() { - memory.set(fp + i, *value)?; - } - fp += private_input.len(); - fp = fp.next_multiple_of(DIMENSION); - - let initial_ap = fp + bytecode.starting_frame_memory; - let initial_ap_vec = - (initial_ap + no_vec_runtime_memory).next_multiple_of(DIMENSION) / DIMENSION; - - let mut pc = 0; - let mut ap = initial_ap; - let mut ap_vec = initial_ap_vec; - - let mut poseidon16_calls = 0; - let mut poseidon24_calls = 0; - let mut dot_product_ext_ext_calls = 0; - let mut dot_product_base_ext_calls = 0; - let mut cpu_cycles = 0; - - let mut last_checkpoint_cpu_cycles = 0; - let mut checkpoint_ap = initial_ap; - let mut checkpoint_ap_vec = ap_vec; - - let mut pcs = Vec::new(); - let mut fps = Vec::new(); - - while pc != bytecode.ending_pc { - if pc >= bytecode.instructions.len() { - return Err(RunnerError::PCOutOfBounds); - } - - pcs.push(pc); - fps.push(fp); - - cpu_cycles += 1; - - for hint in bytecode.hints.get(&pc).unwrap_or(&vec![]) { - hint.execute( - &mut memory, - fp, - &mut ap, - &mut ap_vec, - std_out, - cpu_cycles, - &mut last_checkpoint_cpu_cycles, - &mut checkpoint_ap, - &mut checkpoint_ap_vec, - )?; - } - - bytecode.instructions[pc].execute( - &mut memory, - &mut fp, - &mut pc, - &poseidon_16, - &poseidon_24, - &mut poseidon16_calls, - &mut poseidon24_calls, - &mut dot_product_ext_ext_calls, - &mut dot_product_base_ext_calls, - )?; - } - - debug_assert_eq!(pc, bytecode.ending_pc); - pcs.push(pc); - fps.push(fp); - - if final_execution { - if !std_out.is_empty() { - print!("{std_out}"); - } - let runtime_memory_size = memory.0.len() - (PUBLIC_INPUT_START + public_input.len()); - println!( - "\nBytecode size: {}", - pretty_integer(bytecode.instructions.len()) - ); - println!("Public input size: {}", pretty_integer(public_input.len())); - println!( - "Private input size: {}", - pretty_integer(private_input.len()) - ); - println!("Executed {} instructions", pretty_integer(cpu_cycles)); - println!( - "Runtime memory: {} ({:.2}% vec)", - pretty_integer(runtime_memory_size), - (DIMENSION * (ap_vec - initial_ap_vec)) as f64 / runtime_memory_size as f64 * 100.0 - ); - if poseidon16_calls + poseidon24_calls > 0 { - println!( - "Poseidon2_16 calls: {}, Poseidon2_24 calls: {} (1 poseidon per {} instructions)", - pretty_integer(poseidon16_calls), - pretty_integer(poseidon24_calls), - cpu_cycles / (poseidon16_calls + poseidon24_calls) - ); - } - if dot_product_ext_ext_calls > 0 { - println!( - "DotProductExtExt calls: {}", - pretty_integer(dot_product_ext_ext_calls) - ); - } - if dot_product_base_ext_calls > 0 { - println!( - "DotProductBaseExt calls: {}", - pretty_integer(dot_product_base_ext_calls) - ); - } - let used_memory_cells = memory - .0 - .iter() - .skip(PUBLIC_INPUT_START + public_input.len()) - .filter(|&&x| x.is_some()) - .count(); - println!( - "Memory usage: {:.1}%", - used_memory_cells as f64 / runtime_memory_size as f64 * 100.0 - ); - } - let no_vec_runtime_memory = ap - initial_ap; - Ok(ExecutionResult { - no_vec_runtime_memory, - public_memory_size, - memory, - pcs, - fps, - }) -} diff --git a/crates/zk_vm/src/lib.rs b/crates/zk_vm/src/lib.rs index c1d7b321..473afbc2 100644 --- a/crates/zk_vm/src/lib.rs +++ b/crates/zk_vm/src/lib.rs @@ -1,7 +1,6 @@ use std::ops::Range; -use compiler::{PRECOMPILES, compile_program}; -use vm::execute_bytecode; +use compiler::PRECOMPILES; mod air; mod common; @@ -71,11 +70,6 @@ fn exec_column_groups() -> Vec> { .concat() } -pub fn compile_and_run(program: &str, public_input: &[vm::F], private_input: &[vm::F]) { - let bytecode = compile_program(program); - let _ = execute_bytecode(&bytecode, public_input, private_input); -} - pub trait InAirColumnIndex { fn index_in_air(self) -> usize; } diff --git a/crates/zk_vm/src/prove_execution.rs b/crates/zk_vm/src/prove_execution.rs index 0ebaabc0..3fed3d92 100644 --- a/crates/zk_vm/src/prove_execution.rs +++ b/crates/zk_vm/src/prove_execution.rs @@ -13,10 +13,7 @@ use utils::{ fold_multilinear_in_large_field, pack_extension, padd_with_zero_to_next_power_of_two, to_big_endian_bits, }; -use vm::{ - Bytecode, DIMENSION, EF, F, POSEIDON_16_NULL_HASH_PTR, POSEIDON_24_NULL_HASH_PTR, - execute_bytecode, -}; +use vm::{Bytecode, DIMENSION, EF, F, POSEIDON_16_NULL_HASH_PTR, POSEIDON_24_NULL_HASH_PTR}; use whir_p3::{ dft::EvalsDft, poly::{ @@ -59,7 +56,7 @@ pub fn prove_execution( public_memory_size, memory, } = info_span!("Witness generation").in_scope(|| { - let execution_result = execute_bytecode(bytecode, public_input, private_input); + let execution_result = bytecode.execute(public_input, private_input); get_execution_trace(bytecode, &execution_result) });