Am/test/multi bit noise check

tfhe/src/core_crypto/algorithms/lwe_multi_bit_programmable_bootstrapping.rs

@@ -1084,18 +1084,10 @@ pub fn multi_bit_programmable_bootstrap_lwe_ciphertext<
     );
 
     assert_eq!(
-        input.ciphertext_modulus(),
         output.ciphertext_modulus(),
-        "Mismatched CiphertextModulus between input ({:?}) and output ({:?})",
-        input.ciphertext_modulus(),
-        output.ciphertext_modulus(),
-    );
-
-    assert_eq!(
-        input.ciphertext_modulus(),
         accumulator.ciphertext_modulus(),
-        "Mismatched CiphertextModulus between input ({:?}) and accumulator ({:?})",
-        input.ciphertext_modulus(),
+        "Mismatched CiphertextModulus between output ({:?}) and accumulator ({:?})",
+        output.ciphertext_modulus(),
         accumulator.ciphertext_modulus(),
     );
 
@@ -1795,18 +1787,10 @@ pub fn std_multi_bit_programmable_bootstrap_lwe_ciphertext<
     );
 
     assert_eq!(
-        input.ciphertext_modulus(),
         output.ciphertext_modulus(),
-        "Mismatched CiphertextModulus between input ({:?}) and output ({:?})",
-        input.ciphertext_modulus(),
-        output.ciphertext_modulus(),
-    );
-
-    assert_eq!(
-        input.ciphertext_modulus(),
         accumulator.ciphertext_modulus(),
-        "Mismatched CiphertextModulus between input ({:?}) and accumulator ({:?})",
-        input.ciphertext_modulus(),
+        "Mismatched CiphertextModulus between output ({:?}) and accumulator ({:?})",
+        output.ciphertext_modulus(),
         accumulator.ciphertext_modulus(),
     );
 
@@ -1868,14 +1852,6 @@ pub fn std_multi_bit_f128_blind_rotate_assign<Scalar, InputCont, OutputCont, Key
         multi_bit_bsk.input_lwe_dimension(),
     );
 
-    assert_eq!(
-        input.ciphertext_modulus(),
-        accumulator.ciphertext_modulus(),
-        "Mismatched CiphertextModulus between input ({:?}) and accumulator ({:?})",
-        input.ciphertext_modulus(),
-        accumulator.ciphertext_modulus(),
-    );
-
     let grouping_factor = multi_bit_bsk.grouping_factor();
 
     let lut_poly_size = accumulator.polynomial_size();
@@ -2213,18 +2189,10 @@ pub fn std_multi_bit_programmable_bootstrap_f128_lwe_ciphertext<
     );
 
     assert_eq!(
-        input.ciphertext_modulus(),
         output.ciphertext_modulus(),
-        "Mismatched CiphertextModulus between input ({:?}) and output ({:?})",
-        input.ciphertext_modulus(),
-        output.ciphertext_modulus(),
-    );
-
-    assert_eq!(
-        input.ciphertext_modulus(),
         accumulator.ciphertext_modulus(),
-        "Mismatched CiphertextModulus between input ({:?}) and accumulator ({:?})",
-        input.ciphertext_modulus(),
+        "Mismatched CiphertextModulus between output ({:?}) and accumulator ({:?})",
+        output.ciphertext_modulus(),
         accumulator.ciphertext_modulus(),
     );
 
@@ -2718,3 +2686,224 @@ pub fn multi_bit_programmable_bootstrap_f128_lwe_ciphertext<
 
     extract_lwe_sample_from_glwe_ciphertext(&local_accumulator, output, MonomialDegree(0));
 }
