Skip to content

Commit

Permalink
chore(bench): add packing keyswtich benchmarks
Browse files Browse the repository at this point in the history
  • Loading branch information
soonum committed Oct 24, 2024
1 parent 5db44d0 commit 9d51134
Show file tree
Hide file tree
Showing 4 changed files with 178 additions and 12 deletions.
142 changes: 139 additions & 3 deletions tfhe/benches/core_crypto/ks_bench.rs
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,17 @@ fn benchmark_parameters_32bits() -> Vec<(String, CryptoParametersRecord<u32>)> {
.collect()
}

fn benchmark_compression_parameters() -> Vec<(String, CryptoParametersRecord<u64>)> {
vec![(
COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64.name(),
(
COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
)
.into(),
)]
}

fn keyswitch<Scalar: UnsignedTorus + CastInto<usize> + Serialize>(
criterion: &mut Criterion,
parameters: &[(String, CryptoParametersRecord<Scalar>)],
Expand Down Expand Up @@ -175,6 +186,107 @@ fn keyswitch<Scalar: UnsignedTorus + CastInto<usize> + Serialize>(
}
}

fn packing_keyswitch<Scalar, F>(
criterion: &mut Criterion,
bench_name: &str,
parameters: &[(String, CryptoParametersRecord<Scalar>)],
ks_op: F,
) where
Scalar: UnsignedTorus + CastInto<usize> + Serialize,
F: Fn(
&LwePackingKeyswitchKey<Vec<Scalar>>,
&LweCiphertextList<Vec<Scalar>>,
&mut GlweCiphertext<Vec<Scalar>>,
),
{
let bench_name = format!("core_crypto::{bench_name}");
let mut bench_group = criterion.benchmark_group(&bench_name);

// Create the PRNG
let mut seeder = new_seeder();
let seeder = seeder.as_mut();
let mut encryption_generator =
EncryptionRandomGenerator::<ActivatedRandomGenerator>::new(seeder.seed(), seeder);
let mut secret_generator =
SecretRandomGenerator::<ActivatedRandomGenerator>::new(seeder.seed());

for (name, params) in parameters.iter() {
let lwe_dimension = params.lwe_dimension.unwrap();
let packing_glwe_dimension = params.packing_ks_glwe_dimension.unwrap();
let packing_polynomial_size = params.packing_ks_polynomial_size.unwrap();
let packing_ks_decomp_base_log = params.packing_ks_base_log.unwrap();
let packing_ks_decomp_level_count = params.packing_ks_level.unwrap();
let ciphertext_modulus = params.ciphertext_modulus.unwrap();
let count = params.lwe_per_glwe.unwrap();

let lwe_sk =
allocate_and_generate_new_binary_lwe_secret_key(lwe_dimension, &mut secret_generator);

let glwe_sk = allocate_and_generate_new_binary_glwe_secret_key(
packing_glwe_dimension,
packing_polynomial_size,
&mut secret_generator,
);

let pksk = allocate_and_generate_new_lwe_packing_keyswitch_key(
&lwe_sk,
&glwe_sk,
packing_ks_decomp_base_log,
packing_ks_decomp_level_count,
params.packing_ks_key_noise_distribution.unwrap(),
ciphertext_modulus,
&mut encryption_generator,
);

let mut input_lwe_list = LweCiphertextList::new(
Scalar::ZERO,
lwe_sk.lwe_dimension().to_lwe_size(),
count,
ciphertext_modulus,
);

let plaintext_list = PlaintextList::new(
Scalar::ZERO,
PlaintextCount(input_lwe_list.lwe_ciphertext_count().0),
);

encrypt_lwe_ciphertext_list(
&lwe_sk,
&mut input_lwe_list,
&plaintext_list,
params.lwe_noise_distribution.unwrap(),
&mut encryption_generator,
);

let mut output_glwe = GlweCiphertext::new(
Scalar::ZERO,
glwe_sk.glwe_dimension().to_glwe_size(),
glwe_sk.polynomial_size(),
ciphertext_modulus,
);

let id = format!("{bench_name}::{name}");
{
bench_group.bench_function(&id, |b| {
b.iter(|| {
ks_op(&pksk, &input_lwe_list, &mut output_glwe);
black_box(&mut output_glwe);
})
});
}
let bit_size = (params.message_modulus.unwrap_or(2) as u32).ilog2();
write_to_json(
&id,
*params,
name,
"packing_ks",
&OperatorType::Atomic,
bit_size,
vec![bit_size],
);
}
}

