mirror of
https://github.com/zama-ai/tfhe-rs.git
synced 2026-01-08 22:28:01 -05:00
445 lines
18 KiB
Rust
445 lines
18 KiB
Rust
#[cfg(feature = "gpu")]
|
|
mod cuda {
|
|
use benchmark::params_aliases::*;
|
|
use benchmark::utilities::cuda_integer_utils::cuda_local_streams;
|
|
use benchmark::utilities::{
|
|
cuda_local_keys, get_bench_type, write_to_json, BenchmarkType, OperatorType,
|
|
};
|
|
use criterion::{black_box, criterion_group, Criterion, Throughput};
|
|
use rayon::prelude::*;
|
|
use tfhe::core_crypto::gpu::{get_number_of_gpus, CudaStreams};
|
|
use tfhe::integer::ciphertext::{
|
|
NoiseSquashingCompressionKey, NoiseSquashingCompressionPrivateKey,
|
|
};
|
|
use tfhe::integer::gpu::ciphertext::squashed_noise::CudaSquashedNoiseRadixCiphertext;
|
|
use tfhe::integer::gpu::ciphertext::{
|
|
CudaCompressedSquashedNoiseCiphertextList, CudaUnsignedRadixCiphertext,
|
|
};
|
|
use tfhe::integer::gpu::gen_keys_radix_gpu;
|
|
use tfhe::integer::gpu::list_compression::server_keys::CudaNoiseSquashingCompressionKey;
|
|
use tfhe::integer::gpu::noise_squashing::keys::CudaNoiseSquashingKey;
|
|
use tfhe::integer::noise_squashing::{CompressedNoiseSquashingKey, NoiseSquashingPrivateKey};
|
|
use tfhe::integer::ClientKey;
|
|
use tfhe::keycache::NamedParam;
|
|
use tfhe::shortint::parameters::NoiseSquashingCompressionParameters;
|
|
use tfhe::shortint::PBSParameters;
|
|
|
|
#[derive(Clone)]
|
|
struct BenchConfig {
|
|
param: PBSParameters,
|
|
noise_squashing_compression_parameters: NoiseSquashingCompressionParameters,
|
|
noise_squashing_compression_key: NoiseSquashingCompressionKey,
|
|
compressed_noise_squashing_compression_key: CompressedNoiseSquashingKey,
|
|
bit_size: usize,
|
|
cks: ClientKey,
|
|
}
|
|
|
|
fn get_num_elements_per_gpu(_bit_size: usize) -> usize {
|
|
// 200 elements per GPU seems enough to saturate H100s
|
|
// This is an empirical value and might need to be adjusted in the future
|
|
200
|
|
}
|
|
|
|
fn execute_gpu_glwe_packing_128(c: &mut Criterion, config: BenchConfig) {
|
|
let bench_name = "integer::cuda::128b_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 BenchConfig {
|
|
param,
|
|
noise_squashing_compression_parameters,
|
|
noise_squashing_compression_key,
|
|
compressed_noise_squashing_compression_key,
|
|
bit_size,
|
|
cks,
|
|
} = config;
|
|
|
|
let log_message_modulus = param.message_modulus().0.ilog2() as usize;
|
|
|
|
assert_eq!(bit_size % log_message_modulus, 0);
|
|
let num_blocks = bit_size / log_message_modulus;
|
|
|
|
let bench_id_pack;
|
|
|
|
match get_bench_type() {
|
|
BenchmarkType::Latency => {
|
|
let (_, cuda_sks) = gen_keys_radix_gpu(param, num_blocks, &stream);
|
|
let cuda_noise_squashing_key =
|
|
compressed_noise_squashing_compression_key.decompress_to_cuda(&stream);
|
|
|
|
// Encrypt
|
|
let ct = cks.encrypt_radix(0_u32, num_blocks);
|
|
let d_ct = CudaUnsignedRadixCiphertext::from_radix_ciphertext(&ct, &stream);
|
|
let d_ns_ct = cuda_noise_squashing_key
|
|
.squash_radix_ciphertext_noise(&cuda_sks, &d_ct.ciphertext, &stream)
|
|
.unwrap();
|
|
let cuda_noise_squashing_compression_key =
|
|
CudaNoiseSquashingCompressionKey::from_noise_squashing_compression_key(
|
|
&noise_squashing_compression_key,
|
|
&stream,
|
|
);
|
|
|
|
// Benchmark
|
|
let mut builder = CudaCompressedSquashedNoiseCiphertextList::builder();
|
|
|
|
builder.push(d_ns_ct, &stream);
|
|
|
|
bench_id_pack = format!("{bench_name}::pack_u{bit_size}");
|
|
bench_group.bench_function(&bench_id_pack, |b| {
|
|
b.iter(|| {
|
|
let compressed =
|
|
builder.build(&cuda_noise_squashing_compression_key, &stream);
|
|
|
|
_ = black_box(compressed);
|
|
})
|
|
});
|
|
}
|
|
BenchmarkType::Throughput => {
|
|
let cuda_sks = cuda_local_keys(&cks);
|
|
let num_block =
|
|
(bit_size as f64 / (param.message_modulus().0 as f64).log(2.0)).ceil() as usize;
|
|
let elements = get_num_elements_per_gpu(bit_size) as u64;
|
|
bench_group.throughput(Throughput::Elements(elements));
|
|
|
|
// Encrypt
|
|
let local_streams = cuda_local_streams(num_block, elements as usize);
|
|
|
|
let num_gpus = get_number_of_gpus() as usize;
|
|
|
|
let cuda_compression_key_vec: Vec<CudaNoiseSquashingKey> = (0..num_gpus)
|
|
.into_par_iter()
|
|
.map(|i| {
|
|
let local_stream = &local_streams[i % local_streams.len()];
|
|
compressed_noise_squashing_compression_key.decompress_to_cuda(local_stream)
|
|
})
|
|
.collect();
|
|
let cuda_noise_squashing_compression_key_vec: Vec<
|
|
CudaNoiseSquashingCompressionKey,
|
|
> = (0..num_gpus)
|
|
.into_par_iter()
|
|
.map(|i| {
|
|
let local_stream = &local_streams[i % local_streams.len()];
|
|
CudaNoiseSquashingCompressionKey::from_noise_squashing_compression_key(
|
|
&noise_squashing_compression_key,
|
|
local_stream,
|
|
)
|
|
})
|
|
.collect();
|
|
|
|
// Benchmark
|
|
let builders = (0..elements)
|
|
.into_par_iter()
|
|
.map(|i| {
|
|
let ct = cks.encrypt_radix(0_u32, num_blocks);
|
|
let local_stream = &local_streams[i as usize % local_streams.len()];
|
|
let d_ct =
|
|
CudaUnsignedRadixCiphertext::from_radix_ciphertext(&ct, local_stream);
|
|
let cuda_noise_squashing_key =
|
|
&cuda_compression_key_vec[(i as usize) % num_gpus];
|
|
let cuda_noise_squashing_compression_key =
|
|
&cuda_noise_squashing_compression_key_vec[(i as usize) % num_gpus];
|
|
let d_ns_ct = cuda_noise_squashing_key
|
|
.squash_radix_ciphertext_noise(
|
|
&cuda_sks[(i as usize) % num_gpus],
|
|
&d_ct.ciphertext,
|
|
local_stream,
|
|
)
|
|
.unwrap();
|
|
let mut builder = CudaCompressedSquashedNoiseCiphertextList::builder();
|
|
builder.push(d_ns_ct, local_stream);
|
|
|
|
(builder, cuda_noise_squashing_compression_key, local_stream)
|
|
})
|
|
.collect::<Vec<_>>();
|
|
|
|
bench_id_pack = format!("{bench_name}::throughput::pack_u{bit_size}");
|
|
bench_group.bench_function(&bench_id_pack, |b| {
|
|
b.iter(|| {
|
|
builders.par_iter().for_each(
|
|
|(builder, cuda_noise_squashing_compression_key, local_stream)| {
|
|
builder.build(cuda_noise_squashing_compression_key, local_stream);
|
|
},
|
|
)
|
|
})
|
|
});
|
|
}
|
|
}
|
|
|
|
write_to_json::<u64, _>(
|
|
&bench_id_pack,
|
|
(noise_squashing_compression_parameters, param.into()),
|
|
noise_squashing_compression_parameters.name(),
|
|
"pack",
|
|
&OperatorType::Atomic,
|
|
bit_size as u32,
|
|
vec![param.message_modulus().0.ilog2(); num_blocks],
|
|
);
|
|
|
|
bench_group.finish()
|
|
}
|
|
|
|
fn execute_gpu_glwe_unpacking_128(c: &mut Criterion, config: BenchConfig) {
|
|
let bench_name = "integer::cuda::128b_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 BenchConfig {
|
|
param,
|
|
noise_squashing_compression_parameters,
|
|
noise_squashing_compression_key,
|
|
compressed_noise_squashing_compression_key,
|
|
bit_size,
|
|
cks,
|
|
} = config;
|
|
|
|
let log_message_modulus = param.message_modulus().0.ilog2() as usize;
|
|
|
|
assert_eq!(bit_size % log_message_modulus, 0);
|
|
let num_blocks = bit_size / log_message_modulus;
|
|
|
|
let bench_id_unpack;
|
|
|
|
match get_bench_type() {
|
|
BenchmarkType::Latency => {
|
|
let (_, cuda_sks) = gen_keys_radix_gpu(param, num_blocks, &stream);
|
|
let cuda_noise_squashing_key =
|
|
compressed_noise_squashing_compression_key.decompress_to_cuda(&stream);
|
|
|
|
// Encrypt
|
|
let ct = cks.encrypt_radix(0_u32, num_blocks);
|
|
let d_ct = CudaUnsignedRadixCiphertext::from_radix_ciphertext(&ct, &stream);
|
|
let d_ns_ct = cuda_noise_squashing_key
|
|
.squash_radix_ciphertext_noise(&cuda_sks, &d_ct.ciphertext, &stream)
|
|
.unwrap();
|
|
let cuda_noise_squashing_compression_key =
|
|
CudaNoiseSquashingCompressionKey::from_noise_squashing_compression_key(
|
|
&noise_squashing_compression_key,
|
|
&stream,
|
|
);
|
|
|
|
// Benchmark
|
|
let mut builder = CudaCompressedSquashedNoiseCiphertextList::builder();
|
|
|
|
builder.push(d_ns_ct, &stream);
|
|
|
|
let compressed = builder.build(&cuda_noise_squashing_compression_key, &stream);
|
|
|
|
bench_id_unpack = format!("{bench_name}::unpack_u{bit_size}");
|
|
bench_group.bench_function(&bench_id_unpack, |b| {
|
|
b.iter(|| {
|
|
let unpacked: CudaSquashedNoiseRadixCiphertext =
|
|
compressed.get(0, &stream).unwrap().unwrap();
|
|
|
|
_ = black_box(unpacked);
|
|
})
|
|
});
|
|
}
|
|
BenchmarkType::Throughput => {
|
|
let cuda_sks = cuda_local_keys(&cks);
|
|
let num_block =
|
|
(bit_size as f64 / (param.message_modulus().0 as f64).log(2.0)).ceil() as usize;
|
|
let elements = get_num_elements_per_gpu(bit_size) as u64;
|
|
bench_group.throughput(Throughput::Elements(elements));
|
|
|
|
// Encrypt
|
|
let local_streams = cuda_local_streams(num_block, elements as usize);
|
|
|
|
let num_gpus = get_number_of_gpus() as usize;
|
|
|
|
let cuda_compression_key_vec: Vec<CudaNoiseSquashingKey> = (0..num_gpus)
|
|
.into_par_iter()
|
|
.map(|i| {
|
|
let local_stream = &local_streams[i % local_streams.len()];
|
|
compressed_noise_squashing_compression_key.decompress_to_cuda(local_stream)
|
|
})
|
|
.collect();
|
|
let cuda_noise_squashing_compression_key_vec: Vec<
|
|
CudaNoiseSquashingCompressionKey,
|
|
> = (0..num_gpus)
|
|
.into_par_iter()
|
|
.map(|i| {
|
|
let local_stream = &local_streams[i % local_streams.len()];
|
|
CudaNoiseSquashingCompressionKey::from_noise_squashing_compression_key(
|
|
&noise_squashing_compression_key,
|
|
local_stream,
|
|
)
|
|
})
|
|
.collect();
|
|
|
|
// Benchmark
|
|
let builders = (0..elements)
|
|
.into_par_iter()
|
|
.map(|i| {
|
|
let ct = cks.encrypt_radix(0_u32, num_blocks);
|
|
let local_stream = &local_streams[i as usize % local_streams.len()];
|
|
let d_ct =
|
|
CudaUnsignedRadixCiphertext::from_radix_ciphertext(&ct, local_stream);
|
|
let cuda_noise_squashing_key =
|
|
&cuda_compression_key_vec[(i as usize) % num_gpus];
|
|
let cuda_noise_squashing_compression_key =
|
|
&cuda_noise_squashing_compression_key_vec[(i as usize) % num_gpus];
|
|
let d_ns_ct = cuda_noise_squashing_key
|
|
.squash_radix_ciphertext_noise(
|
|
&cuda_sks[(i as usize) % num_gpus],
|
|
&d_ct.ciphertext,
|
|
local_stream,
|
|
)
|
|
.unwrap();
|
|
let mut builder = CudaCompressedSquashedNoiseCiphertextList::builder();
|
|
builder.push(d_ns_ct, local_stream);
|
|
|
|
(builder, cuda_noise_squashing_compression_key, local_stream)
|
|
})
|
|
.collect::<Vec<_>>();
|
|
|
|
let compressed = builders
|
|
.into_par_iter()
|
|
.map(
|
|
|(builder, cuda_noise_squashing_compression_key, local_stream)| {
|
|
builder.build(cuda_noise_squashing_compression_key, local_stream)
|
|
},
|
|
)
|
|
.collect::<Vec<_>>();
|
|
|
|
bench_id_unpack = format!("{bench_name}::throughput::unpack_u{bit_size}");
|
|
bench_group.bench_function(&bench_id_unpack, |b| {
|
|
b.iter(|| {
|
|
compressed.par_iter().enumerate().for_each(|(i, comp)| {
|
|
let local_stream = &local_streams[i % local_streams.len()];
|
|
|
|
comp.get::<CudaSquashedNoiseRadixCiphertext>(0, local_stream)
|
|
.unwrap()
|
|
.unwrap();
|
|
})
|
|
})
|
|
});
|
|
}
|
|
}
|
|
|
|
write_to_json::<u64, _>(
|
|
&bench_id_unpack,
|
|
(noise_squashing_compression_parameters, param.into()),
|
|
noise_squashing_compression_parameters.name(),
|
|
"unpack",
|
|
&OperatorType::Atomic,
|
|
bit_size as u32,
|
|
vec![param.message_modulus().0.ilog2(); num_blocks],
|
|
);
|
|
|
|
bench_group.finish()
|
|
}
|
|
|
|
fn gpu_glwe_packing_128(c: &mut Criterion) {
|
|
let param = BENCH_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
|
|
let noise_squashing_compression_parameters =
|
|
BENCH_COMP_NOISE_SQUASHING_PARAM_GPU_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
|
|
let noise_squashing_parameters =
|
|
BENCH_NOISE_SQUASHING_PARAM_GPU_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
|
|
|
|
let log_message_modulus = param.message_modulus.0.ilog2() as usize;
|
|
|
|
let cks = ClientKey::new(param);
|
|
|
|
let noise_squashing_compression_private_key =
|
|
NoiseSquashingCompressionPrivateKey::new(noise_squashing_compression_parameters);
|
|
let noise_squashing_private_key = NoiseSquashingPrivateKey::new(noise_squashing_parameters);
|
|
let noise_squashing_compression_key = noise_squashing_private_key
|
|
.new_noise_squashing_compression_key(&noise_squashing_compression_private_key);
|
|
|
|
// Generate and convert compression keys
|
|
let compressed_noise_squashing_compression_key =
|
|
cks.new_compressed_noise_squashing_key(&noise_squashing_private_key);
|
|
|
|
let mut config = BenchConfig {
|
|
param: PBSParameters::PBS(param),
|
|
noise_squashing_compression_key,
|
|
noise_squashing_compression_parameters,
|
|
compressed_noise_squashing_compression_key,
|
|
bit_size: 0,
|
|
cks,
|
|
};
|
|
for bit_size in [
|
|
2,
|
|
8,
|
|
16,
|
|
32,
|
|
64,
|
|
128,
|
|
256,
|
|
noise_squashing_compression_parameters.lwe_per_glwe.0 * log_message_modulus,
|
|
] {
|
|
config.bit_size = bit_size;
|
|
execute_gpu_glwe_packing_128(c, config.clone());
|
|
}
|
|
}
|
|
|
|
fn gpu_glwe_unpacking_128(c: &mut Criterion) {
|
|
let param = BENCH_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
|
|
let noise_squashing_compression_parameters =
|
|
BENCH_COMP_NOISE_SQUASHING_PARAM_GPU_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
|
|
let noise_squashing_parameters =
|
|
BENCH_NOISE_SQUASHING_PARAM_GPU_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
|
|
|
|
let log_message_modulus = param.message_modulus.0.ilog2() as usize;
|
|
|
|
let cks = ClientKey::new(param);
|
|
|
|
let noise_squashing_compression_private_key =
|
|
NoiseSquashingCompressionPrivateKey::new(noise_squashing_compression_parameters);
|
|
let noise_squashing_private_key = NoiseSquashingPrivateKey::new(noise_squashing_parameters);
|
|
let noise_squashing_compression_key = noise_squashing_private_key
|
|
.new_noise_squashing_compression_key(&noise_squashing_compression_private_key);
|
|
|
|
// Generate and convert compression keys
|
|
let compressed_noise_squashing_compression_key =
|
|
cks.new_compressed_noise_squashing_key(&noise_squashing_private_key);
|
|
|
|
let mut config = BenchConfig {
|
|
param: PBSParameters::PBS(param),
|
|
noise_squashing_compression_key,
|
|
noise_squashing_compression_parameters,
|
|
compressed_noise_squashing_compression_key,
|
|
bit_size: 0,
|
|
cks,
|
|
};
|
|
for bit_size in [
|
|
2,
|
|
8,
|
|
16,
|
|
32,
|
|
64,
|
|
128,
|
|
256,
|
|
noise_squashing_compression_parameters.lwe_per_glwe.0 * log_message_modulus,
|
|
] {
|
|
config.bit_size = bit_size;
|
|
execute_gpu_glwe_unpacking_128(c, config.clone());
|
|
}
|
|
}
|
|
|
|
criterion_group!(gpu_glwe_packing_128_2, gpu_glwe_packing_128);
|
|
criterion_group!(gpu_glwe_unpacking_128_2, gpu_glwe_unpacking_128);
|
|
}
|
|
|
|
use criterion::Criterion;
|
|
#[cfg(feature = "gpu")]
|
|
use cuda::gpu_glwe_packing_128_2;
|
|
#[cfg(feature = "gpu")]
|
|
use cuda::gpu_glwe_unpacking_128_2;
|
|
|
|
fn main() {
|
|
#[cfg(feature = "gpu")]
|
|
gpu_glwe_packing_128_2();
|
|
#[cfg(feature = "gpu")]
|
|
gpu_glwe_unpacking_128_2();
|
|
Criterion::default().configure_from_args().final_summary();
|
|
}
|