test(shortint): add compression atomic pattern for noise checks

- noise checks and pfail based on expected noise have been added
- compatible with KS PBS and KS32 PBS

scripts/pfail_estimate.py

@@ -0,0 +1,42 @@
+import scipy.stats as stats
+from scipy.special import erfcinv, erfc
+import math
+
+# utilities
+t = 1 / (2 ** (4 + 2))  # noise bound
+standard_score = lambda p_fail: math.sqrt(2) * erfcinv(p_fail)  # standard score
+
+pfail = lambda z: erfc(z / math.sqrt(2))
+
+# Noise squashing after compression
+# measured_variance = 7.598561171474912e-35
+# variance_after_flood = measured_variance * (2**40 * 100) ** 2
+
+# measured_std_dev = math.sqrt(variance_after_flood)
+
+# New params GPU before MS 128
+# measured_variance = 1.438540449823688e-6
+# Rerand noise
+# measured_variance = 1.4064222454361346e-6
+# measured_variance = 1.408401059719539e-6
+
+# measured_variance = 1.4120971218065554e-6 #KS32
+measured_variance = 1.4150031500067098e-6
+measured_std_dev = math.sqrt(measured_variance)
+
+measured_std_score = t / measured_std_dev
+
+estimated_pfail = pfail(measured_std_score)
+
+print(estimated_pfail, math.log2(estimated_pfail))
+
+
+# Compression encoding for 2_2
+t_compression = 1 / (2 ** (2 + 2))
+measured_variance = 1.0216297411906617e-5
+measured_std_dev = math.sqrt(measured_variance)
+
+measured_std_score = t_compression / measured_std_dev
+
+estimated_pfail = pfail(measured_std_score)
+print(estimated_pfail, math.log2(estimated_pfail))

tfhe/src/core_crypto/algorithms/modulus_switch.rs

@@ -106,6 +106,7 @@ use crate::core_crypto::commons::noise_formulas::noise_simulation::traits::{
     AllocateCenteredBinaryShiftedStandardModSwitchResult, AllocateStandardModSwitchResult,
     CenteredBinaryShiftedStandardModSwitch, StandardModSwitch,
 };
+use crate::core_crypto::entities::glwe_ciphertext::{GlweCiphertext, GlweCiphertextOwned};
 
 impl<Scalar: UnsignedInteger, C: Container<Element = Scalar>> AllocateStandardModSwitchResult
     for LweCiphertext<C>
@@ -206,6 +207,58 @@ impl<
     }
 }
 
