Skip to content

Commit

Permalink
chore(ci): switch gpu tests on push to l40 and deactivate fast pks
Browse files Browse the repository at this point in the history
  • Loading branch information
agnesLeroy committed Jan 6, 2025
1 parent 57a31d1 commit 0175f2c
Show file tree
Hide file tree
Showing 3 changed files with 131 additions and 129 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -31,8 +31,8 @@ __host__ inline bool can_use_pks_fast_path(uint32_t lwe_dimension,
uint32_t polynomial_size,
uint32_t level_count,
uint32_t glwe_dimension) {
// TODO: Generalize to level_count > 1 by transposing the KSK
return level_count == 1;
// TODO: activate it back, fix tests and extend to level_count > 1
return false;
}

// Initialize decomposition by performing rounding
Expand Down
2 changes: 1 addition & 1 deletion ci/slab.toml
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ instance_type = "m6i.4xlarge"
[backend.hyperstack.gpu-test]
environment_name = "canada"
image_name = "Ubuntu Server 22.04 LTS R535 CUDA 12.2"
flavor_name = "n3-RTX-A6000x1"
flavor_name = "n3-L40x1"

[backend.hyperstack.single-h100]
environment_name = "canada"
Expand Down
254 changes: 128 additions & 126 deletions tfhe/src/integer/gpu/ciphertext/compressed_ciphertext_list.rs
Original file line number Diff line number Diff line change
Expand Up @@ -726,130 +726,132 @@ mod tests {
}
}