#[cfg(feature = "gpu")]
mod cuda {
use crate::benchmark_parameters_64bits;
Expand Down Expand Up @@ -402,15 +514,39 @@ mod cuda {

#[cfg(feature = "gpu")]
use cuda::cuda_keyswitch_group;
use tfhe::shortint::parameters::{
COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64, PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
};

pub fn keyswitch_group() {
let mut criterion: Criterion<_> =
(Criterion::default().sample_size(2000)).configure_from_args();
let mut criterion: Criterion<_> = (Criterion::default()
.sample_size(15)
.measurement_time(std::time::Duration::from_secs(60)))
.configure_from_args();
keyswitch(&mut criterion, &benchmark_parameters_64bits());
keyswitch(&mut criterion, &benchmark_parameters_32bits());
}

pub fn packing_keyswitch_group() {
let mut criterion: Criterion<_> = (Criterion::default()
.sample_size(10)
.measurement_time(std::time::Duration::from_secs(30)))
.configure_from_args();
packing_keyswitch(
&mut criterion,
"packing_keyswitch",
&benchmark_compression_parameters(),
keyswitch_lwe_ciphertext_list_and_pack_in_glwe_ciphertext,
);
packing_keyswitch(
&mut criterion,
"par_packing_keyswitch",
&benchmark_compression_parameters(),
par_keyswitch_lwe_ciphertext_list_and_pack_in_glwe_ciphertext,
);
}

#[cfg(not(feature = "gpu"))]
criterion_main!(keyswitch_group);
criterion_main!(keyswitch_group, packing_keyswitch_group);
#[cfg(feature = "gpu")]
criterion_main!(cuda_keyswitch_group);
8 changes: 4 additions & 4 deletions tfhe/benches/integer/glwe_packing_compression.rs
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ fn cpu_glwe_packing(c: &mut Criterion) {

write_to_json::<u64, _>(
&bench_id,
comp_param,
(comp_param, param),
comp_param.name(),
"pack",
&OperatorType::Atomic,
Expand All @@ -80,7 +80,7 @@ fn cpu_glwe_packing(c: &mut Criterion) {

write_to_json::<u64, _>(
&bench_id,
comp_param,
(comp_param, param),
comp_param.name(),
"unpack",
&OperatorType::Atomic,
Expand Down Expand Up @@ -165,7 +165,7 @@ mod cuda {

write_to_json::<u64, _>(
&bench_id,
comp_param,
(comp_param, param),
comp_param.name(),
"pack",
&OperatorType::Atomic,
Expand All @@ -189,7 +189,7 @@ mod cuda {

write_to_json::<u64, _>(
&bench_id,
comp_param,
(comp_param, param),
comp_param.name(),
"unpack",
&OperatorType::Atomic,
Expand Down
32 changes: 29 additions & 3 deletions tfhe/benches/utilities.rs
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ pub mod shortint_utils {
use tfhe::shortint::parameters::{
ShortintKeySwitchingParameters, PARAM_MESSAGE_2_CARRY_2_KS_PBS,
};
use tfhe::shortint::PBSParameters;
use tfhe::shortint::{ClassicPBSParameters, PBSParameters, ShortintParameterSet};

/// An iterator that yields a succession of combinations
/// of parameters and a num_block to achieve a certain bit_size ciphertext
Expand Down Expand Up @@ -140,9 +140,25 @@ pub mod shortint_utils {
}
}

impl From<CompressionParameters> for CryptoParametersRecord<u64> {
fn from(_params: CompressionParameters) -> Self {
impl From<(CompressionParameters, ClassicPBSParameters)> for CryptoParametersRecord<u64> {
fn from((comp_params, pbs_params): (CompressionParameters, ClassicPBSParameters)) -> Self {
let pbs_params = ShortintParameterSet::new_pbs_param_set(pbs_params.into());
let lwe_dimension = pbs_params.encryption_lwe_dimension();
CryptoParametersRecord {
lwe_dimension: Some(lwe_dimension),
br_level: Some(comp_params.br_level),
br_base_log: Some(comp_params.br_base_log),
packing_ks_level: Some(comp_params.packing_ks_level),
packing_ks_base_log: Some(comp_params.packing_ks_base_log),
packing_ks_polynomial_size: Some(comp_params.packing_ks_polynomial_size),
packing_ks_glwe_dimension: Some(comp_params.packing_ks_glwe_dimension),
lwe_per_glwe: Some(comp_params.lwe_per_glwe),
storage_log_modulus: Some(comp_params.storage_log_modulus),
lwe_noise_distribution: Some(pbs_params.encryption_noise_distribution()),
packing_ks_key_noise_distribution: Some(
comp_params.packing_ks_key_noise_distribution,
),
ciphertext_modulus: Some(pbs_params.ciphertext_modulus()),
..Default::default()
}
}
Expand All @@ -157,11 +173,15 @@ pub use shortint_utils::*;
pub struct CryptoParametersRecord<Scalar: UnsignedInteger> {
pub lwe_dimension: Option<LweDimension>,
pub glwe_dimension: Option<GlweDimension>,
pub packing_ks_glwe_dimension: Option<GlweDimension>,
pub polynomial_size: Option<PolynomialSize>,
pub packing_ks_polynomial_size: Option<PolynomialSize>,
#[serde(serialize_with = "CryptoParametersRecord::serialize_distribution")]
pub lwe_noise_distribution: Option<DynamicDistribution<Scalar>>,
#[serde(serialize_with = "CryptoParametersRecord::serialize_distribution")]
pub glwe_noise_distribution: Option<DynamicDistribution<Scalar>>,
#[serde(serialize_with = "CryptoParametersRecord::serialize_distribution")]
pub packing_ks_key_noise_distribution: Option<DynamicDistribution<Scalar>>,
pub pbs_base_log: Option<DecompositionBaseLog>,
pub pbs_level: Option<DecompositionLevelCount>,
pub ks_base_log: Option<DecompositionBaseLog>,
Expand All @@ -171,9 +191,15 @@ pub struct CryptoParametersRecord<Scalar: UnsignedInteger> {
pub pfks_std_dev: Option<StandardDev>,
pub cbs_level: Option<DecompositionLevelCount>,
pub cbs_base_log: Option<DecompositionBaseLog>,
pub br_level: Option<DecompositionLevelCount>,
pub br_base_log: Option<DecompositionBaseLog>,
pub packing_ks_level: Option<DecompositionLevelCount>,
pub packing_ks_base_log: Option<DecompositionBaseLog>,
pub message_modulus: Option<usize>,
pub carry_modulus: Option<usize>,
pub ciphertext_modulus: Option<CiphertextModulus<Scalar>>,
pub lwe_per_glwe: Option<LweCiphertextCount>,
pub storage_log_modulus: Option<CiphertextModulusLog>,
}

impl<Scalar: UnsignedInteger> CryptoParametersRecord<Scalar> {
Expand Down
8 changes: 6 additions & 2 deletions tfhe/examples/utilities/shortint_key_sizes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -230,22 +230,26 @@ fn tuniform_key_set_sizes(results_file: &Path) {

let param_compression = COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
let param_compression_name = stringify!(COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64);
let params_tuple = (
param_compression,
PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
);

let private_compression_key = cks.new_compression_private_key(param_compression);
let (compression_key, decompression_key) =
cks.new_compression_decompression_keys(&private_compression_key);

measure_serialized_size(
&compression_key,
param_compression,
params_tuple,
param_compression_name,
"compression_key",
"CompressionKey",
&mut file,
);
measure_serialized_size(
&decompression_key,
param_compression,
params_tuple,
param_compression_name,
"decompression_key",
"CompressionKey",
Expand Down

0 comments on commit 9d51134

Please sign in to comment.