Skip to content

Commit

Permalink
chore(bench): make compression benchmarks available for database
Browse files Browse the repository at this point in the history
  • Loading branch information
@soonum
soonum committed Sep 13, 2024
1 parent 8314e7d commit 077d572
Show file tree
Hide file tree
Showing 6 changed files with 186 additions and 103 deletions.
6 changes: 6 additions & 0 deletions .github/workflows/benchmark_gpu_integer_full.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,12 @@ jobs:
run: |
make BENCH_OP_FLAVOR=${{ matrix.op_flavor }} bench_${{ matrix.command }}_gpu
# Run these benchmarks only once
- name: Run compression benchmarks with AVX512
if: matrix.op_flavor == 'default' && matrix.command == 'integer'
run: |
make bench_integer_compression_gpu
- name: Parse results
run: |
python3 ./ci/benchmark_parser.py target/criterion ${{ env.RESULTS_FILENAME }} \
Expand Down
6 changes: 6 additions & 0 deletions .github/workflows/benchmark_integer.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,12 @@ jobs:
run: |
make BENCH_OP_FLAVOR=${{ matrix.op_flavor }} bench_${{ matrix.command }}
# Run these benchmarks only once
- name: Run compression benchmarks with AVX512
if: matrix.op_flavor == 'default' && matrix.command == 'integer'
run: |
make bench_integer_compression
- name: Parse results
run: |
python3 ./ci/benchmark_parser.py target/criterion ${{ env.RESULTS_FILENAME }} \
Expand Down
6 changes: 6 additions & 0 deletions Makefile
View file
Original file line number Diff line number Diff line change
Expand Up @@ -893,6 +893,12 @@ bench_integer_gpu: install_rs_check_toolchain
--bench integer-bench \
--features=$(TARGET_ARCH_FEATURE),integer,gpu,internal-keycache,nightly-avx512 -p $(TFHE_SPEC) --

.PHONY: bench_integer_compression # Run benchmarks for unsigned integer compression
bench_integer_compression: install_rs_check_toolchain
RUSTFLAGS="$(RUSTFLAGS)" cargo $(CARGO_RS_CHECK_TOOLCHAIN) bench \
--bench glwe_packing_compression-integer-bench \
--features=$(TARGET_ARCH_FEATURE),integer,internal-keycache,nightly-avx512 -p $(TFHE_SPEC) --

.PHONY: bench_integer_compression_gpu
bench_integer_compression_gpu: install_rs_check_toolchain
RUSTFLAGS="$(RUSTFLAGS)" cargo $(CARGO_RS_CHECK_TOOLCHAIN) bench \
Expand Down
241 changes: 139 additions & 102 deletions tfhe/benches/integer/glwe_packing_compression.rs
View file
Original file line number Diff line number Diff line change
@@ -1,33 +1,24 @@
#[path = "../utilities.rs"]
mod utilities;

use crate::utilities::{write_to_json, OperatorType};
use criterion::{black_box, criterion_group, Criterion};
use tfhe::integer::ciphertext::CompressedCiphertextListBuilder;
use tfhe::integer::{ClientKey, RadixCiphertext};
use tfhe::keycache::NamedParam;
use tfhe::shortint::parameters::list_compression::COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
use tfhe::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;

#[cfg(feature = "gpu")]
use tfhe::core_crypto::gpu::CudaStreams;

#[cfg(feature = "gpu")]
use tfhe::integer::gpu::ciphertext::compressed_ciphertext_list::CudaCompressedCiphertextListBuilder;

#[cfg(feature = "gpu")]
use tfhe::integer::gpu::ciphertext::{CudaRadixCiphertext, CudaUnsignedRadixCiphertext};

#[cfg(feature = "gpu")]
use tfhe::integer::gpu::gen_keys_radix_gpu;

