Skip to content
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.

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

Parallelize test_only segment simulation #498

Merged
merged 2 commits into from
Aug 15, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
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
23 changes: 21 additions & 2 deletions zero_bin/ops/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,15 @@ use std::time::Instant;
#[cfg(not(feature = "test_only"))]
use evm_arithmetization::generation::TrimmedGenerationInputs;
use evm_arithmetization::proof::PublicValues;
use evm_arithmetization::prover::SegmentDataIterator;
use evm_arithmetization::GenerationInputs;
#[cfg(not(feature = "test_only"))]
use paladin::operation::FatalStrategy;
use paladin::{
operation::{FatalError, Monoid, Operation, Result},
registry, RemoteExecute,
};
use proof_gen::types::Field;
use proof_gen::{
proof_gen::{generate_block_proof, generate_segment_agg_proof, generate_transaction_agg_proof},
proof_types::{
Expand All @@ -25,6 +28,22 @@ use zero_bin_common::{debug_utils::save_inputs_to_disk, prover_state::p_state};

registry!();

#[cfg(feature = "test_only")]
#[derive(Deserialize, Serialize, RemoteExecute)]
pub struct BatchTestOnly {}
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is there any reason why you're not using save_inputs_on_error field in this? Even though we technically can't log yet the inputs upon error because the iterator is panicking, we are aiming at changing this logic before the release, which would allow to save inputs even in testing mode.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I assumed we would add it when logging is available but we can have it there already. I also mimicked the error handling of the other operations for the future.


#[cfg(feature = "test_only")]
impl Operation for BatchTestOnly {
type Input = (GenerationInputs, usize);
type Output = ();

fn execute(&self, inputs: Self::Input) -> Result<Self::Output> {
let _ = SegmentDataIterator::<Field>::new(&inputs.0, Some(inputs.1)).collect::<Vec<_>>();

Ok(())
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
let _ = SegmentDataIterator::<Field>::new(&inputs.0, Some(inputs.1)).collect::<Vec<_>>();
Ok(())
simulate_execution_all_segments::<Field>(&inputs.0, inputs.1)

}
}

#[derive(Deserialize, Serialize, RemoteExecute)]
pub struct SegmentProof {
pub save_inputs_on_error: bool,
Expand Down Expand Up @@ -73,8 +92,8 @@ impl Operation for SegmentProof {
type Input = AllData;
type Output = ();

fn execute(&self, _input: Self::Input) -> Result<Self::Output> {
todo!() // currently unused, change or remove
fn execute(&self, _all_data: Self::Input) -> Result<Self::Output> {
Ok(())
}
}

Expand Down
22 changes: 15 additions & 7 deletions zero_bin/prover/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -135,12 +135,14 @@ impl BlockProverInput {
#[cfg(feature = "test_only")]
pub async fn prove(
self,
_runtime: &Runtime,
runtime: &Runtime,
previous: Option<impl Future<Output = Result<GeneratedBlockProof>>>,
prover_config: ProverConfig,
) -> Result<GeneratedBlockProof> {
use evm_arithmetization::prover::testing::simulate_execution_all_segments;
use plonky2::field::goldilocks_field::GoldilocksField;
use std::iter::repeat;

use futures::StreamExt;
use paladin::directive::{Directive, IndexedStream};

let block_number = self.get_block_number();
info!("Testing witness generation for block {block_number}.");
Expand All @@ -152,10 +154,16 @@ impl BlockProverInput {
prover_config.batch_size,
)?;

type F = GoldilocksField;
for txn in txs.into_iter() {
simulate_execution_all_segments::<F>(txn, prover_config.max_cpu_len_log)?;
}
let batch_ops = ops::BatchTestOnly {};

let simulation = Directive::map(
IndexedStream::from(txs.into_iter().zip(repeat(prover_config.max_cpu_len_log))),
&batch_ops,
);

let result = simulation.run(runtime).await?;

result.collect::<Vec<_>>().await;

// Wait for previous block proof
let _prev = match previous {
Expand Down
Loading