#[test]
fn test_gpu_ciphertext_compression_fast_path() {
/// Implement a test only for the storage of ciphertexts
/// using a custom parameter set which is supported by a fast-path
/// packing keyswitch (only for level_count==1)
const COMP_PARAM_CUSTOM_FAST_PATH: CompressionParameters = CompressionParameters {
br_level: DecompositionLevelCount(1),
br_base_log: DecompositionBaseLog(21),
packing_ks_level: DecompositionLevelCount(1),
packing_ks_base_log: DecompositionBaseLog(19),
packing_ks_polynomial_size: PolynomialSize(2048),
packing_ks_glwe_dimension: GlweDimension(1),
lwe_per_glwe: LweCiphertextCount(2048),
storage_log_modulus: CiphertextModulusLog(55),
packing_ks_key_noise_distribution: DynamicDistribution::new_gaussian_from_std_dev(
StandardDev(2.845267479601915e-15),
),
};

const NUM_BLOCKS: usize = 32;

let streams = CudaStreams::new_multi_gpu();

let (radix_cks, sks) = gen_keys_radix_gpu(
PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
NUM_BLOCKS,
&streams,
);
let cks = radix_cks.as_ref();

let private_compression_key = cks.new_compression_private_key(COMP_PARAM_CUSTOM_FAST_PATH);

let (cuda_compression_key, cuda_decompression_key) =
radix_cks.new_cuda_compression_decompression_keys(&private_compression_key, &streams);

const MAX_NB_MESSAGES: usize = 2 * COMP_PARAM_CUSTOM_FAST_PATH.lwe_per_glwe.0 / NUM_BLOCKS;

let mut rng = rand::thread_rng();

let message_modulus: u128 = cks.parameters().message_modulus().0 as u128;

// Hybrid
enum MessageType {
Unsigned(u128),
Signed(i128),
Boolean(bool),
}
for _ in 0..NB_OPERATOR_TESTS {
let mut builder = CudaCompressedCiphertextListBuilder::new();

let nb_messages = rng.gen_range(1..=MAX_NB_MESSAGES as u64);
let mut messages = vec![];
for _ in 0..nb_messages {
let case_selector = rng.gen_range(0..3);
match case_selector {
0 => {
// Unsigned
let modulus = message_modulus.pow(NUM_BLOCKS as u32);
let message = rng.gen::<u128>() % modulus;
let ct = radix_cks.encrypt(message);
let d_ct =
CudaUnsignedRadixCiphertext::from_radix_ciphertext(&ct, &streams);
let d_and_ct = sks.bitand(&d_ct, &d_ct, &streams);
builder.push(d_and_ct, &streams);
messages.push(MessageType::Unsigned(message));
}
1 => {
// Signed
let modulus = message_modulus.pow((NUM_BLOCKS - 1) as u32) as i128;
let message = rng.gen::<i128>() % modulus;
let ct = radix_cks.encrypt_signed(message);
let d_ct =
CudaSignedRadixCiphertext::from_signed_radix_ciphertext(&ct, &streams);
let d_and_ct = sks.bitand(&d_ct, &d_ct, &streams);
builder.push(d_and_ct, &streams);
messages.push(MessageType::Signed(message));
}
_ => {
// Boolean
let message = rng.gen::<i64>() % 2 != 0;
let ct = radix_cks.encrypt_bool(message);
let d_boolean_ct = CudaBooleanBlock::from_boolean_block(&ct, &streams);
let d_ct = d_boolean_ct.0;
let d_and_boolean_ct =
CudaBooleanBlock::from_cuda_radix_ciphertext(d_ct.ciphertext);
builder.push(d_and_boolean_ct, &streams);
messages.push(MessageType::Boolean(message));
}
}
}

let cuda_compressed = builder.build(&cuda_compression_key, &streams);

for (i, val) in messages.iter().enumerate() {
match val {
MessageType::Unsigned(message) => {
let d_decompressed: CudaUnsignedRadixCiphertext = cuda_compressed
.get(i, &cuda_decompression_key, &streams)
.unwrap()
.unwrap();
let decompressed = d_decompressed.to_radix_ciphertext(&streams);
let decrypted: u128 = radix_cks.decrypt(&decompressed);
assert_eq!(decrypted, *message);
}
MessageType::Signed(message) => {
let d_decompressed: CudaSignedRadixCiphertext = cuda_compressed
.get(i, &cuda_decompression_key, &streams)
.unwrap()
.unwrap();
let decompressed = d_decompressed.to_signed_radix_ciphertext(&streams);
let decrypted: i128 = radix_cks.decrypt_signed(&decompressed);
assert_eq!(decrypted, *message);
}
MessageType::Boolean(message) => {
let d_decompressed: CudaBooleanBlock = cuda_compressed
.get(i, &cuda_decompression_key, &streams)
.unwrap()
.unwrap();
let decompressed = d_decompressed.to_boolean_block(&streams);
let decrypted = radix_cks.decrypt_bool(&decompressed);
assert_eq!(decrypted, *message);
}
}
}
}
}
//#[test]
//fn test_gpu_ciphertext_compression_fast_path() {
// /// Implement a test only for the storage of ciphertexts
// /// using a custom parameter set which is supported by a fast-path
// /// packing keyswitch (only for level_count==1)
// const COMP_PARAM_CUSTOM_FAST_PATH: CompressionParameters = CompressionParameters {
// br_level: DecompositionLevelCount(1),
// br_base_log: DecompositionBaseLog(21),
// packing_ks_level: DecompositionLevelCount(1),
// packing_ks_base_log: DecompositionBaseLog(19),
// packing_ks_polynomial_size: PolynomialSize(2048),
// packing_ks_glwe_dimension: GlweDimension(1),
// lwe_per_glwe: LweCiphertextCount(2048),
// storage_log_modulus: CiphertextModulusLog(55),
// packing_ks_key_noise_distribution: DynamicDistribution::new_gaussian_from_std_dev(
// StandardDev(2.845267479601915e-15),
// ),
// };

// const NUM_BLOCKS: usize = 32;

// let streams = CudaStreams::new_multi_gpu();

// let (radix_cks, sks) = gen_keys_radix_gpu(
// PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
// NUM_BLOCKS,
// &streams,
// );
// let cks = radix_cks.as_ref();

// let private_compression_key =
// cks.new_compression_private_key(COMP_PARAM_CUSTOM_FAST_PATH);

