Standardize tracing crate + implement micro-benchmarking (#36)

* add tracing spans * more instrumenting * more tracing * use spans ig * more spans * more spans * commit_T spans 🤔 * nice spans * remove commit_T * Add trace for ppsnark and snark * fix cargo.toml and fmt --------- Co-authored-by: Samuel Burnham <45365069+samuelburnham@users.noreply.github.com>
lurk-lab · Aug 31, 2023 · c360828 · c360828
1 parent b1f971b
commit c360828
Show file tree

Hide file tree

Showing 8 changed files with 51 additions and 36 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -21,8 +21,8 @@ rayon = "1.7"
 rand_core = { version = "0.6", default-features = false }
 rand_chacha = "0.3"
 subtle = "2.5"
-pasta_curves = { git="https://github.com/lurk-lab/pasta_curves", branch="dev", features = ["repr-c", "serde"] }
-neptune = { git="https://github.com/lurk-lab/neptune", branch="dev", default-features = false }
+pasta_curves = { git = "https://github.com/lurk-lab/pasta_curves", branch="dev", features = ["repr-c", "serde"] }
+neptune = { git = "https://github.com/lurk-lab/neptune", branch="dev", default-features = false }
 generic-array = "0.14.4"
 num-bigint = { version = "0.4", features = ["serde", "rand"] }
 num-traits = "0.2"
@@ -34,10 +34,12 @@ byteorder = "1.4.3"
 thiserror = "1.0"
 halo2curves = { version = "0.4.0", features = ["derive_serde"] }
 group = "0.13.0"
-log = "0.4.17"
 abomonation = "0.7.3"
 abomonation_derive = { git = "https://github.com/lurk-lab/abomonation_derive.git" }
 tap = "1.0.1"
+tracing = "0.1.37"
+tracing-texray = "0.2.0"
+tracing-subscriber = "0.3.17"
 
 [target.'cfg(any(target_arch = "x86_64", target_arch = "aarch64"))'.dependencies]
 pasta-msm = { git="https://github.com/lurk-lab/pasta-msm", branch="dev", version = "0.1.4" }
@@ -86,7 +88,7 @@ flamegraph = ["pprof/flamegraph", "pprof/criterion"]
 
 # This is needed to ensure halo2curves, which imports pasta-curves, uses the *same* traits in bn256_grumpkin
 [patch.crates-io]
-pasta_curves = { git="https://github.com/lurk-lab/pasta_curves", branch="dev" }
+pasta_curves = { git = "https://github.com/lurk-lab/pasta_curves", branch = "dev" }
 
 [profile.dev-ci]
 inherits = "dev"

diff --git a/src/lib.rs b/src/lib.rs
@@ -362,6 +362,7 @@ where
 
   /// Create a new `RecursiveSNARK` (or updates the provided `RecursiveSNARK`)
   /// by executing a step of the incremental computation
+  #[tracing::instrument(skip_all, name = "RecursiveSNARK::prove_step")]
   pub fn prove_step(
     &mut self,
     pp: &PublicParams<G1, G2, C1, C2>,
@@ -410,6 +411,7 @@ where
       c_primary,
       pp.ro_consts_circuit_primary.clone(),
     );
+
     let zi_primary = circuit_primary
       .synthesize(&mut cs_primary)
       .map_err(|_| NovaError::SynthesisError)?;

diff --git a/src/nifs.rs b/src/nifs.rs
@@ -31,6 +31,7 @@ impl<G: Group> NIFS<G> {
   /// with the guarantee that the folded witness `W` satisfies the folded instance `U`
   /// if and only if `W1` satisfies `U1` and `W2` satisfies `U2`.
   #[allow(clippy::too_many_arguments)]
+  #[tracing::instrument(skip_all, name = "NIFS::prove")]
   pub fn prove(
     ck: &CommitmentKey<G>,
     ro_consts: &ROConstants<G>,

diff --git a/src/provider/pasta.rs b/src/provider/pasta.rs
@@ -68,6 +68,7 @@ macro_rules! impl_traits {
       type CE = CommitmentEngine<Self>;
 
       #[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))]
+      #[tracing::instrument(skip_all, name = "<_ as Group>::vartime_multiscalar_mul")]
       fn vartime_multiscalar_mul(
         scalars: &[Self::Scalar],
         bases: &[Self::PreprocessedGroupElement],

diff --git a/src/r1cs.rs b/src/r1cs.rs
@@ -293,40 +293,47 @@ impl<G: Group> R1CSShape<G> {
     U2: &R1CSInstance<G>,
     W2: &R1CSWitness<G>,
   ) -> Result<(Vec<G::Scalar>, Commitment<G>), NovaError> {
-    let (AZ_1, BZ_1, CZ_1) = {
+    let (AZ_1, BZ_1, CZ_1) = tracing::info_span!("AZ_1, BZ_1, CZ_1").in_scope(|| {
       let Z1 = [W1.W.clone(), vec![U1.u], U1.X.clone()].concat();
-      self.multiply_vec(&Z1)?
-    };
+      self.multiply_vec(&Z1)
+    })?;
 
-    let (AZ_2, BZ_2, CZ_2) = {
+    let (AZ_2, BZ_2, CZ_2) = tracing::info_span!("AZ_2, BZ_2, CZ_2").in_scope(|| {
       let Z2 = [W2.W.clone(), vec![G::Scalar::ONE], U2.X.clone()].concat();
-      self.multiply_vec(&Z2)?
-    };
-
-    let AZ_1_circ_BZ_2 = (0..AZ_1.len())
-      .into_par_iter()
-      .map(|i| AZ_1[i] * BZ_2[i])
-      .collect::<Vec<G::Scalar>>();
-    let AZ_2_circ_BZ_1 = (0..AZ_2.len())
-      .into_par_iter()
-      .map(|i| AZ_2[i] * BZ_1[i])
-      .collect::<Vec<G::Scalar>>();
-    let u_1_cdot_CZ_2 = (0..CZ_2.len())
-      .into_par_iter()
-      .map(|i| U1.u * CZ_2[i])
-      .collect::<Vec<G::Scalar>>();
-    let u_2_cdot_CZ_1 = (0..CZ_1.len())
-      .into_par_iter()
-      .map(|i| CZ_1[i])
-      .collect::<Vec<G::Scalar>>();
-
-    let T = AZ_1_circ_BZ_2
-      .par_iter()
-      .zip(&AZ_2_circ_BZ_1)
-      .zip(&u_1_cdot_CZ_2)
-      .zip(&u_2_cdot_CZ_1)
-      .map(|(((a, b), c), d)| *a + *b - *c - *d)
-      .collect::<Vec<G::Scalar>>();
+      self.multiply_vec(&Z2)
+    })?;
+
+    // forgive the horror here, but it's for grouping into one span
+    let (AZ_1_circ_BZ_2, AZ_2_circ_BZ_1, u_1_cdot_CZ_2, u_2_cdot_CZ_1) =
+      tracing::info_span!("cross terms").in_scope(|| {
+        let AZ_1_circ_BZ_2 = (0..AZ_1.len())
+          .into_par_iter()
+          .map(|i| AZ_1[i] * BZ_2[i])
+          .collect::<Vec<G::Scalar>>();
+        let AZ_2_circ_BZ_1 = (0..AZ_2.len())
+          .into_par_iter()
+          .map(|i| AZ_2[i] * BZ_1[i])
+          .collect::<Vec<G::Scalar>>();
+        let u_1_cdot_CZ_2 = (0..CZ_2.len())
+          .into_par_iter()
+          .map(|i| U1.u * CZ_2[i])
+          .collect::<Vec<G::Scalar>>();
+        let u_2_cdot_CZ_1 = (0..CZ_1.len())
+          .into_par_iter()
+          .map(|i| CZ_1[i])
+          .collect::<Vec<G::Scalar>>();
+        (AZ_1_circ_BZ_2, AZ_2_circ_BZ_1, u_1_cdot_CZ_2, u_2_cdot_CZ_1)
+      });
+
+    let T = tracing::info_span!("T").in_scope(|| {
+      AZ_1_circ_BZ_2
+        .par_iter()
+        .zip(&AZ_2_circ_BZ_1)
+        .zip(&u_1_cdot_CZ_2)
+        .zip(&u_2_cdot_CZ_1)
+        .map(|(((a, b), c), d)| *a + *b - *c - *d)
+        .collect::<Vec<G::Scalar>>()
+    });
 
     let comm_T = CE::<G>::commit(ck, &T);
 

diff --git a/src/spartan/ppsnark.rs b/src/spartan/ppsnark.rs
@@ -911,6 +911,7 @@ where
   }
 
   /// produces a succinct proof of satisfiability of a `RelaxedR1CS` instance
+  #[tracing::instrument(skip_all, name = "PPSNARK::prove")]
   fn prove(
     ck: &CommitmentKey<G>,
     pk: &Self::ProverKey,

diff --git a/src/spartan/snark.rs b/src/spartan/snark.rs
@@ -99,6 +99,7 @@ where
   }
 
   /// produces a succinct proof of satisfiability of a `RelaxedR1CS` instance
+  #[tracing::instrument(skip_all, name = "SNARK::prove")]
   fn prove(
     ck: &CommitmentKey<G>,
     pk: &Self::ProverKey,

diff --git a/src/supernova/mod.rs b/src/supernova/mod.rs
@@ -20,8 +20,8 @@ use crate::{
 };
 
 use ff::Field;
-use log::debug;
 use serde::{Deserialize, Serialize};
+use tracing::debug;
 
 use crate::bellpepper::{
   r1cs::{NovaShape, NovaWitness},