Merge pull request #841 from mir-protocol/more_timing

More timing for zkEVM proofs
This commit is contained in:
Daniel Lubarov 2022-12-10 23:25:19 -08:00 committed by GitHub
commit 64c385721f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 177 additions and 80 deletions

View File

@ -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

View File

@ -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<F: Field>() -> Vec<CrossTableLookup<F>> {
let mut ctls = vec![ctl_keccak(), ctl_logic(), ctl_memory(), ctl_keccak_sponge()];
// TODO: Some CTLs temporarily disabled while we get them working.

View File

@ -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::<F, D>::COLUMNS]:,
[(); MemoryStark::<F, D>::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::<F, C, D>::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::<F, C, D>::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::<Vec<_>>()
@ -107,66 +117,30 @@ where
challenger.observe_cap(cap);
}
let ctl_data_per_table = cross_table_lookup_data::<F, C, D>(
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::<F, C, D>(
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<F, C, const D: usize>(
all_stark: &AllStark<F, D>,
config: &StarkConfig,
trace_poly_values: [Vec<PolynomialValues<F>>; NUM_TABLES],
trace_commitments: Vec<PolynomialBatch<F, C, D>>,
ctl_data_per_table: [CtlData<F>; NUM_TABLES],
challenger: &mut Challenger<F, C::Hasher>,
timing: &mut TimingTree,
) -> Result<[StarkProof<F, C, D>; NUM_TABLES]>
where
F: RichField + Extendable<D>,
C: GenericConfig<D, F = F>,
[(); C::Hasher::HASH_SIZE]:,
[(); CpuStark::<F, D>::COLUMNS]:,
[(); KeccakStark::<F, D>::COLUMNS]:,
[(); KeccakSpongeStark::<F, D>::COLUMNS]:,
[(); LogicStark::<F, D>::COLUMNS]:,
[(); MemoryStark::<F, D>::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<F, C, S, const D: usize>(
stark: &S,

View File

@ -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<T: Copy> Traces<T> {
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,

View File

@ -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::<F, C, D>(&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,
@ -80,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"));
}

View File

@ -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::<F, C, D>(&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 {
@ -112,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"));
}