plonk::prover: Introduce counter feature for FFTs and MSMs

Co-authored-by: Andrija <[email protected]>

halo2_proofs/Cargo.toml

@@ -50,12 +50,14 @@ group = "0.13"
 pasta_curves = "0.5"
 rand_core = { version = "0.6", default-features = false }
 tracing = "0.1"
+tracing-subscriber = "0.3.16"
 blake2b_simd = "1"
 maybe-rayon = {version = "0.1.0", default-features = false}
 
 # Developer tooling dependencies
 plotters = { version = "0.3.0", default-features = false, optional = true }
 tabbycat = { version = "0.1", features = ["attributes"], optional = true }
+lazy_static = { version = "1", optional = true }
 
 # Legacy circuit compatibility
 halo2_legacy_pdqsort = { version = "0.1.0", optional = true }
@@ -82,6 +84,7 @@ gadget-traces = ["backtrace"]
 sanity-checks = []
 batch = ["rand_core/getrandom"]
 floor-planner-v1-legacy-pdqsort = ["halo2_legacy_pdqsort"]
+counter = ["lazy_static"]
 
 # In-development features
 # See https://zcash.github.io/halo2/dev/features.html

halo2_proofs/src/arithmetic.rs

@@ -118,6 +118,17 @@ fn multiexp_serial<C: CurveAffine>(coeffs: &[C::Scalar], bases: &[C], acc: &mut
 /// Performs a small multi-exponentiation operation.
 /// Uses the double-and-add algorithm with doublings shared across points.
 pub fn small_multiexp<C: CurveAffine>(coeffs: &[C::Scalar], bases: &[C]) -> C::Curve {
+    #[cfg(feature = "counter")]
+    {
+        use crate::MSM_COUNTER;
+        let _ = *MSM_COUNTER
+            .lock()
+            .unwrap()
+            .entry(coeffs.len())
+            .and_modify(|cnt| *cnt += 1)
+            .or_insert(1);
+    }
+
     let coeffs: Vec<_> = coeffs.iter().map(|a| a.to_repr()).collect();
     let mut acc = C::Curve::identity();
 
@@ -145,6 +156,17 @@ pub fn small_multiexp<C: CurveAffine>(coeffs: &[C::Scalar], bases: &[C]) -> C::C
 ///
 /// This will use multithreading if beneficial.
 pub fn best_multiexp<C: CurveAffine>(coeffs: &[C::Scalar], bases: &[C]) -> C::Curve {
+    #[cfg(feature = "counter")]
+    {
+        use crate::MSM_COUNTER;
+        let _ = *MSM_COUNTER
+            .lock()
+            .unwrap()
+            .entry(coeffs.len())
+            .and_modify(|cnt| *cnt += 1)
+            .or_insert(1);
+    }
+
     assert_eq!(coeffs.len(), bases.len());
 
     let num_threads = multicore::current_num_threads();
@@ -184,6 +206,17 @@ pub fn best_multiexp<C: CurveAffine>(coeffs: &[C::Scalar], bases: &[C]) -> C::Cu
 ///
 /// This will use multithreading if beneficial.
 pub fn best_fft<Scalar: Field, G: FftGroup<Scalar>>(a: &mut [G], omega: Scalar, log_n: u32) {
+    #[cfg(feature = "counter")]
+    {
+        use crate::FFT_COUNTER;
+        let _ = *FFT_COUNTER
+            .lock()
+            .unwrap()
+            .entry(a.len())
+            .and_modify(|cnt| *cnt += 1)
+            .or_insert(1);
+    }
+
     fn bitreverse(mut n: usize, l: usize) -> usize {
         let mut r = 0;
         for _ in 0..l {

halo2_proofs/src/lib.rs

@@ -18,3 +18,21 @@ pub mod transcript;
 
 pub mod dev;
 mod helpers;
+
+#[cfg(feature = "counter")]
+extern crate lazy_static;
+
+#[cfg(feature = "counter")]
+use lazy_static::lazy_static;
+
+#[cfg(feature = "counter")]
+use std::sync::Mutex;
+
+#[cfg(feature = "counter")]
+use std::collections::BTreeMap;
+
+#[cfg(feature = "counter")]
+lazy_static! {
+    static ref FFT_COUNTER: Mutex<BTreeMap<usize, usize>> = Mutex::new(BTreeMap::new());
+    static ref MSM_COUNTER: Mutex<BTreeMap<usize, usize>> = Mutex::new(BTreeMap::new());
+}

halo2_proofs/src/plonk/prover.rs

@@ -46,6 +46,16 @@ pub fn create_proof<
     mut rng: R,
     transcript: &mut T,
 ) -> Result<(), Error> {
+    #[cfg(feature = "counter")]
+    {
+        use crate::{FFT_COUNTER, MSM_COUNTER};
+        use std::collections::BTreeMap;
+
+        // reset counters at the beginning of the prove
+        *MSM_COUNTER.lock().unwrap() = BTreeMap::new();
+        *FFT_COUNTER.lock().unwrap() = BTreeMap::new();
+    }
+
     if circuits.len() != instances.len() {
         return Err(Error::InvalidInstances);
     }
@@ -721,6 +731,18 @@ pub fn create_proof<
         // We query the h(X) polynomial at x
         .chain(vanishing.open(x));
 
+    #[cfg(feature = "counter")]
+    {
+        use crate::{FFT_COUNTER, MSM_COUNTER};
+        use std::collections::BTreeMap;
+        tracing::debug!("MSM_COUNTER: {:?}", MSM_COUNTER.lock().unwrap());
+        tracing::debug!("FFT_COUNTER: {:?}", *FFT_COUNTER.lock().unwrap());
+
+        // reset counters at the end of the proving
+        *MSM_COUNTER.lock().unwrap() = BTreeMap::new();
+        *FFT_COUNTER.lock().unwrap() = BTreeMap::new();
+    }
+
     multiopen::create_proof(params, rng, transcript, instances).map_err(|_| Error::Opening)
 }
 

halo2_proofs/tests/plonk_api.rs

@@ -19,6 +19,12 @@ use std::marker::PhantomData;
 
 #[test]
 fn plonk_api() {
+    tracing_subscriber::fmt()
+        .with_max_level(tracing::Level::DEBUG)
+        .with_ansi(false)
+        .without_time()
+        .init();
+
     const K: u32 = 5;
 
     /// This represents an advice column at a certain row in the ConstraintSystem