chore(zk): add benches for zk v2 bounds ghl and cs

zama-ai · Nov 26, 2024 · bb3529f · bb3529f
1 parent 83575c1
commit bb3529f
Show file tree

Hide file tree

Showing 2 changed files with 107 additions and 51 deletions.
diff --git a/tfhe-zk-pok/benches/pke_v2.rs b/tfhe-zk-pok/benches/pke_v2.rs
@@ -1,7 +1,9 @@
 use criterion::{criterion_group, criterion_main, Criterion};
-use tfhe_zk_pok::proofs::pke_v2::{prove, verify};
+use tfhe_zk_pok::proofs::pke_v2::{prove, verify, Bound};
 use tfhe_zk_pok::proofs::ComputeLoad;
-use utils::{init_params_v2, write_to_json, PKEV1_TEST_PARAMS, PKEV2_TEST_PARAMS};
+use utils::{
+    init_params_v2_cs, init_params_v2_ghl, write_to_json, PKEV1_TEST_PARAMS, PKEV2_TEST_PARAMS,
+};
 
 #[path = "./utils.rs"]
 mod utils;
@@ -16,35 +18,36 @@ fn bench_pke_v2_prove(c: &mut Criterion) {
 
     let rng = &mut rand::thread_rng();
 
-    for (params, param_name) in [
-        (PKEV1_TEST_PARAMS, "PKEV1_TEST_PARAMS"),
-        (PKEV2_TEST_PARAMS, "PKEV2_TEST_PARAMS"),
-    ] {
-        let (public_param, public_commit, private_commit, metadata) = init_params_v2(params);
+    for ((params, param_name), load, bound) in itertools::iproduct!(
+        [
+            (PKEV1_TEST_PARAMS, "PKEV1_TEST_PARAMS"),
+            (PKEV2_TEST_PARAMS, "PKEV2_TEST_PARAMS"),
+        ],
+        [ComputeLoad::Proof, ComputeLoad::Verify],
+        [Bound::CS, Bound::GHL]
+    ) {
+        let (public_param, public_commit, private_commit, metadata) = match bound {
+            Bound::GHL => init_params_v2_ghl(params),
+            Bound::CS => init_params_v2_cs(params),
+        };
         let effective_t = params.t >> 1;
         let bits = (params.k as u32) * effective_t.ilog2();
 
-        for load in [ComputeLoad::Proof, ComputeLoad::Verify] {
-            let zk_load = match load {
-                ComputeLoad::Proof => "compute_load_proof",
-                ComputeLoad::Verify => "compute_load_verify",
-            };
-            let bench_id = format!("{bench_name}::{param_name}_{bits}_bits_packed_{zk_load}");
+        let bench_id = format!("{bench_name}::{param_name}_{bits}_bits_packed_{load}_{bound:?}");
 
-            bench_group.bench_function(&bench_id, |b| {
-                b.iter(|| {
-                    prove(
-                        (&public_param, &public_commit),
-                        &private_commit,
-                        &metadata,
-                        load,
-                        rng,
-                    )
-                })
-            });
+        bench_group.bench_function(&bench_id, |b| {
+            b.iter(|| {
+                prove(
+                    (&public_param, &public_commit),
+                    &private_commit,
+                    &metadata,
+                    load,
+                    rng,
+                )
+            })
+        });
 
-            write_to_json(&bench_id, params, param_name, bench_shortname);
-        }
+        write_to_json(&bench_id, params, param_name, bench_shortname);
     }
 }
 
@@ -58,33 +61,38 @@ fn bench_pke_v2_verify(c: &mut Criterion) {
 
     let rng = &mut rand::thread_rng();
 
-    for (params, param_name) in [
-        (PKEV1_TEST_PARAMS, "PKEV1_TEST_PARAMS"),
-        (PKEV2_TEST_PARAMS, "PKEV2_TEST_PARAMS"),
-    ] {
-        let (public_param, public_commit, private_commit, metadata) = init_params_v2(params);
+    for ((params, param_name), load, bound) in itertools::iproduct!(
+        [
+            (PKEV1_TEST_PARAMS, "PKEV1_TEST_PARAMS"),
+            (PKEV2_TEST_PARAMS, "PKEV2_TEST_PARAMS"),
+        ],
+        [ComputeLoad::Proof, ComputeLoad::Verify],
+        [Bound::CS, Bound::GHL]
+    ) {
+        let (public_param, public_commit, private_commit, metadata) = match bound {
+            Bound::GHL => init_params_v2_ghl(params),
+            Bound::CS => init_params_v2_cs(params),
+        };
         let effective_t = params.t >> 1;
         let bits = (params.k as u32) * effective_t.ilog2();
 
-        for load in [ComputeLoad::Proof, ComputeLoad::Verify] {
-            let bench_id = format!("{bench_name}::{param_name}_{bits}_bits_packed_{load}");
+        let bench_id = format!("{bench_name}::{param_name}_{bits}_bits_packed_{load}_{bound:?}");
 
-            let proof = prove(
-                (&public_param, &public_commit),
-                &private_commit,
-                &metadata,
-                load,
-                rng,
-            );
+        let proof = prove(
+            (&public_param, &public_commit),
+            &private_commit,
+            &metadata,
+            load,
+            rng,
+        );
 
-            bench_group.bench_function(&bench_id, |b| {
-                b.iter(|| {
-                    verify(&proof, (&public_param, &public_commit), &metadata).unwrap();
-                })
-            });
+        bench_group.bench_function(&bench_id, |b| {
+            b.iter(|| {
+                verify(&proof, (&public_param, &public_commit), &metadata).unwrap();
+            })
+        });
 
-            write_to_json(&bench_id, params, param_name, bench_shortname);
-        }
+        write_to_json(&bench_id, params, param_name, bench_shortname);
     }
 }
 

diff --git a/tfhe-zk-pok/benches/utils.rs b/tfhe-zk-pok/benches/utils.rs
@@ -9,8 +9,9 @@ use serde::Serialize;
 use tfhe_zk_pok::proofs::pke::{commit, crs_gen, PrivateCommit, PublicCommit, PublicParams};
 
 use tfhe_zk_pok::proofs::pke_v2::{
-    commit as commitv2, crs_gen as crs_genv2, PrivateCommit as PrivateCommitv2,
-    PublicCommit as PublicCommitv2, PublicParams as PublicParamsv2,
+    commit as commitv2, crs_gen_cs as crs_genv2_cs, crs_gen_ghl as crs_genv2_ghl,
+    PrivateCommit as PrivateCommitv2, PublicCommit as PublicCommitv2,
+    PublicParams as PublicParamsv2,
 };
 
 // One of our usecases uses 320 bits of additional metadata
@@ -412,7 +413,7 @@ pub fn init_params_v1(
 }
 
 #[allow(unused)]
-pub fn init_params_v2(
+pub fn init_params_v2_cs(
     test_params: PkeTestParameters,
 ) -> (
     PublicParamsv2<Curve>,
@@ -435,7 +436,54 @@ pub fn init_params_v2(
 
     let ct = testcase.encrypt(test_params);
 
-    let public_param = crs_genv2::<Curve>(d, k, B, q, t, msbs_zero_padding_bit_count, rng);
+    let public_param = crs_genv2_cs::<Curve>(d, k, B, q, t, msbs_zero_padding_bit_count, rng);
+
+    let (public_commit, private_commit) = commitv2(
+        testcase.a.clone(),
+        testcase.b.clone(),
+        ct.c1.clone(),
+        ct.c2.clone(),
+        testcase.r.clone(),
+        testcase.e1.clone(),
+        testcase.m.clone(),
+        testcase.e2.clone(),
+        &public_param,
+        rng,
+    );
+
+    (
+        public_param,
+        public_commit,
+        private_commit,
+        testcase.metadata,
+    )
+}
+
+#[allow(unused)]
+pub fn init_params_v2_ghl(
+    test_params: PkeTestParameters,
+) -> (
+    PublicParamsv2<Curve>,
+    PublicCommitv2<Curve>,
+    PrivateCommitv2<Curve>,
+    [u8; METADATA_LEN],
+) {
+    let PkeTestParameters {
+        d,
+        k,
+        B,
+        q,
+        t,
+        msbs_zero_padding_bit_count,
+    } = test_params;
+
+    let rng = &mut StdRng::seed_from_u64(0);
+
+    let testcase = PkeTestcase::gen(rng, test_params);
+
+    let ct = testcase.encrypt(test_params);
+
+    let public_param = crs_genv2_ghl::<Curve>(d, k, B, q, t, msbs_zero_padding_bit_count, rng);
 
     let (public_commit, private_commit) = commitv2(
         testcase.a.clone(),