// let (cuda_compression_key, cuda_decompression_key) =
// radix_cks.new_cuda_compression_decompression_keys(&private_compression_key, &streams);

// const MAX_NB_MESSAGES: usize = 2 * COMP_PARAM_CUSTOM_FAST_PATH.lwe_per_glwe.0 /
// NUM_BLOCKS;

// let mut rng = rand::thread_rng();

// let message_modulus: u128 = cks.parameters().message_modulus().0 as u128;

// // Hybrid
// enum MessageType {
// Unsigned(u128),
// Signed(i128),
// Boolean(bool),
// }
// for _ in 0..NB_OPERATOR_TESTS {
// let mut builder = CudaCompressedCiphertextListBuilder::new();

// let nb_messages = rng.gen_range(1..=MAX_NB_MESSAGES as u64);
// let mut messages = vec![];
// for _ in 0..nb_messages {
// let case_selector = rng.gen_range(0..3);
// match case_selector {
// 0 => {
// // Unsigned
// let modulus = message_modulus.pow(NUM_BLOCKS as u32);
// let message = rng.gen::<u128>() % modulus;
// let ct = radix_cks.encrypt(message);
// let d_ct =
// CudaUnsignedRadixCiphertext::from_radix_ciphertext(&ct, &streams);
// let d_and_ct = sks.bitand(&d_ct, &d_ct, &streams);
// builder.push(d_and_ct, &streams);
// messages.push(MessageType::Unsigned(message));
// }
// 1 => {
// // Signed
// let modulus = message_modulus.pow((NUM_BLOCKS - 1) as u32) as i128;
// let message = rng.gen::<i128>() % modulus;
// let ct = radix_cks.encrypt_signed(message);
// let d_ct =
// CudaSignedRadixCiphertext::from_signed_radix_ciphertext(&ct,
// &streams); let d_and_ct = sks.bitand(&d_ct, &d_ct, &streams);
// builder.push(d_and_ct, &streams);
// messages.push(MessageType::Signed(message));
// }
// _ => {
// // Boolean
// let message = rng.gen::<i64>() % 2 != 0;
// let ct = radix_cks.encrypt_bool(message);
// let d_boolean_ct = CudaBooleanBlock::from_boolean_block(&ct, &streams);
// let d_ct = d_boolean_ct.0;
// let d_and_boolean_ct =
// CudaBooleanBlock::from_cuda_radix_ciphertext(d_ct.ciphertext);
// builder.push(d_and_boolean_ct, &streams);
// messages.push(MessageType::Boolean(message));
// }
// }
// }

// let cuda_compressed = builder.build(&cuda_compression_key, &streams);

// for (i, val) in messages.iter().enumerate() {
// match val {
// MessageType::Unsigned(message) => {
// let d_decompressed: CudaUnsignedRadixCiphertext = cuda_compressed
// .get(i, &cuda_decompression_key, &streams)
// .unwrap()
// .unwrap();
// let decompressed = d_decompressed.to_radix_ciphertext(&streams);
// let decrypted: u128 = radix_cks.decrypt(&decompressed);
// assert_eq!(decrypted, *message);
// }
// MessageType::Signed(message) => {
// let d_decompressed: CudaSignedRadixCiphertext = cuda_compressed
// .get(i, &cuda_decompression_key, &streams)
// .unwrap()
// .unwrap();
// let decompressed = d_decompressed.to_signed_radix_ciphertext(&streams);
// let decrypted: i128 = radix_cks.decrypt_signed(&decompressed);
// assert_eq!(decrypted, *message);
// }
// MessageType::Boolean(message) => {
// let d_decompressed: CudaBooleanBlock = cuda_compressed
// .get(i, &cuda_decompression_key, &streams)
// .unwrap()
// .unwrap();
// let decompressed = d_decompressed.to_boolean_block(&streams);
// let decrypted = radix_cks.decrypt_bool(&decompressed);
// assert_eq!(decrypted, *message);
// }
// }
// }
// }
//}
}

0 comments on commit 0175f2c

Please sign in to comment.