feat: add noise level checks

This adds the noise-asserts feature, which will make PBS functions do a noise level check. This also adds an extra MaxNoiseLevel parameter to Ciphertext::set_noise_level that is used when the noise-asserts feature is on, to check that the given new-noise level does not exceed the given MaxNoiseLevel. In case of problems, the code will panic By default these checks will also be make in cfg(test)
zama-ai · Nov 14, 2024 · ff11669 · ff11669
1 parent cf56e58
commit ff11669
Show file tree

Hide file tree

Showing 16 changed files with 121 additions and 40 deletions.
diff --git a/tfhe/Cargo.toml b/tfhe/Cargo.toml
@@ -99,6 +99,7 @@ gpu = ["dep:tfhe-cuda-backend"]
 zk-pok = ["dep:tfhe-zk-pok"]
 
 pbs-stats = []
+noise-asserts = []
 
 # Experimental section
 experimental = []

diff --git a/tfhe/src/integer/server_key/comparator.rs b/tfhe/src/integer/server_key/comparator.rs
@@ -182,7 +182,10 @@ impl<'a> Comparator<'a> {
 
         // Here we need the true lwe sub, not the one that comes from shortint.
         crate::core_crypto::algorithms::lwe_ciphertext_sub_assign(&mut lhs.ct, &rhs.ct);
-        lhs.set_noise_level(lhs.noise_level() + rhs.noise_level());
+        lhs.set_noise_level(
+            lhs.noise_level() + rhs.noise_level(),
+            self.server_key.key.max_noise_level,
+        );
         self.server_key
             .key
             .apply_lookup_table_assign(lhs, &self.sign_lut);

diff --git a/tfhe/src/integer/server_key/radix_parallel/scalar_sub.rs b/tfhe/src/integer/server_key/radix_parallel/scalar_sub.rs
@@ -177,7 +177,10 @@ impl ServerKey {
             // And here, it's because shortint sub_assign adds a correcting term,
             // which we do not want here
             crate::core_crypto::algorithms::lwe_ciphertext_sub_assign(&mut lhs_b.ct, &borrow.ct);
-            lhs_b.set_noise_level(lhs_b.noise_level() + borrow.noise_level());
+            lhs_b.set_noise_level(
+                lhs_b.noise_level() + borrow.noise_level(),
+                self.key.max_noise_level,
+            );
 
             borrow.clone_from(lhs_b);
 
@@ -304,7 +307,10 @@ impl ServerKey {
                             &mut block.ct,
                             &simulator.ct,
                         );
-                        block.set_noise_level(block.noise_level() + simulator.noise_level());
+                        block.set_noise_level(
+                            block.noise_level() + simulator.noise_level(),
+                            self.key.max_noise_level,
+                        );
                         self.key.unchecked_scalar_add_assign(block, 1);
                     }
                 });
@@ -334,7 +340,10 @@ impl ServerKey {
                             &mut block.ct,
                             &borrow.ct,
                         );
-                        block.set_noise_level(block.noise_level() + borrow.noise_level());
+                        block.set_noise_level(
+                            block.noise_level() + borrow.noise_level(),
+                            self.key.max_noise_level,
+                        );
 
                         let lut = if i % 2 == 0 {
                             &extract_message_low_block_mut

diff --git a/tfhe/src/integer/server_key/radix_parallel/sub.rs b/tfhe/src/integer/server_key/radix_parallel/sub.rs
@@ -418,7 +418,10 @@ impl ServerKey {
                         &mut block.ct,
                         &simulator.ct,
                     );
-                    block.set_noise_level(block.noise_level() + simulator.noise_level());
+                    block.set_noise_level(
+                        block.noise_level() + simulator.noise_level(),
+                        self.key.max_noise_level,
+                    );
                     self.key.unchecked_scalar_add_assign(block, 1);
                 });
 
@@ -447,7 +450,10 @@ impl ServerKey {
                         &mut block.ct,
                         &borrow.ct,
                     );
