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

feat(prover): Moved setup key generation logic to test harness #1113

Merged
merged 19 commits into from
Feb 19, 2024
21 changes: 6 additions & 15 deletions prover/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

10 changes: 5 additions & 5 deletions prover/prover_fri/tests/basic_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ use zksync_prover_fri_types::{
};
use zksync_types::{basic_fri_types::AggregationRound, L1BatchNumber};
use zksync_vk_setup_data_server_fri::{
keystore::Keystore, setup_data_generator::generate_cpu_base_layer_setup_data,
keystore::Keystore, setup_data_generator::generate_setup_data_common,
};

fn compare_serialized<T: Serialize>(expected: &T, actual: &T) {
Expand Down Expand Up @@ -57,10 +57,10 @@ async fn prover_and_assert_base_layer(
CircuitWrapper::Recursive(_) => anyhow::bail!("Expected base layer circuit"),
};
let keystore = Keystore::default();
let setup_data = Arc::new(
generate_cpu_base_layer_setup_data(&keystore, circuit)
.context("generate_cpu_base_layers_setup_data()")?,
);
let circuit_setup_data =
generate_setup_data_common(&keystore, true, circuit.numeric_circuit_type())
mm-zk marked this conversation as resolved.
Show resolved Hide resolved
.context("generate_cpu_base_layers_setup_data()")?;
let setup_data = Arc::new(circuit_setup_data.into());
let setup_key = ProverServiceDataKey::new(circuit_id, aggregation_round);
let prover_job = ProverJob::new(block_number, expected_proof_id, circuit_wrapper, setup_key);
let artifacts = Prover::prove(
Expand Down
16 changes: 15 additions & 1 deletion prover/vk_setup_data_generator_server_fri/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
#![feature(generic_const_exprs)]
#![feature(allocator_api)]

use serde::{Deserialize, Serialize};
use zkevm_test_harness::compute_setups::CircuitSetupData;
use zksync_prover_fri_types::circuit_definitions::boojum::{
algebraic_props::{round_function::AbsorptionModeOverwrite, sponge::GenericAlgebraicSponge},
cs::{
Expand Down Expand Up @@ -68,6 +68,20 @@ pub type GoldilocksProverSetupData = ProverSetupData<
>,
>;

impl From<CircuitSetupData> for GoldilocksProverSetupData {
fn from(circuit_setup_data: CircuitSetupData) -> Self {
Self {
setup_base: circuit_setup_data.setup_base,
setup: circuit_setup_data.setup,
vk: circuit_setup_data.vk,
setup_tree: circuit_setup_data.setup_tree,
vars_hint: circuit_setup_data.vars_hint,
wits_hint: circuit_setup_data.wits_hint,
finalization_hint: circuit_setup_data.finalization_hint,
}
}
}

#[cfg(feature = "gpu")]
#[derive(Debug, Serialize, Deserialize)]
#[serde(bound = "F: serde::Serialize + serde::de::DeserializeOwned")]
Expand Down
202 changes: 108 additions & 94 deletions prover/vk_setup_data_generator_server_fri/src/main.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
//! Tool to generate different types of keys used by the proving system.
mm-zk marked this conversation as resolved.
Show resolved Hide resolved
//!
//! It can generate verification keys, setup keys, and also commitments.
use std::collections::HashMap;

use anyhow::Context as _;
Expand All @@ -15,10 +18,7 @@ use zkevm_test_harness::{
use zksync_prover_fri_types::circuit_definitions::circuit_definitions::recursion_layer::ZkSyncRecursionLayerStorageType;
use zksync_vk_setup_data_server_fri::{
keystore::Keystore,
setup_data_generator::{
generate_all_cpu_setup_data, generate_all_gpu_setup_data, generate_cpu_setup_data,
generate_gpu_setup_data,
},
setup_data_generator::{CPUSetupDataGenerator, GPUSetupDataGenerator, SetupDataGenerator},
};

mod commitment_generator;
Expand Down Expand Up @@ -80,43 +80,53 @@ enum CircuitSelector {
Basic,
}

#[derive(Debug, Parser)]
struct GeneratorOptions {
circuits_type: CircuitSelector,

/// Specify for which circuit to generate the keys.
#[arg(long)]
numeric_circuit: Option<u8>,

/// If true, then setup keys are not written and only md5 sum is printed.
#[arg(long, default_value = "false")]
dry_run: bool,

#[arg(long)]
path: Option<String>,

#[arg(long)]
setup_path: Option<String>,
}

#[derive(Debug, Subcommand)]
enum Command {
/// Generates verification keys (and finalization hints) for all the basic & leaf circuits.
/// Used for verification.
#[command(name = "generate-vk")]
GenerateVerificationKeys {},
GenerateVerificationKeys {
#[arg(long)]
path: Option<String>,
},
/// Generates setup keys (used by the CPU prover).
#[command(name = "generate-sk")]
GenerateSetupKeys {
circuits_type: CircuitSelector,

/// Specify for which circuit to generate the keys.
#[arg(long)]
numeric_circuit: Option<u8>,

/// If true, then setup keys are not written and only md5 sum is printed.
#[arg(long, default_value = "false")]
dry_run: bool,
#[command(flatten)]
options: GeneratorOptions,
},
/// Generates setup keys (used by the GPU prover).
#[command(name = "generate-sk-gpu")]
GenerateGPUSetupKeys {
circuits_type: CircuitSelector,

/// Specify for which circuit to generate the keys.
#[arg(long)]
numeric_circuit: Option<u8>,

/// If true, then setup keys are not written and only md5 sum is printed.
#[arg(long, default_value = "false")]
dry_run: bool,
#[command(flatten)]
options: GeneratorOptions,
},
/// Generates and updates the commitments - used by the verification contracts.
#[command(name = "update-commitments")]
UpdateCommitments {
#[arg(long, default_value = "true")]
dryrun: bool,
#[arg(long)]
path: Option<String>,
},
}

Expand All @@ -132,6 +142,55 @@ fn print_stats(digests: HashMap<String, String>) -> anyhow::Result<()> {
Ok(())
}

fn keystore_from_optional_path(path: Option<String>, setup_path: Option<String>) -> Keystore {
if let Some(path) = path {
mm-zk marked this conversation as resolved.
Show resolved Hide resolved
return Keystore::new_with_optional_setup_path(path, setup_path);
}
if setup_path.is_some() {
panic!("--setup_path must not be set when --path is not set");
}
Keystore::default()
}

fn generate_setup_keys(
generator: &dyn SetupDataGenerator,
options: &GeneratorOptions,
) -> anyhow::Result<()> {
match options.circuits_type {
CircuitSelector::All => {
let digests = generator.generate_all(options.dry_run)?;
tracing::info!("Setup keys md5(s):");
print_stats(digests)
}
CircuitSelector::Recursive => {
let digest = generator
mm-zk marked this conversation as resolved.
Show resolved Hide resolved
.generate_and_write_setup_data(
false,
options
.numeric_circuit
.expect("--numeric-circuit must be provided"),
options.dry_run,
)
.context("generate_setup_data()")?;
tracing::info!("digest: {:?}", digest);
Ok(())
}
CircuitSelector::Basic => {
let digest = generator
.generate_and_write_setup_data(
true,
options
.numeric_circuit
.expect("--numeric-circuit must be provided"),
options.dry_run,
)
.context("generate_setup_data()")?;
tracing::info!("digest: {:?}", digest);
Ok(())
}
}
}

fn main() -> anyhow::Result<()> {
tracing_subscriber::fmt()
.with_env_filter(
Expand All @@ -142,84 +201,39 @@ fn main() -> anyhow::Result<()> {
.init();

let opt = Cli::parse();
// Setting key store from argument flags will come in next PR.
let keystore = Keystore::default();

match opt.command {
Command::GenerateVerificationKeys {} => {
Command::GenerateVerificationKeys { path } => {
let keystore = keystore_from_optional_path(path, None);
tracing::info!(
"Generating verification keys and storing them inside {:?}",
keystore.get_base_path()
);
generate_vks(&keystore).context("generate_vks()")
}
Command::UpdateCommitments { dryrun } => read_and_update_contract_toml(&keystore, dryrun),

Command::GenerateSetupKeys {
circuits_type,
numeric_circuit,
dry_run,
} => match circuits_type {
CircuitSelector::All => {
let digests = generate_all_cpu_setup_data(&keystore, dry_run)?;
tracing::info!("CPU Setup keys md5(s):");
print_stats(digests)
}
CircuitSelector::Recursive => {
let digest = generate_cpu_setup_data(
&keystore,
false,
numeric_circuit.expect("--numeric-circuit must be provided"),
dry_run,
)
.context("generate_cpu_setup_data()")?;
tracing::info!("digest: {:?}", digest);
Ok(())
}
CircuitSelector::Basic => {
let digest = generate_cpu_setup_data(
&keystore,
true,
numeric_circuit.expect("--numeric-circuit must be provided"),
dry_run,
)
.context("generate_cpu_setup_data()")?;
tracing::info!("digest: {:?}", digest);
Ok(())
}
},
Command::GenerateGPUSetupKeys {
circuits_type,
numeric_circuit,
dry_run,
} => match circuits_type {
CircuitSelector::All => {
let digests = generate_all_gpu_setup_data(&keystore, dry_run)?;
tracing::info!("GPU Setup keys md5(s):");
print_stats(digests)
}
CircuitSelector::Recursive => {
let digest = generate_gpu_setup_data(
&keystore,
false,
numeric_circuit.expect("--numeric-circuit must be provided"),
dry_run,
)
.context("generate_gpu_setup_data()")?;
tracing::info!("digest: {:?}", digest);
Ok(())
}
CircuitSelector::Basic => {
let digest = generate_gpu_setup_data(
&keystore,
true,
numeric_circuit.expect("--numeric-circuit must be provided"),
dry_run,
)
.context("generate_gpu_setup_data()")?;
tracing::info!("digest: {:?}", digest);
Ok(())
}
},
Command::UpdateCommitments { dryrun, path } => {
let keystore = keystore_from_optional_path(path, None);

read_and_update_contract_toml(&keystore, dryrun)
}

Command::GenerateSetupKeys { options } => {
let generator = CPUSetupDataGenerator {
keystore: keystore_from_optional_path(
options.path.clone(),
options.setup_path.clone(),
),
};
generate_setup_keys(&generator, &options)
}
Command::GenerateGPUSetupKeys { options } => {
let generator = GPUSetupDataGenerator {
keystore: keystore_from_optional_path(
options.path.clone(),
options.setup_path.clone(),
),
};
generate_setup_keys(&generator, &options)
}
}
}
Loading
Loading