Merge pull request #3008 from ProvableHQ/update/kary-bench

Improves kary Merkle tree benchmarking

Author: raychu86

synthesizer/benches/kary_merkle_tree.rs

@@ -19,11 +19,11 @@ extern crate criterion;
 use snarkvm_algorithms::snark::varuna::VarunaVersion;
 use snarkvm_circuit::{AleoV0, Eject, Environment, Inject, Mode, collections::kary_merkle_tree::*};
 use snarkvm_console::{
-    algorithms::Sha3_256,
+    algorithms::Poseidon8,
     collections::kary_merkle_tree::KaryMerkleTree,
     network::{
         MainnetV0,
-        prelude::{TestRng, ToBits, Uniform},
+        prelude::{Rng, TestRng, Uniform},
     },
     types::Field,
 };
@@ -34,78 +34,119 @@ use criterion::Criterion;
 type CurrentNetwork = MainnetV0;
 type CurrentAleo = AleoV0;
 
-type NativePathHasher = Sha3_256;
-type NativeLeafHasher = Sha3_256;
-type CircuitPathHasher = snarkvm_circuit::Sha3_256<AleoV0>;
-type CircuitLeafHasher = snarkvm_circuit::Sha3_256<AleoV0>;
+type NativePathHasher = Poseidon8<CurrentNetwork>;
+type NativeLeafHasher = Poseidon8<CurrentNetwork>;
+type CircuitPathHasher = snarkvm_circuit::Poseidon8<AleoV0>;
+type CircuitLeafHasher = snarkvm_circuit::Poseidon8<AleoV0>;
 