-                    block.set_noise_level(block.noise_level() + borrow.noise_level());
+                    block.set_noise_level(
+                        block.noise_level() + borrow.noise_level(),
+                        self.key.max_noise_level,
+                    );
 
                     self.key
                         .apply_lookup_table_assign(block, &message_extract_lut)
@@ -633,7 +639,10 @@ impl ServerKey {
                 &mut lhs_block.ct,
                 &borrow.ct,
             );
-            lhs_block.set_noise_level(lhs_block.noise_level() + borrow.noise_level());
+            lhs_block.set_noise_level(
+                lhs_block.noise_level() + borrow.noise_level(),
+                self.key.max_noise_level,
+            );
             let (msg, new_borrow) = rayon::join(
                 || self.key.message_extract(lhs_block),
                 || self.key.apply_lookup_table(lhs_block, &compute_borrow_lut),

diff --git a/tfhe/src/shortint/ciphertext/common.rs b/tfhe/src/shortint/ciphertext/common.rs
@@ -26,6 +26,10 @@ impl std::error::Error for NotTrivialCiphertextError {}
 pub struct MaxNoiseLevel(usize);
 
 impl MaxNoiseLevel {
+    #[cfg(test)]
+    pub(crate) const MAX: Self = Self(usize::MAX);
+    pub(crate) const UNKNOWN: Self = Self(usize::MAX);
+
     pub const fn new(value: usize) -> Self {
         Self(value)
     }

diff --git a/tfhe/src/shortint/ciphertext/standard.rs b/tfhe/src/shortint/ciphertext/standard.rs
@@ -138,7 +138,13 @@ impl Ciphertext {
         self.noise_level
     }
 
-    pub fn set_noise_level(&mut self, noise_level: NoiseLevel) {
+    #[cfg_attr(any(feature = "noise-asserts", test), track_caller)]
+    pub fn set_noise_level(&mut self, noise_level: NoiseLevel, max_noise_level: MaxNoiseLevel) {
+        if cfg!(feature = "noise-asserts") || cfg!(test) {
+            max_noise_level.validate(noise_level).unwrap()
+        } else {
+            let _ = max_noise_level;
+        }
         self.noise_level = noise_level;
     }
 

diff --git a/tfhe/src/shortint/key_switching_key/mod.rs b/tfhe/src/shortint/key_switching_key/mod.rs
@@ -14,7 +14,7 @@ use crate::shortint::parameters::{
     EncryptionKeyChoice, NoiseLevel, PBSOrder, ShortintKeySwitchingParameters,
 };
 use crate::shortint::server_key::apply_programmable_bootstrap;
-use crate::shortint::{Ciphertext, ClientKey, CompressedServerKey, ServerKey};
+use crate::shortint::{Ciphertext, ClientKey, CompressedServerKey, MaxNoiseLevel, ServerKey};
 use core::cmp::Ordering;
 use rayon::prelude::*;
 use serde::{Deserialize, Serialize};
@@ -520,7 +520,7 @@ impl<'keys> KeySwitchingKeyView<'keys> {
         // TODO: We are outside the standard AP, if we chain keyswitches, we will refresh, which is
         // safer for now. We can likely add an additional flag in shortint to indicate if we
         // want to refresh or not, for now refresh anyways.
-        keyswitched.set_noise_level(NoiseLevel::UNKNOWN);
+        keyswitched.set_noise_level(NoiseLevel::UNKNOWN, MaxNoiseLevel::UNKNOWN);
 
         let cast_rshift = self.key_switching_key_material.cast_rshift;
 
@@ -575,7 +575,10 @@ impl<'keys> KeySwitchingKeyView<'keys> {
                         let wrong_key_ct = keyswitched;
                         let mut correct_key_ct = self.dest_server_key.create_trivial(0);
                         correct_key_ct.degree = wrong_key_ct.degree;
-                        correct_key_ct.set_noise_level(wrong_key_ct.noise_level());
+                        correct_key_ct.set_noise_level(
+                            wrong_key_ct.noise_level(),
+                            self.dest_server_key.max_noise_level,
+                        );
 
                         keyswitch_lwe_ciphertext(
                             &self.dest_server_key.key_switching_key,
@@ -644,7 +647,10 @@ impl<'keys> KeySwitchingKeyView<'keys> {
                                         correct_key_ct.degree = degree_after_keyswitch;
                                     }
                                     // Update the noise as well
-                                    correct_key_ct.set_noise_level(NoiseLevel::NOMINAL);
+                                    correct_key_ct.set_noise_level(
+                                        NoiseLevel::NOMINAL,
+                                        self.dest_server_key.max_noise_level,
+                                    );
                                 });
                             });
                     }
@@ -688,7 +694,10 @@ impl<'keys> KeySwitchingKeyView<'keys> {
                                     );
                                     // Update degree and noise as it's a raw PBS
                                     correct_key_ct.degree = acc.degree;
-                                    correct_key_ct.set_noise_level(NoiseLevel::NOMINAL);
+                                    correct_key_ct.set_noise_level(
+                                        NoiseLevel::NOMINAL,
+                                        self.dest_server_key.max_noise_level,
+                                    );
                                 });
                             });
                     }
