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.

Sign up for GitHub

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: optimize LWG and NWG #2512

Merged
merged 31 commits into from
Aug 2, 2024
Merged

feat: optimize LWG and NWG #2512

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
31 commits
Select commit Hold shift + click to select a range
ff57b38
Move recursive circuits out of AggregationWrapper struct
0xVolosnikov Jul 26, 2024
64fb5b7
Add comments in prover tests
0xVolosnikov Jul 26, 2024
88e610c
Merge branch 'main' into vv-optimize-lwg-nwg
0xVolosnikov Jul 29, 2024
24ef093
Cargo fmt
0xVolosnikov Jul 29, 2024
cc583c7
Process LWG in batches
0xVolosnikov Jul 29, 2024
7c261a0
Move LWG processing in separate thread
0xVolosnikov Jul 29, 2024
a5105f3
Remove separate thread; cleanup
0xVolosnikov Jul 29, 2024
7b369d5
Fix after merge
0xVolosnikov Jul 30, 2024
69c80f8
Parallelize LWG
0xVolosnikov Jul 30, 2024
9b9c52d
Fix semaphore usage
0xVolosnikov Jul 30, 2024
5a3f20d
Merge branch 'main' into vv-optimize-lwg-nwg
0xVolosnikov Jul 30, 2024
89249e9
Add comments to constants
0xVolosnikov Jul 30, 2024
3ed06aa
Fix todos in basic_test
0xVolosnikov Jul 30, 2024
745461e
Cargo fmt
0xVolosnikov Jul 30, 2024
4b32eb0
Simplify LWG
0xVolosnikov Jul 30, 2024
2dbca20
Parallelize NWG as well
0xVolosnikov Jul 30, 2024
b5b9a04
Update deps
0xVolosnikov Aug 2, 2024
fb75fc2
Address comments
0xVolosnikov Aug 2, 2024
a119262
Cargo.lock
0xVolosnikov Aug 2, 2024
8b6add2
Merge remote-tracking branch 'origin/main' into vv-optimize-lwg-nwg
0xVolosnikov Aug 2, 2024
e7e0235
Hotfix
0xVolosnikov Aug 2, 2024
9899892
Cargo fmt
0xVolosnikov Aug 2, 2024
6468ddd
Use max_circuits_in_flight from config
0xVolosnikov Aug 2, 2024
c2a5e72
Fix
0xVolosnikov Aug 2, 2024
e0694a2
Fix lint issues
0xVolosnikov Aug 2, 2024
0f4db7c
Merge remote-tracking branch 'origin/main' into vv-optimize-lwg-nwg
0xVolosnikov Aug 2, 2024
b61d466
Revert unnessesary Cargo.lock changes
0xVolosnikov Aug 2, 2024
f8c3236
Add comments
0xVolosnikov Aug 2, 2024
1219f7f
Fix lint issues in tests
0xVolosnikov Aug 2, 2024
752c541
Revert blocking spawns
0xVolosnikov Aug 2, 2024
7486f2d
Cleanup
0xVolosnikov Aug 2, 2024
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
6 changes: 2 additions & 4 deletions Cargo.lock
View file
Open in desktop

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