-const DEPTH: u8 = 8;
+const DEPTH: u8 = 7;
 const ARITY: u8 = 8;
 
 /// Generates the specified number of random Merkle tree leaves.
 macro_rules! generate_leaves {
-    ($num_leaves:expr, $rng:expr) => {{ (0..$num_leaves).map(|_| Field::<MainnetV0>::rand($rng).to_bits_le()).collect::<Vec<_>>() }};
+    ("bits", $num_leaves:expr, $rng:expr) => {{
+        use snarkvm_console::network::prelude::ToBits;
+        // Generate leaf bits.
+        (0..$num_leaves).map(|_| Field::<CurrentNetwork>::rand($rng).to_bits_le()).collect::<Vec<_>>()
+    }};
+    ("fields", $num_leaves:expr, $rng:expr) => {{
+        use rand::SeedableRng;
+        use rayon::prelude::*;
+        // Generate leaf fields in parallel.
+        (0..$num_leaves)
+            .map(|_| u64::rand($rng))
+            .collect::<Vec<u64>>()
+            .into_par_iter()
+            .map(|seed| vec![Field::<CurrentNetwork>::rand(&mut rand::rngs::StdRng::seed_from_u64(seed))])
+            .collect::<Vec<_>>()
+    }};
 }
 
 fn batch_prove(c: &mut Criterion) {
     let mut rng = TestRng::default();
 
+    // Start the timer.
+    let timer = std::time::Instant::now();
+
     // Initialize the hashers.
-    let native_path_hasher = NativePathHasher::default();
-    let native_leaf_hasher = NativeLeafHasher::default();
-    let circuit_path_hasher = CircuitPathHasher::new();
-    let circuit_leaf_hasher = CircuitLeafHasher::new();
+    let native_path_hasher = NativePathHasher::setup("test").unwrap();
+    let native_leaf_hasher = NativeLeafHasher::setup("test").unwrap();
+    let circuit_path_hasher = CircuitPathHasher::new(Mode::Private, native_path_hasher.clone());
+    let circuit_leaf_hasher = CircuitLeafHasher::new(Mode::Private, native_leaf_hasher.clone());
 
     // Determine the maximum number of leaves.
-    let max_num_leaves = (ARITY as u32).pow(DEPTH as u32);
+    let max_num_leaves = (ARITY as u64).pow(DEPTH as u32);
     // Initialize the leaves.
-    let leaves = generate_leaves!(max_num_leaves, &mut rng);
-
+    let leaves = generate_leaves!("fields", max_num_leaves, &mut rng);
     // Initialize the tree.
     let merkle_tree =
         KaryMerkleTree::<_, _, DEPTH, ARITY>::new(&native_leaf_hasher, &native_path_hasher, &leaves).unwrap();
-    // Initialize the leaf.
-    let merkle_leaf = leaves[0].clone();
-    // Initialize the Merkle path.
-    let merkle_path = merkle_tree.prove(0, &merkle_leaf).unwrap();
 
-    // Start the timer.
+    // Log the current time elapsed.
+    println!(" • Synthesized the Merkle tree in: {} secs", timer.elapsed().as_secs());
     let timer = std::time::Instant::now();
 
-    // Construct the circuit.
-    CurrentAleo::reset();
+    // Construct the assignment closure.
+    let generate_assignment = |rng: &mut TestRng| {
+        // Construct the circuit.
+        CurrentAleo::reset();
+
+        // Select the leaf index to prove.
+        let leaf_index = rng.gen_range(0..max_num_leaves as usize);
+        // Initialize the leaf.
+        let merkle_leaf = leaves[leaf_index].clone();
+        // Initialize the Merkle path.
+        let merkle_path = merkle_tree.prove(leaf_index, &merkle_leaf).unwrap();
 
-    // Initialize the Merkle path circuit.
-    let path = KaryMerklePath::<CurrentAleo, snarkvm_circuit::Sha3_256<CurrentAleo>, DEPTH, ARITY>::new(
-        Mode::Private,
-        merkle_path,
-    );
-    // Initialize the Merkle root.
-    let root = <CircuitPathHasher as PathHash<CurrentAleo>>::Hash::new(Mode::Private, *merkle_tree.root());
-    // Initialize the Merkle leaf.
-    let leaf: Vec<_> = Inject::new(Mode::Private, merkle_leaf);
+        // Uncomment me to enable logging.
+        // println!("\t• Proving leaf index: {leaf_index}");
 
-    // Verify the merkle path.
-    let candidate = path.verify(&circuit_leaf_hasher, &circuit_path_hasher, &root, &leaf);
-    assert!(candidate.eject_value());
+        // Initialize the Merkle path circuit.
+        let path = KaryMerklePath::<CurrentAleo, CircuitPathHasher, DEPTH, ARITY>::new(Mode::Private, merkle_path);
+        // Initialize the Merkle leaf circuit.
+        let leaf: Vec<_> = Inject::new(Mode::Private, merkle_leaf);
+        // Initialize the Merkle root circuit.
+        let root = <CircuitPathHasher as PathHash<CurrentAleo>>::Hash::new(Mode::Private, *merkle_tree.root());
 
-    // Construct the assignment.
-    let assignment = CurrentAleo::eject_assignment_and_reset();
+        // Verify the Merkle path circuit.
+        let candidate = path.verify(&circuit_leaf_hasher, &circuit_path_hasher, &root, &leaf);
+        assert!(candidate.eject_value());
 
+        // Eject the assignment.
+        CurrentAleo::eject_assignment_and_reset()
+    };
+
+    // Synthesize a single assignment.
+    let assignment = generate_assignment(&mut rng);
+
+    // Log the current time elapsed.
     println!(" • Synthesized the circuit in: {} ms", timer.elapsed().as_millis());
+    let timer = std::time::Instant::now();
 
     // Load the universal srs.
     let universal_srs = UniversalSRS::<CurrentNetwork>::load().unwrap();
     // Construct the proving key.
     let (proving_key, _) = universal_srs.to_circuit_key("KaryMerklePathVerification", &assignment).unwrap();
 
+    // Log the current time elapsed.
+    println!(" • Generated the proving key in: {} ms", timer.elapsed().as_millis());
+    println!("\t• Number of public & private variables: {:?}", (assignment.num_public(), assignment.num_private()));
+    println!("\t• Number of constraints: {}", assignment.num_constraints());
+    println!("\t• Number of nonzeros: {:?}", assignment.num_nonzeros());
+
     // Bench the proof construction.
-    for num_assignments in &[1, 2, 4, 8] {
+    for num_assignments in &[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024] {
+        // Reset the timer.
+        let timer = std::time::Instant::now();
+
         // Construct the assignments.
         let assignments =
-            [(proving_key.clone(), (0..*num_assignments).map(|_| assignment.clone()).collect::<Vec<_>>())];
+            [(proving_key.clone(), (0..*num_assignments).map(|_| generate_assignment(&mut rng)).collect::<Vec<_>>())];
+
+        // Log the current time elapsed.
+        println!(" • Generated {num_assignments} assignments in: {} ms", timer.elapsed().as_millis());
 
-        let varuna_version = VarunaVersion::V1;
+        let varuna_version = VarunaVersion::V2;
         c.bench_function(&format!("KaryMerkleTree batch prove {num_assignments} assignments"), |b| {
             b.iter(|| {
                 let _proof =