Skip to content

Commit

Permalink
chore(gpu): downgrade compression conversion tests to become doc tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@pdroalves
pdroalves committed Sep 11, 2024
1 parent 9c979aa commit d01f700
Showing 1 changed file with 156 additions and 153 deletions.
309 changes: 156 additions & 153 deletions tfhe/src/integer/gpu/ciphertext/compressed_ciphertext_list.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,9 +15,8 @@ use crate::integer::gpu::list_compression::server_keys::{
CudaCompressionKey, CudaDecompressionKey, CudaPackedGlweCiphertext,
};
use crate::shortint::ciphertext::{
CompressedCiphertextList as ShortintCompressedCiphertextList, NoiseLevel,
CompressedCiphertextList as ShortintCompressedCiphertextList, Degree, NoiseLevel,
};
use crate::shortint::server_key::generate_lookup_table;
use itertools::Itertools;

pub struct CudaCompressedCiphertextList {
Expand Down Expand Up @@ -59,6 +58,91 @@ impl CudaCompressedCiphertextList {
)
}

/// ```rust
/// use tfhe::core_crypto::gpu::CudaStreams;
/// use tfhe::integer::{BooleanBlock, ClientKey, RadixCiphertext, SignedRadixCiphertext};
/// use tfhe::integer::ciphertext::CompressedCiphertextListBuilder;
/// use tfhe::integer::gpu::ciphertext::boolean_value::CudaBooleanBlock;
/// use tfhe::integer::gpu::ciphertext::compressed_ciphertext_list::CudaCompressedCiphertextListBuilder;
/// use tfhe::integer::gpu::ciphertext::{CudaSignedRadixCiphertext, CudaUnsignedRadixCiphertext};
/// use tfhe::integer::gpu::gen_keys_radix_gpu;
/// use tfhe::shortint::parameters::list_compression::COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64;
/// use tfhe::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64;
///
/// let cks = ClientKey::new(PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64);
///
/// let private_compression_key =
/// cks.new_compression_private_key(COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64);
///
/// let streams = CudaStreams::new_multi_gpu();
///
/// let num_blocks = 32;
/// let (radix_cks, _) = gen_keys_radix_gpu(
/// PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64,
/// num_blocks,
/// &streams,
/// );
/// 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(&streams);
///
/// let compression_key = compressed_compression_key.decompress();
/// let decompression_key = compressed_decompression_key.decompress();
///
/// let ct1 = radix_cks.encrypt(3_u32);
/// let ct2 = radix_cks.encrypt_signed(-2);
/// let ct3 = radix_cks.encrypt_bool(true);
///
/// /// Copy to GPU
/// let d_ct1 = CudaUnsignedRadixCiphertext::from_radix_ciphertext(&ct1, &streams);
/// let d_ct2 = CudaSignedRadixCiphertext::from_signed_radix_ciphertext(&ct2, &streams);
/// let d_ct3 = CudaBooleanBlock::from_boolean_block(&ct3, &streams);
///
/// let cuda_compressed = CudaCompressedCiphertextListBuilder::new()
/// .push(d_ct1, &streams)
/// .push(d_ct2, &streams)
/// .push(d_ct3, &streams)
/// .build(&cuda_compression_key, &streams);
///
/// let reference_compressed = CompressedCiphertextListBuilder::new()
/// .push(ct1)
/// .push(ct2)
/// .push(ct3)
/// .build(&compression_key);
///
/// let converted_compressed = cuda_compressed.to_compressed_ciphertext_list(&streams);
///
/// let decompressed1: RadixCiphertext = converted_compressed
/// .get(0, &decompression_key)
/// .unwrap()
/// .unwrap();
/// let reference_decompressed1 = reference_compressed
/// .get(0, &decompression_key)
/// .unwrap()
/// .unwrap();
/// assert_eq!(decompressed1, reference_decompressed1);
///
/// let decompressed2: SignedRadixCiphertext = converted_compressed
/// .get(1, &decompression_key)
/// .unwrap()
/// .unwrap();
/// let reference_decompressed2 = reference_compressed
/// .get(1, &decompression_key)
/// .unwrap()
/// .unwrap();
/// assert_eq!(decompressed2, reference_decompressed2);
///
/// let decompressed3: BooleanBlock = converted_compressed
/// .get(2, &decompression_key)
/// .unwrap()
/// .unwrap();
/// let reference_decompressed3 = reference_compressed
/// .get(2, &decompression_key)
/// .unwrap()
/// .unwrap();
/// assert_eq!(decompressed3, reference_decompressed3);
/// ```
pub fn to_compressed_ciphertext_list(&self, streams: &CudaStreams) -> CompressedCiphertextList {
let glwe_list = self
.packed_list
Expand Down Expand Up @@ -115,6 +199,70 @@ impl CudaCompressedCiphertextList {
}

impl CompressedCiphertextList {
/// ```rust
/// use tfhe::core_crypto::gpu::CudaStreams;
/// use tfhe::integer::ciphertext::CompressedCiphertextListBuilder;
/// use tfhe::integer::ClientKey;
/// use tfhe::integer::gpu::ciphertext::{CudaSignedRadixCiphertext, CudaUnsignedRadixCiphertext};
/// use tfhe::integer::gpu::ciphertext::boolean_value::CudaBooleanBlock;
/// use tfhe::integer::gpu::gen_keys_radix_gpu;
/// use tfhe::shortint::parameters::list_compression::COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64;
/// use tfhe::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64;
///
/// let cks = ClientKey::new(PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64);
///
/// let private_compression_key =
/// cks.new_compression_private_key(COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64);
///
/// let streams = CudaStreams::new_multi_gpu();
///
/// let num_blocks = 32;
/// let (radix_cks, _) = gen_keys_radix_gpu(
/// PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64,
/// num_blocks,
/// &streams,
/// );
/// let (compressed_compression_key, compressed_decompression_key) =
/// radix_cks.new_compressed_compression_decompression_keys(&private_compression_key);
///
/// let cuda_decompression_key =
/// compressed_decompression_key.decompress_to_cuda(radix_cks.parameters(), &streams);
///
/// let compression_key = compressed_compression_key.decompress();
///
/// let ct1 = radix_cks.encrypt(3_u32);
/// let ct2 = radix_cks.encrypt_signed(-2);
/// let ct3 = radix_cks.encrypt_bool(true);
///
/// let compressed = CompressedCiphertextListBuilder::new()
/// .push(ct1)
/// .push(ct2)
/// .push(ct3)
/// .build(&compression_key);
///
/// let cuda_compressed = compressed.to_cuda_compressed_ciphertext_list(&streams);
///
/// let d_decompressed1 = CudaUnsignedRadixCiphertext {
/// ciphertext: cuda_compressed.get(0, &cuda_decompression_key, &streams),
/// };
/// let decompressed1 = d_decompressed1.to_radix_ciphertext(&streams);
/// let decrypted: u32 = radix_cks.decrypt(&decompressed1);
/// assert_eq!(decrypted, 3_u32);
///
/// let d_decompressed2 = CudaSignedRadixCiphertext {
/// ciphertext: cuda_compressed.get(1, &cuda_decompression_key, &streams),
/// };
/// let decompressed2 = d_decompressed2.to_signed_radix_ciphertext(&streams);
/// let decrypted: i32 = radix_cks.decrypt_signed(&decompressed2);
/// assert_eq!(decrypted, -2);
///
/// let d_decompressed3 = CudaBooleanBlock::from_cuda_radix_ciphertext(
/// cuda_compressed.get(2, &cuda_decompression_key, &streams),
/// );
/// let decompressed3 = d_decompressed3.to_boolean_block(&streams);
/// let decrypted = radix_cks.decrypt_bool(&decompressed3);
/// assert!(decrypted);
/// ```
pub fn to_cuda_compressed_ciphertext_list(
&self,
streams: &CudaStreams,
Expand All @@ -129,20 +277,13 @@ impl CompressedCiphertextList {
let initial_len = first_ct.packed_integers.initial_len;
let bodies_count = first_ct.bodies_count.0;

// To-do: is there a better way to calculate degree?
let carry_extract = generate_lookup_table(
first_ct.glwe_dimension.to_glwe_size(),
first_ct.polynomial_size,
self.packed_list.ciphertext_modulus,
self.packed_list.message_modulus,
self.packed_list.carry_modulus,
|x| x / self.packed_list.message_modulus.0 as u64,
);
let message_modulus = self.packed_list.message_modulus;
let carry_modulus = self.packed_list.carry_modulus;

let first_block_info = CudaBlockInfo {
degree: carry_extract.degree,
message_modulus: self.packed_list.message_modulus,
carry_modulus: self.packed_list.carry_modulus,
degree: Degree::new(message_modulus.0 * carry_modulus.0 - 1),
message_modulus: message_modulus,
carry_modulus: carry_modulus,
pbs_order: self.packed_list.pbs_order,
noise_level: NoiseLevel::NOMINAL,
};
Expand Down Expand Up @@ -270,9 +411,8 @@ impl CudaCompressedCiphertextListBuilder {
#[cfg(test)]
mod tests {
use super::*;
use crate::integer::ciphertext::CompressedCiphertextListBuilder;
use crate::integer::gpu::gen_keys_radix_gpu;
use crate::integer::{BooleanBlock, ClientKey, RadixCiphertext, SignedRadixCiphertext};
use crate::integer::ClientKey;
use crate::shortint::parameters::list_compression::COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64;
use crate::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64;

Expand Down Expand Up @@ -325,143 +465,6 @@ mod tests {
let decrypted: i32 = radix_cks.decrypt_signed(&decompressed2);
assert_eq!(decrypted, -2);

let d_decompressed3 = CudaBooleanBlock::from_cuda_radix_ciphertext(
cuda_compressed.get(2, &cuda_decompression_key, &streams),
);
let decompressed3 = d_decompressed3.to_boolean_block(&streams);
let decrypted = radix_cks.decrypt_bool(&decompressed3);
assert!(decrypted);
}
}
#[test]
fn test_gpu_compressed_ciphertext_conversion_to_cpu() {
let cks = ClientKey::new(PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64);

let private_compression_key =
cks.new_compression_private_key(COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64);

let streams = CudaStreams::new_multi_gpu();

let num_blocks = 32;
let (radix_cks, _) = gen_keys_radix_gpu(
PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64,
num_blocks,
&streams,
);
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(&streams);

let compression_key = compressed_compression_key.decompress();
let decompression_key = compressed_decompression_key.decompress();

for _ in 0..NB_TESTS {
let ct1 = radix_cks.encrypt(3_u32);
let ct2 = radix_cks.encrypt_signed(-2);
let ct3 = radix_cks.encrypt_bool(true);

// Copy to GPU
let d_ct1 = CudaUnsignedRadixCiphertext::from_radix_ciphertext(&ct1, &streams);
let d_ct2 = CudaSignedRadixCiphertext::from_signed_radix_ciphertext(&ct2, &streams);
let d_ct3 = CudaBooleanBlock::from_boolean_block(&ct3, &streams);

let cuda_compressed = CudaCompressedCiphertextListBuilder::new()
.push(d_ct1, &streams)
.push(d_ct2, &streams)
.push(d_ct3, &streams)
.build(&cuda_compression_key, &streams);

let reference_compressed = CompressedCiphertextListBuilder::new()
.push(ct1)
.push(ct2)
.push(ct3)
.build(&compression_key);

let converted_compressed = cuda_compressed.to_compressed_ciphertext_list(&streams);

let decompressed1: RadixCiphertext = converted_compressed
.get(0, &decompression_key)
.unwrap()
.unwrap();
let reference_decompressed1 = reference_compressed
.get(0, &decompression_key)
.unwrap()
.unwrap();
assert_eq!(decompressed1, reference_decompressed1);

let decompressed2: SignedRadixCiphertext = converted_compressed
.get(1, &decompression_key)
.unwrap()
.unwrap();
let reference_decompressed2 = reference_compressed
.get(1, &decompression_key)
.unwrap()
.unwrap();
assert_eq!(decompressed2, reference_decompressed2);

let decompressed3: BooleanBlock = converted_compressed
.get(2, &decompression_key)
.unwrap()
.unwrap();
let reference_decompressed3 = reference_compressed
.get(2, &decompression_key)
.unwrap()
.unwrap();
assert_eq!(decompressed3, reference_decompressed3);
}
}

#[test]
fn test_gpu_compressed_ciphertext_conversion_to_gpu() {
let cks = ClientKey::new(PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64);

let private_compression_key =
cks.new_compression_private_key(COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64);

let streams = CudaStreams::new_multi_gpu();

let num_blocks = 32;
let (radix_cks, _) = gen_keys_radix_gpu(
PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M64,
num_blocks,
&streams,
);
let (compressed_compression_key, compressed_decompression_key) =
radix_cks.new_compressed_compression_decompression_keys(&private_compression_key);

let cuda_decompression_key =
compressed_decompression_key.decompress_to_cuda(radix_cks.parameters(), &streams);

let compression_key = compressed_compression_key.decompress();

for _ in 0..NB_TESTS {
let ct1 = radix_cks.encrypt(3_u32);
let ct2 = radix_cks.encrypt_signed(-2);
let ct3 = radix_cks.encrypt_bool(true);

let compressed = CompressedCiphertextListBuilder::new()
.push(ct1)
.push(ct2)
.push(ct3)
.build(&compression_key);

let cuda_compressed = compressed.to_cuda_compressed_ciphertext_list(&streams);

let d_decompressed1 = CudaUnsignedRadixCiphertext {
ciphertext: cuda_compressed.get(0, &cuda_decompression_key, &streams),
};
let decompressed1 = d_decompressed1.to_radix_ciphertext(&streams);
let decrypted: u32 = radix_cks.decrypt(&decompressed1);
assert_eq!(decrypted, 3_u32);

let d_decompressed2 = CudaSignedRadixCiphertext {
ciphertext: cuda_compressed.get(1, &cuda_decompression_key, &streams),
};
let decompressed2 = d_decompressed2.to_signed_radix_ciphertext(&streams);
let decrypted: i32 = radix_cks.decrypt_signed(&decompressed2);
assert_eq!(decrypted, -2);

let d_decompressed3 = CudaBooleanBlock::from_cuda_radix_ciphertext(
cuda_compressed.get(2, &cuda_decompression_key, &streams),
);
Expand Down

0 comments on commit d01f700

Please sign in to comment.