fix(shortint): add ciphertext_modulus_after_packing_ks to compression parameters

mayeul-zama · mayeul-zama · commit fd554e757566 · 2025-01-08T09:44:58.000+01:00
diff --git a/tfhe/src/integer/gpu/client_key/radix.rs b/tfhe/src/integer/gpu/client_key/radix.rs
@@ -43,7 +43,7 @@ impl RadixClientKey {
                 params.packing_ks_base_log,
                 params.packing_ks_level,
                 params.packing_ks_key_noise_distribution,
-                self.parameters().ciphertext_modulus(),
+                params.ciphertext_modulus_after_packing_ks,
                 &mut engine.encryption_generator,
             )
         });
diff --git a/tfhe/src/shortint/backward_compatibility/parameters/list_compression.rs b/tfhe/src/shortint/backward_compatibility/parameters/list_compression.rs
@@ -1,8 +1,59 @@
-use tfhe_versionable::VersionsDispatch;
-
 use super::parameters::list_compression::CompressionParameters;
+use super::CiphertextModulus;
+use crate::core_crypto::prelude::{CiphertextModulusLog, LweCiphertextCount};
+use crate::shortint::parameters::{
+    DecompositionBaseLog, DecompositionLevelCount, DynamicDistribution, GlweDimension,
+    PolynomialSize,
+};
+use std::convert::Infallible;
+use tfhe_versionable::{Upgrade, Version, VersionsDispatch};
+
+#[derive(Version)]
+pub struct CompressionParametersV0 {
+    pub br_level: DecompositionLevelCount,
+    pub br_base_log: DecompositionBaseLog,
+    pub packing_ks_level: DecompositionLevelCount,
+    pub packing_ks_base_log: DecompositionBaseLog,
+    pub packing_ks_polynomial_size: PolynomialSize,
+    pub packing_ks_glwe_dimension: GlweDimension,
+    pub lwe_per_glwe: LweCiphertextCount,
+    pub storage_log_modulus: CiphertextModulusLog,
+    pub packing_ks_key_noise_distribution: DynamicDistribution<u64>,
+}
+
+impl Upgrade<CompressionParameters> for CompressionParametersV0 {
+    type Error = Infallible;
+
+    fn upgrade(self) -> Result<CompressionParameters, Self::Error> {
+        let Self {
+            br_level,
+            br_base_log,
+            packing_ks_level,
+            packing_ks_base_log,
+            packing_ks_polynomial_size,
+            packing_ks_glwe_dimension,
+            lwe_per_glwe,
+            storage_log_modulus,
+            packing_ks_key_noise_distribution,
+        } = self;
+
+        Ok(CompressionParameters {
+            br_level,
+            br_base_log,
+            packing_ks_level,
+            packing_ks_base_log,
+            packing_ks_polynomial_size,
+            packing_ks_glwe_dimension,
+            lwe_per_glwe,
+            storage_log_modulus,
+            packing_ks_key_noise_distribution,
+            ciphertext_modulus_after_packing_ks: CiphertextModulus::new_native(),
+        })
+    }
+}
 
 #[derive(VersionsDispatch)]
 pub enum CompressionParametersVersions {
-    V0(CompressionParameters),
+    V0(CompressionParametersV0),
+    V1(CompressionParameters),
 }
diff --git a/tfhe/src/shortint/list_compression/compressed_server_keys.rs b/tfhe/src/shortint/list_compression/compressed_server_keys.rs
@@ -103,7 +103,7 @@ impl ClientKey {
                 params.packing_ks_base_log,
                 params.packing_ks_level,
                 params.packing_ks_key_noise_distribution,
-                self.parameters.ciphertext_modulus(),
+                params.ciphertext_modulus_after_packing_ks,
                 &mut engine.seeder,
             )
         });
diff --git a/tfhe/src/shortint/list_compression/compression.rs b/tfhe/src/shortint/list_compression/compression.rs
@@ -25,7 +25,9 @@ impl CompressionKey {
         let lwe_pksk = &self.packing_key_switching_key;
 
         let polynomial_size = lwe_pksk.output_polynomial_size();
-        let ciphertext_modulus = lwe_pksk.ciphertext_modulus();
+
+        let out_ciphertext_modulus = lwe_pksk.ciphertext_modulus();
+
         let glwe_size = lwe_pksk.output_glwe_size();
         let lwe_size = lwe_pksk.input_key_lwe_dimension().to_lwe_size();
 
@@ -43,6 +45,7 @@ impl CompressionKey {
         let message_modulus = first_ct.message_modulus;
         let carry_modulus = first_ct.carry_modulus;
         let pbs_order = first_ct.pbs_order;
+        let in_ciphertext_modulus = first_ct.ct.ciphertext_modulus();
 
         assert!(
             message_modulus.0 <= carry_modulus.0,
@@ -86,6 +89,12 @@ impl CompressionKey {
                         "All ciphertexts do not have the same pbs order"
                     );
 
+                    assert_eq!(
+                        in_ciphertext_modulus,
+                        ct.ct.ciphertext_modulus(),
+                        "All ciphertexts do not have the same ciphertext modulus"
+                    );
+
                     let mut ct = ct.clone();
                     let max_noise_level =
                         MaxNoiseLevel::new((ct.noise_level() * message_modulus.0).get());
@@ -94,12 +103,12 @@ impl CompressionKey {
                     list.extend(ct.ct.as_ref());
                 }
 
-                let list = LweCiphertextList::from_container(list, lwe_size, ciphertext_modulus);
+                let list = LweCiphertextList::from_container(list, lwe_size, in_ciphertext_modulus);
 
                 let bodies_count = LweCiphertextCount(ct_list.len());
 
                 let mut out =
-                    GlweCiphertext::new(0, glwe_size, polynomial_size, ciphertext_modulus);
+                    GlweCiphertext::new(0, glwe_size, polynomial_size, out_ciphertext_modulus);
 
                 par_keyswitch_lwe_ciphertext_list_and_pack_in_glwe_ciphertext(
                     lwe_pksk, &list, &mut out,
@@ -120,7 +129,7 @@ impl CompressionKey {
             pbs_order,
             lwe_per_glwe,
             count,
-            ciphertext_modulus,
+            ciphertext_modulus: out_ciphertext_modulus,
         }
     }
 }
diff --git a/tfhe/src/shortint/list_compression/server_keys.rs b/tfhe/src/shortint/list_compression/server_keys.rs
@@ -65,7 +65,7 @@ impl ClientKey {
                 params.packing_ks_base_log,
                 params.packing_ks_level,
                 params.packing_ks_key_noise_distribution,
-                self.parameters.ciphertext_modulus(),
+                params.ciphertext_modulus_after_packing_ks,
                 &mut engine.encryption_generator,
             )
         });