+impl<Scalar: UnsignedInteger, C: Container<Element = Scalar>> AllocateStandardModSwitchResult
+    for GlweCiphertext<C>
+{
+    type Output = GlweCiphertextOwned<Scalar>;
+    type SideResources = ();
+
+    fn allocate_standard_mod_switch_result(
+        &self,
+        _side_resources: &mut Self::SideResources,
+    ) -> Self::Output {
+        // We will mod switch but we keep the current modulus as the noise is interesting in the
+        // context of the input modulus
+        Self::Output::new(
+            Scalar::ZERO,
+            self.glwe_size(),
+            self.polynomial_size(),
+            self.ciphertext_modulus(),
+        )
+    }
+}
+
+impl<
+        Scalar: UnsignedInteger,
+        InputCont: Container<Element = Scalar>,
+        OutputCont: ContainerMut<Element = Scalar>,
+    > StandardModSwitch<GlweCiphertext<OutputCont>> for GlweCiphertext<InputCont>
+{
+    type SideResources = ();
+
+    fn standard_mod_switch(
+        &self,
+        output_modulus_log: CiphertextModulusLog,
+        output: &mut GlweCiphertext<OutputCont>,
+        _side_resources: &mut Self::SideResources,
+    ) {
+        assert!(self
+            .ciphertext_modulus()
+            .is_compatible_with_native_modulus());
+        assert_eq!(self.glwe_size(), output.glwe_size());
+        assert_eq!(self.polynomial_size(), output.polynomial_size());
+        // Mod switched but the noise is to be interpreted with respect to the input modulus, as
+        // strictly the operation adding the noise is the rounding under the original modulus
+        assert_eq!(self.ciphertext_modulus(), output.ciphertext_modulus());
+
+        for (inp, out) in self.as_ref().iter().zip(output.as_mut().iter_mut()) {
+            let msed = modulus_switch(*inp, output_modulus_log);
+            // Shift in MSBs to match the power of 2 encoding in core
+            *out = msed << (Scalar::BITS - output_modulus_log.0);
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

tfhe/src/core_crypto/commons/noise_formulas/noise_simulation/lwe_packing_keyswitch.rs

@@ -8,11 +8,11 @@ use crate::core_crypto::commons::noise_formulas::noise_simulation::traits::{
 };
 use crate::core_crypto::commons::noise_formulas::noise_simulation::{
     NoiseSimulationGlwe, NoiseSimulationLwe, NoiseSimulationModulus,
+    NoiseSimulationNoiseDistribution, NoiseSimulationNoiseDistributionKind,
 };
 use crate::core_crypto::commons::numeric::UnsignedInteger;
 use crate::core_crypto::commons::parameters::{
-    DecompositionBaseLog, DecompositionLevelCount, DynamicDistribution, GlweSize, LweDimension,
+    DecompositionBaseLog, DecompositionLevelCount, GlweDimension, LweDimension, PolynomialSize,
-    PolynomialSize,
 };
 use crate::core_crypto::commons::traits::container::Container;
 use crate::core_crypto::entities::lwe_packing_keyswitch_key::LwePackingKeyswitchKey;
@@ -22,9 +22,9 @@ pub struct NoiseSimulationLwePackingKeyswitchKey {
     input_lwe_dimension: LweDimension,
     decomp_base_log: DecompositionBaseLog,
     decomp_level_count: DecompositionLevelCount,
-    output_glwe_size: GlweSize,
+    output_glwe_dimension: GlweDimension,
     output_polynomial_size: PolynomialSize,
-    noise_distribution: DynamicDistribution<u128>,
+    noise_distribution: NoiseSimulationNoiseDistribution,
     modulus: NoiseSimulationModulus,
 }
 
@@ -33,16 +33,16 @@ impl NoiseSimulationLwePackingKeyswitchKey {
         input_lwe_dimension: LweDimension,
         decomp_base_log: DecompositionBaseLog,
         decomp_level_count: DecompositionLevelCount,
-        output_glwe_size: GlweSize,
+        output_glwe_dimension: GlweDimension,
         output_polynomial_size: PolynomialSize,
-        noise_distribution: DynamicDistribution<u128>,
+        noise_distribution: NoiseSimulationNoiseDistribution,
         modulus: NoiseSimulationModulus,
     ) -> Self {
         Self {
             input_lwe_dimension,
             decomp_base_log,
             decomp_level_count,
-            output_glwe_size,
+            output_glwe_dimension,
             output_polynomial_size,
             noise_distribution,
             modulus,
@@ -57,7 +57,7 @@ impl NoiseSimulationLwePackingKeyswitchKey {
             input_lwe_dimension,
             decomp_base_log,
             decomp_level_count,
-            output_glwe_size,
+            output_glwe_dimension,
             output_polynomial_size,
             noise_distribution: _,
             modulus,
@@ -66,7 +66,7 @@ impl NoiseSimulationLwePackingKeyswitchKey {
         let pksk_input_lwe_dimension = pksk.input_key_lwe_dimension();
         let pksk_decomp_base_log = pksk.decomposition_base_log();
         let pksk_decomp_level_count = pksk.decomposition_level_count();
-        let pksk_output_glwe_size = pksk.output_glwe_size();
+        let pksk_output_glwe_dimension = pksk.output_key_glwe_dimension();
         let pksk_output_polynomial_size = pksk.output_key_polynomial_size();
         let pksk_modulus =
             NoiseSimulationModulus::from_ciphertext_modulus(pksk.ciphertext_modulus());
@@ -74,7 +74,7 @@ impl NoiseSimulationLwePackingKeyswitchKey {
         input_lwe_dimension == pksk_input_lwe_dimension
             && decomp_base_log == pksk_decomp_base_log
             && decomp_level_count == pksk_decomp_level_count
-            && output_glwe_size == pksk_output_glwe_size
+            && output_glwe_dimension == pksk_output_glwe_dimension
             && output_polynomial_size == pksk_output_polynomial_size
             && modulus == pksk_modulus
     }
@@ -91,15 +91,15 @@ impl NoiseSimulationLwePackingKeyswitchKey {
         self.decomp_level_count
     }
 
-    pub fn output_glwe_size(&self) -> GlweSize {
+    pub fn output_glwe_dimension(&self) -> GlweDimension {
-        self.output_glwe_size
+        self.output_glwe_dimension
     }
 
     pub fn output_polynomial_size(&self) -> PolynomialSize {
         self.output_polynomial_size
     }
 
-    pub fn noise_distribution(&self) -> DynamicDistribution<u128> {
+    pub fn noise_distribution(&self) -> NoiseSimulationNoiseDistribution {
         self.noise_distribution
     }
 
@@ -117,7 +117,7 @@ impl AllocateLwePackingKeyswitchResult for NoiseSimulationLwePackingKeyswitchKey
         _side_resources: &mut Self::SideResources,
     ) -> Self::Output {
         Self::Output::new(
-            self.output_glwe_size().to_glwe_dimension(),
+            self.output_glwe_dimension(),
             self.output_polynomial_size(),
             Variance(f64::NAN),
             self.modulus,
@@ -137,43 +137,44 @@ impl LwePackingKeyswitch<[&NoiseSimulationLwe], NoiseSimulationGlwe>
         _side_resources: &mut Self::SideResources,
     ) {
         let mut input_iter = input.iter();
-        let input = input_iter.next().unwrap();
+        let first_input = input_iter.next().unwrap();
 
-        let mut lwe_to_pack = 1;
+        // Check first input is compatible with us
+        assert_eq!(first_input.lwe_dimension(), self.input_lwe_dimension());
+        // Check all inputs are the same as first input
+        assert!(input_iter.all(|x| x == first_input));
 
-        assert!(input_iter.inspect(|_| lwe_to_pack += 1).all(|x| x == input));
+        let lwe_to_pack = input.len() as f64;
 
-        assert_eq!(input.lwe_dimension(), self.input_lwe_dimension());
+        let packing_ks_additive_var = match self.noise_distribution().kind() {
-
+            NoiseSimulationNoiseDistributionKind::Gaussian => {
-        let packing_ks_additive_var = match self.noise_distribution() {
-            DynamicDistribution::Gaussian(_) => {
                 packing_keyswitch_additive_variance_132_bits_security_gaussian(
                     self.input_lwe_dimension(),
-                    self.output_glwe_size().to_glwe_dimension(),
+                    self.output_glwe_dimension(),
                     self.output_polynomial_size(),
                     self.decomp_base_log(),
                     self.decomp_level_count(),
-                    lwe_to_pack.into(),
+                    lwe_to_pack,
                     self.modulus().as_f64(),
                 )
             }
-            DynamicDistribution::TUniform(_) => {
+            NoiseSimulationNoiseDistributionKind::TUniform => {
                 packing_keyswitch_additive_variance_132_bits_security_tuniform(
                     self.input_lwe_dimension(),
-                    self.output_glwe_size().to_glwe_dimension(),
+                    self.output_glwe_dimension(),
                     self.output_polynomial_size(),
                     self.decomp_base_log(),
                     self.decomp_level_count(),
-                    lwe_to_pack.into(),
+                    lwe_to_pack,
                     self.modulus().as_f64(),
                 )
             }
         };
 
         *output = NoiseSimulationGlwe::new(
-            self.output_glwe_size().to_glwe_dimension(),
+            self.output_glwe_dimension(),
             self.output_polynomial_size(),
-            Variance(input.variance().0 + packing_ks_additive_var.0),
+            Variance(first_input.variance().0 + packing_ks_additive_var.0),
             self.modulus(),
         );
     }

tfhe/src/core_crypto/commons/noise_formulas/noise_simulation/mod.rs

@@ -13,13 +13,14 @@ pub use lwe_programmable_bootstrap::{
 
 use crate::core_crypto::commons::ciphertext_modulus::CiphertextModulus;
 use crate::core_crypto::commons::dispersion::Variance;
+use crate::core_crypto::commons::math::random::DynamicDistribution;
 use crate::core_crypto::commons::noise_formulas::noise_simulation::traits::{
     AllocateLweBootstrapResult, AllocateLweMultiBitBlindRotateResult, LweUncorrelatedAdd,
     LweUncorrelatedSub, ScalarMul, ScalarMulAssign,
 };
 use crate::core_crypto::commons::numeric::{CastInto, UnsignedInteger};
 use crate::core_crypto::commons::parameters::{
-    CiphertextModulusLog, GlweDimension, LweDimension, PolynomialSize,
+    CiphertextModulusLog, GlweDimension, GlweSize, LweDimension, PolynomialSize,
 };
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -71,6 +72,40 @@ impl NoiseSimulationModulus {
     }
 }
 
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub enum NoiseSimulationNoiseDistribution {
+    U32(DynamicDistribution<u32>),
+    U64(DynamicDistribution<u64>),
+    U128(DynamicDistribution<u128>),
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum NoiseSimulationNoiseDistributionKind {
+    Gaussian,
+    TUniform,
+}
+
+impl NoiseSimulationNoiseDistribution {
+    pub fn kind(&self) -> NoiseSimulationNoiseDistributionKind {
+        match self {
+            Self::U32(dynamic_distribution) => dynamic_distribution.into(),
+            Self::U64(dynamic_distribution) => dynamic_distribution.into(),
+            Self::U128(dynamic_distribution) => dynamic_distribution.into(),
+        }
+    }
+}
+
+impl<Scalar: UnsignedInteger> From<&DynamicDistribution<Scalar>>
+    for NoiseSimulationNoiseDistributionKind
+{
+    fn from(value: &DynamicDistribution<Scalar>) -> Self {
+        match value {
+            DynamicDistribution::Gaussian(_) => Self::Gaussian,
+            DynamicDistribution::TUniform(_) => Self::TUniform,
+        }
+    }
+}
+
 // Avoids fields to be public/accessible in the noise_simulation module to make sure all functions
 // use constructors
 mod simulation_ciphertexts {
@@ -211,6 +246,10 @@ mod simulation_ciphertexts {
             self.glwe_dimension
         }
 
+        pub fn glwe_size(&self) -> GlweSize {
+            self.glwe_dimension().to_glwe_size()
+        }
+
         pub fn polynomial_size(&self) -> PolynomialSize {
             self.polynomial_size
         }

tfhe/src/core_crypto/commons/noise_formulas/noise_simulation/modulus_switch.rs

@@ -8,7 +8,7 @@ use crate::core_crypto::commons::noise_formulas::noise_simulation::traits::{
     StandardModSwitch,
 };
 use crate::core_crypto::commons::noise_formulas::noise_simulation::{
-    NoiseSimulationLwe, NoiseSimulationModulus,
+    NoiseSimulationGlwe, NoiseSimulationLwe, NoiseSimulationModulus,
 };
 use crate::core_crypto::commons::parameters::{CiphertextModulusLog, LweBskGroupingFactor};
 
@@ -160,3 +160,56 @@ impl CenteredBinaryShiftedStandardModSwitch<Self> for NoiseSimulationLwe {
         );
     }
 }
+
+impl AllocateStandardModSwitchResult for NoiseSimulationGlwe {
+    type Output = Self;
+    type SideResources = ();
+
+    fn allocate_standard_mod_switch_result(
+        &self,
+        _side_resources: &mut Self::SideResources,
+    ) -> Self::Output {
+        Self::Output::new(
+            self.glwe_dimension(),
+            self.polynomial_size(),
+            self.variance_per_occupied_slot(),
+            self.modulus(),
+        )
+    }
+}
+
+impl StandardModSwitch<Self> for NoiseSimulationGlwe {
+    type SideResources = ();
+
+    fn standard_mod_switch(
+        &self,
+        output_modulus_log: CiphertextModulusLog,
+        output: &mut Self,
+        _side_resources: &mut Self::SideResources,
+    ) {
+        let simulation_after_mod_switch_modulus =
+            NoiseSimulationModulus::from_ciphertext_modulus_log(output_modulus_log);
+
+        let input_modulus_f64 = self.modulus().as_f64();
+        let output_modulus_f64 = simulation_after_mod_switch_modulus.as_f64();
+
+        assert!(output_modulus_f64 < input_modulus_f64);
+
+        let mod_switch_additive_variance = modulus_switch_additive_variance(
+            self.glwe_dimension()
+                .to_equivalent_lwe_dimension(self.polynomial_size()),
+            input_modulus_f64,
+            output_modulus_f64,
+        );
+
+        *output = Self::new(
+            self.glwe_dimension(),
+            self.polynomial_size(),
+            Variance(self.variance_per_occupied_slot().0 + mod_switch_additive_variance.0),
+            // Mod switched but the noise is to be interpreted with respect to the input modulus,
+            // as strictly the operation adding the noise is the rounding under the
+            // original modulus
+            self.modulus(),
+        )
+    }
+}

tfhe/src/shortint/client_key/atomic_pattern/standard.rs

@@ -251,7 +251,21 @@ impl StandardAtomicPatternClientKey {
         &self,
         private_compression_key: &CompressionPrivateKeys,
     ) -> CompressionKey {
-        private_compression_key.new_compression_key(&self.glwe_secret_key, self.parameters())
+        ShortintEngine::with_thread_local_mut(|engine| {
+            self.new_compression_key_with_engine(private_compression_key, engine)
+        })
+    }
+
+    pub(crate) fn new_compression_key_with_engine(
+        &self,
+        private_compression_key: &CompressionPrivateKeys,
+        engine: &mut ShortintEngine,
+    ) -> CompressionKey {
+        private_compression_key.new_compression_key_with_engine(
+            &self.glwe_secret_key,
+            self.parameters(),
+            engine,
+        )
     }
 
     pub fn new_compressed_compression_key(

tfhe/src/shortint/list_compression/private_key.rs

@@ -39,6 +39,17 @@ impl CompressionPrivateKeys {
         &self,
         glwe_secret_key: &GlweSecretKey<Vec<u64>>,
         pbs_params: ShortintParameterSet,
+    ) -> CompressionKey {
+        ShortintEngine::with_thread_local_mut(|engine| {
+            self.new_compression_key_with_engine(glwe_secret_key, pbs_params, engine)
+        })
+    }
+
+    pub(crate) fn new_compression_key_with_engine(
+        &self,
+        glwe_secret_key: &GlweSecretKey<Vec<u64>>,
+        pbs_params: ShortintParameterSet,
+        engine: &mut ShortintEngine,
     ) -> CompressionKey {
         assert_eq!(
             pbs_params.encryption_key_choice(),
@@ -57,8 +68,7 @@ impl CompressionPrivateKeys {
             "Compression parameters say to store more bits than useful"
         );
 
-        let packing_key_switching_key = ShortintEngine::with_thread_local_mut(|engine| {
+        let packing_key_switching_key = allocate_and_generate_new_lwe_packing_keyswitch_key(
-            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(),
@@ -66,8 +76,7 @@ impl CompressionPrivateKeys {
             compression_params.packing_ks_key_noise_distribution(),
             pbs_params.ciphertext_modulus(),
             &mut engine.encryption_generator,
-            )
+        );
-        });
 
         CompressionKey {
             packing_key_switching_key,
@@ -75,6 +84,7 @@ impl CompressionPrivateKeys {
             storage_log_modulus: compression_params.storage_log_modulus(),
         }
     }
+
     pub(crate) fn new_compressed_compression_key(
         &self,
         glwe_secret_key: &GlweSecretKey<Vec<u64>>,

tfhe/src/shortint/list_compression/server_keys.rs

@@ -78,6 +78,14 @@ impl ClientKey {
             )
     }
 
+    pub fn new_compression_key(
+        &self,
+        private_compression_key: &CompressionPrivateKeys,
+    ) -> CompressionKey {
+        self.atomic_pattern
+            .new_compression_key(private_compression_key)
+    }
+
     pub fn new_compression_decompression_keys(
         &self,
         private_compression_key: &CompressionPrivateKeys,

tfhe/src/shortint/server_key/tests/noise_distribution/br_dp_packingks_ms.rs

@@ -0,0 +1,735 @@
+use super::utils::noise_simulation::*;
+use super::utils::traits::*;
+use super::utils::{
+    expected_pfail_for_precision, mean_and_variance_check, normality_check, pfail_check,
+    precision_with_padding, DecryptionAndNoiseResult, NoiseSample, PfailAndPrecision,
+    PfailTestMeta, PfailTestResult,
+};
+use super::{should_run_short_pfail_tests_debug, should_use_single_key_debug};
+use crate::shortint::atomic_pattern::AtomicPattern;
+use crate::shortint::ciphertext::{Ciphertext, Degree, NoiseLevel};
+use crate::shortint::client_key::atomic_pattern::AtomicPatternClientKey;
+use crate::shortint::client_key::ClientKey;
+use crate::shortint::engine::ShortintEngine;
+use crate::shortint::list_compression::{CompressionKey, CompressionPrivateKeys};
+use crate::shortint::parameters::test_params::{
+    TEST_COMP_PARAM_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::parameters::{
+    AtomicPatternParameters, CarryModulus, CiphertextModulusLog, CompressionParameters,
+    MessageModulus, PBSParameters, Variance,
+};
+use crate::shortint::server_key::ServerKey;
+use crate::shortint::{PaddingBit, ShortintEncoding};
+use rayon::prelude::*;
+
+#[allow(clippy::too_many_arguments)]
+#[allow(clippy::type_complexity)]
+pub fn br_dp_packing_ks_ms<
+    InputCt,
+    PBSResult,
+    PBSKey,
+    Accumulator,
+    DPScalar,
+    DPResult,
+    PackingKsk,
+    PackingKsResult,
+    MsResult,
+    Resources,
+>(
+    input: Vec<InputCt>,
+    bsk: &PBSKey,
+    accumulator: &Accumulator,
+    scalar: DPScalar,
+    packing_ksk: &PackingKsk,
+    storage_modulus_log: CiphertextModulusLog,
+    side_resources: &mut [Resources],
+) -> (
+    Vec<(InputCt, PBSResult, DPResult)>,
+    PackingKsResult,
+    MsResult,
+)
+where
+    Accumulator: AllocateLweBootstrapResult<Output = PBSResult, SideResources = Resources> + Sync,
+    PBSKey:
+        LweClassicFftBootstrap<InputCt, PBSResult, Accumulator, SideResources = Resources> + Sync,
+    PBSResult: ScalarMul<DPScalar, Output = DPResult, SideResources = Resources> + Send,
+    PackingKsk: AllocateLwePackingKeyswitchResult<Output = PackingKsResult, SideResources = Resources>
+        + for<'a> LwePackingKeyswitch<[&'a DPResult], PackingKsResult, SideResources = Resources>,
+    PackingKsResult: AllocateStandardModSwitchResult<Output = MsResult, SideResources = Resources>
+        + StandardModSwitch<MsResult, SideResources = Resources>,
+    InputCt: Send,
+    DPResult: Send,
+    DPScalar: Copy + Send + Sync,
+    Resources: Send,
+{
+    let res: Vec<_> = input
+        .into_par_iter()
+        .zip(side_resources.par_iter_mut())
+        .map(|(input, side_resources)| {
+            let mut pbs_result = accumulator.allocate_lwe_bootstrap_result(side_resources);
+            bsk.lwe_classic_fft_pbs(&input, &mut pbs_result, accumulator, side_resources);
+            let after_dp = pbs_result.scalar_mul(scalar, side_resources);
+
+            (input, pbs_result, after_dp)
+        })
+        .collect();
+
+    let after_dp: Vec<_> = res
+        .iter()
+        .map(|(_input, _pbs_result, after_dp)| after_dp)
+        .collect();
+
+    let mut packing_result =
+        packing_ksk.allocate_lwe_packing_keyswitch_result(&mut side_resources[0]);
+    packing_ksk.keyswitch_lwes_and_pack_in_glwe(
+        after_dp.as_slice(),
+        &mut packing_result,
+        &mut side_resources[0],
+    );
+
+    let mut ms_result = packing_result.allocate_standard_mod_switch_result(&mut side_resources[0]);
+    packing_result.standard_mod_switch(storage_modulus_log, &mut ms_result, &mut side_resources[0]);
+
+    (res, packing_result, ms_result)
+}
+
+fn sanity_check_encrypt_br_dp_packing_ks_ms<P>(params: P, comp_params: CompressionParameters)
+where
+    P: Into<AtomicPatternParameters>,
+{
+    let params: AtomicPatternParameters = params.into();
+    let cks = ClientKey::new(params);
+    let sks = ServerKey::new(&cks);
+    let compression_private_key = cks.new_compression_private_key(comp_params);
+    let compression_key = cks.new_compression_key(&compression_private_key);
+
+    let lwe_per_glwe = compression_key.lwe_per_glwe;
+    // The multiplication done in the compression is made to move the message up at the top of the
+    // carry space, multiplying by the carry modulus achieves that
+    let dp_scalar = params.carry_modulus().0;
+    let br_input_modulus_log = sks.br_input_modulus_log();
+    let storage_modulus_log = compression_key.storage_log_modulus;
+
+    let id_lut = sks.generate_lookup_table(|x| x);
+
+    let input_zeros: Vec<_> = (0..lwe_per_glwe.0)
+        .map(|_| cks.encrypt_noiseless_pbs_input_dyn_lwe(br_input_modulus_log, 0))
+        .collect();
+    let mut side_resources = vec![(); input_zeros.len()];
+
+    let (before_packing, _after_packing, mut after_ms) = br_dp_packing_ks_ms(
+        input_zeros,
+        &sks,
+        &id_lut,
+        dp_scalar,
+        &compression_key,
+        storage_modulus_log,
+        &mut side_resources,
+    );
+
+    let compression_inputs: Vec<_> = before_packing
+        .into_iter()
+        .map(|(_input, pbs_result, _dp_result)| {
+            Ciphertext::new(
+                pbs_result.into_lwe_64(),
+                Degree::new(sks.message_modulus.0 - 1),
+                NoiseLevel::NOMINAL,
+                sks.message_modulus,
+                sks.carry_modulus,
+                sks.atomic_pattern.kind(),
+            )
+        })
+        .collect();
+
+    let compressed = compression_key.compress_ciphertexts_into_list(&compression_inputs);
+
+    let underlying_glwes = compressed.modulus_switched_glwe_ciphertext_list;
+
+    assert_eq!(underlying_glwes.len(), 1);
+
+    let extracted = underlying_glwes[0].extract();
+
+    // Bodies that were not filled are discarded
+    after_ms.get_mut_body().as_mut()[lwe_per_glwe.0..].fill(0);
+
+    assert_eq!(after_ms.as_view(), extracted.as_view());
+}
+
+#[test]
+fn test_sanity_check_encrypt_br_dp_packing_ks_ms_test_param_message_2_carry_2_ks_pbs_tuniform_2m128(
+) {
+    sanity_check_encrypt_br_dp_packing_ks_ms(
+        TEST_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
+        TEST_COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
+    )
+}
+
+#[test]
+fn test_sanity_check_encrypt_br_dp_packing_ks_ms_test_param_message_2_carry_2_ks32_pbs_tuniform_2m128(
+) {
+    sanity_check_encrypt_br_dp_packing_ks_ms(
+        TEST_PARAM_MESSAGE_2_CARRY_2_KS32_PBS_TUNIFORM_2M128,
+        TEST_COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
+    )
+}
+
+#[allow(clippy::type_complexity)]
+fn encrypt_br_dp_packing_ks_ms_inner_helper(
+    params: AtomicPatternParameters,
+    comp_params: CompressionParameters,
+    single_cks: &ClientKey,
+    single_sks: &ServerKey,
+    single_compression_private_key: &CompressionPrivateKeys,
+    single_compression_key: &CompressionKey,
+    msg: u64,
+) -> (
+    Vec<(
+        DecryptionAndNoiseResult,
+        DecryptionAndNoiseResult,
+        DecryptionAndNoiseResult,
+    )>,
+    Vec<DecryptionAndNoiseResult>,
+    Vec<DecryptionAndNoiseResult>,
+) {
+    let mut engine = ShortintEngine::new();
+    let thread_cks;
+    let thread_sks;
+    let thread_compression_private_key;
+    let thread_compression_key;
+    let (cks, sks, compression_private_key, compression_key) = if should_use_single_key_debug() {
+        (
+            single_cks,
+            single_sks,
+            single_compression_private_key,
+            single_compression_key,
+        )
+    } else {
+        thread_cks = engine.new_client_key(params);
+        thread_sks = engine.new_server_key(&thread_cks);
+
+        thread_compression_private_key =
+            thread_cks.new_compression_private_key_with_engine(comp_params, &mut engine);
+        thread_compression_key = thread_cks.new_compression_key(&thread_compression_private_key);
+
+        (
+            &thread_cks,
+            &thread_sks,
+            &thread_compression_private_key,
+            &thread_compression_key,
+        )
+    };
+
+    let br_input_modulus_log = sks.br_input_modulus_log();
+    let lwe_per_glwe = compression_key.lwe_per_glwe;
+
+    let input_zeros: Vec<_> = (0..lwe_per_glwe.0)
+        .map(|_| {
+            cks.encrypt_noiseless_pbs_input_dyn_lwe_with_engine(
+                br_input_modulus_log,
+                msg,
+                &mut engine,
+            )
+        })
+        .collect();
+
+    let id_lut = sks.generate_lookup_table(|x| x);
+    let mut side_resources = vec![(); input_zeros.len()];
+    let dp_scalar = params.carry_modulus().0;
+    let storage_modulus_log = compression_key.storage_log_modulus;
+
+    let (before_packing, after_packing, after_ms) = br_dp_packing_ks_ms(
+        input_zeros,
+        sks,
+        &id_lut,
+        dp_scalar,
+        compression_key,
+        storage_modulus_log,
+        &mut side_resources,
+    );
+
+    let compute_large_lwe_secret_key = cks.encryption_key();
+    let compression_glwe_secret_key = &compression_private_key.post_packing_ks_key;
+
+    let compute_encoding = sks.encoding(PaddingBit::Yes);
+    let compression_encoding = ShortintEncoding {
+        carry_modulus: CarryModulus(1),
+        ..compute_encoding
+    };
+
+    (
+        before_packing
+            .into_iter()
+            .map(|(input, pbs_result, dp_result)| {
+                (
+                    match &cks.atomic_pattern {
+                        AtomicPatternClientKey::Standard(standard_atomic_pattern_client_key) => {
+                            DecryptionAndNoiseResult::new_from_lwe(
+                                input.as_ref_64(),
+                                &standard_atomic_pattern_client_key.lwe_secret_key,
+                                msg,
+                                &compute_encoding,
+                            )
+                        }
+                        AtomicPatternClientKey::KeySwitch32(ks32_atomic_pattern_client_key) => {
+                            let ks32_params = ks32_atomic_pattern_client_key.parameters;
+                            let compute_encoding_32 = ShortintEncoding {
+                                ciphertext_modulus: ks32_params.post_keyswitch_ciphertext_modulus,
+                                message_modulus: ks32_params.message_modulus,
+                                carry_modulus: ks32_params.carry_modulus,
+                                padding_bit: PaddingBit::Yes,
+                            };
+
+                            DecryptionAndNoiseResult::new_from_lwe(
+                                input.as_ref_32(),
+                                &ks32_atomic_pattern_client_key.lwe_secret_key,
+                                msg.try_into().unwrap(),
+                                &compute_encoding_32,
+                            )
+                        }
+                    },
+                    DecryptionAndNoiseResult::new_from_lwe(
+                        pbs_result.as_ref_64(),
+                        &compute_large_lwe_secret_key,
+                        msg,
+                        &compute_encoding,
+                    ),
+                    DecryptionAndNoiseResult::new_from_lwe(
+                        dp_result.as_ref_64(),
+                        &compute_large_lwe_secret_key,
+                        msg,
+                        &compression_encoding,
+                    ),
+                )
+            })
+            .collect(),
+        DecryptionAndNoiseResult::new_from_glwe(
+            &after_packing,
+            compression_glwe_secret_key,
+            compression_private_key.params.lwe_per_glwe(),
+            msg,
+            &compression_encoding,
+        ),
+        DecryptionAndNoiseResult::new_from_glwe(
+            &after_ms,
+            compression_glwe_secret_key,
+            compression_private_key.params.lwe_per_glwe(),
+            msg,
+            &compression_encoding,
+        ),
+    )
+}
+
+#[allow(clippy::type_complexity)]
+fn encrypt_br_dp_packing_ks_ms_noise_helper(
+    params: AtomicPatternParameters,
+    comp_params: CompressionParameters,
+    single_cks: &ClientKey,
+    single_sks: &ServerKey,
+    single_compression_private_key: &CompressionPrivateKeys,
+    single_compression_key: &CompressionKey,
+    msg: u64,
+) -> (
+    Vec<(NoiseSample, NoiseSample, NoiseSample)>,
+    Vec<NoiseSample>,
+    Vec<NoiseSample>,
+) {
+    let (before_packing, after_packing, after_ms) = encrypt_br_dp_packing_ks_ms_inner_helper(
+        params,
+        comp_params,
+        single_cks,
+        single_sks,
+        single_compression_private_key,
+        single_compression_key,
+        msg,
+    );
+
+    (
+        before_packing
+            .into_iter()
+            .map(|(input, after_pbs, after_dp)| {
+                (
+                    input
+                        .get_noise_if_decryption_was_correct()
+                        .expect("Decryption Failed"),
+                    after_pbs
+                        .get_noise_if_decryption_was_correct()
+                        .expect("Decryption Failed"),
+                    after_dp
+                        .get_noise_if_decryption_was_correct()
+                        .expect("Decryption Failed"),
+                )
+            })
+            .collect(),
+        after_packing
+            .into_iter()
+            .map(|x| {
+                x.get_noise_if_decryption_was_correct()
+                    .expect("Decryption Failed")
+            })
+            .collect(),
+        after_ms
+            .into_iter()
+            .map(|x| {
+                x.get_noise_if_decryption_was_correct()
+                    .expect("Decryption Failed")
+            })
+            .collect(),
+    )
+}
+
+#[allow(clippy::type_complexity)]
+fn encrypt_br_dp_packing_ks_ms_pfail_helper(
+    params: AtomicPatternParameters,
+    comp_params: CompressionParameters,
+    single_cks: &ClientKey,
+    single_sks: &ServerKey,
+    single_compression_private_key: &CompressionPrivateKeys,
+    single_compression_key: &CompressionKey,
+    msg: u64,
+) -> Vec<DecryptionAndNoiseResult> {
+    let (_before_packing, _after_packing, after_ms) = encrypt_br_dp_packing_ks_ms_inner_helper(
+        params,
+        comp_params,
+        single_cks,
+        single_sks,
+        single_compression_private_key,
+        single_compression_key,
+        msg,
+    );
+
+    after_ms
+}
+
+fn noise_check_encrypt_br_dp_packing_ks_ms_noise<P>(params: P, comp_params: CompressionParameters)
+where
+    P: Into<AtomicPatternParameters>,
+{
+    let params: AtomicPatternParameters = params.into();
+    let cks = ClientKey::new(params);
+    let sks = ServerKey::new(&cks);
+    let compression_private_key = cks.new_compression_private_key(comp_params);
+    let compression_key = cks.new_compression_key(&compression_private_key);
+
+    let noise_simulation_bsk =
+        NoiseSimulationLweFourierBsk::new_from_atomic_pattern_parameters(params);
+    let noise_simulation_packing_key =
+        NoiseSimulationLwePackingKeyswitchKey::new_from_comp_parameters(params, comp_params);
+
+    assert!(noise_simulation_bsk.matches_actual_shortint_server_key(&sks));
+    assert!(noise_simulation_packing_key.matches_actual_shortint_comp_key(&compression_key));
+
+    // The multiplication done in the compression is made to move the message up at the top of the
+    // carry space, multiplying by the carry modulus achieves that
+    let dp_scalar = params.carry_modulus().0;
+
+    let noise_simulation_accumulator = NoiseSimulationGlwe::new(
+        noise_simulation_bsk.output_glwe_size().to_glwe_dimension(),
+        noise_simulation_bsk.output_polynomial_size(),
+        Variance(0.0),
+        noise_simulation_bsk.modulus(),
+    );
+
+    let lwe_per_glwe = compression_key.lwe_per_glwe;
+    let storage_modulus_log = compression_key.storage_log_modulus;
+    let br_input_modulus_log = sks.br_input_modulus_log();
+
+    let (_before_packing_sim, _after_packing_sim, after_ms_sim) = {
+        let noise_simulation = NoiseSimulationLwe::new(
+            cks.parameters().lwe_dimension(),
+            Variance(0.0),
+            NoiseSimulationModulus::from_ciphertext_modulus(cks.parameters().ciphertext_modulus()),
+        );
+        br_dp_packing_ks_ms(
+            vec![noise_simulation; lwe_per_glwe.0],
+            &noise_simulation_bsk,
+            &noise_simulation_accumulator,
+            dp_scalar,
+            &noise_simulation_packing_key,
+            storage_modulus_log,
+            &mut vec![(); lwe_per_glwe.0],
+        )
+    };
+
+    let input_zeros: Vec<_> = (0..lwe_per_glwe.0)
+        .map(|_| cks.encrypt_noiseless_pbs_input_dyn_lwe(br_input_modulus_log, 0))
+        .collect();
+    let id_lut = sks.generate_lookup_table(|x| x);
+    let mut side_resources = vec![(); input_zeros.len()];
+
+    // Check that the circuit is correct with respect to core implementation, i.e. does not crash on
+    // dimension checks
+    let (expected_glwe_size_out, expected_polynomial_size_out, expected_modulus_f64_out) = {
+        let (_before_packing_sim, _after_packing, after_ms) = br_dp_packing_ks_ms(
+            input_zeros,
+            &sks,
+            &id_lut,
+            dp_scalar,
+            &compression_key,
+            storage_modulus_log,
+            &mut side_resources,
+        );
+
+        (
+            after_ms.glwe_size(),
+            after_ms.polynomial_size(),
+            after_ms.ciphertext_modulus().raw_modulus_float(),
+        )
+    };
+
+    assert_eq!(after_ms_sim.glwe_size(), expected_glwe_size_out);
+    assert_eq!(after_ms_sim.polynomial_size(), expected_polynomial_size_out);
+    assert_eq!(after_ms_sim.modulus().as_f64(), expected_modulus_f64_out);
+
+    let cleartext_modulus = params.message_modulus().0 * params.carry_modulus().0;
+    let mut noise_samples_before_ms = vec![];
+    let mut noise_samples_after_ms = vec![];
+
+    let sample_count_per_msg = 1000usize;
+
+    for _ in 0..cleartext_modulus {
+        let (current_noise_samples_before_ms, current_noise_samples_after_ms): (Vec<_>, Vec<_>) =
+            (0..sample_count_per_msg)
+                .into_par_iter()
+                .map(|_| {
+                    let (_before_packing, after_packing, after_ms) =
+                        encrypt_br_dp_packing_ks_ms_noise_helper(
+                            params,
+                            comp_params,
+                            &cks,
+                            &sks,
+                            &compression_private_key,
+                            &compression_key,
+                            0,
+                        );
+                    (after_packing, after_ms)
+                })
+                .flatten()
+                .unzip();
+
+        noise_samples_before_ms
+            .extend(current_noise_samples_before_ms.into_iter().map(|x| x.value));
+        noise_samples_after_ms.extend(current_noise_samples_after_ms.into_iter().map(|x| x.value));
+    }
+
+    let before_ms_normality = normality_check(&noise_samples_before_ms, "before ms", 0.01);
+
+    let after_ms_is_ok = mean_and_variance_check(
+        &noise_samples_after_ms,
+        "after_ms",
+        0.0,
+        after_ms_sim.variance_per_occupied_slot(),
+        comp_params.packing_ks_key_noise_distribution(),
+        after_ms_sim
+            .glwe_dimension()
+            .to_equivalent_lwe_dimension(after_ms_sim.polynomial_size()),
+        after_ms_sim.modulus().as_f64(),
+    );
+
+    assert!(before_ms_normality.null_hypothesis_is_valid && after_ms_is_ok);
+}
+
+#[test]
+fn test_noise_check_encrypt_br_dp_packing_ks_ms_noise_test_param_message_2_carry_2_ks_pbs_tuniform_2m128(
+) {
+    noise_check_encrypt_br_dp_packing_ks_ms_noise(
+        TEST_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
+        TEST_COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
+    )
+}
+
+#[test]
+fn test_noise_check_encrypt_br_dp_packing_ks_ms_noise_test_param_message_2_carry_2_ks32_pbs_tuniform_2m128(
+) {
+    noise_check_encrypt_br_dp_packing_ks_ms_noise(
+        TEST_PARAM_MESSAGE_2_CARRY_2_KS32_PBS_TUNIFORM_2M128,
+        TEST_COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
+    )
+}
+
+fn noise_check_encrypt_br_dp_packing_ks_ms_pfail<P>(params: P, comp_params: CompressionParameters)
+where
+    P: Into<AtomicPatternParameters>,
+{
+    let (pfail_test_meta, params) = {
+        let mut params: AtomicPatternParameters = params.into();
+
+        let original_message_modulus = params.message_modulus();
+        let original_carry_modulus = params.carry_modulus();
+
+        // For now only allow 2_2 parameters, and see later for heuristics to use
+        assert_eq!(original_message_modulus.0, 4);
+        assert_eq!(original_carry_modulus.0, 4);
+
+        let noise_simulation_bsk =
+            NoiseSimulationLweFourierBsk::new_from_atomic_pattern_parameters(params);
+        let noise_simulation_packing_key =
+            NoiseSimulationLwePackingKeyswitchKey::new_from_comp_parameters(params, comp_params);
+
+        // The multiplication done in the compression is made to move the message up at the top of
+        // the carry space, multiplying by the carry modulus achieves that
+        let dp_scalar = params.carry_modulus().0;
+
+        let noise_simulation_accumulator = NoiseSimulationGlwe::new(
+            noise_simulation_bsk.output_glwe_size().to_glwe_dimension(),
+            noise_simulation_bsk.output_polynomial_size(),
+            Variance(0.0),
+            noise_simulation_bsk.modulus(),
+        );
+
+        let lwe_per_glwe = comp_params.lwe_per_glwe();
+        let storage_modulus_log = comp_params.storage_log_modulus();
+
+        let (_before_packing_sim, _after_packing_sim, after_ms_sim) = {
+            let noise_simulation = NoiseSimulationLwe::new(
+                params.lwe_dimension(),
+                Variance(0.0),
+                NoiseSimulationModulus::from_ciphertext_modulus(params.ciphertext_modulus()),
+            );
+            br_dp_packing_ks_ms(
+                vec![noise_simulation; lwe_per_glwe.0],
+                &noise_simulation_bsk,
+                &noise_simulation_accumulator,
+                dp_scalar,
+                &noise_simulation_packing_key,
+                storage_modulus_log,
+                &mut vec![(); lwe_per_glwe.0],
+            )
+        };
+
+        let expected_variance_after_storage = after_ms_sim.variance_per_occupied_slot();
+
+        let compression_carry_mod = CarryModulus(1);
+        let compression_message_mod = original_message_modulus;
+        let compression_precision_with_padding =
+            precision_with_padding(compression_message_mod, compression_carry_mod);
+        let expected_pfail_for_storage = expected_pfail_for_precision(
+            compression_precision_with_padding,
+            expected_variance_after_storage,
+        );
+
+        let original_pfail_and_precision = PfailAndPrecision::new(
+            expected_pfail_for_storage,
+            compression_message_mod,
+            compression_carry_mod,
+        );
+
+        let updated_message_mod = MessageModulus(1 << 6);
+        let updated_carry_mod = compression_carry_mod;
+
+        let updated_precision_with_padding =
+            precision_with_padding(updated_message_mod, updated_carry_mod);
+
+        let new_expected_pfail_for_storage = expected_pfail_for_precision(
+            updated_precision_with_padding,
+            expected_variance_after_storage,
+        );
+
+        let new_expected_pfail_and_precision = PfailAndPrecision::new(
+            new_expected_pfail_for_storage,
+            updated_message_mod,
+            updated_carry_mod,
+        );
+
+        // Here we update the message modulus only:
+        // - because the message modulus matches for the compression encoding and compute encoding
+        // - so that the carry modulus stays the same and we apply the same dot product as normal
+        //   for 2_2
+        // - so that the effective encoding after the storage is the one we used to evaluate the
+        //   pfail
+        // TODO: do something about this
+        match &mut params {
+            AtomicPatternParameters::Standard(pbsparameters) => match pbsparameters {
+                PBSParameters::PBS(classic_pbsparameters) => {
+                    classic_pbsparameters.message_modulus = updated_message_mod
+                }
+                PBSParameters::MultiBitPBS(multi_bit_pbsparameters) => {
+                    multi_bit_pbsparameters.message_modulus = updated_message_mod
+                }
+            },
+            AtomicPatternParameters::KeySwitch32(key_switch32_pbsparameters) => {
+                key_switch32_pbsparameters.message_modulus = updated_message_mod
+            }
+        }
+
+        let pfail_test_meta = if should_run_short_pfail_tests_debug() {
+            // To have the same amount of keys generated as the case where a single run is a single
+            // sample
+            let expected_fails = 200 * lwe_per_glwe.0 as u32;
+            PfailTestMeta::new_with_desired_expected_fails(
+                original_pfail_and_precision,
+                new_expected_pfail_and_precision,
+                expected_fails,
+            )
+        } else {
+            // To guarantee 1_000_000 keysets are generated
+            let total_runs = 1_000_000 * lwe_per_glwe.0 as u32;
+            PfailTestMeta::new_with_total_runs(
+                original_pfail_and_precision,
+                new_expected_pfail_and_precision,
+                total_runs,
+            )
+        };
+
+        (pfail_test_meta, params)
+    };
+
+    let cks = ClientKey::new(params);
+    let sks = ServerKey::new(&cks);
+    let compression_private_key = cks.new_compression_private_key(comp_params);
+    let compression_key = cks.new_compression_key(&compression_private_key);
+
+    let lwe_per_glwe = compression_key.lwe_per_glwe;
+
+    let total_runs_for_expected_fails = pfail_test_meta
+        .total_runs_for_expected_fails()
+        .div_ceil(lwe_per_glwe.0.try_into().unwrap());
+
+    println!(
+        "Actual runs with {} samples per run: {total_runs_for_expected_fails}",
+        lwe_per_glwe.0
+    );
+
+    let measured_fails: f64 = (0..total_runs_for_expected_fails)
+        .into_par_iter()
+        .map(|_| {
+            let after_ms_decryption_result = encrypt_br_dp_packing_ks_ms_pfail_helper(
+                params,
+                comp_params,
+                &cks,
+                &sks,
+                &compression_private_key,
+                &compression_key,
+                0,
+            );
+            after_ms_decryption_result
+                .into_iter()
+                .map(|x| x.failure_as_f64())
+                .sum::<f64>()
+        })
+        .sum();
+
+    let test_result = PfailTestResult { measured_fails };
+
+    pfail_check(&pfail_test_meta, test_result);
+}
+
+#[test]
+fn test_noise_check_encrypt_br_dp_packing_ks_ms_pfail_test_param_message_2_carry_2_ks_pbs_tuniform_2m128(
+) {
+    noise_check_encrypt_br_dp_packing_ks_ms_pfail(
+        TEST_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
+        TEST_COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
+    )
+}
+
+#[test]
+fn test_noise_check_encrypt_br_dp_packing_ks_ms_pfail_test_param_message_2_carry_2_ks32_pbs_tuniform_2m128(
+) {
+    noise_check_encrypt_br_dp_packing_ks_ms_pfail(
+        TEST_PARAM_MESSAGE_2_CARRY_2_KS32_PBS_TUNIFORM_2M128,
+        TEST_COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
+    )
+}

tfhe/src/shortint/server_key/tests/noise_distribution/br_rerand_dp_ks_ms.rs

@@ -233,7 +233,7 @@ fn encrypt_decomp_br_rerand_dp_ks_any_ms_inner_helper(
     let br_input_modulus_log = sks.br_input_modulus_log();
     let modulus_switch_config = sks.noise_simulation_modulus_switch_config();
 
-    let ct = comp_private_key.encrypt_noiseless_decompression_input_dyn_lwe(cks, 0, &mut engine);
+    let ct = comp_private_key.encrypt_noiseless_decompression_input_dyn_lwe(cks, msg, &mut engine);
 
     let cpk_ct_zero_rerand = {
         let compact_list = cpk.encrypt_iter_with_modulus_with_engine(

tfhe/src/shortint/server_key/tests/noise_distribution/mod.rs

@@ -1,4 +1,5 @@
 pub(crate) mod br_dp_ks_ms;
+pub(crate) mod br_dp_packingks_ms;
 pub(crate) mod br_rerand_dp_ks_ms;
 pub(crate) mod cpk_ks_ms;
 pub(crate) mod dp_ks_ms;

tfhe/src/shortint/server_key/tests/noise_distribution/utils/mod.rs

@@ -35,6 +35,11 @@ use crate::shortint::parameters::{
     AtomicPatternParameters, CarryModulus, MessageModulus, PBSParameters,
 };
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct PrecisionWithPadding {
+    value: u32,
+}
+
 pub fn normality_check(
     noise_samples: &[f64],
     check_location: &str,
@@ -187,7 +192,7 @@ pub fn encrypt_new_noiseless_lwe<
 #[derive(Clone, Copy)]
 pub struct PfailAndPrecision {
     pfail: f64,
-    precision_with_padding: u32,
+    precision_with_padding: PrecisionWithPadding,
 }
 
 impl PfailAndPrecision {
@@ -215,7 +220,7 @@ impl PfailAndPrecision {
         self.pfail
     }
 
-    pub fn precision_with_padding(&self) -> u32 {
+    pub fn precision_with_padding(&self) -> PrecisionWithPadding {
         self.precision_with_padding
     }
 }
@@ -310,9 +315,9 @@ pub fn pfail_check(pfail_test_meta: &PfailTestMeta, pfail_test_result: PfailTest
     println!("expected_pfail={expected_pfail}");
 
     let equivalent_measured_pfail = equivalent_pfail_gaussian_noise(
-        new_precision_with_padding,
+        new_precision_with_padding.value,
         measured_pfail,
-        original_precision_with_padding,
+        original_precision_with_padding.value,
     );
 
     println!("equivalent_measured_pfail={equivalent_measured_pfail}");
@@ -336,14 +341,14 @@ pub fn pfail_check(pfail_test_meta: &PfailTestMeta, pfail_test_result: PfailTest
         println!("pfail_upper_bound={pfail_upper_bound}");
 
         let equivalent_pfail_lower_bound = equivalent_pfail_gaussian_noise(
-            new_precision_with_padding,
+            new_precision_with_padding.value,
             pfail_lower_bound,
-            original_precision_with_padding,
+            original_precision_with_padding.value,
         );
         let equivalent_pfail_upper_bound = equivalent_pfail_gaussian_noise(
-            new_precision_with_padding,
+            new_precision_with_padding.value,
             pfail_upper_bound,
-            original_precision_with_padding,
+            original_precision_with_padding.value,
         );
 
         println!("equivalent_pfail_lower_bound={equivalent_pfail_lower_bound}");
@@ -529,9 +534,9 @@ pub fn update_ap_params_for_pfail(
     }
 
     let new_expected_pfail = equivalent_pfail_gaussian_noise(
-        orig_pfail_and_precision.precision_with_padding(),
+        orig_pfail_and_precision.precision_with_padding().value,
         orig_pfail_and_precision.pfail(),
-        precision_with_padding(ap_params.message_modulus(), ap_params.carry_modulus()),
+        precision_with_padding(ap_params.message_modulus(), ap_params.carry_modulus()).value,
     );
     let new_expected_log2_pfail = new_expected_pfail.log2();
 
@@ -560,8 +565,41 @@ pub fn update_ap_params_for_pfail(
     (orig_pfail_and_precision, new_expected_pfail)
 }
 
-pub fn precision_with_padding(msg_mod: MessageModulus, carr_mod: CarryModulus) -> u32 {
+pub fn precision_with_padding(
-    let cleartext_modulus = msg_mod.0 * carr_mod.0;
+    msg_mod: MessageModulus,
+    carry_mod: CarryModulus,
+) -> PrecisionWithPadding {
+    let cleartext_modulus = msg_mod.0 * carry_mod.0;
     assert!(cleartext_modulus.is_power_of_two());
-    cleartext_modulus.ilog2() + 1
+    PrecisionWithPadding {
+        value: cleartext_modulus.ilog2() + 1,
+    }
+}
+
+pub fn expected_pfail_for_precision(
+    precision_with_padding: PrecisionWithPadding,
+    variance: Variance,
+) -> f64 {
+    // The additional 1 is to guarantee proper decryption
+    let precision_for_proper_decryption: i32 =
+        (precision_with_padding.value + 1).try_into().unwrap();
+    let correctness_threshold = 2.0f64.powi(-precision_for_proper_decryption);
+
+    let measured_std_dev = variance.get_standard_dev().0;
+    let measured_std_score = correctness_threshold / measured_std_dev;
+
+    statrs::function::erf::erfc(measured_std_score / core::f64::consts::SQRT_2)
+}
+
+#[test]
+fn test_expected_pfail_for_ci_run_filter() {
+    // Practical check on a compression-like scenario, of interest because pfail is known to be very
+    // low
+    let precision_with_padding = precision_with_padding(MessageModulus(1 << 2), CarryModulus(1));
+    let theoretical_variance = Variance(1.0216297411906617e-5);
+
+    assert_eq!(
+        expected_pfail_for_precision(precision_with_padding, theoretical_variance).log2(),
+        -280.4295428516361
+    );
 }

tfhe/src/shortint/server_key/tests/noise_distribution/utils/noise_simulation.rs

@@ -30,7 +30,7 @@ use crate::shortint::key_switching_key::{
     KeySwitchingKeyDestinationAtomicPattern, KeySwitchingKeyView,
 };
 use crate::shortint::list_compression::{
-    CompressionPrivateKeys, DecompressionKey, NoiseSquashingCompressionKey,
+    CompressionKey, CompressionPrivateKeys, DecompressionKey, NoiseSquashingCompressionKey,
 };
 use crate::shortint::noise_squashing::atomic_pattern::AtomicPatternNoiseSquashingKey;
 use crate::shortint::noise_squashing::{
@@ -142,6 +142,33 @@ impl DynLwe {
             Self::U128(lwe_ciphertext) => lwe_ciphertext.as_view(),
         }
     }
+
+    #[track_caller]
+    pub fn as_ref_32(&self) -> &LweCiphertextOwned<u32> {
+        match self {
+            Self::U32(lwe_ciphertext) => lwe_ciphertext,
+            Self::U64(_) => panic!("Tried getting a u64 LweCiphertext as u32."),
+            Self::U128(_) => panic!("Tried getting a u128 LweCiphertext as u32."),
+        }
+    }
+
+    #[track_caller]
+    pub fn as_ref_64(&self) -> &LweCiphertextOwned<u64> {
+        match self {
+            Self::U32(_) => panic!("Tried getting a u32 LweCiphertext as u64."),
+            Self::U64(lwe_ciphertext) => lwe_ciphertext,
+            Self::U128(_) => panic!("Tried getting a u128 LweCiphertext as u64."),
+        }
+    }
+
+    #[track_caller]
+    pub fn as_ref_128(&self) -> &LweCiphertextOwned<u128> {
+        match self {
+            Self::U32(_) => panic!("Tried getting a u32 LweCiphertext as u128."),
+            Self::U64(_) => panic!("Tried getting a u64 LweCiphertext as u128."),
+            Self::U128(lwe_ciphertext) => lwe_ciphertext,
+        }
+    }
 }
 
 impl<Scalar: CastInto<u32> + CastInto<u64> + CastInto<u128>> ScalarMul<Scalar> for DynLwe {
@@ -320,6 +347,17 @@ impl ClientKey {
         &self,
         modulus_log: CiphertextModulusLog,
         msg: u64,
+    ) -> DynLwe {
+        ShortintEngine::with_thread_local_mut(|engine| {
+            self.encrypt_noiseless_pbs_input_dyn_lwe_with_engine(modulus_log, msg, engine)
+        })
+    }
+
+    pub fn encrypt_noiseless_pbs_input_dyn_lwe_with_engine(
+        &self,
+        modulus_log: CiphertextModulusLog,
+        msg: u64,
+        engine: &mut ShortintEngine,
     ) -> DynLwe {
         match &self.atomic_pattern {
             AtomicPatternClientKey::Standard(standard_atomic_pattern_client_key) => {
@@ -331,7 +369,6 @@ impl ClientKey {
                     padding_bit: PaddingBit::Yes,
                 };
 
-                ShortintEngine::with_thread_local_mut(|engine| {
                 DynLwe::U64(encrypt_new_noiseless_lwe(
                     &standard_atomic_pattern_client_key.lwe_secret_key,
                     CiphertextModulus::try_new_power_of_2(modulus_log.0).unwrap(),
@@ -339,7 +376,6 @@ impl ClientKey {
                     &encoding,
                     &mut engine.encryption_generator,
                 ))
-                })
             }
             AtomicPatternClientKey::KeySwitch32(ks32_atomic_pattern_client_key) => {
                 let params = ks32_atomic_pattern_client_key.parameters;
@@ -350,7 +386,6 @@ impl ClientKey {
                     padding_bit: PaddingBit::Yes,
                 };
 
-                ShortintEngine::with_thread_local_mut(|engine| {
                 DynLwe::U32(encrypt_new_noiseless_lwe(
                     &ks32_atomic_pattern_client_key.lwe_secret_key,
                     CiphertextModulus::try_new_power_of_2(modulus_log.0).unwrap(),
@@ -358,7 +393,6 @@ impl ClientKey {
                     &encoding,
                     &mut engine.encryption_generator,
                 ))
-                })
             }
         }
     }
@@ -1351,6 +1385,35 @@ impl LweKeyswitch<DynLwe, DynLwe> for KeySwitchingKeyView<'_> {
     }
 }
 
+impl AllocateLwePackingKeyswitchResult for CompressionKey {
+    type Output = GlweCiphertextOwned<u64>;
+    type SideResources = ();
+
+    fn allocate_lwe_packing_keyswitch_result(
+        &self,
+        side_resources: &mut Self::SideResources,
+    ) -> Self::Output {
+        self.packing_key_switching_key
+            .allocate_lwe_packing_keyswitch_result(side_resources)
+    }
+}
+
+impl LwePackingKeyswitch<[&DynLwe], GlweCiphertextOwned<u64>> for CompressionKey {
+    type SideResources = ();
+
+    fn keyswitch_lwes_and_pack_in_glwe(
+        &self,
+        input: &[&DynLwe],
+        output: &mut GlweCiphertextOwned<u64>,
+        side_resources: &mut Self::SideResources,
+    ) {
+        let input: Vec<_> = input.iter().map(|x| x.as_ref_64()).collect();
+
+        self.packing_key_switching_key
+            .keyswitch_lwes_and_pack_in_glwe(&input, output, side_resources);
+    }
+}
+
 impl LweClassicFftBootstrap<DynLwe, DynLwe, LookupTable<Vec<u64>>> for DecompressionKey {
     type SideResources = ();
 
@@ -1688,16 +1751,41 @@ impl NoiseSimulationLwePackingKeyswitchKey {
             squashing_lwe_dim,
             noise_squashing_compression_params.packing_ks_base_log,
             noise_squashing_compression_params.packing_ks_level,
-            noise_squashing_compression_params
+            noise_squashing_compression_params.packing_ks_glwe_dimension,
-                .packing_ks_glwe_dimension
-                .to_glwe_size(),
             noise_squashing_compression_params.packing_ks_polynomial_size,
+            NoiseSimulationNoiseDistribution::U128(
                 noise_squashing_compression_params.packing_ks_key_noise_distribution,
+            ),
             NoiseSimulationModulus::from_ciphertext_modulus(
                 noise_squashing_compression_params.ciphertext_modulus,
             ),
         )
     }
+
+    pub fn new_from_comp_parameters(
+        params: AtomicPatternParameters,
+        compression_params: CompressionParameters,
+    ) -> Self {
+        let params_big_lwe_dim = params
+            .glwe_dimension()
+            .to_equivalent_lwe_dimension(params.polynomial_size());
+
+        Self::new(
+            params_big_lwe_dim,
+            compression_params.packing_ks_base_log(),
+            compression_params.packing_ks_level(),
+            compression_params.packing_ks_glwe_dimension(),
+            compression_params.packing_ks_polynomial_size(),
+            NoiseSimulationNoiseDistribution::U64(
+                compression_params.packing_ks_key_noise_distribution(),
+            ),
+            NoiseSimulationModulus::from_ciphertext_modulus(params.ciphertext_modulus()),
+        )
+    }
+
+    pub fn matches_actual_shortint_comp_key(&self, comp_key: &CompressionKey) -> bool {
+        self.matches_actual_pksk(&comp_key.packing_key_switching_key)
+    }
 }
 
 impl NoiseSimulationLweKeyswitchKey {