Merge pull request #2177 from ljedrz/perf/fewer_bhp_allocs

Fewer BHP allocations
ProvableHQ · Nov 24, 2023 · b93ab37 · b93ab37
2 parents b93295c + 941a065
commit b93ab37
Show file tree

Hide file tree

Showing 5 changed files with 39 additions and 23 deletions.
diff --git a/algorithms/src/snark/varuna/ahp/matrices.rs b/algorithms/src/snark/varuna/ahp/matrices.rs
@@ -27,7 +27,6 @@ use snarkvm_fields::{Field, PrimeField};
 use snarkvm_utilities::{cfg_iter, cfg_iter_mut, serialize::*};
 
 use anyhow::{anyhow, ensure, Result};
-use itertools::Itertools;
 use std::collections::BTreeMap;
 
 #[cfg(not(feature = "serial"))]

diff --git a/circuit/algorithms/src/bhp/hash_uncompressed.rs b/circuit/algorithms/src/bhp/hash_uncompressed.rs
@@ -38,31 +38,32 @@ impl<E: Environment, const NUM_WINDOWS: u8, const WINDOW_SIZE: u8> HashUncompres
         // Initialize a variable to store the hash from the current iteration.
         let mut digest = Group::zero();
 
+        // Prepare a reusable vector for the preimage.
+        let mut preimage = Vec::with_capacity(num_hasher_bits);
+
         // Compute the hash of the input.
         for (i, input_bits) in input.chunks(max_input_bits_per_iteration).enumerate() {
             // Determine if this is the first iteration.
-            let preimage = match i == 0 {
+            match i == 0 {
                 // Construct the first iteration as: [ 0...0 || DOMAIN || LENGTH(INPUT) || INPUT[0..BLOCK_SIZE] ].
                 true => {
                     // Initialize a vector for the hash preimage.
-                    let mut preimage = Vec::with_capacity(num_hasher_bits);
                     preimage.extend(self.domain.clone());
                     U64::constant(console::U64::new(input.len() as u64)).write_bits_le(&mut preimage);
                     preimage.extend_from_slice(input_bits);
-                    preimage
                 }
                 // Construct the subsequent iterations as: [ PREVIOUS_HASH[0..DATA_BITS] || INPUT[I * BLOCK_SIZE..(I + 1) * BLOCK_SIZE] ].
                 false => {
                     // Initialize a vector for the hash preimage.
-                    let mut preimage = Vec::with_capacity(num_hasher_bits);
                     digest.to_x_coordinate().write_bits_le(&mut preimage);
                     preimage.truncate(num_data_bits);
                     preimage.extend_from_slice(input_bits);
-                    preimage
                 }
-            };
+            }
             // Hash the preimage for this iteration.
             digest = self.hasher.hash_uncompressed(&preimage);
+            // Clear the preimage vector for the next iteration.
+            preimage.clear();
         }
 
         digest

diff --git a/circuit/algorithms/src/bhp/hasher/hash_uncompressed.rs b/circuit/algorithms/src/bhp/hasher/hash_uncompressed.rs
@@ -14,6 +14,8 @@
 
 use super::*;
 
+use std::borrow::Cow;
+
 impl<E: Environment, const NUM_WINDOWS: u8, const WINDOW_SIZE: u8> HashUncompressed
     for BHPHasher<E, NUM_WINDOWS, WINDOW_SIZE>
 {
@@ -31,18 +33,27 @@ impl<E: Environment, const NUM_WINDOWS: u8, const WINDOW_SIZE: u8> HashUncompres
         }
 
         // Ensure the input size is within the parameter size.
-        let mut input = input.to_vec();
-        match input.len() <= Self::MAX_BITS {
+        let input = match input.len() <= Self::MAX_BITS {
             true => {
                 // Pad the input to a multiple of `BHP_CHUNK_SIZE` for hashing.
                 if input.len() % BHP_CHUNK_SIZE != 0 {
+                    // Compute the number of padding bits.
                     let padding = BHP_CHUNK_SIZE - (input.len() % BHP_CHUNK_SIZE);
-                    input.resize(input.len() + padding, Boolean::constant(false));
-                    assert_eq!(input.len() % BHP_CHUNK_SIZE, 0, "Input must be a multiple of {BHP_CHUNK_SIZE}");
+                    // Pad the input with `false` bits.
+                    let mut padded_input = Vec::with_capacity(input.len() + padding);
+                    padded_input.extend_from_slice(input);
+                    padded_input.resize(input.len() + padding, Boolean::constant(false));
+                    // Ensure the input is a multiple of `BHP_CHUNK_SIZE`.
+                    assert_eq!(padded_input.len() % BHP_CHUNK_SIZE, 0, "Input must be a multiple of {BHP_CHUNK_SIZE}");
+                    // Return the padded input.
+                    Cow::Owned(padded_input)
+                } else {
+                    // Return the input as a borrowed slice.
+                    Cow::Borrowed(input)
                 }
             }
             false => E::halt(format!("Inputs to this BHP cannot exceed {} bits", Self::MAX_BITS)),
-        }
+        };
 
         // Declare the 1 constant field element.
         let one = Field::one();

