diff --git a/halo2_proofs/Cargo.toml b/halo2_proofs/Cargo.toml index a1982bf538..fa752f61c6 100644 --- a/halo2_proofs/Cargo.toml +++ b/halo2_proofs/Cargo.toml @@ -1,12 +1,7 @@ [package] name = "halo2_proofs" version = "0.2.0" -authors = [ - "Sean Bowe ", - "Ying Tong Lai ", - "Daira Hopwood ", - "Jack Grigg ", -] +authors = ["Sean Bowe ", "Ying Tong Lai ", "Daira Hopwood ", "Jack Grigg "] edition = "2021" rust-version = "1.56.1" description = """ @@ -55,7 +50,7 @@ ff = "0.12" group = "0.12" halo2curves = { path = "../arithmetic/curves" } rand = "0.8" -rand_core = { version = "0.6", default-features = false} +rand_core = { version = "0.6", default-features = false } tracing = "0.1" blake2b_simd = "1" rustc-hash = "1.1.0" @@ -63,6 +58,10 @@ sha3 = "0.9.1" ark-std = { version = "0.3.0", features = ["print-trace"], optional = true } serde = { version = "1.0", default-features = false, features = ["derive"] } bincode = "1.3.3" +maybe-rayon = { version = "0.1.0", default-features = false } +itertools = "0.10" +tokio = { version = "1.33", features = ["full"] } +serde_json = { version = "1.0", default-features = false } # Developer tooling dependencies plotters = { version = "0.3.0", optional = true } @@ -80,8 +79,10 @@ rand_chacha = "0.3.1" getrandom = { version = "0.2", features = ["js"] } [features] -default = ["batch"] +default = ["batch", "multicore"] +multicore = ["maybe-rayon/threads"] dev-graph = ["plotters", "tabbycat"] +test-dev-graph = ["dev-graph", "plotters/bitmap_backend", "plotters/bitmap_encoder", "plotters/ttf"] gadget-traces = ["backtrace"] sanity-checks = [] batch = ["rand/getrandom"] diff --git a/halo2_proofs/examples/serialization.rs b/halo2_proofs/examples/serialization.rs index fb10d8246e..721301c515 100644 --- a/halo2_proofs/examples/serialization.rs +++ b/halo2_proofs/examples/serialization.rs @@ -126,91 +126,121 @@ impl Circuit for StandardPlonk { } } -fn main() -> std::io::Result<()> { +#[tokio::main(flavor = "multi_thread", worker_threads = 24)] +async fn main() -> std::io::Result<()> { let k = 22; - let circuit = StandardPlonk(Fr::random(OsRng)); - let params = ParamsKZG::::setup(k, OsRng); - let vk = keygen_vk(¶ms, &circuit).expect("vk should not fail"); - let pk = keygen_pk(¶ms, vk, &circuit).expect("pk should not fail"); - - for buf_size in [1024, 8 * 1024, 1024 * 1024, 1024 * 1024 * 1024] { - println!("buf_size: {buf_size}"); - // Using halo2_proofs serde implementation - let f = File::create("serialization-test.pk")?; - let mut writer = BufWriter::with_capacity(buf_size, f); - let start = std::time::Instant::now(); - pk.write(&mut writer, SerdeFormat::RawBytes)?; - writer.flush().unwrap(); - println!("SerdeFormat::RawBytes pk write time: {:?}", start.elapsed()); - - let f = File::open("serialization-test.pk")?; - let mut reader = BufReader::with_capacity(buf_size, f); - let start = std::time::Instant::now(); - let pk = - ProvingKey::::read::<_, StandardPlonk>(&mut reader, SerdeFormat::RawBytes) - .unwrap(); - println!("SerdeFormat::RawBytes pk read time: {:?}", start.elapsed()); - - let metadata = fs::metadata("serialization-test.pk")?; - let file_size = metadata.len(); - println!("The size of the file is {} bytes", file_size); - std::fs::remove_file("serialization-test.pk")?; - - // Using bincode - let f = File::create("serialization-test.pk")?; - let mut writer = BufWriter::with_capacity(buf_size, f); - let start = std::time::Instant::now(); - bincode::serialize_into(&mut writer, &pk).unwrap(); - writer.flush().unwrap(); - println!("bincode pk write time: {:?}", start.elapsed()); - - let f = File::open("serialization-test.pk").unwrap(); - let mut reader = BufReader::with_capacity(buf_size, f); - let start = std::time::Instant::now(); - let pk: ProvingKey = bincode::deserialize_from(&mut reader).unwrap(); - println!("bincode pk read time: {:?}", start.elapsed()); - - let metadata = fs::metadata("serialization-test.pk")?; - let file_size = metadata.len(); - println!("The size of the file is {} bytes", file_size); - std::fs::remove_file("serialization-test.pk").unwrap(); - - let instances: &[&[Fr]] = &[&[circuit.clone().0]]; - let mut transcript = Blake2bWrite::<_, _, Challenge255<_>>::init(vec![]); - create_proof::< - KZGCommitmentScheme, - ProverGWC<'_, Bn256>, - Challenge255, - _, - Blake2bWrite, G1Affine, Challenge255<_>>, - _, - >( - ¶ms, - &pk, - &[circuit.clone()], - &[instances], - OsRng, - &mut transcript, - ) - .expect("prover should not fail"); - let proof = transcript.finalize(); - - let strategy = SingleStrategy::new(¶ms); - let mut transcript = Blake2bRead::<_, _, Challenge255<_>>::init(&proof[..]); - assert!(verify_proof::< - KZGCommitmentScheme, - VerifierGWC<'_, Bn256>, - Challenge255, - Blake2bRead<&[u8], G1Affine, Challenge255>, - SingleStrategy<'_, Bn256>, - >( - ¶ms, - pk.get_vk(), - strategy, - &[instances], - &mut transcript - ) - .is_ok()); - } + + let buf_size = 1024 * 1024; + + // let pk_path = "/home/ubuntu/playground/serialization-test.pk"; + // let f = File::open(pk_path)?; + // let mut reader = BufReader::with_capacity(buf_size, f); + // let start = std::time::Instant::now(); + // let pk = ProvingKey::::read::<_, StandardPlonk>(&mut reader, SerdeFormat::RawBytes) + // .unwrap(); + // println!("SerdeFormat::RawBytes pk read time: {:?}", start.elapsed()); + + // let pk_folder = "/home/ubuntu/playground/serialization-test/"; + let pk_folder = "/mnt/ramdisk/serialization-test"; + // let start = std::time::Instant::now(); + // pk.multi_thread_write(pk_folder, SerdeFormat::RawBytes)?; + // println!( + // "SerdeFormat::RawBytes pk multi thread write time: {:?}", + // start.elapsed() + // ); + + let start = std::time::Instant::now(); + ProvingKey::::multi_thread_read::(pk_folder, SerdeFormat::RawBytes) + .await?; + println!( + "SerdeFormat::RawBytes pk multi thread read time: {:?}", + start.elapsed() + ); + Ok(()) + // let circuit = StandardPlonk(Fr::random(OsRng)); + // let params = ParamsKZG::::setup(k, OsRng); + // let vk = keygen_vk(¶ms, &circuit).expect("vk should not fail"); + // let pk = keygen_pk(¶ms, vk, &circuit).expect("pk should not fail"); + + // for buf_size in [1024, 8 * 1024, 1024 * 1024, 1024 * 1024 * 1024] { + // println!("buf_size: {buf_size}"); + // // Using halo2_proofs serde implementation + // let f = File::create("serialization-test.pk")?; + // let mut writer = BufWriter::with_capacity(buf_size, f); + // let start = std::time::Instant::now(); + // pk.write(&mut writer, SerdeFormat::RawBytes)?; + // writer.flush().unwrap(); + // println!("SerdeFormat::RawBytes pk write time: {:?}", start.elapsed()); + + // let f = File::open("serialization-test.pk")?; + // let mut reader = BufReader::with_capacity(buf_size, f); + // let start = std::time::Instant::now(); + // let pk = + // ProvingKey::::read::<_, StandardPlonk>(&mut reader, SerdeFormat::RawBytes) + // .unwrap(); + // println!("SerdeFormat::RawBytes pk read time: {:?}", start.elapsed()); + + // let metadata = fs::metadata("serialization-test.pk")?; + // let file_size = metadata.len(); + // println!("The size of the file is {} bytes", file_size); + // std::fs::remove_file("serialization-test.pk")?; + + // // Using bincode + // let f = File::create("serialization-test.pk")?; + // let mut writer = BufWriter::with_capacity(buf_size, f); + // let start = std::time::Instant::now(); + // bincode::serialize_into(&mut writer, &pk).unwrap(); + // writer.flush().unwrap(); + // println!("bincode pk write time: {:?}", start.elapsed()); + + // let f = File::open("serialization-test.pk").unwrap(); + // let mut reader = BufReader::with_capacity(buf_size, f); + // let start = std::time::Instant::now(); + // let pk: ProvingKey = bincode::deserialize_from(&mut reader).unwrap(); + // println!("bincode pk read time: {:?}", start.elapsed()); + + // let metadata = fs::metadata("serialization-test.pk")?; + // let file_size = metadata.len(); + // println!("The size of the file is {} bytes", file_size); + // std::fs::remove_file("serialization-test.pk").unwrap(); + + // let instances: &[&[Fr]] = &[&[circuit.clone().0]]; + // let mut transcript = Blake2bWrite::<_, _, Challenge255<_>>::init(vec![]); + // create_proof::< + // KZGCommitmentScheme, + // ProverGWC<'_, Bn256>, + // Challenge255, + // _, + // Blake2bWrite, G1Affine, Challenge255<_>>, + // _, + // >( + // ¶ms, + // &pk, + // &[circuit.clone()], + // &[instances], + // OsRng, + // &mut transcript, + // ) + // .expect("prover should not fail"); + // let proof = transcript.finalize(); + + // let strategy = SingleStrategy::new(¶ms); + // let mut transcript = Blake2bRead::<_, _, Challenge255<_>>::init(&proof[..]); + // assert!(verify_proof::< + // KZGCommitmentScheme, + // VerifierGWC<'_, Bn256>, + // Challenge255, + // Blake2bRead<&[u8], G1Affine, Challenge255>, + // SingleStrategy<'_, Bn256>, + // >( + // ¶ms, + // pk.get_vk(), + // strategy, + // &[instances], + // &mut transcript + // ) + // .is_ok()); + // } + // Ok(()) } diff --git a/halo2_proofs/src/helpers.rs b/halo2_proofs/src/helpers.rs index 8eba57359b..37484e0ad1 100644 --- a/halo2_proofs/src/helpers.rs +++ b/halo2_proofs/src/helpers.rs @@ -1,7 +1,15 @@ use crate::poly::Polynomial; use ff::PrimeField; use halo2curves::{pairing::Engine, serde::SerdeObject, CurveAffine}; -use std::io; +use itertools::Itertools; +use maybe_rayon::prelude::{ + IndexedParallelIterator, IntoParallelIterator, IntoParallelRefIterator, ParallelIterator, +}; +use std::{ + fs::File, + io::{self, BufReader, BufWriter}, + path::Path, +}; /// This enum specifies how various types are serialized and deserialized. #[derive(Clone, Copy, Debug)] @@ -125,6 +133,35 @@ pub(crate) fn read_polynomial_vec( .collect() } +/// Reads a vector of polynomials from buffer +pub(crate) async fn multi_thread_read_polynomial_vec< + F: SerdePrimeField, + B: Send + Sync + 'static, +>( + pk_prefix_path: impl AsRef, + format: SerdeFormat, + n: usize, +) -> Vec> { + const BUFFER_SIZE: usize = 1024 * 1024; + let join_handles = (0..n) + .map(|i| { + let mut poly_path = pk_prefix_path + .as_ref() + .clone() + .to_path_buf() + .into_os_string(); + poly_path.push(format!("_{i}")); + let mut reader = BufReader::with_capacity(BUFFER_SIZE, File::open(poly_path).unwrap()); + tokio::spawn(async move { Polynomial::::read(&mut reader, format) }) + }) + .collect_vec(); + let mut ret = Vec::with_capacity(join_handles.len()); + for join_handle in join_handles { + ret.push(join_handle.await.unwrap()); + } + ret +} + /// Writes a slice of polynomials to buffer pub(crate) fn write_polynomial_slice( slice: &[Polynomial], @@ -139,6 +176,36 @@ pub(crate) fn write_polynomial_slice( } } +/// Writes a slice of polynomials to buffer +pub(crate) fn multi_thread_write_polynomial_slice( + slice: &[Polynomial], + pk_prefix_path: impl AsRef, + format: SerdeFormat, +) { + const BUFFER_SIZE: usize = 1024 * 1024; + let poly_path = slice + .iter() + .enumerate() + .map(|(i, _)| { + let mut poly_path = pk_prefix_path + .as_ref() + .clone() + .to_path_buf() + .into_os_string(); + poly_path.push(format!("_{i}")); + poly_path + }) + .collect_vec(); + slice + .par_iter() + .zip_eq(poly_path.par_iter()) + .for_each(|(poly, poly_path)| { + let mut writer = + BufWriter::with_capacity(BUFFER_SIZE, File::create(poly_path).unwrap()); + poly.write(&mut writer, format); + }); +} + /// Gets the total number of bytes of a slice of polynomials, assuming all polynomials are the same length pub(crate) fn polynomial_slice_byte_length(slice: &[Polynomial]) -> usize { let field_len = F::default().to_repr().as_ref().len(); diff --git a/halo2_proofs/src/plonk.rs b/halo2_proofs/src/plonk.rs index 344f98f6a5..e681170998 100644 --- a/halo2_proofs/src/plonk.rs +++ b/halo2_proofs/src/plonk.rs @@ -8,10 +8,12 @@ use blake2b_simd::Params as Blake2bParams; use ff::PrimeField; use group::ff::Field; +use rayon::prelude::{IntoParallelRefIterator, IntoParallelRefMutIterator, ParallelIterator}; use serde::{Deserialize, Serialize}; use crate::arithmetic::{CurveAffine, FieldExt}; use crate::helpers::{ + multi_thread_read_polynomial_vec, multi_thread_write_polynomial_slice, polynomial_slice_byte_length, read_polynomial_vec, write_polynomial_slice, SerdeCurveAffine, SerdePrimeField, }; @@ -45,7 +47,7 @@ pub use verifier::*; use evaluation::Evaluator; use std::env::var; -use std::io; +use std::io::{self, BufReader}; use std::time::Instant; /// Temp @@ -149,6 +151,10 @@ pub fn log_info(msg: String) { } } +use std::fs::File; +use std::io::BufWriter; +use std::path::{Path, PathBuf}; + /// This is a verifying key which allows for the verification of proofs for a /// particular circuit. #[derive(Clone, Debug, Serialize, Deserialize)] @@ -387,6 +393,13 @@ pub struct ProvingKey { ev: Evaluator, } +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct ProvingKeyMetadata { + pub num_fixed_values: usize, + pub num_fixed_polys: usize, + pub num_fixed_cosets: usize, +} + impl ProvingKey { /// Get the underlying [`VerifyingKey`]. pub fn get_vk(&self) -> &VerifyingKey { @@ -432,6 +445,56 @@ where Ok(()) } + pub fn multi_thread_write( + &self, + path: impl AsRef, + format: SerdeFormat, + ) -> io::Result<()> { + const BUFFER_SIZE: usize = 1024 * 1024; + let create_path = |filename: &str| -> PathBuf { + let mut p = path.as_ref().to_path_buf(); + p.push(filename); + p + }; + let create_writer = |filename: &str| -> BufWriter { + let p = create_path(filename); + BufWriter::with_capacity(BUFFER_SIZE, File::create(p).unwrap()) + }; + let metadata = ProvingKeyMetadata { + num_fixed_values: self.fixed_values.len(), + num_fixed_polys: self.fixed_polys.len(), + num_fixed_cosets: self.fixed_cosets.len(), + }; + serde_json::to_writer(create_writer("metadata"), &metadata)?; + + let mut vk_writer = create_writer("vk"); + self.vk.write(&mut vk_writer, format)?; + + let l0_writer = create_writer("l0"); + let l_last_writer = create_writer("l_last"); + let l_active_writer = create_writer("l_active"); + + let mut v = [ + (l0_writer, &self.l0), + (l_last_writer, &self.l_last), + (l_active_writer, &self.l_last), + ]; + v.par_iter_mut().for_each(|(writer, p)| { + p.write(writer, format); + }); + + let values_path_prefix = create_path("values"); + multi_thread_write_polynomial_slice(&self.fixed_values, values_path_prefix, format); + let polys_path_prefix = create_path("polys"); + multi_thread_write_polynomial_slice(&self.fixed_polys, polys_path_prefix, format); + let cosets_path_prefix = create_path("cosets"); + multi_thread_write_polynomial_slice(&self.fixed_cosets, cosets_path_prefix, format); + + let perm_path = create_path("perm"); + self.permutation.multi_thread_write(perm_path, format); + Ok(()) + } + /// Reads a proving key from a buffer. /// Does so by reading verification key first, and then deserializing the rest of the file into the remaining proving key data. /// @@ -469,6 +532,82 @@ where }) } + pub async fn multi_thread_read>( + path: impl AsRef, + format: SerdeFormat, + ) -> io::Result { + const BUFFER_SIZE: usize = 1024 * 1024; + let create_path = |filename: &str| -> PathBuf { + let mut p = path.as_ref().to_path_buf(); + p.push(filename); + p + }; + let create_reader = |filename: &str| -> BufReader { + let p = create_path(filename); + BufReader::with_capacity(BUFFER_SIZE, File::open(p).unwrap()) + }; + let metadata: ProvingKeyMetadata = serde_json::from_reader(create_reader("metadata"))?; + + let vk_path = create_path("vk"); + let mut reader = BufReader::with_capacity(BUFFER_SIZE, File::open(vk_path).unwrap()); + let vk_promise = tokio::spawn(async move { + VerifyingKey::::read::<_, ConcreteCircuit>(&mut reader, format).unwrap() + }); + + let l0_path = create_path("l0"); + let mut reader = BufReader::with_capacity(BUFFER_SIZE, File::open(l0_path).unwrap()); + let l0_promise = tokio::spawn(async move { Polynomial::read(&mut reader, format) }); + + let l_last_path = create_path("l_last"); + let mut reader = BufReader::with_capacity(BUFFER_SIZE, File::open(l_last_path).unwrap()); + let l_last_promise = tokio::spawn(async move { Polynomial::read(&mut reader, format) }); + + let l_active_row_path = create_path("l_active"); + let mut reader = + BufReader::with_capacity(BUFFER_SIZE, File::open(l_active_row_path).unwrap()); + let l_active_row_promise = + tokio::spawn(async move { Polynomial::read(&mut reader, format) }); + + let values_path = create_path("values"); + let values_n = metadata.num_fixed_values; + let values_promise = tokio::spawn(multi_thread_read_polynomial_vec( + values_path, + format, + values_n, + )); + + let polys_path = create_path("polys"); + let polys_n = metadata.num_fixed_polys; + let polys_promise = tokio::spawn(multi_thread_read_polynomial_vec( + polys_path, format, polys_n, + )); + + let cosets_path = create_path("cosets"); + let cosets_n = metadata.num_fixed_cosets; + let cosets_promise = tokio::spawn(multi_thread_read_polynomial_vec( + cosets_path, + format, + cosets_n, + )); + + let perm_path = create_path("perm"); + let permutation_promise = + tokio::spawn(permutation::ProvingKey::async_read(perm_path, format)); + let vk = vk_promise.await?; + let ev = Evaluator::new(vk.cs()); + Ok(Self { + vk, + l0: l0_promise.await?, + l_last: l_last_promise.await?, + l_active_row: l_active_row_promise.await?, + fixed_values: values_promise.await?, + fixed_polys: polys_promise.await?, + fixed_cosets: cosets_promise.await?, + permutation: permutation_promise.await?, + ev, + }) + } + /// Writes a proving key to a vector of bytes using [`Self::write`]. pub fn to_bytes(&self, format: SerdeFormat) -> Vec { let mut bytes = Vec::::with_capacity(self.bytes_length()); diff --git a/halo2_proofs/src/plonk/permutation.rs b/halo2_proofs/src/plonk/permutation.rs index 255c12a956..fb549d0ca5 100644 --- a/halo2_proofs/src/plonk/permutation.rs +++ b/halo2_proofs/src/plonk/permutation.rs @@ -4,6 +4,7 @@ use super::circuit::{Any, Column}; use crate::{ arithmetic::CurveAffine, helpers::{ + multi_thread_read_polynomial_vec, multi_thread_write_polynomial_slice, polynomial_slice_byte_length, read_polynomial_vec, write_polynomial_slice, SerdeCurveAffine, SerdePrimeField, }, @@ -19,7 +20,7 @@ pub(crate) mod verifier; pub use keygen::Assembly; use serde::{Deserialize, Serialize}; -use std::io; +use std::{fs::File, io, path::Path}; /// A permutation argument. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -135,6 +136,13 @@ pub(crate) struct ProvingKey { pub(super) cosets: Vec>, } +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct ProvingKeyMetadata { + pub num_perms: usize, + pub num_polys: usize, + pub num_cosets: usize, +} + impl ProvingKey where C::Scalar: SerdePrimeField, @@ -150,6 +158,41 @@ where cosets, } } + /// Async version of read + pub(super) async fn async_read(prefix: impl AsRef, format: SerdeFormat) -> Self { + let os_str = prefix.as_ref().clone().to_path_buf().into_os_string(); + let mut metadata_os_str = os_str.clone(); + metadata_os_str.push("_metadata"); + let metadata: ProvingKeyMetadata = + serde_json::from_reader(File::open(metadata_os_str).unwrap()).unwrap(); + + let mut perms_os_str = os_str.clone(); + perms_os_str.push("_perm"); + let permutations_promise = tokio::spawn(multi_thread_read_polynomial_vec( + perms_os_str, + format, + metadata.num_perms, + )); + let mut polys_os_str = os_str.clone(); + polys_os_str.push("_polys"); + let polys_promise = tokio::spawn(multi_thread_read_polynomial_vec( + polys_os_str, + format, + metadata.num_polys, + )); + let mut cosets_os_str = os_str.clone(); + cosets_os_str.push("_cosets"); + let cosets_promise = tokio::spawn(multi_thread_read_polynomial_vec( + cosets_os_str, + format, + metadata.num_cosets, + )); + ProvingKey { + permutations: permutations_promise.await.unwrap(), + polys: polys_promise.await.unwrap(), + cosets: cosets_promise.await.unwrap(), + } + } /// Writes proving key for a single permutation argument to buffer using `Polynomial::write`. pub(super) fn write(&self, writer: &mut W, format: SerdeFormat) { @@ -157,6 +200,29 @@ where write_polynomial_slice(&self.polys, writer, format); write_polynomial_slice(&self.cosets, writer, format); } + /// Multi thread write + pub(super) fn multi_thread_write(&self, prefix: impl AsRef, format: SerdeFormat) { + let os_str = prefix.as_ref().clone().to_path_buf().into_os_string(); + let metadata = ProvingKeyMetadata { + num_perms: self.permutations.len(), + num_polys: self.polys.len(), + num_cosets: self.cosets.len(), + }; + let mut metadata_os_str = os_str.clone(); + metadata_os_str.push("_metadata"); + serde_json::to_writer(File::create(metadata_os_str).unwrap(), &metadata).unwrap(); + + let permutations = &self.permutations; + let mut perm_os_str = os_str.clone(); + perm_os_str.push("_perm"); + multi_thread_write_polynomial_slice(permutations, perm_os_str, format); + let mut polys_os_str = os_str.clone(); + polys_os_str.push("_polys"); + multi_thread_write_polynomial_slice(&self.polys, polys_os_str, format); + let mut cosets_os_str = os_str.clone(); + cosets_os_str.push("_cosets"); + multi_thread_write_polynomial_slice(&self.cosets, cosets_os_str, format); + } } impl ProvingKey { diff --git a/halo2_proofs/tests/plonk_api.rs b/halo2_proofs/tests/plonk_api.rs index 5f0919f086..679163c7e7 100644 --- a/halo2_proofs/tests/plonk_api.rs +++ b/halo2_proofs/tests/plonk_api.rs @@ -19,8 +19,11 @@ use halo2_proofs::transcript::{ Blake2bRead, Blake2bWrite, Challenge255, EncodedChallenge, TranscriptReadBuffer, TranscriptWriterBuffer, }; +use halo2_proofs::SerdeFormat; use rand_core::{OsRng, RngCore}; +use std::fs::File; use std::hash::Hash; +use std::io; use std::marker::PhantomData; #[test]