+
+// ============== Noise measurement trait implementations ============== //
+use crate::core_crypto::commons::noise_formulas::noise_simulation::traits::{
+    AllocateMultiBitModSwitchResult, LweMultiBitFft128BlindRotate, LweMultiBitFft128Bootstrap,
+    LweMultiBitFftBlindRotate, LweMultiBitFftBootstrap, MultiBitModSwitch,
+};
+
+impl<
+        Scalar: UnsignedInteger + CastInto<usize> + CastFrom<usize>,
+        C: Container<Element = Scalar> + Sync,
+    > AllocateMultiBitModSwitchResult for LweCiphertext<C>
+{
+    type Output = StandardMultiBitModulusSwitchedCt<Scalar, Vec<Scalar>>;
+    type SideResources = ();
+
+    fn allocate_multi_bit_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 {
+            input: LweCiphertextOwned::from_container(
+                self.as_ref().to_vec(),
+                self.ciphertext_modulus(),
+            ),
+            // Placeholder values for those as they will be filled during mod switch
+            // Choose defaults that should crash things
+            grouping_factor: LweBskGroupingFactor(usize::MAX),
+            log_modulus: CiphertextModulusLog(usize::MAX),
+        }
+    }
+}
+
+impl<
+        Scalar: UnsignedInteger + CastInto<usize> + CastFrom<usize>,
+        C: Container<Element = Scalar>,
+        OutCont: ContainerMut<Element = Scalar> + Sync,
+    > MultiBitModSwitch<StandardMultiBitModulusSwitchedCt<Scalar, OutCont>> for LweCiphertext<C>
+{
+    type SideResources = ();
+
+    fn multi_bit_mod_switch(
+        &self,
+        grouping_factor: LweBskGroupingFactor,
+        output_modulus_log: CiphertextModulusLog,
+        output: &mut StandardMultiBitModulusSwitchedCt<Scalar, OutCont>,
+        _side_resources: &mut Self::SideResources,
+    ) {
+        assert_eq!(output.input.ciphertext_modulus(), self.ciphertext_modulus());
+
+        let StandardMultiBitModulusSwitchedCt {
+            input: output_input,
+            grouping_factor: output_grouping_factor,
+            log_modulus: output_log_modulus,
+        } = output;
+
+        *output_grouping_factor = grouping_factor;
+        *output_log_modulus = output_modulus_log;
+        output_input.as_mut().copy_from_slice(self.as_ref());
+
+        // Nothing to do given the mod switch will be lazily evaluated by the following multi bit
+        // blind rotation
+    }
+}
+
+impl<
+        InputScalar: UnsignedInteger + CastInto<usize> + CastFrom<usize>,
+        InputCont: Container<Element = InputScalar> + Sync,
+        OutputScalar: UnsignedTorus + Sync,
+        OutputCont: ContainerMut<Element = OutputScalar>,
+        AccCont: Container<Element = OutputScalar>,
+        KeyCont: Container<Element = c64> + Sync,
+    >
+    LweMultiBitFftBlindRotate<
+        StandardMultiBitModulusSwitchedCt<InputScalar, InputCont>,
+        LweCiphertext<OutputCont>,
+        GlweCiphertext<AccCont>,
+    > for FourierLweMultiBitBootstrapKey<KeyCont>
+{
+    type SideResources = (ThreadCount, bool);
+
+    fn lwe_multi_bit_fft_blind_rotate(
+        &self,
+        input: &StandardMultiBitModulusSwitchedCt<InputScalar, InputCont>,
+        output: &mut LweCiphertext<OutputCont>,
+        accumulator: &GlweCiphertext<AccCont>,
+        side_resources: &mut Self::SideResources,
+    ) {
+        let (thread_count, deterministic_execution) = *side_resources;
+
+        let mut local_accumulator = GlweCiphertext::from_container(
+            accumulator.as_ref().to_vec(),
+            accumulator.polynomial_size(),
+            accumulator.ciphertext_modulus(),
+        );
+
+        multi_bit_blind_rotate_assign(
+            input,
+            &mut local_accumulator,
+            self,
+            thread_count,
+            deterministic_execution,
+        );
+
+        extract_lwe_sample_from_glwe_ciphertext(&local_accumulator, output, MonomialDegree(0));
+    }
+}
+
+impl<
+        InputScalar: UnsignedInteger + CastInto<usize> + CastFrom<usize>,
+        InputCont: Container<Element = InputScalar> + Sync,
+        OutputScalar: UnsignedTorus + Sync,
+        OutputCont: ContainerMut<Element = OutputScalar>,
+        AccCont: Container<Element = OutputScalar>,
+        KeyCont: Container<Element = f64> + Sync + Split,
+    >
+    LweMultiBitFft128BlindRotate<
+        StandardMultiBitModulusSwitchedCt<InputScalar, InputCont>,
+        LweCiphertext<OutputCont>,
+        GlweCiphertext<AccCont>,
+    > for Fourier128LweMultiBitBootstrapKey<KeyCont>
+{
+    type SideResources = ThreadCount;
+
+    fn lwe_multi_bit_fft_128_blind_rotate(
+        &self,
+        input: &StandardMultiBitModulusSwitchedCt<InputScalar, InputCont>,
+        output: &mut LweCiphertext<OutputCont>,
+        accumulator: &GlweCiphertext<AccCont>,
+        side_resources: &mut Self::SideResources,
+    ) {
+        let thread_count = *side_resources;
+
+        let mut local_accumulator = GlweCiphertext::from_container(
+            accumulator.as_ref().to_vec(),
+            accumulator.polynomial_size(),
+            accumulator.ciphertext_modulus(),
+        );
+
+        multi_bit_f128_deterministic_blind_rotate_assign(
+            input,
+            &mut local_accumulator,
+            self,
+            thread_count,
+        );
+
+        extract_lwe_sample_from_glwe_ciphertext(&local_accumulator, output, MonomialDegree(0));
+    }
+}
+
+impl<
+        Scalar: UnsignedTorus + CastInto<usize> + CastFrom<usize>,
+        InputCont: Container<Element = Scalar> + Sync,
+        OutputCont: ContainerMut<Element = Scalar>,
+        AccCont: Container<Element = Scalar>,
+        KeyCont: Container<Element = c64> + Sync,
+    >
+    LweMultiBitFftBootstrap<
+        LweCiphertext<InputCont>,
+        LweCiphertext<OutputCont>,
+        GlweCiphertext<AccCont>,
+    > for FourierLweMultiBitBootstrapKey<KeyCont>
+{
+    type SideResources = (ThreadCount, bool);
+
+    fn lwe_multi_bit_fft_bootstrap(
+        &self,
+        input: &LweCiphertext<InputCont>,
+        output: &mut LweCiphertext<OutputCont>,
+        accumulator: &GlweCiphertext<AccCont>,
+        side_resources: &mut Self::SideResources,
+    ) {
+        let (thread_count, deterministic_execution) = *side_resources;
+
+        multi_bit_programmable_bootstrap_lwe_ciphertext(
+            input,
+            output,
+            accumulator,
+            self,
+            thread_count,
+            deterministic_execution,
+        );
+    }
+}
+
+impl<
+        InputScalar: UnsignedTorus + CastInto<usize> + CastFrom<usize>,
+        InputCont: Container<Element = InputScalar> + Sync,
+        OutputScalar: UnsignedTorus + Sync,
+        OutputCont: ContainerMut<Element = OutputScalar>,
+        AccCont: Container<Element = OutputScalar>,
+        KeyCont: Container<Element = f64> + Sync + Split,
+    >
+    LweMultiBitFft128Bootstrap<
+        LweCiphertext<InputCont>,
+        LweCiphertext<OutputCont>,
+        GlweCiphertext<AccCont>,
+    > for Fourier128LweMultiBitBootstrapKey<KeyCont>
+{
+    type SideResources = (ThreadCount, bool);
+
+    fn lwe_multi_bit_fft_128_bootstrap(
+        &self,
+        input: &LweCiphertext<InputCont>,
+        output: &mut LweCiphertext<OutputCont>,
+        accumulator: &GlweCiphertext<AccCont>,
+        side_resources: &mut Self::SideResources,
+    ) {
+        let (thread_count, deterministic_execution) = *side_resources;
+
+        multi_bit_programmable_bootstrap_f128_lwe_ciphertext(
+            input,
+            output,
+            accumulator,
+            self,
+            thread_count,
+            deterministic_execution,
+        );
+    }
+}

