diff --git a/Cargo.toml b/Cargo.toml index 235fc2f..a0bfdfc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -37,6 +37,7 @@ ark-ff = { version = "0.3.0", default-features = false, features = ["parallel", ark-groth16 = { git = "https://github.com/arkworks-rs/groth16", rev = "765817f", features = ["parallel"] } ark-relations = { version = "0.3.0", default-features = false } ark-std = { version = "0.3.0", default-features = false, features = ["parallel"] } +ark-zkey = { path = "crates/ark-zkey" } color-eyre = "0.6" criterion = { version = "0.3", optional = true, features = ["async_tokio"] } hex = "0.4.0" @@ -73,6 +74,7 @@ tiny-keccak = "2.0.2" tracing-test = "0.2" [build-dependencies] +ark-zkey = { path = "crates/ark-zkey" } color-eyre = "0.6" enumset = "1.0.8" reqwest = { version = "0.11", features = ["blocking"] } diff --git a/build.rs b/build.rs index 4bf31d8..eda45f7 100644 --- a/build.rs +++ b/build.rs @@ -6,6 +6,8 @@ use std::{ extern crate reqwest; +use ark_zkey; + const SEMAPHORE_FILES_PATH: &str = "semaphore_files"; const SEMAPHORE_DOWNLOAD_URL: &str = "https://www.trusted-setup-pse.org/semaphore"; @@ -44,6 +46,23 @@ fn semaphore_file_path(file_name: &str, depth: usize) -> PathBuf { .join(file_name) } +fn create_arkzkey(path: PathBuf) -> Result { + let ark_zkey_path = path.join("-arkzkey"); + + let (original_proving_key, original_constraint_matrices) = + ark_zkey::read_proving_key_and_matrices_from_zkey( + path.to_str().expect("Failed to convert path."), + )?; + + ark_zkey::convert_zkey( + original_proving_key, + original_constraint_matrices, + &ark_zkey_path.to_str().unwrap(), + )?; + + Ok(ark_zkey_path) +} + fn build_circuit(depth: usize) -> Result<()> { let base_path = Path::new(SEMAPHORE_FILES_PATH); if !base_path.exists() { @@ -63,15 +82,20 @@ fn build_circuit(depth: usize) -> Result<()> { let download_url = format!("{SEMAPHORE_DOWNLOAD_URL}/{depth_str}/{filename}.{extension}"); let path = Path::new(&depth_subfolder).join(format!("{filename}.{extension}")); download_and_store_binary(&download_url, &path)?; + create_arkzkey(path)?; } // Compute absolute paths let zkey_file = absolute(semaphore_file_path("semaphore.zkey", depth))?; - let graph_file = absolute(Path::new("graphs") - .join(depth.to_string()) - .join("graph.bin"))?; + let arkzkey_file = absolute(semaphore_file_path("semaphore.zkey-arkzkey", depth))?; + let graph_file = absolute( + Path::new("graphs") + .join(depth.to_string()) + .join("graph.bin"), + )?; assert!(zkey_file.exists()); + assert!(arkzkey_file.exists()); assert!(graph_file.exists()); // Export generated paths @@ -80,6 +104,11 @@ fn build_circuit(depth: usize) -> Result<()> { depth, zkey_file.display() ); + println!( + "cargo:rustc-env=BUILD_RS_ARKZKEY_FILE_{}={}", + depth, + arkzkey_file.display() + ); println!( "cargo:rustc-env=BUILD_RS_GRAPH_FILE_{}={}", depth, diff --git a/crates/ark-zkey/.gitignore b/crates/ark-zkey/.gitignore new file mode 100644 index 0000000..6985cf1 --- /dev/null +++ b/crates/ark-zkey/.gitignore @@ -0,0 +1,14 @@ +# Generated by Cargo +# will have compiled files and executables +debug/ +target/ + +# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries +# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html +Cargo.lock + +# These are backup files generated by rustfmt +**/*.rs.bk + +# MSVC Windows builds of rustc generate these, which store debugging information +*.pdb diff --git a/crates/ark-zkey/Cargo.toml b/crates/ark-zkey/Cargo.toml new file mode 100644 index 0000000..d327801 --- /dev/null +++ b/crates/ark-zkey/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "ark-zkey" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +# XXX: Shouldn't be necessary, but this way we stay consistent with wasmer version and fix +# error[E0432]: unresolved import `wasmer` error +# (likely due to other packages) +[patch.crates-io] +# NOTE: Forked wasmer to work around memory limits +# See https://github.com/wasmerio/wasmer/commit/09c7070 +wasmer = { git = "https://github.com/oskarth/wasmer.git", rev = "09c7070" } + +[dependencies] +color-eyre = "0.6" +memmap2 = "0.9" +flame = "0.2" +flamer = "0.5" + +ark-serialize = { version = "=0.4.1", features = ["derive"] } +ark-bn254 = { version = "=0.4.0" } +ark-groth16 = { version = "=0.4.0" } +ark-circom = { git = "https://github.com/arkworks-rs/circom-compat.git" } +ark-relations = { version = "=0.4.0" } +ark-ff = { version = "=0.4.1" } +ark-ec = { version = "=0.4.1" } diff --git a/crates/ark-zkey/README.md b/crates/ark-zkey/README.md new file mode 100644 index 0000000..067d853 --- /dev/null +++ b/crates/ark-zkey/README.md @@ -0,0 +1,79 @@ +# ark-zkey + +Library to read `zkey` faster by serializing to `arkworks` friendly format. + +See https://github.com/oskarth/mopro/issues/25 for context. + +## To generate arkzkey + +Hacky, but the way we generate `arkzkey` now is by running the corresponding test. + +Note that we also neeed to change the const `ZKEY_BYTES` above. + +E.g.: + +``` +cargo test multiplier2 --release -- --nocapture +cargo test keccak256 --release -- --nocapture +cargo test rsa --release -- --nocapture +``` + +Will take corresponding `zkey` and put `arkzkey`` in same folder. + +## Multiplier + +NOTE: Need to change const ZKEY here + +`cargo test multiplier2 --release -- --nocapture` + +``` +running 1 test +[build] Processing zkey data... +[build] Time to process zkey data: 3.513041ms +[build] Serializing proving key and constraint matrices +[build] Time to serialize proving key and constraint matrices: 42ns +[build] Writing arkzkey to: ../mopro-core/examples/circom/multiplier2/target/multiplier2_final.arkzkey +[build] Time to write arkzkey: 1.884875ms +Reading arkzkey from: ../mopro-core/examples/circom/multiplier2/target/multiplier2_final.arkzkey + +Time to open arkzkey file: 18.084µs +Time to mmap arkzkey: 8.542µs +Time to deserialize proving key: 305.75µs +Time to deserialize matrices: 5µs +Time to read arkzkey: 348.083µs +test tests::test_multiplier2_serialization_deserialization ... ok +``` + +Naive test: `cargo test naive --release -- --nocapture` (with right zkey constant). + +**Result: `350µs` vs naive `3.3ms`** + +## Keccak + +NOTE: Need to change const ZKEY here + +`cargo test keccak256 --release -- --nocapture` + +``` +[build] Processing zkey data... +test tests::test_keccak256_serialization_deserialization has been running for over 60 seconds +[build]Time to process zkey data: 158.753181958s +[build] Serializing proving key and constraint matrices +[build] Time to serialize proving key and constraint matrices: 42ns +[build] Writing arkzkey to: ../mopro-core/examples/circom/keccak256/target/keccak256_256_test_final.arkzkey +[build] Time to write arkzkey: 16.204274125s +Reading arkzkey from: ../mopro-core/examples/circom/keccak256/target/keccak256_256_test_final.arkzkey +Time to open arkzkey file: 51.75µs +Time to mmap arkzkey: 17.25µs +Time to deserialize proving key: 18.323550083s +Time to deserialize matrices: 46.935792ms +Time to read arkzkey: 18.3730695s +test tests::test_keccak256_serialization_deserialization ... ok +``` + +Vs naive: + +`[build] Time to process zkey data: 158.753181958s` + + +**Result: 18s vs 158s** diff --git a/crates/ark-zkey/src/lib.rs b/crates/ark-zkey/src/lib.rs new file mode 100644 index 0000000..65d23af --- /dev/null +++ b/crates/ark-zkey/src/lib.rs @@ -0,0 +1,329 @@ +use ark_bn254::{Bn254, Fr}; +use ark_circom::read_zkey; +//use ark_ec::pairing::Pairing; +use ark_ff::Field; +use ark_groth16::ProvingKey; +//use ark_groth16::VerifyingKey; +use ark_relations::r1cs::ConstraintMatrices; +use ark_serialize::{CanonicalDeserialize, CanonicalSerialize}; +use color_eyre::eyre::{Result, WrapErr}; +use memmap2::Mmap; +use std::fs::File; +use std::io::Cursor; +//use std::io::{Read,self}; +use std::io::BufReader; +use std::path::PathBuf; +use std::time::Instant; + +#[derive(CanonicalSerialize, CanonicalDeserialize, Clone, Debug, PartialEq)] +pub struct SerializableProvingKey(pub ProvingKey); + +#[derive(CanonicalSerialize, CanonicalDeserialize, Clone, Debug, PartialEq)] +pub struct SerializableMatrix { + pub data: Vec>, +} + +#[derive(CanonicalSerialize, CanonicalDeserialize, Clone, Debug, PartialEq)] +pub struct SerializableConstraintMatrices { + pub num_instance_variables: usize, + pub num_witness_variables: usize, + pub num_constraints: usize, + pub a_num_non_zero: usize, + pub b_num_non_zero: usize, + pub c_num_non_zero: usize, + pub a: SerializableMatrix, + pub b: SerializableMatrix, + pub c: SerializableMatrix, +} + +impl From>> for SerializableMatrix { + fn from(matrix: Vec>) -> Self { + SerializableMatrix { data: matrix } + } +} + +impl From> for Vec> { + fn from(serializable_matrix: SerializableMatrix) -> Self { + serializable_matrix.data + } +} + +pub fn serialize_proving_key(pk: &SerializableProvingKey) -> Vec { + let mut serialized_data = Vec::new(); + pk.serialize_compressed(&mut serialized_data) + .expect("Serialization failed"); + serialized_data +} + +pub fn deserialize_proving_key(data: Vec) -> SerializableProvingKey { + SerializableProvingKey::deserialize_compressed_unchecked(&mut &data[..]) + .expect("Deserialization failed") +} + +// takes bytes from .zkey files +const ZKEY_DATA: &[u8] = &[0u8]; + +pub fn read_proving_key_and_matrices( +) -> Result<(SerializableProvingKey, SerializableConstraintMatrices)> { + println!("[build] Processing zkey data..."); + let now = Instant::now(); + + let mut cursor = Cursor::new(ZKEY_DATA); + + let (proving_key, matrices) = read_zkey(&mut cursor).wrap_err("Failed to process zkey data")?; + println!("[build] Time to process zkey data: {:?}", now.elapsed()); + + println!("[build] Serializing proving key and constraint matrices"); + let now = Instant::now(); + let serializable_proving_key = SerializableProvingKey(proving_key); + let serializable_constrain_matrices = SerializableConstraintMatrices { + num_instance_variables: matrices.num_instance_variables, + num_witness_variables: matrices.num_witness_variables, + num_constraints: matrices.num_constraints, + a_num_non_zero: matrices.a_num_non_zero, + b_num_non_zero: matrices.b_num_non_zero, + c_num_non_zero: matrices.c_num_non_zero, + a: SerializableMatrix { data: matrices.a }, + b: SerializableMatrix { data: matrices.b }, + c: SerializableMatrix { data: matrices.c }, + }; + println!( + "[build] Time to serialize proving key and constraint matrices: {:?}", + now.elapsed() + ); + + Ok((serializable_proving_key, serializable_constrain_matrices)) +} + +pub fn read_arkzkey( + arkzkey_path: &str, +) -> Result<(SerializableProvingKey, SerializableConstraintMatrices)> { + let now = std::time::Instant::now(); + let arkzkey_file_path = PathBuf::from(arkzkey_path); + let arkzkey_file = File::open(arkzkey_file_path).wrap_err("Failed to open arkzkey file")?; + println!("Time to open arkzkey file: {:?}", now.elapsed()); + + //let mut buf_reader = BufReader::new(arkzkey_file); + + // Using mmap + let now = std::time::Instant::now(); + let mmap = unsafe { Mmap::map(&arkzkey_file)? }; + let mut cursor = std::io::Cursor::new(mmap); + println!("Time to mmap arkzkey: {:?}", now.elapsed()); + + // Was &mut buf_reader + let now = std::time::Instant::now(); + let proving_key = SerializableProvingKey::deserialize_compressed_unchecked(&mut cursor) + .wrap_err("Failed to deserialize proving key")?; + println!("Time to deserialize proving key: {:?}", now.elapsed()); + + let now = std::time::Instant::now(); + let constraint_matrices = + SerializableConstraintMatrices::deserialize_compressed_unchecked(&mut cursor) + .wrap_err("Failed to deserialize constraint matrices")?; + println!("Time to deserialize matrices: {:?}", now.elapsed()); + + Ok((proving_key, constraint_matrices)) +} + +// TODO: Return ProvingKey, ConstraintMatrices? +pub fn read_arkzkey_from_bytes( + arkzkey_bytes: &[u8], +) -> Result<(ProvingKey, ConstraintMatrices)> { + let mut cursor = std::io::Cursor::new(arkzkey_bytes); + + let now = std::time::Instant::now(); + let serialized_proving_key = + SerializableProvingKey::deserialize_compressed_unchecked(&mut cursor) + .wrap_err("Failed to deserialize proving key")?; + println!("Time to deserialize proving key: {:?}", now.elapsed()); + + let now = std::time::Instant::now(); + let serialized_constraint_matrices = + SerializableConstraintMatrices::deserialize_compressed_unchecked(&mut cursor) + .wrap_err("Failed to deserialize constraint matrices")?; + println!("Time to deserialize matrices: {:?}", now.elapsed()); + + // Get on right form for API + let proving_key: ProvingKey = serialized_proving_key.0; + let constraint_matrices: ConstraintMatrices = ConstraintMatrices { + num_instance_variables: serialized_constraint_matrices.num_instance_variables, + num_witness_variables: serialized_constraint_matrices.num_witness_variables, + num_constraints: serialized_constraint_matrices.num_constraints, + a_num_non_zero: serialized_constraint_matrices.a_num_non_zero, + b_num_non_zero: serialized_constraint_matrices.b_num_non_zero, + c_num_non_zero: serialized_constraint_matrices.c_num_non_zero, + a: serialized_constraint_matrices.a.data, + b: serialized_constraint_matrices.b.data, + c: serialized_constraint_matrices.c.data, + }; + + Ok((proving_key, constraint_matrices)) +} + +pub fn read_proving_key_and_matrices_from_zkey( + zkey_path: &str, +) -> Result<(SerializableProvingKey, SerializableConstraintMatrices)> { + println!("Reading zkey from: {}", zkey_path); + let now = Instant::now(); + let zkey_file_path = PathBuf::from(zkey_path); + let zkey_file = File::open(zkey_file_path).wrap_err("Failed to open zkey file")?; + + let mut buf_reader = BufReader::new(zkey_file); + + let (proving_key, matrices) = + read_zkey(&mut buf_reader).wrap_err("Failed to read zkey file")?; + println!("Time to read zkey: {:?}", now.elapsed()); + + println!("Serializing proving key and constraint matrices"); + let now = Instant::now(); + let serializable_proving_key = SerializableProvingKey(proving_key); + let serializable_constrain_matrices = SerializableConstraintMatrices { + num_instance_variables: matrices.num_instance_variables, + num_witness_variables: matrices.num_witness_variables, + num_constraints: matrices.num_constraints, + a_num_non_zero: matrices.a_num_non_zero, + b_num_non_zero: matrices.b_num_non_zero, + c_num_non_zero: matrices.c_num_non_zero, + a: SerializableMatrix { data: matrices.a }, + b: SerializableMatrix { data: matrices.b }, + c: SerializableMatrix { data: matrices.c }, + }; + println!( + "Time to serialize proving key and constraint matrices: {:?}", + now.elapsed() + ); + + Ok((serializable_proving_key, serializable_constrain_matrices)) +} + +pub fn convert_zkey( + proving_key: SerializableProvingKey, + constraint_matrices: SerializableConstraintMatrices, + arkzkey_path: &str, +) -> Result<()> { + let arkzkey_file_path = PathBuf::from(arkzkey_path); + + let serialized_path = PathBuf::from(arkzkey_file_path); + + let mut file = + File::create(&serialized_path).wrap_err("Failed to create serialized proving key file")?; + + proving_key + .serialize_compressed(&mut file) + .wrap_err("Failed to serialize proving key")?; + + constraint_matrices + .serialize_compressed(&mut file) + .wrap_err("Failed to serialize constraint matrices")?; + + Ok(()) +} + +// fn read_zkey_with_mmap(zkey_path: &str) -> Result<(ProvingKey, ConstraintMatrices)> { +// let file = File::open(zkey_path)?; + +// let mmap = unsafe { Mmap::map(&file)? }; + +// let cursor = Cursor::new(&mmap); +// let (proving_key, matrices) = read_zkey(&mut cursor.clone())?; + +// Ok((proving_key, matrices)) +// } + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Instant; + + fn circuit_naive_read() -> Result<()> { + println!("Reading zkey..."); + let now = Instant::now(); + + let mut cursor = Cursor::new(ZKEY_DATA); + + let (_proving_key, _matrices) = + read_zkey(&mut cursor).wrap_err("Failed to process zkey data")?; + println!("Time to read zkey data: {:?}", now.elapsed()); + + Ok(()) + } + + fn test_circuit_serialization_deserialization(dir: &str, circuit: &str) -> Result<()> { + let _zkey_path = format!("{}/target/{}_final.zkey", dir, circuit); + let arkzkey_path = format!("{}/target/{}_final.arkzkey", dir, circuit); + + let (original_proving_key, original_constraint_matrices) = read_proving_key_and_matrices()?; + + println!("[build] Writing arkzkey to: {}", arkzkey_path); + let now = Instant::now(); + convert_zkey( + original_proving_key.clone(), + original_constraint_matrices.clone(), + &arkzkey_path, + )?; + println!("[build] Time to write arkzkey: {:?}", now.elapsed()); + + println!("Reading arkzkey from: {}", arkzkey_path); + let now = Instant::now(); + let (deserialized_proving_key, deserialized_constraint_matrices) = + read_arkzkey(&arkzkey_path)?; + println!("Time to read arkzkey: {:?}", now.elapsed()); + + assert_eq!( + original_proving_key, deserialized_proving_key, + "Original and deserialized proving keys do not match" + ); + + assert_eq!( + original_constraint_matrices, deserialized_constraint_matrices, + "Original and deserialized constraint matrices do not match" + ); + + //flame::dump_html(&mut std::fs::File::create("flame-graph.html").unwrap()).unwrap(); + + Ok(()) + } + + #[test] + fn test_multiplier2_serialization_deserialization() -> Result<()> { + test_circuit_serialization_deserialization( + "../mopro-core/examples/circom/multiplier2", + "multiplier2", + ) + } + + #[test] + fn test_keccak256_serialization_deserialization() -> Result<()> { + test_circuit_serialization_deserialization( + "../mopro-core/examples/circom/keccak256", + "keccak256_256_test", + ) + } + + #[test] + fn test_rsa_serialization_deserialization() -> Result<()> { + test_circuit_serialization_deserialization("../mopro-core/examples/circom/rsa", "main") + } + + // XXX: We do include_bytes for zkey data, so need to manually change this + #[test] + fn test_circuit_naive_read() -> Result<()> { + circuit_naive_read() + } + + // #[test] + // fn test_read_arkzkey_from_bytes() -> Result<()> { + // const ARKZKEY_BYTES: &[u8] = include_bytes!( + // "../../mopro-core/examples/circom/keccak256/target/keccak256_256_test_final.arkzkey" + // ); + + // println!("Reading arkzkey from bytes (keccak)"); + // let now = Instant::now(); + // let (_deserialized_proving_key, _deserialized_constraint_matrices) = + // read_arkzkey_from_bytes(ARKZKEY_BYTES)?; + // println!("Time to read arkzkey: {:?}", now.elapsed()); + + // Ok(()) + // } +} diff --git a/src/circuit.rs b/src/circuit.rs index bac96b1..5bdfff2 100644 --- a/src/circuit.rs +++ b/src/circuit.rs @@ -6,11 +6,12 @@ use core::include_bytes; use once_cell::sync::Lazy; use std::io::Cursor; +use ark_zkey; use semaphore_depth_config::{get_depth_index, get_supported_depth_count}; use semaphore_depth_macros::array_for_depths; const ZKEY_BYTES: [&[u8]; get_supported_depth_count()] = - array_for_depths!(|depth| include_bytes!(env!(concat!("BUILD_RS_ZKEY_FILE_", depth)))); + array_for_depths!(|depth| include_bytes!(env!(concat!("BUILD_RS_ARKZKEY_FILE_", depth)))); const GRAPH_BYTES: [&[u8]; get_supported_depth_count()] = array_for_depths!(|depth| include_bytes!(env!(concat!("BUILD_RS_GRAPH_FILE_", depth)))); @@ -18,7 +19,7 @@ const GRAPH_BYTES: [&[u8]; get_supported_depth_count()] = static ZKEY: [Lazy<(ProvingKey, ConstraintMatrices)>; get_supported_depth_count()] = array_for_depths!(|depth| Lazy::new(|| { let mut reader = Cursor::new(ZKEY_BYTES[get_depth_index(depth).unwrap()]); - read_zkey(&mut reader).expect("zkey should be valid") + ark_zkey::read_arkzkey(&mut reader).expect("zkey should be valid") })); #[must_use] @@ -31,4 +32,4 @@ pub fn zkey(depth: usize) -> &'static (ProvingKey, ConstraintMatrices pub fn graph(depth: usize) -> &'static [u8] { let index = get_depth_index(depth).unwrap_or_else(|| panic!("depth {depth} is not supported")); &GRAPH_BYTES[index] -} \ No newline at end of file +}