@@ -116,6 +116,7 @@ pub struct CompressionConformanceParameters {
     pub packing_ks_base_log: DecompositionBaseLog,
     pub packing_ks_polynomial_size: PolynomialSize,
     pub packing_ks_glwe_dimension: GlweDimension,
+    pub ciphertext_modulus_after_packing_ks: CiphertextModulus<u64>,
     pub lwe_per_glwe: LweCiphertextCount,
     pub storage_log_modulus: CiphertextModulusLog,
     pub uncompressed_polynomial_size: PolynomialSize,
@@ -132,6 +133,8 @@ impl From<(PBSParameters, CompressionParameters)> for CompressionConformancePara
             packing_ks_base_log: compression_params.packing_ks_base_log,
             packing_ks_polynomial_size: compression_params.packing_ks_polynomial_size,
             packing_ks_glwe_dimension: compression_params.packing_ks_glwe_dimension,
+            ciphertext_modulus_after_packing_ks: compression_params
+                .ciphertext_modulus_after_packing_ks,
             lwe_per_glwe: compression_params.lwe_per_glwe,
             storage_log_modulus: compression_params.storage_log_modulus,
             uncompressed_polynomial_size: pbs_params.polynomial_size(),
@@ -159,7 +162,7 @@ impl ParameterSetConformant for CompressionKey {
                 .to_equivalent_lwe_dimension(parameter_set.uncompressed_polynomial_size),
             output_glwe_size: parameter_set.packing_ks_glwe_dimension.to_glwe_size(),
             output_polynomial_size: parameter_set.packing_ks_polynomial_size,
-            ciphertext_modulus: parameter_set.cipherext_modulus,
+            ciphertext_modulus: parameter_set.ciphertext_modulus_after_packing_ks,
         };
 
         packing_key_switching_key.is_conformant(&params)
diff --git a/tfhe/src/shortint/parameters/list_compression.rs b/tfhe/src/shortint/parameters/list_compression.rs
@@ -1,5 +1,6 @@
 use tfhe_versionable::Versionize;
 
+use super::CiphertextModulus;
 use crate::core_crypto::prelude::{CiphertextModulusLog, LweCiphertextCount};
 use crate::shortint::backward_compatibility::parameters::list_compression::CompressionParametersVersions;
 use crate::shortint::parameters::{
@@ -20,6 +21,7 @@ pub struct CompressionParameters {
     pub lwe_per_glwe: LweCiphertextCount,
     pub storage_log_modulus: CiphertextModulusLog,
     pub packing_ks_key_noise_distribution: DynamicDistribution<u64>,
+    pub ciphertext_modulus_after_packing_ks: CiphertextModulus,
 }
 
 pub const COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64: CompressionParameters =
@@ -36,4 +38,5 @@ pub const V0_11_COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64: CompressionPa
         lwe_per_glwe: LweCiphertextCount(256),
         storage_log_modulus: CiphertextModulusLog(12),
         packing_ks_key_noise_distribution: DynamicDistribution::new_t_uniform(43),
+        ciphertext_modulus_after_packing_ks: CiphertextModulus::new_native(),
     };

Original file line number	Diff line number	Diff line change
`@@ -43,7 +43,7 @@ impl RadixClientKey {`
`43`	`43`	`params.packing_ks_base_log,`
`44`	`44`	`params.packing_ks_level,`
`45`	`45`	`params.packing_ks_key_noise_distribution,`
`46`		`- self.parameters().ciphertext_modulus(),`
	`46`	`+ params.ciphertext_modulus_after_packing_ks,`
`47`	`47`	`&mut engine.encryption_generator,`
`48`	`48`	`)`
`49`	`49`	`});`
Original file line number	Diff line number	Diff line change
`@@ -103,7 +103,7 @@ impl ClientKey {`
`103`	`103`	`params.packing_ks_base_log,`
`104`	`104`	`params.packing_ks_level,`
`105`	`105`	`params.packing_ks_key_noise_distribution,`
`106`		`- self.parameters.ciphertext_modulus(),`
	`106`	`+ params.ciphertext_modulus_after_packing_ks,`
`107`	`107`	`&mut engine.seeder,`
`108`	`108`	`)`
`109`	`109`	`});`