@@ -750,7 +759,10 @@ impl<'keys> KeySwitchingKeyView<'keys> {
                                         );
                                         correct_key_ct.degree = new_degree;
                                     }
-                                    correct_key_ct.set_noise_level(NoiseLevel::NOMINAL);
+                                    correct_key_ct.set_noise_level(
+                                        NoiseLevel::NOMINAL,
+                                        self.dest_server_key.max_noise_level,
+                                    );
                                 });
                             });
                     }

diff --git a/tfhe/src/shortint/list_compression/compression.rs b/tfhe/src/shortint/list_compression/compression.rs
@@ -11,7 +11,7 @@ use crate::shortint::parameters::NoiseLevel;
 use crate::shortint::server_key::{
     apply_programmable_bootstrap, generate_lookup_table, unchecked_scalar_mul_assign,
 };
-use crate::shortint::{Ciphertext, CiphertextModulus};
+use crate::shortint::{Ciphertext, CiphertextModulus, MaxNoiseLevel};
 use rayon::iter::ParallelIterator;
 use rayon::slice::ParallelSlice;
 
@@ -87,8 +87,9 @@ impl CompressionKey {
                     );
 
                     let mut ct = ct.clone();
-
-                    unchecked_scalar_mul_assign(&mut ct, message_modulus.0 as u8);
+                    let max_noise_level =
+                        MaxNoiseLevel::new((ct.noise_level() * message_modulus.0).get());
+                    unchecked_scalar_mul_assign(&mut ct, message_modulus.0 as u8, max_noise_level);
 
                     list.extend(ct.ct.as_ref());
                 }

diff --git a/tfhe/src/shortint/server_key/add.rs b/tfhe/src/shortint/server_key/add.rs
@@ -2,7 +2,7 @@ use super::{CiphertextNoiseDegree, SmartCleaningOperation};
 use crate::core_crypto::algorithms::*;
 use crate::shortint::ciphertext::Degree;
 use crate::shortint::server_key::CheckError;