diff --git a/console/algorithms/src/bhp/hash_uncompressed.rs b/console/algorithms/src/bhp/hash_uncompressed.rs
@@ -38,31 +38,31 @@ impl<E: Environment, const NUM_WINDOWS: u8, const WINDOW_SIZE: u8> HashUncompres
         // Initialize a variable to store the hash from the current iteration.
         let mut digest = Group::<E>::zero();
 
+        // Prepare a reusable vector for the preimage.
+        let mut preimage = Vec::with_capacity(num_hasher_bits);
+
         // Compute the hash of the input.
         for (i, input_bits) in input.chunks(max_input_bits_per_iteration).enumerate() {
             // Determine if this is the first iteration.
-            let preimage = match i == 0 {
+            match i == 0 {
                 // Construct the first iteration as: [ 0...0 || DOMAIN || LENGTH(INPUT) || INPUT[0..BLOCK_SIZE] ].
                 true => {
                     // Initialize a vector for the hash preimage.
-                    let mut preimage = Vec::with_capacity(num_hasher_bits);
                     preimage.extend(&self.domain);
                     (input.len() as u64).write_bits_le(&mut preimage);
                     preimage.extend(input_bits);
-                    preimage
                 }
                 // Construct the subsequent iterations as: [ PREVIOUS_HASH[0..DATA_BITS] || INPUT[I * BLOCK_SIZE..(I + 1) * BLOCK_SIZE] ].
                 false => {
                     // Initialize a vector for the hash preimage.
-                    let mut preimage = Vec::with_capacity(num_hasher_bits);
                     digest.to_x_coordinate().write_bits_le(&mut preimage);
                     preimage.truncate(num_data_bits);
                     preimage.extend(input_bits);
-                    preimage
                 }
-            };
+            }
             // Hash the preimage for this iteration.
             digest = self.hasher.hash_uncompressed(&preimage)?;
+            preimage.clear();
         }
 
         Ok(digest)

diff --git a/console/algorithms/src/bhp/hasher/hash_uncompressed.rs b/console/algorithms/src/bhp/hasher/hash_uncompressed.rs
@@ -14,6 +14,8 @@
 
 use super::*;
 
+use std::borrow::Cow;
+
 impl<E: Environment, const NUM_WINDOWS: u8, const WINDOW_SIZE: u8> HashUncompressed
     for BHPHasher<E, NUM_WINDOWS, WINDOW_SIZE>
 {
@@ -36,12 +38,15 @@ impl<E: Environment, const NUM_WINDOWS: u8, const WINDOW_SIZE: u8> HashUncompres
         );
 
         // Pad the input to a multiple of `BHP_CHUNK_SIZE` for hashing.
-        let mut input = input.to_vec();
-        if input.len() % BHP_CHUNK_SIZE != 0 {
+        let input = if input.len() % BHP_CHUNK_SIZE != 0 {
             let padding = BHP_CHUNK_SIZE - (input.len() % BHP_CHUNK_SIZE);
-            input.resize(input.len() + padding, false);
-            ensure!((input.len() % BHP_CHUNK_SIZE) == 0, "Input must be a multiple of {BHP_CHUNK_SIZE}");
-        }
+            let mut padded_input = vec![false; input.len() + padding];
+            padded_input[..input.len()].copy_from_slice(input);
+            ensure!((padded_input.len() % BHP_CHUNK_SIZE) == 0, "Input must be a multiple of {BHP_CHUNK_SIZE}");
+            Cow::Owned(padded_input)
+        } else {
+            Cow::Borrowed(input)
+        };
 
         // Compute sum of h_i^{sum of (1-2*c_{i,j,2})*(1+c_{i,j,0}+2*c_{i,j,1})*2^{4*(j-1)} for all j in segment}
         // for all i. Described in section 5.4.1.7 in the Zcash protocol specification.