tfhe/src/core_crypto/algorithms/lwe_programmable_bootstrapping/mod.rs

@@ -77,7 +77,9 @@ where
 }
 
 // ============== Noise measurement trait implementations ============== //
-use crate::core_crypto::commons::noise_formulas::noise_simulation::traits::AllocateLweBootstrapResult;
+use crate::core_crypto::commons::noise_formulas::noise_simulation::traits::{
+    AllocateLweBootstrapResult, AllocateLweMultiBitBlindRotateResult,
+};
 
 impl<Scalar: UnsignedInteger, AccCont: Container<Element = Scalar>> AllocateLweBootstrapResult
     for GlweCiphertext<AccCont>
@@ -100,3 +102,25 @@ impl<Scalar: UnsignedInteger, AccCont: Container<Element = Scalar>> AllocateLweB
         )
     }
 }
+
+impl<Scalar: UnsignedInteger, AccCont: Container<Element = Scalar>>
+    AllocateLweMultiBitBlindRotateResult for GlweCiphertext<AccCont>
+{
+    type Output = LweCiphertextOwned<Scalar>;
+    type SideResources = ();
+
+    fn allocate_lwe_multi_bit_blind_rotate_result(
+        &self,
+        _side_resources: &mut Self::SideResources,
+    ) -> Self::Output {
+        let glwe_dim = self.glwe_size().to_glwe_dimension();
+        let poly_size = self.polynomial_size();
+        let equivalent_lwe_dim = glwe_dim.to_equivalent_lwe_dimension(poly_size);
+
+        LweCiphertext::new(
+            Scalar::ZERO,
+            equivalent_lwe_dim.to_lwe_size(),
+            self.ciphertext_modulus(),
+        )
+    }
+}