2 changes: 1 addition & 1 deletion Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -200,7 +200,7 @@ circuit_sequencer_api_1_3_3 = { package = "circuit_sequencer_api", version = "=0
circuit_sequencer_api_1_4_0 = { package = "circuit_sequencer_api", version = "=0.140.0" }
circuit_sequencer_api_1_4_1 = { package = "circuit_sequencer_api", version = "=0.141.0" }
circuit_sequencer_api_1_4_2 = { package = "circuit_sequencer_api", version = "=0.142.0" }
circuit_sequencer_api_1_5_0 = { package = "circuit_sequencer_api", version = "=0.150.2-rc.1" }
circuit_sequencer_api_1_5_0 = { git = "https://github.com/matter-labs/era-zkevm_test_harness", branch="vv-optimize-lwg-nwg", package = "circuit_sequencer_api"}
crypto_codegen = { package = "zksync_solidity_vk_codegen", version = "=0.1.0" }
kzg = { package = "zksync_kzg", version = "=0.150.2-rc.1" }
zk_evm = { version = "=0.133.0" }
Expand Down
51 changes: 36 additions & 15 deletions prover/Cargo.lock
View file
Open in desktop

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

6 changes: 3 additions & 3 deletions prover/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,9 +56,9 @@ tracing-subscriber = { version = "0.3" }
vise = "0.1.0"

# Proving dependencies
circuit_definitions = "=0.150.2-rc.1"
circuit_sequencer_api = "=0.150.2-rc.1"
zkevm_test_harness = "=0.150.2-rc.1"
circuit_definitions = { git = "https://github.com/matter-labs/era-zkevm_test_harness", branch="vv-optimize-lwg-nwg", package = "circuit_definitions"}
circuit_sequencer_api = { git = "https://github.com/matter-labs/era-zkevm_test_harness", branch="vv-optimize-lwg-nwg", package = "circuit_sequencer_api"}
zkevm_test_harness = { git = "https://github.com/matter-labs/era-zkevm_test_harness", branch="vv-optimize-lwg-nwg"}

# GPU proving dependencies
wrapper_prover = { package = "zksync-wrapper-prover", version = "=0.140.0-gpu-wrapper.0" }
Expand Down
14 changes: 6 additions & 8 deletions prover/crates/bin/witness_generator/src/leaf_aggregation.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,11 +46,8 @@ use crate::{
pub struct LeafAggregationArtifacts {
circuit_id: u8,
block_number: L1BatchNumber,
pub aggregations: Vec<(
u64,
RecursionQueueSimulator<GoldilocksField>,
ZkSyncRecursiveLayerCircuit,
)>,
pub aggregations: Vec<(u64, RecursionQueueSimulator<GoldilocksField>)>,
pub recursive_circuits: Vec<ZkSyncRecursiveLayerCircuit>,
#[allow(dead_code)]
closed_form_inputs: Vec<ZkSyncBaseLayerClosedFormInput<GoldilocksField>>,
}
Expand Down Expand Up @@ -268,7 +265,7 @@ pub fn process_leaf_aggregation_job(
job.closed_form_inputs.0,
);
let leaf_params = (circuit_id, job.leaf_params);
let (aggregations, closed_form_inputs) =
let (aggregations, recursive_circuits, closed_form_inputs) =
create_leaf_witnesses(subsets, job.proofs, job.base_vk, leaf_params);
WITNESS_GENERATOR_METRICS.witness_generation_time[&AggregationRound::LeafAggregation.into()]
.observe(started_at.elapsed());
Expand All @@ -284,6 +281,7 @@ pub fn process_leaf_aggregation_job(
circuit_id,
block_number: job.block_number,
aggregations,
recursive_circuits,
closed_form_inputs,
}
}
Expand Down Expand Up @@ -385,13 +383,13 @@ async fn save_artifacts(
artifacts.block_number,
get_recursive_layer_circuit_id_for_base_layer(artifacts.circuit_id),
0,
artifacts.aggregations.clone(),
artifacts.aggregations,
object_store,
)
.await;
let circuit_ids_and_urls = save_recursive_layer_prover_input_artifacts(
artifacts.block_number,
artifacts.aggregations,
artifacts.recursive_circuits,
AggregationRound::LeafAggregation,
0,
object_store,
Expand Down
20 changes: 7 additions & 13 deletions prover/crates/bin/witness_generator/src/node_aggregation.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,11 +41,8 @@ pub struct NodeAggregationArtifacts {
circuit_id: u8,
block_number: L1BatchNumber,
depth: u16,
pub next_aggregations: Vec<(
u64,
RecursionQueueSimulator<GoldilocksField>,
ZkSyncRecursiveLayerCircuit,
)>,
pub next_aggregations: Vec<(u64, RecursionQueueSimulator<GoldilocksField>)>,
pub recursive_circuits: Vec<ZkSyncRecursiveLayerCircuit>,
}

#[derive(Debug)]
Expand All @@ -59,11 +56,7 @@ pub struct NodeAggregationWitnessGeneratorJob {
circuit_id: u8,
block_number: L1BatchNumber,
depth: u16,
aggregations: Vec<(
u64,
RecursionQueueSimulator<GoldilocksField>,
ZkSyncRecursiveLayerCircuit,
)>,
aggregations: Vec<(u64, RecursionQueueSimulator<GoldilocksField>)>,
proofs: Vec<ZkSyncRecursionLayerProof>,
leaf_vk: ZkSyncRecursionLayerVerificationKey,
node_vk: ZkSyncRecursionLayerVerificationKey,
Expand Down Expand Up @@ -109,7 +102,7 @@ impl NodeAggregationWitnessGenerator {
0 => job.leaf_vk,
_ => job.node_vk,
};
let next_aggregations = create_node_witnesses(
let (next_aggregations, recursive_circuits) = create_node_witnesses(
job.aggregations,
job.proofs,
vk,
Expand All @@ -134,6 +127,7 @@ impl NodeAggregationWitnessGenerator {
block_number: job.block_number,
depth: job.depth + 1,
next_aggregations,
recursive_circuits,
}
}
}
Expand Down Expand Up @@ -375,13 +369,13 @@ async fn save_artifacts(
artifacts.block_number,
artifacts.circuit_id,
artifacts.depth,
artifacts.next_aggregations.clone(),
artifacts.next_aggregations,
object_store,
)
.await;
let circuit_ids_and_urls = save_recursive_layer_prover_input_artifacts(
artifacts.block_number,
artifacts.next_aggregations,
artifacts.recursive_circuits,
AggregationRound::NodeAggregation,
artifacts.depth,
object_store,
Expand Down
24 changes: 5 additions & 19 deletions prover/crates/bin/witness_generator/src/utils.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,13 +78,7 @@ impl StoredObject for ClosedFormInputWrapper {
}

#[derive(serde::Serialize, serde::Deserialize)]
pub struct AggregationWrapper(
pub Vec<(
u64,
RecursionQueueSimulator<GoldilocksField>,
ZkSyncRecursiveLayerCircuit,
)>,
);
pub struct AggregationWrapper(pub Vec<(u64, RecursionQueueSimulator<GoldilocksField>)>);
Copy link
Contributor

Choose a reason for hiding this comment

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

Nit: Would be nice to turn those into some DTO structs.


impl StoredObject for AggregationWrapper {
const BUCKET: Bucket = Bucket::NodeAggregationWitnessJobsFri;
Expand Down Expand Up @@ -145,18 +139,14 @@ pub async fn save_circuit(

pub async fn save_recursive_layer_prover_input_artifacts(
block_number: L1BatchNumber,
aggregations: Vec<(
u64,
RecursionQueueSimulator<GoldilocksField>,
ZkSyncRecursiveLayerCircuit,
)>,
recursive_circuits: Vec<ZkSyncRecursiveLayerCircuit>,
aggregation_round: AggregationRound,
depth: u16,
object_store: &dyn ObjectStore,
base_layer_circuit_id: Option<u8>,
) -> Vec<(u8, String)> {
let mut ids_and_urls = Vec::with_capacity(aggregations.len());
for (sequence_number, (_, _, circuit)) in aggregations.into_iter().enumerate() {
let mut ids_and_urls = Vec::with_capacity(recursive_circuits.len());
for (sequence_number, circuit) in recursive_circuits.into_iter().enumerate() {
let circuit_id = base_layer_circuit_id.unwrap_or_else(|| circuit.numeric_circuit_type());
let circuit_key = FriCircuitKey {
block_number,
Expand All @@ -178,11 +168,7 @@ pub async fn save_node_aggregations_artifacts(
block_number: L1BatchNumber,
circuit_id: u8,
depth: u16,
aggregations: Vec<(
u64,
RecursionQueueSimulator<GoldilocksField>,
ZkSyncRecursiveLayerCircuit,
)>,
aggregations: Vec<(u64, RecursionQueueSimulator<GoldilocksField>)>,
object_store: &dyn ObjectStore,
) -> String {
let key = AggregationsKey {
Expand Down
4 changes: 2 additions & 2 deletions prover/crates/bin/witness_generator/tests/basic_test.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ async fn test_leaf_witness_gen() {
.unwrap();

let artifacts = LeafAggregationWitnessGenerator::process_job_sync(job, Instant::now());
let aggregations = AggregationWrapper(artifacts.aggregations);
let aggregations = AggregationWrapper(artifacts.aggregations); // TODO check circuits as well
compare_serialized(&expected_aggregation, &aggregations);
}

Expand Down Expand Up @@ -103,6 +103,6 @@ async fn test_node_witness_gen() {
.unwrap();

let artifacts = NodeAggregationWitnessGenerator::process_job_sync(job, Instant::now());
let aggregations = AggregationWrapper(artifacts.next_aggregations);
let aggregations = AggregationWrapper(artifacts.next_aggregations); // TODO check circuits as well
compare_serialized(&expected_aggregation, &aggregations);
}
Loading