Skip to content

Commit

Permalink
chore(zk): add bench zk v1 vs v2
Browse files Browse the repository at this point in the history
  • Loading branch information
@nsarlin-zama
nsarlin-zama committed Jan 8, 2025
1 parent adf52ac commit 2353e85
Show file tree
Hide file tree
Showing 3 changed files with 77 additions and 25 deletions.
72 changes: 49 additions & 23 deletions tfhe/benches/integer/zk_pke.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,14 @@ use tfhe::integer::parameters::IntegerCompactCiphertextListExpansionMode;
use tfhe::integer::{ClientKey, CompactPrivateKey, CompactPublicKey, ServerKey};
use tfhe::keycache::NamedParam;
use tfhe::shortint::parameters::classic::tuniform::p_fail_2_minus_64::ks_pbs::PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
use tfhe::shortint::parameters::compact_public_key_only::p_fail_2_minus_64::ks_pbs::V0_11_PARAM_PKE_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
use tfhe::shortint::parameters::key_switching::p_fail_2_minus_64::ks_pbs::V0_11_PARAM_KEYSWITCH_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
use tfhe::shortint::parameters::compact_public_key_only::p_fail_2_minus_64::ks_pbs::{
V0_11_PARAM_PKE_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
V0_11_PARAM_PKE_TO_SMALL_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64_ZKV1,
};
use tfhe::shortint::parameters::key_switching::p_fail_2_minus_64::ks_pbs::{
V0_11_PARAM_KEYSWITCH_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
V0_11_PARAM_KEYSWITCH_PKE_TO_SMALL_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64_ZKV1,
};
use tfhe::shortint::parameters::PBSParameters;
use tfhe::zk::{CompactPkeCrs, ZkComputeLoad};
use utilities::{write_to_json, OperatorType};
Expand All @@ -33,11 +39,18 @@ fn pke_zk_proof(c: &mut Criterion) {
.sample_size(15)
.measurement_time(std::time::Duration::from_secs(60));

for (param_pke, _param_casting, param_fhe) in [(
V0_11_PARAM_PKE_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
V0_11_PARAM_KEYSWITCH_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
)] {
for (param_pke, _param_casting, param_fhe) in [
(
V0_11_PARAM_PKE_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
V0_11_PARAM_KEYSWITCH_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
),
(
V0_11_PARAM_PKE_TO_SMALL_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64_ZKV1,
V0_11_PARAM_KEYSWITCH_PKE_TO_SMALL_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64_ZKV1,
PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
),
] {
let param_name = param_fhe.name();
let param_name = param_name.as_str();
let cks = ClientKey::new(param_fhe);
Expand All @@ -53,6 +66,8 @@ fn pke_zk_proof(c: &mut Criterion) {
let mut rng = rand::thread_rng();
metadata.fill_with(|| rng.gen());

let zk_vers = param_pke.zk_scheme;

for bits in [64usize, 640, 1280, 4096] {
assert_eq!(bits % 64, 0);
// Packing, so we take the message and carry modulus to compute our block count
Expand Down Expand Up @@ -81,8 +96,9 @@ fn pke_zk_proof(c: &mut Criterion) {

match BENCH_TYPE.get().unwrap() {
BenchmarkType::Latency => {
bench_id =
format!("{bench_name}::{param_name}_{bits}_bits_packed_{zk_load}");
bench_id = format!(
"{bench_name}::{param_name}_{bits}_bits_packed_{zk_load}_ZK{zk_vers:?}"
);
bench_group.bench_function(&bench_id, |b| {
let input_msg = rng.gen::<u64>();
let messages = vec![input_msg; fhe_uint_count];
Expand All @@ -100,7 +116,7 @@ fn pke_zk_proof(c: &mut Criterion) {
bench_group.throughput(Throughput::Elements(elements));

bench_id = format!(
"{bench_name}::throughput::{param_name}_{bits}_bits_packed_{zk_load}"
"{bench_name}::throughput::{param_name}_{bits}_bits_packed_{zk_load}_ZK{zk_vers:?}"
);
bench_group.bench_function(&bench_id, |b| {
let messages = (0..elements)
Expand Down Expand Up @@ -155,11 +171,18 @@ fn pke_zk_verify(c: &mut Criterion, results_file: &Path) {
.open(results_file)
.expect("cannot open results file");

for (param_pke, param_casting, param_fhe) in [(
V0_11_PARAM_PKE_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
V0_11_PARAM_KEYSWITCH_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
)] {
for (param_pke, param_casting, param_fhe) in [
(
V0_11_PARAM_PKE_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
V0_11_PARAM_KEYSWITCH_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
),
(
V0_11_PARAM_PKE_TO_SMALL_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64_ZKV1,
V0_11_PARAM_KEYSWITCH_PKE_TO_SMALL_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64_ZKV1,
PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
),
] {
let param_name = param_fhe.name();
let param_name = param_name.as_str();
let cks = ClientKey::new(param_fhe);
Expand All @@ -174,6 +197,8 @@ fn pke_zk_verify(c: &mut Criterion, results_file: &Path) {
let mut rng = rand::thread_rng();
metadata.fill_with(|| rng.gen());

let zk_vers = param_pke.zk_scheme;

for bits in [64usize, 640, 1280, 4096] {
assert_eq!(bits % 64, 0);
// Packing, so we take the message and carry modulus to compute our block count
Expand All @@ -199,7 +224,7 @@ fn pke_zk_verify(c: &mut Criterion, results_file: &Path) {

println!("CRS size: {}", crs_data.len());

let test_name = format!("zk::crs_sizes::{param_name}_{bits}_bits_packed");
let test_name = format!("zk::crs_sizes::{param_name}_{bits}_bits_packed_ZK{zk_vers:?}");

write_result(&mut file, &test_name, crs_data.len());
write_to_json::<u64, _>(
Expand All @@ -223,10 +248,11 @@ fn pke_zk_verify(c: &mut Criterion, results_file: &Path) {

match BENCH_TYPE.get().unwrap() {
BenchmarkType::Latency => {
bench_id_verify =
format!("{bench_name}::{param_name}_{bits}_bits_packed_{zk_load}");
bench_id_verify = format!(
"{bench_name}::{param_name}_{bits}_bits_packed_{zk_load}_ZK{zk_vers:?}"
);
bench_id_verify_and_expand = format!(
"{bench_name}_and_expand::{param_name}_{bits}_bits_packed_{zk_load}"
"{bench_name}_and_expand::{param_name}_{bits}_bits_packed_{zk_load}_ZK{zk_vers:?}"
);

let input_msg = rng.gen::<u64>();
Expand All @@ -246,7 +272,7 @@ fn pke_zk_verify(c: &mut Criterion, results_file: &Path) {
);

let test_name = format!(
"zk::proven_list_size::{param_name}_{bits}_bits_packed_{zk_load}"
"zk::proven_list_size::{param_name}_{bits}_bits_packed_{zk_load}_ZK{zk_vers:?}"
);

write_result(
Expand All @@ -268,7 +294,7 @@ fn pke_zk_verify(c: &mut Criterion, results_file: &Path) {
println!("proof size: {}", ct1.proof_size());

let test_name =
format!("zk::proof_sizes::{param_name}_{bits}_bits_packed_{zk_load}");
format!("zk::proof_sizes::{param_name}_{bits}_bits_packed_{zk_load}_ZK{zk_vers:?}");

write_result(&mut file, &test_name, proof_size);
write_to_json::<u64, _>(
Expand Down Expand Up @@ -308,10 +334,10 @@ fn pke_zk_verify(c: &mut Criterion, results_file: &Path) {
bench_group.throughput(Throughput::Elements(elements));

bench_id_verify = format!(
"{bench_name}::throughput::{param_name}_{bits}_bits_packed_{zk_load}"
"{bench_name}::throughput::{param_name}_{bits}_bits_packed_{zk_load}_ZK{zk_vers:?}"
);
bench_id_verify_and_expand = format!(
"{bench_name}_and_expand::{param_name}_{bits}_bits_packed_{zk_load}"
"{bench_name}_and_expand::{param_name}_{bits}_bits_packed_{zk_load}_ZK{zk_vers:?}"
);

println!("Generating proven ciphertexts list ({zk_load})... ");
Expand Down
15 changes: 13 additions & 2 deletions tfhe/src/js_on_wasm_api/shortint.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,14 @@ use crate::core_crypto::commons::math::random::Seed;
use crate::core_crypto::prelude::DefaultRandomGenerator;
use crate::js_on_wasm_api::js_high_level_api::into_js_error;
use crate::shortint::parameters::classic::compact_pk::*;
use crate::shortint::parameters::compact_public_key_only::p_fail_2_minus_64::ks_pbs::V0_11_PARAM_PKE_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
use crate::shortint::parameters::key_switching::p_fail_2_minus_64::ks_pbs::V0_11_PARAM_KEYSWITCH_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
use crate::shortint::parameters::compact_public_key_only::p_fail_2_minus_64::ks_pbs::{
V0_11_PARAM_PKE_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
V0_11_PARAM_PKE_TO_SMALL_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64_ZKV1,
};
use crate::shortint::parameters::key_switching::p_fail_2_minus_64::ks_pbs::{
V0_11_PARAM_KEYSWITCH_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
V0_11_PARAM_KEYSWITCH_PKE_TO_SMALL_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64_ZKV1,
};
use crate::shortint::parameters::*;
use std::panic::set_hook;
use wasm_bindgen::prelude::*;
Expand Down Expand Up @@ -216,6 +222,7 @@ pub struct ShortintNoiseDistribution(
#[allow(non_camel_case_types)]
pub enum ShortintCompactPublicKeyEncryptionParametersName {
SHORTINT_PARAM_PKE_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
SHORTINT_PARAM_PKE_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64_ZKV1,
}

#[wasm_bindgen]
Expand All @@ -227,6 +234,10 @@ impl ShortintCompactPublicKeyEncryptionParameters {
ShortintCompactPublicKeyEncryptionParametersName::SHORTINT_PARAM_PKE_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64 => Self {
compact_pke_params: V0_11_PARAM_PKE_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
casting_parameters: V0_11_PARAM_KEYSWITCH_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
},
ShortintCompactPublicKeyEncryptionParametersName::SHORTINT_PARAM_PKE_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64_ZKV1 => Self {
compact_pke_params: V0_11_PARAM_PKE_TO_SMALL_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64_ZKV1,
casting_parameters: V0_11_PARAM_KEYSWITCH_PKE_TO_SMALL_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64_ZKV1,
}
}
}
Expand Down
15 changes: 15 additions & 0 deletions tfhe/web_wasm_parallel_tests/worker.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -645,6 +645,7 @@ async function compressedServerKeyBenchMessage2Carry2() {
async function compactPublicKeyZeroKnowledgeBench() {
let params_to_bench = [
{
zk_scheme: "ZKV2",
name: shortint_params_name(
ShortintParametersName.PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
),
Expand All @@ -655,6 +656,18 @@ async function compactPublicKeyZeroKnowledgeBench() {
ShortintCompactPublicKeyEncryptionParametersName.SHORTINT_PARAM_PKE_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
),
},
{
zk_scheme: "ZKV1",
name: shortint_params_name(
ShortintParametersName.PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
),
block_params: new ShortintParameters(
ShortintParametersName.PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
),
casting_params: new ShortintCompactPublicKeyEncryptionParameters(
ShortintCompactPublicKeyEncryptionParametersName.SHORTINT_PARAM_PKE_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64_ZKV1,
),
},
];

let bench_results = {};
Expand Down Expand Up @@ -718,6 +731,8 @@ async function compactPublicKeyZeroKnowledgeBench() {
const mean = timing / bench_loops;
const common_bench_str =
"compact_fhe_uint_proven_encryption_" +
params.zk_scheme +
"_" +
encrypt_count * 64 +
"_bits_packed_" +
load_to_str[loadChoice];
Expand Down

0 comments on commit 2353e85

Please sign in to comment.