From 6f841678a57b67099f3a37747b83131482acd6db Mon Sep 17 00:00:00 2001 From: Daniel Lubarov Date: Fri, 9 Dec 2022 12:51:42 -0800 Subject: [PATCH 1/3] More timing for zkEVM proofs --- evm/src/all_stark.rs | 12 ++ evm/src/prover.rs | 203 ++++++++++++++++++++---------- evm/src/witness/traces.rs | 31 +++-- evm/tests/empty_txn_list.rs | 3 +- evm/tests/transfer_to_new_addr.rs | 3 +- 5 files changed, 174 insertions(+), 78 deletions(-) diff --git a/evm/src/all_stark.rs b/evm/src/all_stark.rs index 22f7c123..62a6c2cc 100644 --- a/evm/src/all_stark.rs +++ b/evm/src/all_stark.rs @@ -76,6 +76,18 @@ pub enum Table { pub(crate) const NUM_TABLES: usize = Table::Memory as usize + 1; +impl Table { + pub(crate) fn all() -> [Self; NUM_TABLES] { + [ + Self::Cpu, + Self::Keccak, + Self::KeccakSponge, + Self::Logic, + Self::Memory, + ] + } +} + pub(crate) fn all_cross_table_lookups() -> Vec> { let mut ctls = vec![ctl_keccak(), ctl_logic(), ctl_memory(), ctl_keccak_sponge()]; // TODO: Some CTLs temporarily disabled while we get them working. diff --git a/evm/src/prover.rs b/evm/src/prover.rs index c99e1873..55b57437 100644 --- a/evm/src/prover.rs +++ b/evm/src/prover.rs @@ -1,6 +1,7 @@ use std::any::type_name; use anyhow::{ensure, Result}; +use itertools::Itertools; use maybe_rayon::*; use plonky2::field::extension::Extendable; use plonky2::field::packable::Packable; @@ -53,7 +54,11 @@ where [(); LogicStark::::COLUMNS]:, [(); MemoryStark::::COLUMNS]:, { - let (traces, public_values) = generate_traces(all_stark, inputs, config, timing); + let (traces, public_values) = timed!( + timing, + "generate all traces", + generate_traces(all_stark, inputs, config, timing) + ); prove_with_traces(all_stark, config, traces, public_values, timing) } @@ -80,19 +85,24 @@ where let trace_commitments = timed!( timing, - "compute trace commitments", + "compute all trace commitments", trace_poly_values .iter() - .map(|trace| { - PolynomialBatch::::from_values( - // TODO: Cloning this isn't great; consider having `from_values` accept a reference, - // or having `compute_permutation_z_polys` read trace values from the `PolynomialBatch`. - trace.clone(), - rate_bits, - false, - cap_height, + .zip_eq(Table::all()) + .map(|(trace, table)| { + timed!( timing, - None, + &format!("compute trace commitment for {:?}", table), + PolynomialBatch::::from_values( + // TODO: Cloning this isn't great; consider having `from_values` accept a reference, + // or having `compute_permutation_z_polys` read trace values from the `PolynomialBatch`. + trace.clone(), + rate_bits, + false, + cap_height, + timing, + None, + ) ) }) .collect::>() @@ -107,66 +117,30 @@ where challenger.observe_cap(cap); } - let ctl_data_per_table = cross_table_lookup_data::( - config, - &trace_poly_values, - &all_stark.cross_table_lookups, - &mut challenger, + let ctl_data_per_table = timed!( + timing, + "compute CTL data", + cross_table_lookup_data::( + config, + &trace_poly_values, + &all_stark.cross_table_lookups, + &mut challenger, + ) ); - let cpu_proof = prove_single_table( - &all_stark.cpu_stark, - config, - &trace_poly_values[Table::Cpu as usize], - &trace_commitments[Table::Cpu as usize], - &ctl_data_per_table[Table::Cpu as usize], - &mut challenger, + let stark_proofs = timed!( timing, - )?; - let keccak_proof = prove_single_table( - &all_stark.keccak_stark, - config, - &trace_poly_values[Table::Keccak as usize], - &trace_commitments[Table::Keccak as usize], - &ctl_data_per_table[Table::Keccak as usize], - &mut challenger, - timing, - )?; - let keccak_sponge_proof = prove_single_table( - &all_stark.keccak_sponge_stark, - config, - &trace_poly_values[Table::KeccakSponge as usize], - &trace_commitments[Table::KeccakSponge as usize], - &ctl_data_per_table[Table::KeccakSponge as usize], - &mut challenger, - timing, - )?; - let logic_proof = prove_single_table( - &all_stark.logic_stark, - config, - &trace_poly_values[Table::Logic as usize], - &trace_commitments[Table::Logic as usize], - &ctl_data_per_table[Table::Logic as usize], - &mut challenger, - timing, - )?; - let memory_proof = prove_single_table( - &all_stark.memory_stark, - config, - &trace_poly_values[Table::Memory as usize], - &trace_commitments[Table::Memory as usize], - &ctl_data_per_table[Table::Memory as usize], - &mut challenger, - timing, - )?; - - let stark_proofs = [ - cpu_proof, - keccak_proof, - keccak_sponge_proof, - logic_proof, - memory_proof, - ]; + "compute all proofs given commitments", + prove_with_commitments( + all_stark, + config, + trace_poly_values, + trace_commitments, + ctl_data_per_table, + &mut challenger, + timing + )? + ); Ok(AllProof { stark_proofs, @@ -174,6 +148,99 @@ where }) } +fn prove_with_commitments( + all_stark: &AllStark, + config: &StarkConfig, + trace_poly_values: [Vec>; NUM_TABLES], + trace_commitments: Vec>, + ctl_data_per_table: [CtlData; NUM_TABLES], + challenger: &mut Challenger, + timing: &mut TimingTree, +) -> Result<[StarkProof; NUM_TABLES]> +where + F: RichField + Extendable, + C: GenericConfig, + [(); C::Hasher::HASH_SIZE]:, + [(); CpuStark::::COLUMNS]:, + [(); KeccakStark::::COLUMNS]:, + [(); KeccakSpongeStark::::COLUMNS]:, + [(); LogicStark::::COLUMNS]:, + [(); MemoryStark::::COLUMNS]:, +{ + let cpu_proof = timed!( + timing, + "prove CPU STARK", + prove_single_table( + &all_stark.cpu_stark, + config, + &trace_poly_values[Table::Cpu as usize], + &trace_commitments[Table::Cpu as usize], + &ctl_data_per_table[Table::Cpu as usize], + challenger, + timing, + )? + ); + let keccak_proof = timed!( + timing, + "prove Keccak STARK", + prove_single_table( + &all_stark.keccak_stark, + config, + &trace_poly_values[Table::Keccak as usize], + &trace_commitments[Table::Keccak as usize], + &ctl_data_per_table[Table::Keccak as usize], + challenger, + timing, + )? + ); + let keccak_sponge_proof = timed!( + timing, + "prove Keccak sponge STARK", + prove_single_table( + &all_stark.keccak_sponge_stark, + config, + &trace_poly_values[Table::KeccakSponge as usize], + &trace_commitments[Table::KeccakSponge as usize], + &ctl_data_per_table[Table::KeccakSponge as usize], + challenger, + timing, + )? + ); + let logic_proof = timed!( + timing, + "prove logic STARK", + prove_single_table( + &all_stark.logic_stark, + config, + &trace_poly_values[Table::Logic as usize], + &trace_commitments[Table::Logic as usize], + &ctl_data_per_table[Table::Logic as usize], + challenger, + timing, + )? + ); + let memory_proof = timed!( + timing, + "prove memory STARK", + prove_single_table( + &all_stark.memory_stark, + config, + &trace_poly_values[Table::Memory as usize], + &trace_commitments[Table::Memory as usize], + &ctl_data_per_table[Table::Memory as usize], + challenger, + timing, + )? + ); + Ok([ + cpu_proof, + keccak_proof, + keccak_sponge_proof, + logic_proof, + memory_proof, + ]) +} + /// Compute proof for a single STARK table. pub(crate) fn prove_single_table( stark: &S, diff --git a/evm/src/witness/traces.rs b/evm/src/witness/traces.rs index 41b654fb..c904d2e5 100644 --- a/evm/src/witness/traces.rs +++ b/evm/src/witness/traces.rs @@ -4,6 +4,7 @@ use itertools::Itertools; use plonky2::field::extension::Extendable; use plonky2::field::polynomial::PolynomialValues; use plonky2::hash::hash_types::RichField; +use plonky2::timed; use plonky2::util::timing::TimingTree; use crate::all_stark::{AllStark, NUM_TABLES}; @@ -131,18 +132,32 @@ impl Traces { let cpu_rows = cpu.into_iter().map(|x| x.into()).collect(); let cpu_trace = trace_rows_to_poly_values(cpu_rows); - let keccak_trace = + let keccak_trace = timed!( + timing, + "generate Keccak trace", all_stark .keccak_stark - .generate_trace(keccak_inputs, cap_elements, timing); - let keccak_sponge_trace = + .generate_trace(keccak_inputs, cap_elements, timing) + ); + let keccak_sponge_trace = timed!( + timing, + "generate Keccak sponge trace", all_stark .keccak_sponge_stark - .generate_trace(keccak_sponge_ops, cap_elements, timing); - let logic_trace = all_stark - .logic_stark - .generate_trace(logic_ops, cap_elements, timing); - let memory_trace = all_stark.memory_stark.generate_trace(memory_ops, timing); + .generate_trace(keccak_sponge_ops, cap_elements, timing) + ); + let logic_trace = timed!( + timing, + "generate logic trace", + all_stark + .logic_stark + .generate_trace(logic_ops, cap_elements, timing) + ); + let memory_trace = timed!( + timing, + "generate memory trace", + all_stark.memory_stark.generate_trace(memory_ops, timing) + ); [ cpu_trace, diff --git a/evm/tests/empty_txn_list.rs b/evm/tests/empty_txn_list.rs index abeef644..e7de7b40 100644 --- a/evm/tests/empty_txn_list.rs +++ b/evm/tests/empty_txn_list.rs @@ -1,4 +1,5 @@ use std::collections::HashMap; +use std::time::Duration; use env_logger::{try_init_from_env, Env, DEFAULT_FILTER_ENV}; use eth_trie_utils::partial_trie::PartialTrie; @@ -49,7 +50,7 @@ fn test_empty_txn_list() -> anyhow::Result<()> { let mut timing = TimingTree::new("prove", log::Level::Debug); let proof = prove::(&all_stark, &config, inputs, &mut timing)?; - timing.print(); + timing.filter(Duration::from_millis(100)).print(); assert_eq!( proof.public_values.trie_roots_before.state_root, diff --git a/evm/tests/transfer_to_new_addr.rs b/evm/tests/transfer_to_new_addr.rs index 88d129a3..fc48f80f 100644 --- a/evm/tests/transfer_to_new_addr.rs +++ b/evm/tests/transfer_to_new_addr.rs @@ -1,4 +1,5 @@ use std::collections::HashMap; +use std::time::Duration; use env_logger::{try_init_from_env, Env, DEFAULT_FILTER_ENV}; use eth_trie_utils::partial_trie::{Nibbles, PartialTrie}; @@ -68,7 +69,7 @@ fn test_simple_transfer() -> anyhow::Result<()> { let mut timing = TimingTree::new("prove", log::Level::Debug); let proof = prove::(&all_stark, &config, inputs, &mut timing)?; - timing.print(); + timing.filter(Duration::from_millis(100)).print(); let expected_state_trie_after = { let sender_account_after = AccountRlp { From 569cd058a0a27611d09b38561782fe2342acc99b Mon Sep 17 00:00:00 2001 From: Daniel Lubarov Date: Fri, 9 Dec 2022 13:20:33 -0800 Subject: [PATCH 2/3] log level --- evm/tests/empty_txn_list.rs | 2 +- evm/tests/transfer_to_new_addr.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/evm/tests/empty_txn_list.rs b/evm/tests/empty_txn_list.rs index e7de7b40..aa7b60b9 100644 --- a/evm/tests/empty_txn_list.rs +++ b/evm/tests/empty_txn_list.rs @@ -81,5 +81,5 @@ fn test_empty_txn_list() -> anyhow::Result<()> { } fn init_logger() { - let _ = try_init_from_env(Env::default().filter_or(DEFAULT_FILTER_ENV, "debug")); + let _ = try_init_from_env(Env::default().filter_or(DEFAULT_FILTER_ENV, "info")); } diff --git a/evm/tests/transfer_to_new_addr.rs b/evm/tests/transfer_to_new_addr.rs index fc48f80f..351506da 100644 --- a/evm/tests/transfer_to_new_addr.rs +++ b/evm/tests/transfer_to_new_addr.rs @@ -113,5 +113,5 @@ fn eth_to_wei(eth: U256) -> U256 { } fn init_logger() { - let _ = try_init_from_env(Env::default().filter_or(DEFAULT_FILTER_ENV, "debug")); + let _ = try_init_from_env(Env::default().filter_or(DEFAULT_FILTER_ENV, "info")); } From 94b73e87bc8d0faa6e9e1127543939ae919522f4 Mon Sep 17 00:00:00 2001 From: Daniel Lubarov Date: Fri, 9 Dec 2022 14:48:00 -0800 Subject: [PATCH 3/3] backtraces --- .github/workflows/continuous-integration-workflow.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/continuous-integration-workflow.yml b/.github/workflows/continuous-integration-workflow.yml index 48848b73..0e3e5977 100644 --- a/.github/workflows/continuous-integration-workflow.yml +++ b/.github/workflows/continuous-integration-workflow.yml @@ -46,6 +46,7 @@ jobs: env: RUSTFLAGS: -Copt-level=3 -Cdebug-assertions -Coverflow-checks=y -Cdebuginfo=0 -Cprefer-dynamic=y CARGO_INCREMENTAL: 1 + RUST_BACKTRACE: 1 lints: name: Formatting and Clippy