Skip to content

Commit

Permalink
feat(prover): Moved setup key generation logic to test harness (matte…
Browse files Browse the repository at this point in the history
…r-labs#1113)

## What ❔

* Moved setup key generation logic to test harness
* Allow key generator to pass the different directories when generating
keys.
* Moved the setup key generation for CPU and GPU under a common trait.

## Why ❔

* Setup key generation logic should live in test harness (where it is
closer to actual circuit logic)
* Passing different directories when generating keys, allow for more
flexibility and easier tooling integrations.
  • Loading branch information
mm-zk authored Feb 19, 2024
1 parent 1471615 commit 469ab06
Show file tree
Hide file tree
Showing 6 changed files with 294 additions and 676 deletions.
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())
.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.
//!
//! 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 {
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
.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

0 comments on commit 469ab06

Please sign in to comment.