github/tfhe-rs
1
1
Fork 0
mirror of https://github.com/zama-ai/tfhe-rs.git synced 2026-01-09 22:57:59 -05:00
Code Issues Packages Projects Releases Wiki Activity

feat(shortint): add ct compression for the ks32 atomic pattern

Browse Source
This commit is contained in:
Nicolas Sarlin
2025-10-01 18:27:23 +02:00
committed by Nicolas Sarlin
parent f9e876730a
commit aefec1fe64
9 changed files with 371 additions and 166 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
tfhe/src/high_level_api/compressed_ciphertext_list.rs
View File

@@ -957,13 +957,18 @@ pub mod gpu {
mod tests {
use crate::prelude::*;
use crate::safe_serialization::{safe_deserialize, safe_serialize};
#[cfg(not(feature = "gpu"))]
use crate::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS32_PBS_TUNIFORM_2M128;
#[cfg(feature = "gpu")]
use crate::shortint::parameters::{
COMP_PARAM_GPU_MULTI_BIT_GROUP_4_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
PARAM_GPU_MULTI_BIT_GROUP_4_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
};
use crate::shortint::parameters::{
COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
};
use crate::shortint::PBSParameters;
use crate::shortint::AtomicPatternParameters;
#[cfg(feature = "gpu")]
use crate::GpuIndex;
use crate::{
@@ -974,26 +979,25 @@ mod tests {
#[test]
fn test_compressed_ct_list_cpu_gpu() {
for (params, comp_params) in [
if cfg!(not(feature = "gpu")) {
(
PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128.into(),
COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
)
} else {
(
crate::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128
.into(),
crate::shortint::parameters::COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
)
},
(
PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128.into(),
COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
),
#[cfg(not(feature = "gpu"))]
(
PARAM_MESSAGE_2_CARRY_2_KS32_PBS_TUNIFORM_2M128.into(),
COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
),
#[cfg(feature = "gpu")]
(
PARAM_GPU_MULTI_BIT_GROUP_4_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128.into(),
COMP_PARAM_GPU_MULTI_BIT_GROUP_4_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
),
] {
let config = crate::ConfigBuilder::with_custom_parameters::<PBSParameters>(params)
.enable_compression(comp_params)
.build();
let config =
crate::ConfigBuilder::with_custom_parameters::<AtomicPatternParameters>(params)
.enable_compression(comp_params)
.build();
let ck = crate::ClientKey::generate(config);
let sk = crate::CompressedServerKey::new(&ck);
@@ -1202,9 +1206,10 @@ mod tests {
COMP_PARAM_GPU_MULTI_BIT_GROUP_4_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
),
] {
let config = crate::ConfigBuilder::with_custom_parameters::<PBSParameters>(params)
.enable_compression(comp_params)
.build();
let config =
crate::ConfigBuilder::with_custom_parameters::<AtomicPatternParameters>(params)
.enable_compression(comp_params)
.build();
let ck = crate::ClientKey::generate(config);
let sk = crate::CompressedServerKey::new(&ck);

5
tfhe/src/integer/ciphertext/compressed_ciphertext_list.rs
View File

@@ -225,6 +225,7 @@ mod tests {
TEST_COMP_PARAM_GPU_MULTI_BIT_GROUP_4_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
TEST_COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
TEST_PARAM_GPU_MULTI_BIT_GROUP_4_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
TEST_PARAM_MESSAGE_2_CARRY_2_KS32_PBS_TUNIFORM_2M128,
TEST_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
};
use crate::shortint::ShortintParameterSet;
@@ -265,6 +266,10 @@ mod tests {
TEST_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128.into(),
TEST_COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
),
(
TEST_PARAM_MESSAGE_2_CARRY_2_KS32_PBS_TUNIFORM_2M128.into(),
TEST_COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
),
(
TEST_PARAM_GPU_MULTI_BIT_GROUP_4_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128.into(),
TEST_COMP_PARAM_GPU_MULTI_BIT_GROUP_4_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,

2
tfhe/src/integer/mod.rs
View File

@@ -175,7 +175,7 @@ where
gen_keys_inner(shortint_parameters_set)
} else {
keycache::KEY_CACHE
.get_from_params(shortint_parameters_set.pbs_parameters().unwrap(), key_kind)
.get_from_params(shortint_parameters_set.ap_parameters().unwrap(), key_kind)
}
}
#[cfg(all(not(test), not(feature = "internal-keycache")))]

50
tfhe/src/shortint/client_key/atomic_pattern/ks32.rs
View File

@@ -8,7 +8,13 @@ use crate::core_crypto::prelude::{
use crate::shortint::backward_compatibility::client_key::atomic_pattern::KS32AtomicPatternClientKeyVersions;
use crate::shortint::client_key::{GlweSecretKeyOwned, LweSecretKeyOwned, LweSecretKeyView};
use crate::shortint::engine::ShortintEngine;
use crate::shortint::parameters::{DynamicDistribution, KeySwitch32PBSParameters};
use crate::shortint::list_compression::{
CompressedCompressionKey, CompressedDecompressionKey, CompressionKey, CompressionPrivateKeys,
DecompressionKey,
};
use crate::shortint::parameters::{
CompressionParameters, DynamicDistribution, KeySwitch32PBSParameters,
};
use crate::shortint::{AtomicPatternKind, ShortintParameterSet};
use super::EncryptionAtomicPattern;
@@ -146,6 +152,48 @@ impl KS32AtomicPatternClientKey {
pub fn small_lwe_secret_key(&self) -> LweSecretKeyView<'_, u32> {
self.lwe_secret_key.as_view()
}
pub fn new_compression_key(
&self,
private_compression_key: &CompressionPrivateKeys,
) -> CompressionKey {
private_compression_key.new_compression_key(&self.glwe_secret_key, self.parameters())
}
pub fn new_compressed_compression_key(
&self,
private_compression_key: &CompressionPrivateKeys,
) -> CompressedCompressionKey {
private_compression_key
.new_compressed_compression_key(&self.glwe_secret_key, self.parameters())
}
pub fn new_decompression_key(
&self,
private_compression_key: &CompressionPrivateKeys,
) -> DecompressionKey {
private_compression_key.new_decompression_key(&self.glwe_secret_key, self.parameters())
}
pub fn new_decompression_key_with_params(
&self,
private_compression_key: &CompressionPrivateKeys,
compression_params: CompressionParameters,
) -> DecompressionKey {
private_compression_key.new_decompression_key_with_params(
&self.glwe_secret_key,
self.parameters(),
compression_params,
)
}
pub fn new_compressed_decompression_key(
&self,
private_compression_key: &CompressionPrivateKeys,
) -> CompressedDecompressionKey {
private_compression_key
.new_compressed_decompression_key(&self.glwe_secret_key, self.parameters())
}
}
impl EncryptionAtomicPattern for KS32AtomicPatternClientKey {

70
tfhe/src/shortint/client_key/atomic_pattern/mod.rs
View File

@@ -6,7 +6,11 @@ use tfhe_versionable::Versionize;
use crate::shortint::backward_compatibility::client_key::atomic_pattern::AtomicPatternClientKeyVersions;
use crate::shortint::engine::ShortintEngine;
use crate::shortint::parameters::DynamicDistribution;
use crate::shortint::list_compression::{
CompressedCompressionKey, CompressedDecompressionKey, CompressionKey, CompressionPrivateKeys,
DecompressionKey,
};
use crate::shortint::parameters::{CompressionParameters, DynamicDistribution};
use crate::shortint::{AtomicPatternKind, AtomicPatternParameters, ShortintParameterSet};
use super::{LweSecretKeyOwned, LweSecretKeyView};
@@ -102,6 +106,70 @@ impl AtomicPatternClientKey {
)),
}
}
pub fn new_compression_key(
&self,
private_compression_key: &CompressionPrivateKeys,
) -> CompressionKey {
match self {
Self::Standard(std_cks) => std_cks.new_compression_key(private_compression_key),
Self::KeySwitch32(ks32_cks) => ks32_cks.new_compression_key(private_compression_key),
}
}
pub fn new_compressed_compression_key(
&self,
private_compression_key: &CompressionPrivateKeys,
) -> CompressedCompressionKey {
match self {
Self::Standard(std_cks) => {
std_cks.new_compressed_compression_key(private_compression_key)
}
Self::KeySwitch32(ks32_cks) => {
ks32_cks.new_compressed_compression_key(private_compression_key)
}
}
}
pub fn new_decompression_key(
&self,
private_compression_key: &CompressionPrivateKeys,
) -> DecompressionKey {
match self {
Self::Standard(std_cks) => std_cks.new_decompression_key(private_compression_key),
Self::KeySwitch32(ks32_cks) => ks32_cks.new_decompression_key(private_compression_key),
}
}
/// Create a decompression key with different parameters than the one in the secret key.
///
/// This allows for example to compress using cpu parameters and decompress with gpu parameters
pub fn new_decompression_key_with_params(
&self,
private_compression_key: &CompressionPrivateKeys,
compression_params: CompressionParameters,
) -> DecompressionKey {
match self {
Self::Standard(std_cks) => std_cks
.new_decompression_key_with_params(private_compression_key, compression_params),
Self::KeySwitch32(ks32_cks) => ks32_cks
.new_decompression_key_with_params(private_compression_key, compression_params),
}
}
pub fn new_compressed_decompression_key(
&self,
private_compression_key: &CompressionPrivateKeys,
) -> CompressedDecompressionKey {
match self {
Self::Standard(std_cks) => {
std_cks.new_compressed_decompression_key(private_compression_key)
}
Self::KeySwitch32(ks32_cks) => {
ks32_cks.new_compressed_decompression_key(private_compression_key)
}
}
}
}
impl EncryptionAtomicPattern for AtomicPatternClientKey {

50
tfhe/src/shortint/client_key/atomic_pattern/standard.rs
View File

@@ -8,7 +8,13 @@ use crate::core_crypto::prelude::{
use crate::shortint::backward_compatibility::client_key::atomic_pattern::StandardAtomicPatternClientKeyVersions;
use crate::shortint::client_key::{GlweSecretKeyOwned, LweSecretKeyOwned, LweSecretKeyView};
use crate::shortint::engine::ShortintEngine;
use crate::shortint::parameters::{DynamicDistribution, ShortintKeySwitchingParameters};
use crate::shortint::list_compression::{
CompressedCompressionKey, CompressedDecompressionKey, CompressionKey, CompressionPrivateKeys,
DecompressionKey,
};
use crate::shortint::parameters::{
CompressionParameters, DynamicDistribution, ShortintKeySwitchingParameters,
};
use crate::shortint::{
AtomicPatternKind, EncryptionKeyChoice, PBSParameters, ShortintParameterSet, WopbsParameters,
};
@@ -237,6 +243,48 @@ impl StandardAtomicPatternClientKey {
),
}
}
pub fn new_compression_key(
&self,
private_compression_key: &CompressionPrivateKeys,
) -> CompressionKey {
private_compression_key.new_compression_key(&self.glwe_secret_key, self.parameters())
}
pub fn new_compressed_compression_key(
&self,
private_compression_key: &CompressionPrivateKeys,
) -> CompressedCompressionKey {
private_compression_key
.new_compressed_compression_key(&self.glwe_secret_key, self.parameters())
}
pub fn new_decompression_key(
&self,
private_compression_key: &CompressionPrivateKeys,
) -> DecompressionKey {
private_compression_key.new_decompression_key(&self.glwe_secret_key, self.parameters())
}
pub fn new_decompression_key_with_params(
&self,
private_compression_key: &CompressionPrivateKeys,
compression_params: CompressionParameters,
) -> DecompressionKey {
private_compression_key.new_decompression_key_with_params(
&self.glwe_secret_key,
self.parameters(),
compression_params,
)
}
pub fn new_compressed_decompression_key(
&self,
private_compression_key: &CompressionPrivateKeys,
) -> CompressedDecompressionKey {
private_compression_key
.new_compressed_decompression_key(&self.glwe_secret_key, self.parameters())
}
}
impl EncryptionAtomicPattern for StandardAtomicPatternClientKey {

63
tfhe/src/shortint/list_compression/compressed_server_keys.rs
View File

@@ -6,8 +6,6 @@ use super::{
use crate::conformance::ParameterSetConformant;
use crate::core_crypto::fft_impl::fft64::crypto::bootstrap::LweBootstrapKeyConformanceParams;
use crate::core_crypto::prelude::{
allocate_and_generate_new_seeded_lwe_packing_keyswitch_key,
par_allocate_and_generate_new_seeded_lwe_bootstrap_key,
par_convert_standard_lwe_bootstrap_key_to_fourier, CiphertextModulusLog,
FourierLweBootstrapKey, LweCiphertextCount, LwePackingKeyswitchKeyConformanceParams,
SeededLweBootstrapKeyOwned, SeededLwePackingKeyswitchKey,
@@ -16,10 +14,7 @@ use crate::shortint::backward_compatibility::list_compression::{
CompressedCompressionKeyVersions, CompressedDecompressionKeyVersions,
CompressedNoiseSquashingCompressionKeyVersions,
};
use crate::shortint::client_key::atomic_pattern::AtomicPatternClientKey;
use crate::shortint::client_key::ClientKey;
use crate::shortint::engine::ShortintEngine;
use crate::shortint::EncryptionKeyChoice;
use serde::{Deserialize, Serialize};
use std::fmt::Debug;
use tfhe_versionable::Versionize;
@@ -84,58 +79,12 @@ impl ClientKey {
&self,
private_compression_key: &CompressionPrivateKeys,
) -> (CompressedCompressionKey, CompressedDecompressionKey) {
let AtomicPatternClientKey::Standard(std_cks) = &self.atomic_pattern else {
panic!("Only the standard atomic pattern supports compression")
};
let pbs_params = std_cks.parameters;
assert_eq!(
pbs_params.encryption_key_choice(),
EncryptionKeyChoice::Big,
"Compression is only compatible with ciphertext in post PBS dimension"
);
let compression_params = &private_compression_key.params;
let packing_key_switching_key = ShortintEngine::with_thread_local_mut(|engine| {
allocate_and_generate_new_seeded_lwe_packing_keyswitch_key(
&std_cks.large_lwe_secret_key(),
&private_compression_key.post_packing_ks_key,
compression_params.packing_ks_base_log,
compression_params.packing_ks_level,
compression_params.packing_ks_key_noise_distribution,
pbs_params.ciphertext_modulus(),
&mut engine.seeder,
)
});
let glwe_compression_key = CompressedCompressionKey {
packing_key_switching_key,
lwe_per_glwe: compression_params.lwe_per_glwe,
storage_log_modulus: compression_params.storage_log_modulus,
};
let blind_rotate_key = ShortintEngine::with_thread_local_mut(|engine| {
par_allocate_and_generate_new_seeded_lwe_bootstrap_key(
&private_compression_key
.post_packing_ks_key
.as_lwe_secret_key(),
&std_cks.glwe_secret_key,
compression_params.br_base_log,
compression_params.br_level,
pbs_params.glwe_noise_distribution(),
pbs_params.ciphertext_modulus(),
&mut engine.seeder,
)
});
let glwe_decompression_key = CompressedDecompressionKey {
blind_rotate_key,
lwe_per_glwe: compression_params.lwe_per_glwe,
};
(glwe_compression_key, glwe_decompression_key)
(
self.atomic_pattern
.new_compressed_compression_key(private_compression_key),
self.atomic_pattern
.new_compressed_decompression_key(private_compression_key),
)
}
}

161
tfhe/src/shortint/list_compression/private_key.rs
View File

@@ -1,7 +1,10 @@
use tfhe_versionable::Versionize;
use crate::core_crypto::prelude::{
allocate_and_generate_new_binary_glwe_secret_key, GlweSecretKeyOwned,
allocate_and_generate_new_binary_glwe_secret_key,
allocate_and_generate_new_lwe_packing_keyswitch_key,
allocate_and_generate_new_seeded_lwe_packing_keyswitch_key,
par_allocate_and_generate_new_seeded_lwe_bootstrap_key, GlweSecretKey, GlweSecretKeyOwned,
};
use crate::shortint::backward_compatibility::list_compression::{
CompressionPrivateKeysVersions, NoiseSquashingCompressionPrivateKeyVersions,
@@ -11,9 +14,13 @@ use crate::shortint::client_key::ClientKey;
use crate::shortint::engine::ShortintEngine;
use crate::shortint::noise_squashing::NoiseSquashingPrivateKeyView;
use crate::shortint::parameters::{CompressionParameters, NoiseSquashingCompressionParameters};
use crate::shortint::EncryptionKeyChoice;
use crate::shortint::{EncryptionKeyChoice, ShortintParameterSet};
use std::fmt::Debug;
use super::{
CompressedCompressionKey, CompressedDecompressionKey, CompressionKey, DecompressionKey,
};
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, Versionize)]
#[versionize(CompressionPrivateKeysVersions)]
pub struct CompressionPrivateKeys {
@@ -21,16 +28,158 @@ pub struct CompressionPrivateKeys {
pub params: CompressionParameters,
}
impl CompressionPrivateKeys {
pub(crate) fn new_compression_key(
&self,
glwe_secret_key: &GlweSecretKey<Vec<u64>>,
pbs_params: ShortintParameterSet,
) -> CompressionKey {
assert_eq!(
pbs_params.encryption_key_choice(),
EncryptionKeyChoice::Big,
"Compression is only compatible with ciphertext in post PBS dimension"
);
let compression_params = &self.params;
assert!(
compression_params.storage_log_modulus.0
<= pbs_params
.polynomial_size()
.to_blind_rotation_input_modulus_log()
.0,
"Compression parameters say to store more bits than useful"
);
let packing_key_switching_key = ShortintEngine::with_thread_local_mut(|engine| {
allocate_and_generate_new_lwe_packing_keyswitch_key(
&glwe_secret_key.as_lwe_secret_key(),
&self.post_packing_ks_key,
compression_params.packing_ks_base_log,
compression_params.packing_ks_level,
compression_params.packing_ks_key_noise_distribution,
pbs_params.ciphertext_modulus(),
&mut engine.encryption_generator,
)
});
CompressionKey {
packing_key_switching_key,
lwe_per_glwe: compression_params.lwe_per_glwe,
storage_log_modulus: compression_params.storage_log_modulus,
}
}
pub(crate) fn new_compressed_compression_key(
&self,
glwe_secret_key: &GlweSecretKey<Vec<u64>>,
pbs_params: ShortintParameterSet,
) -> CompressedCompressionKey {
assert_eq!(
pbs_params.encryption_key_choice(),
EncryptionKeyChoice::Big,
"Compression is only compatible with ciphertext in post PBS dimension"
);
let compression_params = &self.params;
let packing_key_switching_key = ShortintEngine::with_thread_local_mut(|engine| {
allocate_and_generate_new_seeded_lwe_packing_keyswitch_key(
&glwe_secret_key.as_lwe_secret_key(),
&self.post_packing_ks_key,
compression_params.packing_ks_base_log,
compression_params.packing_ks_level,
compression_params.packing_ks_key_noise_distribution,
pbs_params.ciphertext_modulus(),
&mut engine.seeder,
)
});
CompressedCompressionKey {
packing_key_switching_key,
lwe_per_glwe: compression_params.lwe_per_glwe,
storage_log_modulus: compression_params.storage_log_modulus,
}
}
pub(crate) fn new_decompression_key(
&self,
glwe_secret_key: &GlweSecretKey<Vec<u64>>,
pbs_params: ShortintParameterSet,
) -> DecompressionKey {
self.new_decompression_key_with_params(glwe_secret_key, pbs_params, self.params)
}
/// Create a decompression key with different parameters than the one in the secret key.
///
/// This allows for example to compress using cpu parameters and decompress with gpu parameters
pub(crate) fn new_decompression_key_with_params(
&self,
glwe_secret_key: &GlweSecretKey<Vec<u64>>,
pbs_params: ShortintParameterSet,
compression_params: CompressionParameters,
) -> DecompressionKey {
assert_eq!(
pbs_params.encryption_key_choice(),
EncryptionKeyChoice::Big,
"Compression is only compatible with ciphertext in post PBS dimension"
);
let blind_rotate_key = ShortintEngine::with_thread_local_mut(|engine| {
engine.new_classic_bootstrapping_key(
&self.post_packing_ks_key.as_lwe_secret_key(),
glwe_secret_key,
pbs_params.glwe_noise_distribution(),
compression_params.br_base_log,
compression_params.br_level,
pbs_params.ciphertext_modulus(),
)
});
DecompressionKey {
blind_rotate_key,
lwe_per_glwe: compression_params.lwe_per_glwe,
}
}
pub(crate) fn new_compressed_decompression_key(
&self,
glwe_secret_key: &GlweSecretKey<Vec<u64>>,
pbs_params: ShortintParameterSet,
) -> CompressedDecompressionKey {
assert_eq!(
pbs_params.encryption_key_choice(),
EncryptionKeyChoice::Big,
"Compression is only compatible with ciphertext in post PBS dimension"
);
let compression_params = &self.params;
let blind_rotate_key = ShortintEngine::with_thread_local_mut(|engine| {
par_allocate_and_generate_new_seeded_lwe_bootstrap_key(
&self.post_packing_ks_key.as_lwe_secret_key(),
glwe_secret_key,
compression_params.br_base_log,
compression_params.br_level,
pbs_params.glwe_noise_distribution(),
pbs_params.ciphertext_modulus(),
&mut engine.seeder,
)
});
CompressedDecompressionKey {
blind_rotate_key,
lwe_per_glwe: compression_params.lwe_per_glwe,
}
}
}
impl ClientKey {
pub fn new_compression_private_key(
&self,
params: CompressionParameters,
) -> CompressionPrivateKeys {
assert_eq!(
self.parameters()
.pbs_parameters()
.unwrap()
.encryption_key_choice(),
self.parameters().encryption_key_choice(),
EncryptionKeyChoice::Big,
"Compression is only compatible with ciphertext in post PBS dimension"
);

91
tfhe/src/shortint/list_compression/server_keys.rs
View File

@@ -7,7 +7,6 @@ use crate::shortint::atomic_pattern::AtomicPatternParameters;
use crate::shortint::backward_compatibility::list_compression::{
CompressionKeyVersions, DecompressionKeyVersions, NoiseSquashingCompressionKeyVersions,
};
use crate::shortint::client_key::atomic_pattern::AtomicPatternClientKey;
use crate::shortint::client_key::ClientKey;
use crate::shortint::engine::ShortintEngine;
use crate::shortint::list_compression::CompressedNoiseSquashingCompressionKey;
@@ -16,7 +15,6 @@ use crate::shortint::parameters::{
CompressionParameters, NoiseSquashingCompressionParameters, NoiseSquashingParameters,
PolynomialSize,
};
use crate::shortint::EncryptionKeyChoice;
use serde::{Deserialize, Serialize};
use std::fmt::Debug;
use tfhe_versionable::Versionize;
@@ -47,93 +45,28 @@ impl DecompressionKey {
}
impl ClientKey {
pub(crate) fn new_decompression_key_with_params(
/// Create a decompression key with different parameters than the one in the secret key.
///
/// This allows for example to compress using cpu parameters and decompress with gpu parameters
pub fn new_decompression_key_with_params(
&self,
private_compression_key: &CompressionPrivateKeys,
compression_params: CompressionParameters,
) -> DecompressionKey {
let AtomicPatternClientKey::Standard(std_cks) = &self.atomic_pattern else {
panic!("Only the standard atomic pattern supports compression")
};
let pbs_params = std_cks.parameters;
assert_eq!(
pbs_params.encryption_key_choice(),
EncryptionKeyChoice::Big,
"Compression is only compatible with ciphertext in post PBS dimension"
);
let blind_rotate_key = ShortintEngine::with_thread_local_mut(|engine| {
engine.new_classic_bootstrapping_key(
&private_compression_key
.post_packing_ks_key
.as_lwe_secret_key(),
&std_cks.glwe_secret_key,
pbs_params.glwe_noise_distribution(),
compression_params.br_base_log,
compression_params.br_level,
pbs_params.ciphertext_modulus(),
)
});
DecompressionKey {
blind_rotate_key,
lwe_per_glwe: compression_params.lwe_per_glwe,
}
self.atomic_pattern
.new_decompression_key_with_params(private_compression_key, compression_params)
}
pub fn new_compression_decompression_keys(
&self,
private_compression_key: &CompressionPrivateKeys,
) -> (CompressionKey, DecompressionKey) {
let AtomicPatternClientKey::Standard(std_cks) = &self.atomic_pattern else {
panic!("Only the standard atomic pattern supports compression")
};
let pbs_params = std_cks.parameters;
assert_eq!(
pbs_params.encryption_key_choice(),
EncryptionKeyChoice::Big,
"Compression is only compatible with ciphertext in post PBS dimension"
);
let compression_params = &private_compression_key.params;
assert!(
compression_params.storage_log_modulus.0
<= pbs_params
.polynomial_size()
.to_blind_rotation_input_modulus_log()
.0,
"Compression parameters say to store more bits than useful"
);
let packing_key_switching_key = ShortintEngine::with_thread_local_mut(|engine| {
allocate_and_generate_new_lwe_packing_keyswitch_key(
&std_cks.large_lwe_secret_key(),
&private_compression_key.post_packing_ks_key,
compression_params.packing_ks_base_log,
compression_params.packing_ks_level,
compression_params.packing_ks_key_noise_distribution,
pbs_params.ciphertext_modulus(),
&mut engine.encryption_generator,
)
});
let glwe_compression_key = CompressionKey {
packing_key_switching_key,
lwe_per_glwe: compression_params.lwe_per_glwe,
storage_log_modulus: compression_params.storage_log_modulus,
};
let glwe_decompression_key = self.new_decompression_key_with_params(
private_compression_key,
private_compression_key.params,
);
(glwe_compression_key, glwe_decompression_key)
(
self.atomic_pattern
.new_compression_key(private_compression_key),
self.atomic_pattern
.new_decompression_key(private_compression_key),
)
}
}