tfhe/src/core_crypto/algorithms/test/noise_distribution/lwe_multi_bit_programmable_bootstrapping_noise.rs

@@ -3,6 +3,7 @@ use crate::core_crypto::commons::noise_formulas::lwe_multi_bit_programmable_boot
     multi_bit_pbs_variance_132_bits_security_gaussian_gf_3_fft_mul,
     multi_bit_pbs_variance_132_bits_security_tuniform_gf_3_fft_mul,
 };
+use crate::core_crypto::commons::noise_formulas::noise_simulation::PBS_FFT_64_MANTISSA_SIZE;
 use crate::core_crypto::commons::noise_formulas::secure_noise::{
     minimal_lwe_variance_for_132_bits_security_gaussian,
     minimal_lwe_variance_for_132_bits_security_tuniform,
@@ -48,6 +49,7 @@ where
                 polynomial_size,
                 pbs_decomposition_base_log,
                 pbs_decomposition_level_count,
+                PBS_FFT_64_MANTISSA_SIZE,
                 modulus_as_f64,
             )
         }
@@ -58,6 +60,7 @@ where
                 polynomial_size,
                 pbs_decomposition_base_log,
                 pbs_decomposition_level_count,
+                PBS_FFT_64_MANTISSA_SIZE,
                 modulus_as_f64,
             )
         }

tfhe/src/core_crypto/algorithms/test/noise_distribution/lwe_programmable_bootstrapping_noise.rs

@@ -1,5 +1,6 @@
 use super::*;
 use crate::core_crypto::commons::noise_formulas::lwe_programmable_bootstrap::pbs_variance_132_bits_security_gaussian_fft_mul;
+use crate::core_crypto::commons::noise_formulas::noise_simulation::PBS_FFT_64_MANTISSA_SIZE;
 use crate::core_crypto::commons::noise_formulas::secure_noise::minimal_lwe_variance_for_132_bits_security_gaussian;
 use crate::core_crypto::commons::test_tools::{torus_modular_diff, variance};
 use rayon::prelude::*;
@@ -38,6 +39,7 @@ where
         polynomial_size,
         pbs_decomposition_base_log,
         pbs_decomposition_level_count,
+        PBS_FFT_64_MANTISSA_SIZE,
         modulus_as_f64,
     );