Skip to content

Commit

Permalink
chore(bench): benchmark oprf function against all available precisions
Browse files Browse the repository at this point in the history
  • Loading branch information
@soonum
soonum committed Jul 12, 2024
1 parent 7b96f55 commit a2c4d31
Show file tree
Hide file tree
Showing 4 changed files with 60 additions and 151 deletions.
56 changes: 1 addition & 55 deletions tfhe/benches/integer/bench.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,17 +4,14 @@ mod oprf;
#[path = "../utilities.rs"]
mod utilities;

use crate::utilities::{write_to_json, EnvConfig, OperatorType};
use crate::utilities::{write_to_json, EnvConfig, OperatorType, ParamsAndNumBlocksIter};
use criterion::{criterion_group, Criterion};
use itertools::iproduct;
use rand::prelude::*;
use std::env;
use std::vec::IntoIter;
use tfhe::integer::keycache::KEY_CACHE;
use tfhe::integer::prelude::*;
use tfhe::integer::{IntegerKeyKind, RadixCiphertext, RadixClientKey, ServerKey, U256};
use tfhe::keycache::NamedParam;
use tfhe::shortint::parameters::*;

/// The type used to hold scalar values
/// It must be as big as the largest bit size tested
Expand All @@ -27,57 +24,6 @@ fn gen_random_u256(rng: &mut ThreadRng) -> U256 {
tfhe::integer::U256::from((clearlow, clearhigh))
}

/// An iterator that yields a succession of combinations
/// of parameters and a num_block to achieve a certain bit_size ciphertext
/// in radix decomposition
struct ParamsAndNumBlocksIter {
params_and_bit_sizes:
itertools::Product<IntoIter<tfhe::shortint::PBSParameters>, IntoIter<usize>>,
}

impl Default for ParamsAndNumBlocksIter {
fn default() -> Self {
let env_config = EnvConfig::new();

if env_config.is_multi_bit {
#[cfg(feature = "gpu")]
let params = vec![PARAM_GPU_MULTI_BIT_MESSAGE_2_CARRY_2_GROUP_3_KS_PBS.into()];
#[cfg(not(feature = "gpu"))]
let params = vec![PARAM_MULTI_BIT_MESSAGE_2_CARRY_2_GROUP_2_KS_PBS.into()];

let params_and_bit_sizes = iproduct!(params, env_config.bit_sizes());
Self {
params_and_bit_sizes,
}
} else {
// FIXME One set of parameter is tested since we want to benchmark only quickest
// operations.
let params = vec![
PARAM_MESSAGE_2_CARRY_2_KS_PBS.into(),
// PARAM_MESSAGE_3_CARRY_3_KS_PBS.into(),
// PARAM_MESSAGE_4_CARRY_4_KS_PBS.into(),
];

let params_and_bit_sizes = iproduct!(params, env_config.bit_sizes());
Self {
params_and_bit_sizes,
}
}
}
}

impl Iterator for ParamsAndNumBlocksIter {
type Item = (tfhe::shortint::PBSParameters, usize, usize);

fn next(&mut self) -> Option<Self::Item> {
let (param, bit_size) = self.params_and_bit_sizes.next()?;
let num_block =
(bit_size as f64 / (param.message_modulus().0 as f64).log(2.0)).ceil() as usize;

Some((param, num_block, bit_size))
}
}

