feat(shortint): allow the KS32 parameters to have non native KSK modulus

Author: IceTDrinker

tfhe/src/shortint/atomic_pattern/ks32.rs

@@ -83,7 +83,7 @@ impl KS32AtomicPatternServerKey {
             params.ks_base_log(),
             params.ks_level(),
             pbs_params.lwe_noise_distribution(),
-            CoreCiphertextModulus::new_native(), // Does it make sense to parametrize this ?
+            pbs_params.post_keyswitch_ciphertext_modulus(),
             &mut engine.encryption_generator,
         );
 

tfhe/src/shortint/atomic_pattern/standard.rs

@@ -153,10 +153,8 @@ impl AtomicPattern for StandardAtomicPatternServerKey {
 
     fn apply_lookup_table_assign(&self, ct: &mut Ciphertext, acc: &LookupTableOwned) {
         ShortintEngine::with_thread_local_mut(|engine| {
-            let (mut ciphertext_buffer, buffers) = engine.get_buffers(
-                self.intermediate_lwe_dimension(),
-                CiphertextModulus::new_native(),
-            );
+            let (mut ciphertext_buffer, buffers) =
+                engine.get_buffers(self.intermediate_lwe_dimension(), self.ciphertext_modulus());
 
             match self.pbs_order {
                 PBSOrder::KeyswitchBootstrap => {
@@ -349,10 +347,8 @@ impl StandardAtomicPatternServerKey {
         let mut acc = lut.acc.clone();
 
         ShortintEngine::with_thread_local_mut(|engine| {
-            let (mut ciphertext_buffer, buffers) = engine.get_buffers(
-                self.intermediate_lwe_dimension(),
-                CiphertextModulus::new_native(),
-            );
+            let (mut ciphertext_buffer, buffers) =
+                engine.get_buffers(self.intermediate_lwe_dimension(), self.ciphertext_modulus());
 
             // Compute a key switch
             keyswitch_lwe_ciphertext(&self.key_switching_key, &ct.ct, &mut ciphertext_buffer);

tfhe/src/shortint/engine/server_side.rs

@@ -11,8 +11,7 @@ use crate::shortint::atomic_pattern::AtomicPatternServerKey;
 use crate::shortint::ciphertext::MaxDegree;
 use crate::shortint::client_key::secret_encryption_key::SecretEncryptionKeyView;
 use crate::shortint::parameters::{
-    CoreCiphertextModulus, EncryptionKeyChoice, KeySwitch32PBSParameters,
-    ShortintKeySwitchingParameters,
+    EncryptionKeyChoice, KeySwitch32PBSParameters, ShortintKeySwitchingParameters,
 };
 use crate::shortint::server_key::{
     CompressedModulusSwitchNoiseReductionKey, ModulusSwitchNoiseReductionKey,
@@ -109,7 +108,7 @@ impl ShortintEngine {
                     modulus_switch_noise_reduction_params,
                     in_key,
                     self,
-                    CoreCiphertextModulus::new_native(),
+                    pbs_params.post_keyswitch_ciphertext_modulus,
                     pbs_params.lwe_noise_distribution,
                 )
             });

tfhe/src/shortint/parameters/ks32.rs

@@ -2,19 +2,19 @@ use serde::{Deserialize, Serialize};
 use tfhe_versionable::Versionize;
 
 pub use crate::core_crypto::commons::parameters::{
-    CiphertextModulus as CoreCiphertextModulus, CiphertextModulusLog, DecompositionBaseLog,
-    DecompositionLevelCount, DynamicDistribution, EncryptionKeyChoice, GlweDimension,
-    LweBskGroupingFactor, LweCiphertextCount, LweDimension, NoiseEstimationMeasureBound,
-    PolynomialSize, RSigmaFactor,
+    CiphertextModulusLog, DecompositionBaseLog, DecompositionLevelCount, DynamicDistribution,
+    EncryptionKeyChoice, GlweDimension, LweBskGroupingFactor, LweCiphertextCount, LweDimension,
+    NoiseEstimationMeasureBound, PolynomialSize, RSigmaFactor,
 };
 use crate::core_crypto::prelude::{
     LweCiphertextConformanceParams, LweKeyswitchKeyConformanceParams, MsDecompressionType,
 };
 use crate::shortint::backward_compatibility::parameters::KeySwitch32PBSParametersVersions;
 
 use super::{
-    AtomicPatternKind, CarryModulus, CiphertextConformanceParams, CiphertextModulus, Degree,
-    MaxNoiseLevel, MessageModulus, ModulusSwitchNoiseReductionParams, NoiseLevel,
+    AtomicPatternKind, CarryModulus, CiphertextConformanceParams, CiphertextModulus,
+    CiphertextModulus32, Degree, MaxNoiseLevel, MessageModulus, ModulusSwitchNoiseReductionParams,
+    NoiseLevel,
 };
 
 /// A set of cryptographic parameters used with the atomic pattern
@@ -35,6 +35,7 @@ pub struct KeySwitch32PBSParameters {
     pub carry_modulus: CarryModulus,
     pub max_noise_level: MaxNoiseLevel,
     pub log2_p_fail: f64,
+    pub post_keyswitch_ciphertext_modulus: CiphertextModulus32,
     pub ciphertext_modulus: CiphertextModulus,
     pub modulus_switch_noise_reduction_params: Option<ModulusSwitchNoiseReductionParams>,
 }
@@ -49,8 +50,7 @@ impl From<&KeySwitch32PBSParameters> for LweKeyswitchKeyConformanceParams<u32> {
             input_lwe_dimension: value
                 .glwe_dimension()
                 .to_equivalent_lwe_dimension(value.polynomial_size()),
-            // For the moment we only handle the native u32 modulus for the KSK
-            ciphertext_modulus: CoreCiphertextModulus::new_native(),
+            ciphertext_modulus: value.post_keyswitch_ciphertext_modulus(),
         }
     }
 }
@@ -104,6 +104,10 @@ impl KeySwitch32PBSParameters {
         self.max_noise_level
     }
 
+    pub const fn post_keyswitch_ciphertext_modulus(&self) -> CiphertextModulus32 {
+        self.post_keyswitch_ciphertext_modulus
+    }
+
     pub const fn ciphertext_modulus(&self) -> CiphertextModulus {
         self.ciphertext_modulus
     }

tfhe/src/shortint/parameters/mod.rs

@@ -100,6 +100,7 @@ pub struct CarryModulus(pub u64);
 
 /// Determines in what ring computations are made
 pub type CiphertextModulus = CoreCiphertextModulus<u64>;
+pub type CiphertextModulus32 = CoreCiphertextModulus<u32>;
 
 impl From<&PBSConformanceParams> for LweBootstrapKeyConformanceParams<u64> {
     fn from(value: &PBSConformanceParams) -> Self {

tfhe/src/shortint/parameters/v1_1/ks32/tuniform/p_fail_2_minus_128/ks_pbs/mod.rs

@@ -1,7 +1,7 @@
 use crate::core_crypto::prelude::DynamicDistribution;
 use crate::shortint::parameters::{
-    KeySwitch32PBSParameters, LweCiphertextCount, ModulusSwitchNoiseReductionParams,
-    NoiseEstimationMeasureBound, RSigmaFactor, Variance,
+    CiphertextModulus32, KeySwitch32PBSParameters, LweCiphertextCount,
+    ModulusSwitchNoiseReductionParams, NoiseEstimationMeasureBound, RSigmaFactor, Variance,
 };
 use crate::shortint::prelude::{
     DecompositionBaseLog, DecompositionLevelCount, GlweDimension, LweDimension, PolynomialSize,
@@ -25,6 +25,7 @@ pub const V1_1_PARAM_MESSAGE_2_CARRY_2_KS32_PBS_TUNIFORM_2M128: KeySwitch32PBSPa
         carry_modulus: CarryModulus(4),
         max_noise_level: MaxNoiseLevel::new(5),
         log2_p_fail: -129.358380844,
+        post_keyswitch_ciphertext_modulus: CiphertextModulus32::new_native(),
         ciphertext_modulus: CiphertextModulus::new_native(),
         modulus_switch_noise_reduction_params: Some(ModulusSwitchNoiseReductionParams {
             modulus_switch_zeros_count: LweCiphertextCount(1449),