From ab43d6179d1c0950151f34fc6f6e08ec577ae274 Mon Sep 17 00:00:00 2001 From: saitima Date: Mon, 13 Jan 2025 18:30:53 +0300 Subject: [PATCH] pinned memory for fflonk combined monomials --- crates/proof-compression/src/chain.rs | 130 ++++++++++++++---- .../proof-compression/src/proof_system/crs.rs | 31 ++--- .../src/proof_system/fflonk.rs | 16 ++- .../src/proof_system/interface.rs | 3 +- .../src/proof_system/plonk.rs | 36 ++--- .../src/step/snark_wrapper.rs | 4 + crates/proof-compression/src/test.rs | 31 +---- 7 files changed, 154 insertions(+), 97 deletions(-) diff --git a/crates/proof-compression/src/chain.rs b/crates/proof-compression/src/chain.rs index 3a60aa5..fffeaa7 100644 --- a/crates/proof-compression/src/chain.rs +++ b/crates/proof-compression/src/chain.rs @@ -1,10 +1,74 @@ -use circuit_definitions::circuit_definitions::aux_layer::compression_modes::{ - CompressionMode1, CompressionMode1ForWrapper, CompressionMode2, CompressionMode3, - CompressionMode4, CompressionMode5ForWrapper, +use circuit_definitions::circuit_definitions::aux_layer::{ + compression_modes::{ + CompressionMode1, CompressionMode1ForWrapper, CompressionMode2, CompressionMode3, + CompressionMode4, CompressionMode5ForWrapper, + }, + ZkSyncCompressionForWrapperProof, ZkSyncCompressionLayerProof, }; +use fflonk::FflonkSnarkVerifierCircuitProof; use super::*; +pub trait ProofStorage { + fn get_scheduler_proof(&self) -> SchedulerProof; + fn save_compression_layer_proof(&mut self, circuit_id: u8, proof: ZkSyncCompressionLayerProof); + fn save_compression_wrapper_proof( + &mut self, + circuit_id: u8, + proof: ZkSyncCompressionForWrapperProof, + ); + fn save_plonk_proof(&mut self, proof: PlonkSnarkVerifierCircuitProof); + fn save_fflonk_proof(&mut self, proof: FflonkSnarkVerifierCircuitProof); +} + +pub struct SimpleProofStorage { + compression_layer_storage: std::collections::HashMap, + compression_wrapper_storage: std::collections::HashMap, + fflonk: Option, + plonk: Option, +} + +impl SimpleProofStorage { + pub fn new() -> Self { + Self { + compression_layer_storage: std::collections::HashMap::new(), + compression_wrapper_storage: std::collections::HashMap::new(), + fflonk: None, + plonk: None, + } + } +} + +impl ProofStorage for SimpleProofStorage { + fn get_scheduler_proof(&self) -> SchedulerProof { + let scheduler_proof_file = + std::fs::File::open("./data/scheduler_recursive_proof.json").unwrap(); + let scheduler_proof: circuit_definitions::circuit_definitions::recursion_layer::ZkSyncRecursionLayerProof = + serde_json::from_reader(&scheduler_proof_file).unwrap(); + let scheduler_proof = scheduler_proof.into_inner(); + scheduler_proof + } + fn save_compression_layer_proof(&mut self, circuit_id: u8, proof: ZkSyncCompressionLayerProof) { + self.compression_layer_storage.insert(circuit_id, proof); + } + + fn save_compression_wrapper_proof( + &mut self, + circuit_id: u8, + proof: ZkSyncCompressionForWrapperProof, + ) { + self.compression_wrapper_storage.insert(circuit_id, proof); + } + + fn save_plonk_proof(&mut self, proof: PlonkSnarkVerifierCircuitProof) { + self.plonk = Some(proof); + } + + fn save_fflonk_proof(&mut self, proof: FflonkSnarkVerifierCircuitProof) { + self.fflonk = Some(proof); + } +} + pub enum SnarkWrapper { Plonk, FFfonk, @@ -20,61 +84,72 @@ pub type SchedulerProof = franklin_crypto::boojum::cs::implementations::proof::P GoldilocksExt2, >; -pub fn run_proof_chain( - input_proof: SchedulerProof, +pub fn run_proof_chain( snark_wrapper: SnarkWrapper, blob_storage: &BS, -) -> SnarkWrapperProof -where + proof_storage: &mut PS, +) where BS: BlobStorage, + PS: ProofStorage, { match snark_wrapper { - SnarkWrapper::Plonk => { - let proof = run_proof_chain_with_plonk(input_proof, blob_storage); - SnarkWrapperProof::Plonk(proof) - } - SnarkWrapper::FFfonk => { - let proof = run_proof_chain_with_fflonk(input_proof, blob_storage); - SnarkWrapperProof::FFfonk(proof) - } + SnarkWrapper::Plonk => run_proof_chain_with_plonk(blob_storage, proof_storage), + SnarkWrapper::FFfonk => run_proof_chain_with_fflonk(blob_storage, proof_storage), } } -pub fn run_proof_chain_with_fflonk( - input_proof: SchedulerProof, - blob_storage: &BS, -) -> FflonkSnarkVerifierCircuitProof +pub fn run_proof_chain_with_fflonk(blob_storage: &BS, proof_storage: &mut PS) where BS: BlobStorage, + PS: ProofStorage, { let context_manager = SimpleContextManager::new(); let start = std::time::Instant::now(); + ::run_pre_initialization_tasks(); let compact_raw_crs = ::load_compact_raw_crs(blob_storage); let fflonk_precomputation = FflonkSnarkWrapper::get_precomputation(blob_storage); + let input_proof = proof_storage.get_scheduler_proof(); + let next_proof = CompressionMode1::prove_compression_step(input_proof, blob_storage, &context_manager); + let compression_proof_1 = ZkSyncCompressionLayerProof::from_inner(1, next_proof.clone()); + proof_storage.save_compression_layer_proof(1, compression_proof_1.clone()); + let next_proof = CompressionMode2::prove_compression_step::<_, SimpleContextManager>( next_proof, blob_storage, &context_manager, ); + let compression_proof_2 = ZkSyncCompressionLayerProof::from_inner(2, next_proof.clone()); + proof_storage.save_compression_layer_proof(2, compression_proof_2.clone()); + let next_proof = CompressionMode3::prove_compression_step::<_, SimpleContextManager>( next_proof, blob_storage, &context_manager, ); + let compression_proof_3 = ZkSyncCompressionLayerProof::from_inner(3, next_proof.clone()); + proof_storage.save_compression_layer_proof(3, compression_proof_3.clone()); + let next_proof = CompressionMode4::prove_compression_step::<_, SimpleContextManager>( next_proof, blob_storage, &context_manager, ); + let compression_proof_4 = ZkSyncCompressionLayerProof::from_inner(4, next_proof.clone()); + proof_storage.save_compression_layer_proof(4, compression_proof_4.clone()); + let next_proof = CompressionMode5ForWrapper::prove_compression_step::<_, SimpleContextManager>( next_proof, blob_storage, &context_manager, ); + proof_storage.save_compression_wrapper_proof( + 5, + ZkSyncCompressionForWrapperProof::from_inner(5, next_proof.clone()), + ); println!( "Proving entire compression chain took {}s", start.elapsed().as_secs() @@ -90,7 +165,7 @@ where "Proving entire chain with snark wrapper took {}s", start.elapsed().as_secs() ); - final_proof + proof_storage.save_fflonk_proof(final_proof); } pub fn precompute_proof_chain_with_fflonk(blob_storage: &BS) @@ -98,6 +173,7 @@ where BS: BlobStorageExt, { let context_manager = SimpleContextManager::new(); + ::run_pre_initialization_tasks(); let compact_raw_crs = ::load_compact_raw_crs(blob_storage); @@ -125,12 +201,10 @@ where ); } -pub fn run_proof_chain_with_plonk( - input_proof: SchedulerProof, - blob_storage: &BS, -) -> PlonkSnarkVerifierCircuitProof +pub fn run_proof_chain_with_plonk(blob_storage: &BS, proof_storage: &mut PS) where BS: BlobStorage, + PS: ProofStorage, { let context_manager = SimpleContextManager::new(); let start = std::time::Instant::now(); @@ -138,11 +212,17 @@ where ::load_compact_raw_crs(blob_storage); let plonk_precomputation = PlonkSnarkWrapper::get_precomputation(blob_storage); + let input_proof = proof_storage.get_scheduler_proof(); + let next_proof = CompressionMode1ForWrapper::prove_compression_step( input_proof, blob_storage, &context_manager, ); + proof_storage.save_compression_wrapper_proof( + 1, + ZkSyncCompressionForWrapperProof::from_inner(1, next_proof.clone()), + ); let final_proof = PlonkSnarkWrapper::prove_snark_wrapper_step( compact_raw_crs, @@ -155,7 +235,7 @@ where "Entire compression chain with plonk took {}s", start.elapsed().as_secs() ); - final_proof + proof_storage.save_plonk_proof(final_proof); } pub fn precompute_proof_chain_with_plonk(blob_storage: &BS) diff --git a/crates/proof-compression/src/proof_system/crs.rs b/crates/proof-compression/src/proof_system/crs.rs index e63292f..8cf0cec 100644 --- a/crates/proof-compression/src/proof_system/crs.rs +++ b/crates/proof-compression/src/proof_system/crs.rs @@ -1,7 +1,7 @@ use super::*; -use ::fflonk::bellman::{ - self, - kate_commitment::{Crs, CrsForMonomialForm}, +use ::fflonk::{ + bellman::kate_commitment::{Crs, CrsForMonomialForm}, + hardcoded_g2_bases, }; use bellman::{bn256::Bn256, CurveAffine, Engine, Field, PrimeField}; use byteorder::{BigEndian, ReadBytesExt}; @@ -10,13 +10,11 @@ use gpu_prover::ManagerConfigs; pub(crate) fn write_crs_into_raw_compact_form( original_crs: &Crs, mut dst_raw_compact_crs: W, - num_points: usize, ) -> std::io::Result<()> { use bellman::CurveAffine; - assert!(num_points <= original_crs.g1_bases.len()); use bellman::{PrimeField, PrimeFieldRepr}; use byteorder::{BigEndian, WriteBytesExt}; - assert!(num_points < u32::MAX as usize); + let num_points = original_crs.g1_bases.len(); dst_raw_compact_crs.write_u32::(num_points as u32)?; for g1_base in original_crs.g1_bases.iter() { let (x, y) = g1_base.as_xy(); @@ -39,6 +37,7 @@ pub(crate) fn read_crs_from_raw_compact_form std::io::Result> { + // requested number of bases can be smaller than the available bases use byteorder::{BigEndian, ReadBytesExt}; let actual_num_points = src_raw_compact_crs.read_u32::()? as usize; assert!(num_g1_points <= actual_num_points as usize); @@ -51,16 +50,9 @@ pub(crate) fn read_crs_from_raw_compact_form(), - ); - src_raw_compact_crs.read_exact(buf)?; - } + + let g2_bases = hardcoded_g2_bases::().to_vec_in(A::default()); + Ok(Crs::<_, CrsForMonomialForm, A>::new_in(g1_bases, g2_bases)) } @@ -73,7 +65,8 @@ pub(crate) fn create_compact_raw_crs(dst: W) { .max() .unwrap(); let original_crs = make_crs_from_ignition_transcripts(num_points); - write_crs_into_raw_compact_form(&original_crs, dst, num_points).unwrap(); + assert_eq!(original_crs.g1_bases.len(), num_points); + write_crs_into_raw_compact_form(&original_crs, dst).unwrap(); } fn make_crs_from_ignition_transcripts(num_points: usize) -> Crs { @@ -277,3 +270,7 @@ fn create_crs_from_ignition_transcript + ?Sized>( Ok(new) } + +pub(crate) fn hardcoded_canonical_g2_bases() -> [bellman::bn256::G2Affine; 2] { + ::fflonk::hardcoded_g2_bases::() +} diff --git a/crates/proof-compression/src/proof_system/fflonk.rs b/crates/proof-compression/src/proof_system/fflonk.rs index c7af650..a442b88 100644 --- a/crates/proof-compression/src/proof_system/fflonk.rs +++ b/crates/proof-compression/src/proof_system/fflonk.rs @@ -38,7 +38,7 @@ impl ProofSystemDefinition for FflonkSnarkWrapper { type Proof = FflonkSnarkVerifierCircuitProof; type VK = FflonkSnarkVerifierCircuitVK; type FinalizationHint = usize; - // type Allocator = GlobalHost; // TODO need global host with preallocated host memory + // Pinned memory with small allocations is expensive e.g Assembly storage type Allocator = std::alloc::Global; type ProvingAssembly = FflonkAssembly; type Transcript = RollingKeccakTranscript; @@ -69,6 +69,12 @@ impl SnarkWrapperProofSystem for FflonkSnarkWrapper { CrsForMonomialForm, Self::Allocator, >; + + fn pre_init() { + let domain_size = ::fflonk::fflonk::L1_VERIFIER_DOMAIN_SIZE_LOG; + Self::Context::init_pinned_memory(domain_size).unwrap(); + } + fn load_compact_raw_crs(src: R) -> Self::CRS { let domain_size = 1 << ::fflonk::fflonk_cpu::L1_VERIFIER_DOMAIN_SIZE_LOG; let num_g1_bases_for_crs = ::fflonk::fflonk_cpu::MAX_COMBINED_DEGREE_FACTOR * domain_size; @@ -78,7 +84,11 @@ impl SnarkWrapperProofSystem for FflonkSnarkWrapper { fn init_context(compact_raw_crs: AsyncHandler) -> Self::Context { let compact_raw_crs = compact_raw_crs.wait(); let domain_size = 1 << ::fflonk::fflonk_cpu::L1_VERIFIER_DOMAIN_SIZE_LOG; - let context = Self::Context::init_from_preloaded_crs(domain_size, compact_raw_crs).unwrap(); + let context = DeviceContextWithSingleDevice::init_from_preloaded_crs::( + domain_size, + compact_raw_crs, + ) + .unwrap(); context } @@ -120,7 +130,7 @@ impl SnarkWrapperProofSystem for FflonkSnarkWrapper { fn prove_from_witnesses( _: AsyncHandler, - _: Vec, + _: Self::ExternalWitnessData, _: AsyncHandler, _: Self::FinalizationHint, ) -> Self::Proof { diff --git a/crates/proof-compression/src/proof_system/interface.rs b/crates/proof-compression/src/proof_system/interface.rs index 119ec1f..724f76b 100644 --- a/crates/proof-compression/src/proof_system/interface.rs +++ b/crates/proof-compression/src/proof_system/interface.rs @@ -79,6 +79,7 @@ pub(crate) trait SnarkWrapperProofSystem: ProofSystemDefinition { type Circuit; type Context: Send + Sync + 'static; type CRS: Send + Sync + 'static; + fn pre_init(); fn init_context(crs: AsyncHandler) -> Self::Context; fn load_compact_raw_crs(src: R) -> Self::CRS; fn synthesize_for_proving(circuit: Self::Circuit) -> Self::ProvingAssembly; @@ -91,7 +92,7 @@ pub(crate) trait SnarkWrapperProofSystem: ProofSystemDefinition { fn prove_from_witnesses( _: AsyncHandler, - _: Vec, + _: Self::ExternalWitnessData, _: AsyncHandler, _: Self::FinalizationHint, ) -> Self::Proof; diff --git a/crates/proof-compression/src/proof_system/plonk.rs b/crates/proof-compression/src/proof_system/plonk.rs index 2fcbd52..69bcc71 100644 --- a/crates/proof-compression/src/proof_system/plonk.rs +++ b/crates/proof-compression/src/proof_system/plonk.rs @@ -70,10 +70,7 @@ impl ProofSystemDefinition for PlonkSnarkWrapper { type Proof = PlonkSnarkVerifierCircuitProof; type VK = PlonkSnarkVerifierCircuitVK; type FinalizationHint = usize; - #[cfg(feature = "allocator")] type Allocator = gpu_prover::cuda_bindings::CudaAllocator; - #[cfg(not(feature = "allocator"))] - type Allocator = std::alloc::Global; type ProvingAssembly = PlonkAssembly; type Transcript = RollingKeccakTranscript; fn take_witnesses( @@ -98,16 +95,17 @@ impl ProofSystemDefinition for PlonkSnarkWrapper { impl SnarkWrapperProofSystem for PlonkSnarkWrapper { type Circuit = PlonkSnarkVerifierCircuit; - type Context = ( - UnsafePlonkProverDeviceMemoryManagerWrapper, - Vec, - ); + type Context = UnsafePlonkProverDeviceMemoryManagerWrapper; type CRS = bellman::kate_commitment::Crs< bellman::compact_bn256::Bn256, CrsForMonomialForm, Self::Allocator, >; + fn pre_init() { + // TODO: initialize static pinned memory + } + fn load_compact_raw_crs(src: R) -> Self::CRS { let num_g1_points_for_crs = 1 << PlonkProverDeviceMemoryManagerConfig::FULL_SLOT_SIZE_LOG; read_crs_from_raw_compact_form(src, num_g1_points_for_crs).unwrap() @@ -117,12 +115,8 @@ impl SnarkWrapperProofSystem for PlonkSnarkWrapper { let device_ids: Vec<_> = (0..::NUM_GPUS).collect(); let compact_raw_crs = compact_raw_crs.wait(); - let g2_bases = compact_raw_crs.g2_monomial_bases.as_ref().to_vec(); let manager = DeviceMemoryManager::init(&device_ids, &compact_raw_crs.g1_bases).unwrap(); - ( - UnsafePlonkProverDeviceMemoryManagerWrapper(manager), - g2_bases, - ) + UnsafePlonkProverDeviceMemoryManagerWrapper(manager) } fn synthesize_for_proving(circuit: Self::Circuit) -> Self::ProvingAssembly { @@ -146,7 +140,7 @@ impl SnarkWrapperProofSystem for PlonkSnarkWrapper { assert!(domain_size.is_power_of_two()); assert_eq!(domain_size, finalization_hint); - let (ctx, _) = ctx.wait(); + let ctx = ctx.wait(); let mut ctx = ctx.into_inner(); let mut precomputation = precomputation.wait().into_inner(); let worker = bellman::worker::Worker::new(); @@ -197,28 +191,18 @@ impl SnarkWrapperProofSystemExt for PlonkSnarkWrapper { assert!(domain_size.is_power_of_two()); assert_eq!(domain_size, hardcoded_finalization_hint); - let (ctx, compact_g2_bases) = ctx.wait(); + let ctx = ctx.wait(); let mut ctx = ctx.into_inner(); let worker = bellman::worker::Worker::new(); - #[cfg(feature = "allocator")] let mut precomputation = AsyncSetup::::allocate(hardcoded_finalization_hint); - #[cfg(not(feature = "allocator"))] - let mut precomputation = AsyncSetup::allocate(hardcoded_finalization_hint); precomputation .generate_from_assembly(&worker, &setup_assembly, &mut ctx) .unwrap(); + let hardcoded_g2_bases = hardcoded_canonical_g2_bases(); let mut dummy_crs = Crs::::dummy_crs(1); - let mut g2_bases = vec![]; - use bellman::CurveAffine; - for base in compact_g2_bases.into_iter() { - let (x, y) = base.into_xy_unchecked(); - let base = bellman::bn256::G2Affine::from_xy_unchecked(x, y); - g2_bases.push(base); - } - - dummy_crs.g2_monomial_bases = std::sync::Arc::new(g2_bases); + dummy_crs.g2_monomial_bases = std::sync::Arc::new(hardcoded_g2_bases.to_vec()); let vk = gpu_prover::compute_vk_from_assembly::< _, _, diff --git a/crates/proof-compression/src/step/snark_wrapper.rs b/crates/proof-compression/src/step/snark_wrapper.rs index ca3faec..cf465b5 100644 --- a/crates/proof-compression/src/step/snark_wrapper.rs +++ b/crates/proof-compression/src/step/snark_wrapper.rs @@ -115,6 +115,10 @@ pub(crate) trait SnarkWrapperStep: SnarkWrapperProofSystem { AsyncHandler::spawn(f) } + fn run_pre_initialization_tasks() { + Self::pre_init(); + } + fn prove_snark_wrapper_step( ctx_config: AsyncHandler, precomputation: AsyncHandler, diff --git a/crates/proof-compression/src/test.rs b/crates/proof-compression/src/test.rs index ca303ff..bc9f1cf 100644 --- a/crates/proof-compression/src/test.rs +++ b/crates/proof-compression/src/test.rs @@ -1,38 +1,19 @@ use super::*; -use circuit_definitions::circuit_definitions::recursion_layer::ZkSyncRecursionLayerProof; - #[test] fn test_proof_chain_with_fflonk() { - let scheduler_proof_file = - std::fs::File::open("./data/scheduler_recursive_proof.json").unwrap(); - let scheduler_proof: ZkSyncRecursionLayerProof = - serde_json::from_reader(&scheduler_proof_file).unwrap(); - let scheduler_proof = scheduler_proof.into_inner(); - let simple_blob_storage = FileSystemBlobStorage; - let proof = run_proof_chain_with_fflonk(scheduler_proof, &simple_blob_storage); - - let proof_file_path = format!("./data/final_fflonk_proof.json"); - let proof_file = std::fs::File::create(&proof_file_path).unwrap(); - serde_json::to_writer(proof_file, &proof).unwrap(); - println!("Final fflonk snark wrapper roof saved at {proof_file_path}"); + let mut simple_proof_storage = SimpleProofStorage::new(); + run_proof_chain_with_fflonk(&simple_blob_storage, &mut simple_proof_storage); + println!("Final fflonk snark wrapper proof has successfully generated"); } #[test] fn test_proof_chain_with_plonk() { - let scheduler_proof_file = - std::fs::File::open("./data/scheduler_recursive_proof.json").unwrap(); - let scheduler_proof: ZkSyncRecursionLayerProof = - serde_json::from_reader(&scheduler_proof_file).unwrap(); - let scheduler_proof = scheduler_proof.into_inner(); - let simple_blob_storage = FileSystemBlobStorage; - let proof = run_proof_chain_with_plonk(scheduler_proof, &simple_blob_storage); + let mut simple_proof_storage = SimpleProofStorage::new(); + run_proof_chain_with_plonk(&simple_blob_storage, &mut simple_proof_storage); - let proof_file_path = format!("./data/final_plonk_proof.json"); - let proof_file = std::fs::File::create(&proof_file_path).unwrap(); - serde_json::to_writer(proof_file, &proof).unwrap(); - println!("Final plonk snark wrapper roof saved at {proof_file_path}"); + println!("Final plonk snark wrapper proof has successfully generated"); } #[test]