-use crate::shortint::{Ciphertext, ServerKey};
+use crate::shortint::{Ciphertext, MaxNoiseLevel, ServerKey};
 
 impl ServerKey {
     /// Compute homomorphically an addition between two ciphertexts encrypting integer values.
@@ -226,7 +226,7 @@ impl ServerKey {
     /// assert_eq!(msg + msg, two);
     /// ```
     pub fn unchecked_add_assign(&self, ct_left: &mut Ciphertext, ct_right: &Ciphertext) {
-        unchecked_add_assign(ct_left, ct_right);
+        unchecked_add_assign(ct_left, ct_right, self.max_noise_level);
     }
 
     /// Verify if ct_left and ct_right can be added together.
@@ -515,8 +515,15 @@ impl ServerKey {
     }
 }
 
-pub(crate) fn unchecked_add_assign(ct_left: &mut Ciphertext, ct_right: &Ciphertext) {
+pub(crate) fn unchecked_add_assign(
+    ct_left: &mut Ciphertext,
+    ct_right: &Ciphertext,
+    max_noise_level: MaxNoiseLevel,
+) {
     lwe_ciphertext_add_assign(&mut ct_left.ct, &ct_right.ct);
     ct_left.degree = Degree::new(ct_left.degree.get() + ct_right.degree.get());
-    ct_left.set_noise_level(ct_left.noise_level() + ct_right.noise_level());
+    ct_left.set_noise_level(
+        ct_left.noise_level() + ct_right.noise_level(),
+        max_noise_level,
+    );
 }
diff --git a/tfhe/src/shortint/server_key/bivariate_pbs.rs b/tfhe/src/shortint/server_key/bivariate_pbs.rs
@@ -158,7 +158,7 @@ impl ServerKey {
 
         self.unchecked_scalar_mul_assign(ct_left, acc.ct_right_modulus.0 as u8);
 
-        unchecked_add_assign(ct_left, ct_right);
+        unchecked_add_assign(ct_left, ct_right, self.max_noise_level);
 
         // Compute the PBS
         self.apply_lookup_table_assign(ct_left, &acc.acc);

diff --git a/tfhe/src/shortint/server_key/mod.rs b/tfhe/src/shortint/server_key/mod.rs
@@ -805,6 +805,15 @@ impl ServerKey {
             return;
         }
 
+        if cfg!(any(feature = "noise-asserts", test)) && ct.noise_level() != NoiseLevel::UNKNOWN {
+            assert!(
+                self.max_noise_level.validate(ct.noise_level).is_ok(),
+                "{:?} is above the maximum ({:?})",
+                ct.noise_level,
+                self.max_noise_level
+            );
+        }
+
         ShortintEngine::with_thread_local_mut(|engine| {
             let (mut ciphertext_buffers, buffers) = engine.get_buffers(self);
             match self.pbs_order {
@@ -842,7 +851,7 @@ impl ServerKey {
         });
 
         ct.degree = acc.degree;
-        ct.set_noise_level(NoiseLevel::NOMINAL);
+        ct.set_noise_level(NoiseLevel::NOMINAL, self.max_noise_level);
     }
 
     /// Compute a keyswitch and programmable bootstrap applying several functions on an input
@@ -913,6 +922,14 @@ impl ServerKey {
         ct: &Ciphertext,
         acc: &ManyLookupTableOwned,
     ) -> Vec<Ciphertext> {
+        if cfg!(any(feature = "noise-asserts", test)) && ct.noise_level() != NoiseLevel::UNKNOWN {
+            assert!(
+                self.max_noise_level.validate(ct.noise_level).is_ok(),
+                "{:?} is above the maximum ({:?})",
+                ct.noise_level,
+                self.max_noise_level
+            );
+        }
         match self.pbs_order {
             PBSOrder::KeyswitchBootstrap => self.keyswitch_programmable_bootstrap_many_lut(ct, acc),
             PBSOrder::BootstrapKeyswitch => self.programmable_bootstrap_keyswitch_many_lut(ct, acc),
@@ -1185,7 +1202,7 @@ impl ServerKey {
         trivially_encrypt_lwe_ciphertext(&mut ct.ct, encoded);
 
         ct.degree = Degree::new(modular_value);
-        ct.set_noise_level(NoiseLevel::ZERO);
+        ct.set_noise_level(NoiseLevel::ZERO, self.max_noise_level);
     }
 
     pub fn bootstrapping_key_size_elements(&self) -> usize {
@@ -1339,7 +1356,7 @@ impl ServerKey {
             );
 
             output_shortint_ct.degree = *output_degree;
-            output_shortint_ct.set_noise_level(NoiseLevel::NOMINAL);
+            output_shortint_ct.set_noise_level(NoiseLevel::NOMINAL, self.max_noise_level);
             outputs.push(output_shortint_ct);
         }
 
@@ -1390,7 +1407,7 @@ impl ServerKey {
             );
 
             output_shortint_ct.degree = *output_degree;
-            output_shortint_ct.set_noise_level(NoiseLevel::NOMINAL);
+            output_shortint_ct.set_noise_level(NoiseLevel::NOMINAL, self.max_noise_level);
             outputs.push(output_shortint_ct);
         }
 

diff --git a/tfhe/src/shortint/server_key/mul.rs b/tfhe/src/shortint/server_key/mul.rs
@@ -221,7 +221,7 @@ impl ServerKey {
         // ct1 + ct2
         let mut ct_add = ct1.clone();
 
-        unchecked_add_assign(&mut ct_add, ct2);
+        unchecked_add_assign(&mut ct_add, ct2, self.max_noise_level);
 
         // ct1 - ct2
         let (mut ct_sub, z) = self.unchecked_sub_with_correcting_term(ct1, ct2);

diff --git a/tfhe/src/shortint/server_key/scalar_mul.rs b/tfhe/src/shortint/server_key/scalar_mul.rs
@@ -3,7 +3,7 @@ use crate::core_crypto::algorithms::*;
 use crate::core_crypto::entities::*;
 use crate::shortint::ciphertext::Degree;
 use crate::shortint::server_key::CheckError;
-use crate::shortint::{Ciphertext, ServerKey};
+use crate::shortint::{Ciphertext, MaxNoiseLevel, ServerKey};
 
 impl ServerKey {
     /// Compute homomorphically a multiplication of a ciphertext by a scalar.
@@ -203,7 +203,7 @@ impl ServerKey {
     /// assert_eq!(3, clear);
     /// ```
     pub fn unchecked_scalar_mul_assign(&self, ct: &mut Ciphertext, scalar: u8) {
-        unchecked_scalar_mul_assign(ct, scalar);
+        unchecked_scalar_mul_assign(ct, scalar, self.max_noise_level);
     }
 
     /// Multiply one ciphertext with a scalar in the case the carry space cannot fit the product
@@ -516,8 +516,12 @@ impl ServerKey {
     }
 }
 
-pub(crate) fn unchecked_scalar_mul_assign(ct: &mut Ciphertext, scalar: u8) {
-    ct.set_noise_level(ct.noise_level() * scalar as usize);
+pub(crate) fn unchecked_scalar_mul_assign(
+    ct: &mut Ciphertext,
+    scalar: u8,
+    max_noise_level: MaxNoiseLevel,
+) {
+    ct.set_noise_level(ct.noise_level() * scalar as usize, max_noise_level);
     ct.degree = Degree::new(ct.degree.get() * scalar as usize);
 
     match scalar {

diff --git a/tfhe/src/shortint/server_key/shift.rs b/tfhe/src/shortint/server_key/shift.rs
@@ -537,7 +537,7 @@ impl ServerKey {
     /// ```
     pub fn unchecked_scalar_left_shift_assign(&self, ct: &mut Ciphertext, shift: u8) {
         let scalar = 1_u8 << shift;
-        unchecked_scalar_mul_assign(ct, scalar);
+        unchecked_scalar_mul_assign(ct, scalar, self.max_noise_level);
     }
 
     /// Checks if the left shift operation can be applied.

diff --git a/tfhe/src/shortint/server_key/sub.rs b/tfhe/src/shortint/server_key/sub.rs
@@ -401,7 +401,10 @@ impl ServerKey {
 
         lwe_ciphertext_add_assign(&mut ct_left.ct, &neg_right.ct);
 
-        ct_left.set_noise_level(ct_left.noise_level() + ct_right.noise_level());
+        ct_left.set_noise_level(
+            ct_left.noise_level() + ct_right.noise_level(),
+            self.max_noise_level,
+        );
         ct_left.degree = Degree::new(ct_left.degree.get() + z as usize);
 
         z