fn cpu_glwe_packing(c: &mut Criterion) {
let param = PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;

let comp_param = COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;

let bench_name = "integer_packing_compression";

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

let param = PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
let comp_param = COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;

let cks = ClientKey::new(param);

let private_compression_key = cks.new_compression_private_key(comp_param);
Expand All @@ -37,7 +28,7 @@ fn cpu_glwe_packing(c: &mut Criterion) {

let log_message_modulus = param.message_modulus.0.ilog2() as usize;

for num_bits in [
for bit_size in [
8,
16,
32,
Expand All @@ -46,26 +37,38 @@ fn cpu_glwe_packing(c: &mut Criterion) {
256,
comp_param.lwe_per_glwe.0 * log_message_modulus,
] {
assert_eq!(num_bits % log_message_modulus, 0);
let num_blocks = num_bits / log_message_modulus;
assert_eq!(bit_size % log_message_modulus, 0);
let num_blocks = bit_size / log_message_modulus;

let ct = cks.encrypt_radix(0_u32, num_blocks);

let mut builder = CompressedCiphertextListBuilder::new();

builder.push(ct);

bench_group.bench_function(format!("pack_u{num_bits}"), |b| {
let bench_id = format!("{bench_name}::pack_u{bit_size}");
bench_group.bench_function(&bench_id, |b| {
b.iter(|| {
let compressed = builder.build(&compression_key);

_ = black_box(compressed);
})
});

write_to_json::<u64, _>(
&bench_id,
comp_param,
comp_param.name(),
"pack",
&OperatorType::Atomic,
bit_size as u32,
vec![param.message_modulus.0.ilog2(); num_blocks],
);

let compressed = builder.build(&compression_key);

bench_group.bench_function(format!("unpack_u{num_bits}"), |b| {
let bench_id = format!("{bench_name}::unpack_u{bit_size}");
bench_group.bench_function(&bench_id, |b| {
b.iter(|| {
let unpacked: RadixCiphertext =
compressed.get(0, &decompression_key).unwrap().unwrap();
Expand All @@ -74,94 +77,128 @@ fn cpu_glwe_packing(c: &mut Criterion) {
})
});

bench_group.bench_function(format!("pack_unpack_u{num_bits}"), |b| {
b.iter(|| {
let compressed = builder.build(&compression_key);

let unpacked: RadixCiphertext =
compressed.get(0, &decompression_key).unwrap().unwrap();

_ = black_box(unpacked);
})
});
write_to_json::<u64, _>(
&bench_id,
comp_param,
comp_param.name(),
"unpack",
&OperatorType::Atomic,
bit_size as u32,
vec![param.message_modulus.0.ilog2(); num_blocks],
);
}

bench_group.finish()
}

#[cfg(feature = "gpu")]
fn gpu_glwe_packing(c: &mut Criterion) {
let bench_name = "integer_cuda_packing_compression";
let mut bench_group = c.benchmark_group(bench_name);
bench_group
.sample_size(15)
.measurement_time(std::time::Duration::from_secs(30));

let stream = CudaStreams::new_multi_gpu();

let param = PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
let comp_param = COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;

let log_message_modulus = param.message_modulus.0.ilog2() as usize;

for bit_size in [
8,
16,
32,
64,
128,
256,
comp_param.lwe_per_glwe.0 * log_message_modulus,
] {
assert_eq!(bit_size % log_message_modulus, 0);
let num_blocks = bit_size / log_message_modulus;

// Generate private compression key
let cks = ClientKey::new(param);
let private_compression_key = cks.new_compression_private_key(comp_param);

// Generate and convert compression keys
let (radix_cks, _) = gen_keys_radix_gpu(param, num_blocks, &stream);
let (compressed_compression_key, compressed_decompression_key) =
radix_cks.new_compressed_compression_decompression_keys(&private_compression_key);
let cuda_compression_key = compressed_compression_key.decompress_to_cuda(&stream);
let cuda_decompression_key =
compressed_decompression_key.decompress_to_cuda(radix_cks.parameters(), &stream);

// Encrypt
let ct = cks.encrypt_radix(0_u32, num_blocks);
let d_ct = CudaUnsignedRadixCiphertext::from_radix_ciphertext(&ct, &stream);

// Benchmark
let mut builder = CudaCompressedCiphertextListBuilder::new();

builder.push(d_ct, &stream);

let bench_id = format!("pack_u{bit_size}");
bench_group.bench_function(&bench_id, |b| {
b.iter(|| {
let compressed = builder.build(&cuda_compression_key, &stream);

_ = black_box(compressed);
})
});

let compressed = builder.build(&cuda_compression_key, &stream);

let bench_id = format!("unpack_u{bit_size}");
bench_group.bench_function(&bench_id, |b| {
b.iter(|| {
let unpacked: CudaRadixCiphertext =
compressed.get(0, &cuda_decompression_key, &stream);

_ = black_box(unpacked);
})
});
mod cuda {
use super::*;
use tfhe::core_crypto::gpu::CudaStreams;
use tfhe::integer::gpu::ciphertext::compressed_ciphertext_list::CudaCompressedCiphertextListBuilder;
use tfhe::integer::gpu::ciphertext::{CudaRadixCiphertext, CudaUnsignedRadixCiphertext};
use tfhe::integer::gpu::gen_keys_radix_gpu;

fn gpu_glwe_packing(c: &mut Criterion) {
let bench_name = "integer::cuda::packing_compression";
let mut bench_group = c.benchmark_group(bench_name);
bench_group
.sample_size(15)
.measurement_time(std::time::Duration::from_secs(30));

let stream = CudaStreams::new_multi_gpu();

let param = PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
let comp_param = COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;

let log_message_modulus = param.message_modulus.0.ilog2() as usize;

for bit_size in [
8,
16,
32,
64,
128,
256,
comp_param.lwe_per_glwe.0 * log_message_modulus,
] {
assert_eq!(bit_size % log_message_modulus, 0);
let num_blocks = bit_size / log_message_modulus;

// Generate private compression key
let cks = ClientKey::new(param);
let private_compression_key = cks.new_compression_private_key(comp_param);

// Generate and convert compression keys
let (radix_cks, _) = gen_keys_radix_gpu(param, num_blocks, &stream);
let (compressed_compression_key, compressed_decompression_key) =
radix_cks.new_compressed_compression_decompression_keys(&private_compression_key);
let cuda_compression_key = compressed_compression_key.decompress_to_cuda(&stream);
let cuda_decompression_key =
compressed_decompression_key.decompress_to_cuda(radix_cks.parameters(), &stream);

// Encrypt
let ct = cks.encrypt_radix(0_u32, num_blocks);
let d_ct = CudaUnsignedRadixCiphertext::from_radix_ciphertext(&ct, &stream);

// Benchmark
let mut builder = CudaCompressedCiphertextListBuilder::new();

builder.push(d_ct, &stream);

let bench_id = format!("{bench_name}::pack_u{bit_size}");
bench_group.bench_function(&bench_id, |b| {
b.iter(|| {
let compressed = builder.build(&cuda_compression_key, &stream);

_ = black_box(compressed);
})
});

write_to_json::<u64, _>(
&bench_id,
comp_param,
comp_param.name(),
"pack",
&OperatorType::Atomic,
bit_size as u32,
vec![param.message_modulus.0.ilog2(); num_blocks],
);

let compressed = builder.build(&cuda_compression_key, &stream);

let bench_id = format!("{bench_name}::unpack_u{bit_size}");
bench_group.bench_function(&bench_id, |b| {
b.iter(|| {
let unpacked: CudaRadixCiphertext =
compressed.get(0, &cuda_decompression_key, &stream);

_ = black_box(unpacked);
})
});

write_to_json::<u64, _>(
&bench_id,
comp_param,
comp_param.name(),
"unpack",
&OperatorType::Atomic,
bit_size as u32,
vec![param.message_modulus.0.ilog2(); num_blocks],
);
}

bench_group.finish()
}

criterion_group!(gpu_glwe_packing2, gpu_glwe_packing);
}

#[cfg(feature = "gpu")]
criterion_group!(gpu_glwe_packing2, gpu_glwe_packing);
criterion_group!(cpu_glwe_packing2, cpu_glwe_packing);

#[cfg(feature = "gpu")]
use cuda::gpu_glwe_packing2;

fn main() {
#[cfg(feature = "gpu")]
gpu_glwe_packing2();
Expand Down
28 changes: 28 additions & 0 deletions tfhe/src/shortint/keycache.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ use crate::shortint::parameters::classic::tuniform::p_fail_2_minus_64::ks_pbs::P
use crate::shortint::parameters::coverage_parameters::*;
use crate::shortint::parameters::key_switching::p_fail_2_minus_64::ks_pbs::PARAM_KEYSWITCH_1_1_KS_PBS_TO_2_2_KS_PBS;
use crate::shortint::parameters::key_switching::*;
use crate::shortint::parameters::list_compression::*;
use crate::shortint::parameters::multi_bit::*;
use crate::shortint::parameters::parameters_wopbs::*;
use crate::shortint::parameters::*;
Expand Down Expand Up @@ -286,6 +287,33 @@ impl NamedParam for ShortintKeySwitchingParameters {
}
}

impl NamedParam for CompressionParameters {
fn name(&self) -> String {
named_params_impl!(expose
COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64,
COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64
);
named_params_impl!(
{
*self;
Self
} == (COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64)
);

named_params_impl!(
{
*self;
Self
} == (COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64)
);

format!(
"COMP_PARAM_CUSTOM_BR_LEVEL_{}_NOISE_DISTRIB_{}",
self.br_level.0, self.packing_ks_key_noise_distribution
)
}
}

impl From<PBSParameters> for (ClientKey, ServerKey) {
fn from(param: PBSParameters) -> Self {
let param_set = ShortintParameterSet::from(param);
Expand Down
2 changes: 1 addition & 1 deletion tfhe/src/shortint/parameters/list_compression.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ use crate::shortint::parameters::{
};
use std::fmt::Debug;

#[derive(Copy, Clone, Debug, serde::Serialize, serde::Deserialize, Versionize)]
#[derive(Copy, Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize, Versionize)]
#[versionize(CompressionParametersVersions)]
pub struct CompressionParameters {
pub br_level: DecompositionLevelCount,
Expand Down

0 comments on commit 077d572

Please sign in to comment.