/// Base function to bench a server key function that is a binary operation, input ciphertexts will
/// contain non zero carries
fn bench_server_key_binary_function_dirty_inputs<F>(
Expand Down
41 changes: 4 additions & 37 deletions tfhe/benches/integer/oprf.rs
View file
Original file line number Diff line number Diff line change
@@ -1,50 +1,17 @@
use crate::utilities::{write_to_json, OperatorType};
use crate::utilities::{write_to_json, OperatorType, ParamsAndNumBlocksIter};
use concrete_csprng::seeders::Seed;
use criterion::{black_box, Criterion};
use itertools::iproduct;
use std::vec::IntoIter;
use tfhe::integer::keycache::KEY_CACHE;
use tfhe::integer::IntegerKeyKind;
use tfhe::keycache::NamedParam;
use tfhe::shortint::parameters::*;

/// An iterator that yields a succession of combinations
/// of parameters and a num_block to achieve a certain bit_size ciphertext
/// in radix decomposition
struct ParamsAndNumBlocksIter {
params_and_bit_sizes: itertools::Product<IntoIter<PBSParameters>, IntoIter<u64>>,
}

impl Default for ParamsAndNumBlocksIter {
fn default() -> Self {
let params = vec![PARAM_MESSAGE_2_CARRY_2_KS_PBS.into()];
let bit_sizes = vec![64];
let params_and_bit_sizes = iproduct!(params, bit_sizes);
Self {
params_and_bit_sizes,
}
}
}

impl Iterator for ParamsAndNumBlocksIter {
type Item = (PBSParameters, u64, u64);

fn next(&mut self) -> Option<Self::Item> {
let (param, bit_size) = self.params_and_bit_sizes.next()?;
let num_block =
(bit_size as f64 / (param.message_modulus().0 as f64).log(2.0)).ceil() as u64;

Some((param, num_block, bit_size))
}
}

pub fn unsigned_oprf(c: &mut Criterion) {
let bench_name = "integer::unsigned_oprf";

let mut bench_group = c.benchmark_group(bench_name);
bench_group
.sample_size(15)
.measurement_time(std::time::Duration::from_secs(60));
.measurement_time(std::time::Duration::from_secs(30));

for (param, num_block, bit_size) in ParamsAndNumBlocksIter::default() {
let (_, sk) = KEY_CACHE.get_from_params(param, IntegerKeyKind::Radix);
Expand All @@ -54,8 +21,8 @@ pub fn unsigned_oprf(c: &mut Criterion) {
b.iter(|| {
_ = black_box(sk.par_generate_oblivious_pseudo_random_unsigned_integer(
Seed(0),
bit_size,
num_block,
bit_size as u64,
num_block as u64,
));
})
});
Expand Down
60 changes: 1 addition & 59 deletions tfhe/benches/integer/signed_bench.rs
View file
Original file line number Diff line number Diff line change
@@ -1,21 +1,14 @@
#[path = "../utilities.rs"]
mod utilities;

use crate::utilities::{write_to_json, EnvConfig, OperatorType};
use crate::utilities::{write_to_json, EnvConfig, OperatorType, ParamsAndNumBlocksIter};
use criterion::{criterion_group, Criterion};
use itertools::iproduct;
use rand::prelude::*;
use std::env;
use std::vec::IntoIter;
use tfhe::integer::keycache::KEY_CACHE;
use tfhe::integer::prelude::*;
use tfhe::integer::{IntegerKeyKind, RadixCiphertext, ServerKey, SignedRadixCiphertext, I256};
use tfhe::keycache::NamedParam;
#[cfg(feature = "gpu")]
use tfhe::shortint::parameters::PARAM_GPU_MULTI_BIT_MESSAGE_2_CARRY_2_GROUP_3_KS_PBS;
use tfhe::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS_PBS;
#[cfg(not(feature = "gpu"))]
use tfhe::shortint::parameters::PARAM_MULTI_BIT_MESSAGE_2_CARRY_2_GROUP_2_KS_PBS;

fn gen_random_i256(rng: &mut ThreadRng) -> I256 {
let clearlow = rng.gen::<u128>();
Expand All @@ -24,57 +17,6 @@ fn gen_random_i256(rng: &mut ThreadRng) -> I256 {
tfhe::integer::I256::from((clearlow, clearhigh))
}

/// An iterator that yields a succession of combinations
/// of parameters and a num_block to achieve a certain bit_size ciphertext
/// in radix decomposition
struct ParamsAndNumBlocksIter {
params_and_bit_sizes:
itertools::Product<IntoIter<tfhe::shortint::PBSParameters>, IntoIter<usize>>,
}

impl Default for ParamsAndNumBlocksIter {
fn default() -> Self {
let env_config = EnvConfig::new();

if env_config.is_multi_bit {
#[cfg(feature = "gpu")]
let params = vec![PARAM_GPU_MULTI_BIT_MESSAGE_2_CARRY_2_GROUP_3_KS_PBS.into()];
#[cfg(not(feature = "gpu"))]
let params = vec![PARAM_MULTI_BIT_MESSAGE_2_CARRY_2_GROUP_2_KS_PBS.into()];

let params_and_bit_sizes = iproduct!(params, env_config.bit_sizes());
Self {
params_and_bit_sizes,
}
} else {
// FIXME One set of parameter is tested since we want to benchmark only quickest
// operations.
let params = vec![
PARAM_MESSAGE_2_CARRY_2_KS_PBS.into(),
// PARAM_MESSAGE_3_CARRY_3_KS_PBS.into(),
// PARAM_MESSAGE_4_CARRY_4_KS_PBS.into(),
];

let params_and_bit_sizes = iproduct!(params, env_config.bit_sizes());
Self {
params_and_bit_sizes,
}
}
}
}

impl Iterator for ParamsAndNumBlocksIter {
type Item = (tfhe::shortint::PBSParameters, usize, usize);

fn next(&mut self) -> Option<Self::Item> {
let (param, bit_size) = self.params_and_bit_sizes.next()?;
let num_block =
(bit_size as f64 / param.message_modulus().0.ilog2() as f64).ceil() as usize;

Some((param, num_block, bit_size))
}
}

/// Base function to bench a server key function that is a binary operation, input ciphertext will
/// contain only zero carries
fn bench_server_key_signed_binary_function_clean_inputs<F>(
Expand Down
54 changes: 54 additions & 0 deletions tfhe/benches/utilities.rs
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,18 @@
use itertools::iproduct;
use serde::Serialize;
use std::path::PathBuf;
use std::vec::IntoIter;
use std::{env, fs};
#[cfg(feature = "boolean")]
use tfhe::boolean::parameters::BooleanParameters;
use tfhe::core_crypto::prelude::*;
#[cfg(feature = "shortint")]
use tfhe::shortint::parameters::ShortintKeySwitchingParameters;
#[cfg(feature = "gpu")]
use tfhe::shortint::parameters::PARAM_GPU_MULTI_BIT_MESSAGE_2_CARRY_2_GROUP_3_KS_PBS;
use tfhe::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS_PBS;
#[cfg(not(feature = "gpu"))]
use tfhe::shortint::parameters::PARAM_MULTI_BIT_MESSAGE_2_CARRY_2_GROUP_2_KS_PBS;
#[cfg(feature = "shortint")]
use tfhe::shortint::PBSParameters;

Expand Down Expand Up @@ -268,6 +275,53 @@ impl EnvConfig {
}
}

/// An iterator that yields a succession of combinations
/// of parameters and a num_block to achieve a certain bit_size ciphertext
/// in radix decomposition
pub struct ParamsAndNumBlocksIter {
params_and_bit_sizes:
itertools::Product<IntoIter<tfhe::shortint::PBSParameters>, IntoIter<usize>>,
}

impl Default for ParamsAndNumBlocksIter {
fn default() -> Self {
let env_config = EnvConfig::new();

if env_config.is_multi_bit {
#[cfg(feature = "gpu")]
let params = vec![PARAM_GPU_MULTI_BIT_MESSAGE_2_CARRY_2_GROUP_3_KS_PBS.into()];
#[cfg(not(feature = "gpu"))]
let params = vec![PARAM_MULTI_BIT_MESSAGE_2_CARRY_2_GROUP_2_KS_PBS.into()];

let params_and_bit_sizes = iproduct!(params, env_config.bit_sizes());
Self {
params_and_bit_sizes,
}
} else {
// FIXME One set of parameter is tested since we want to benchmark only quickest
// operations.
let params = vec![PARAM_MESSAGE_2_CARRY_2_KS_PBS.into()];

let params_and_bit_sizes = iproduct!(params, env_config.bit_sizes());
Self {
params_and_bit_sizes,
}
}
}
}

impl Iterator for ParamsAndNumBlocksIter {
type Item = (tfhe::shortint::PBSParameters, usize, usize);

fn next(&mut self) -> Option<Self::Item> {
let (param, bit_size) = self.params_and_bit_sizes.next()?;
let num_block =
(bit_size as f64 / (param.message_modulus().0 as f64).log(2.0)).ceil() as usize;

Some((param, num_block, bit_size))
}
}

// Empty main to please clippy.
#[allow(dead_code)]
pub fn main() {}

0 comments on commit a2c4d31

Please sign in to comment.