Skip to content

Commit 12e7a71

Browse files
tmontaigutmontaigu
committed
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)
1 parent 46cf465 commit 12e7a71

File tree

16 files changed

+107
-48
lines changed

16 files changed

+107
-48
lines changed

‎tfhe/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,7 @@ gpu = ["dep:tfhe-cuda-backend"]
102102
zk-pok = ["dep:tfhe-zk-pok"]
103103

104104
pbs-stats = []
105+
noise-asserts = []
105106

106107
# Experimental section
107108
experimental = []

‎tfhe/src/integer/server_key/comparator.rs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -182,7 +182,10 @@ impl<'a> Comparator<'a> {
182182

183183
// Here we need the true lwe sub, not the one that comes from shortint.
184184
crate::core_crypto::algorithms::lwe_ciphertext_sub_assign(&mut lhs.ct, &rhs.ct);
185-
lhs.set_noise_level(lhs.noise_level() + rhs.noise_level());
185+
lhs.set_noise_level(
186+
lhs.noise_level() + rhs.noise_level(),
187+
self.server_key.key.max_noise_level,
188+
);
186189
self.server_key
187190
.key
188191
.apply_lookup_table_assign(lhs, &self.sign_lut);

‎tfhe/src/integer/server_key/radix_parallel/scalar_sub.rs

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -177,7 +177,10 @@ impl ServerKey {
177177
// And here, it's because shortint sub_assign adds a correcting term,
178178
// which we do not want here
179179
crate::core_crypto::algorithms::lwe_ciphertext_sub_assign(&mut lhs_b.ct, &borrow.ct);
180-
lhs_b.set_noise_level(lhs_b.noise_level() + borrow.noise_level());
180+
lhs_b.set_noise_level(
181+
lhs_b.noise_level() + borrow.noise_level(),
182+
self.key.max_noise_level,
183+
);
181184

182185
borrow.clone_from(lhs_b);
183186

@@ -304,7 +307,10 @@ impl ServerKey {
304307
&mut block.ct,
305308
&simulator.ct,
306309
);
307-
block.set_noise_level(block.noise_level() + simulator.noise_level());
310+
block.set_noise_level(
311+
block.noise_level() + simulator.noise_level(),
312+
self.key.max_noise_level,
313+
);
308314
self.key.unchecked_scalar_add_assign(block, 1);
309315
}
310316
});
@@ -334,7 +340,10 @@ impl ServerKey {
334340
&mut block.ct,
335341
&borrow.ct,
336342
);
337-
block.set_noise_level(block.noise_level() + borrow.noise_level());
343+
block.set_noise_level(
344+
block.noise_level() + borrow.noise_level(),
345+
self.key.max_noise_level,
346+
);
338347

339348
let lut = if i % 2 == 0 {
340349
&extract_message_low_block_mut

‎tfhe/src/integer/server_key/radix_parallel/sub.rs

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -418,7 +418,10 @@ impl ServerKey {
418418
&mut block.ct,
419419
&simulator.ct,
420420
);
421-
block.set_noise_level(block.noise_level() + simulator.noise_level());
421+
block.set_noise_level(
422+
block.noise_level() + simulator.noise_level(),
423+
self.key.max_noise_level,
424+
);
422425
self.key.unchecked_scalar_add_assign(block, 1);
423426
});
424427

@@ -447,7 +450,10 @@ impl ServerKey {
447450
&mut block.ct,
448451
&borrow.ct,
449452
);
450-
block.set_noise_level(block.noise_level() + borrow.noise_level());
453+
block.set_noise_level(
454+
block.noise_level() + borrow.noise_level(),
455+
self.key.max_noise_level,
456+
);
451457

452458
self.key
453459
.apply_lookup_table_assign(block, &message_extract_lut)
@@ -633,7 +639,10 @@ impl ServerKey {
633639
&mut lhs_block.ct,
634640
&borrow.ct,
635641
);
636-
lhs_block.set_noise_level(lhs_block.noise_level() + borrow.noise_level());
642+
lhs_block.set_noise_level(
643+
lhs_block.noise_level() + borrow.noise_level(),
644+
self.key.max_noise_level,
645+
);
637646
let (msg, new_borrow) = rayon::join(
638647
|| self.key.message_extract(lhs_block),
639648
|| self.key.apply_lookup_table(lhs_block, &compute_borrow_lut),

‎tfhe/src/shortint/ciphertext/common.rs

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@ impl std::error::Error for NotTrivialCiphertextError {}
2626
pub struct MaxNoiseLevel(usize);
2727

2828
impl MaxNoiseLevel {
29+
pub(crate) const UNKNOWN: Self = Self(usize::MAX);
30+
2931
pub const fn new(value: usize) -> Self {
3032
Self(value)
3133
}
@@ -67,11 +69,8 @@ pub struct NoiseLevel(usize);
6769
impl NoiseLevel {
6870
pub const NOMINAL: Self = Self(1);
6971
pub const ZERO: Self = Self(0);
70-
// To force a refresh no matter the tolerance of the server key, useful for serialization update
71-
// for formats which did not have noise levels saved
72-
pub const MAX: Self = Self(usize::MAX);
7372
// As a safety measure the unknown noise level is set to the max value
74-
pub const UNKNOWN: Self = Self::MAX;
73+
pub const UNKNOWN: Self = Self(usize::MAX);
7574
}
7675

7776
impl NoiseLevel {
@@ -273,16 +272,16 @@ mod tests {
273272

274273
let mut rng = thread_rng();
275274

276-
assert_eq!(NoiseLevel::UNKNOWN, NoiseLevel::MAX);
275+
assert_eq!(NoiseLevel::UNKNOWN.0, usize::MAX);
277276

278-
let max_noise_level = NoiseLevel::MAX;
277+
let max_noise_level = NoiseLevel::UNKNOWN;
279278
let random_addend = rng.gen::<usize>();
280279
let add = max_noise_level + NoiseLevel(random_addend);
281-
assert_eq!(add, NoiseLevel::MAX);
280+
assert_eq!(add, NoiseLevel::UNKNOWN);
282281

283282
let random_positive_multiplier = rng.gen_range(1usize..=usize::MAX);
284283
let mul = max_noise_level * random_positive_multiplier;
285-
assert_eq!(mul, NoiseLevel::MAX);
284+
assert_eq!(mul, NoiseLevel::UNKNOWN);
286285
}
287286

288287
#[test]

‎tfhe/src/shortint/ciphertext/standard.rs

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,13 @@ impl Ciphertext {
140140
self.noise_level
141141
}
142142

143-
pub fn set_noise_level(&mut self, noise_level: NoiseLevel) {
143+
#[cfg_attr(any(feature = "noise-asserts", test), track_caller)]
144+
pub fn set_noise_level(&mut self, noise_level: NoiseLevel, max_noise_level: MaxNoiseLevel) {
145+
if cfg!(feature = "noise-asserts") || cfg!(test) {
146+
max_noise_level.validate(noise_level).unwrap()
147+
} else {
148+
let _ = max_noise_level;
149+
}
144150
self.noise_level = noise_level;
145151
}
146152

‎tfhe/src/shortint/key_switching_key/mod.rs

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ use crate::shortint::parameters::{
1414
EncryptionKeyChoice, NoiseLevel, PBSOrder, ShortintKeySwitchingParameters,
1515
};
1616
use crate::shortint::server_key::apply_programmable_bootstrap;
17-
use crate::shortint::{Ciphertext, ClientKey, CompressedServerKey, ServerKey};
17+
use crate::shortint::{Ciphertext, ClientKey, CompressedServerKey, MaxNoiseLevel, ServerKey};
1818
use core::cmp::Ordering;
1919
use rayon::prelude::*;
2020
use serde::{Deserialize, Serialize};
@@ -520,7 +520,7 @@ impl<'keys> KeySwitchingKeyView<'keys> {
520520
// TODO: We are outside the standard AP, if we chain keyswitches, we will refresh, which is
521521
// safer for now. We can likely add an additional flag in shortint to indicate if we
522522
// want to refresh or not, for now refresh anyways.
523-
keyswitched.set_noise_level(NoiseLevel::UNKNOWN);
523+
keyswitched.set_noise_level(NoiseLevel::UNKNOWN, MaxNoiseLevel::UNKNOWN);
524524

525525
let cast_rshift = self.key_switching_key_material.cast_rshift;
526526

@@ -575,7 +575,10 @@ impl<'keys> KeySwitchingKeyView<'keys> {
575575
let wrong_key_ct = keyswitched;
576576
let mut correct_key_ct = self.dest_server_key.create_trivial(0);
577577
correct_key_ct.degree = wrong_key_ct.degree;
578-
correct_key_ct.set_noise_level(wrong_key_ct.noise_level());
578+
correct_key_ct.set_noise_level(
579+
wrong_key_ct.noise_level(),
580+
self.dest_server_key.max_noise_level,
581+
);
579582

580583
keyswitch_lwe_ciphertext(
581584
&self.dest_server_key.key_switching_key,
@@ -644,7 +647,10 @@ impl<'keys> KeySwitchingKeyView<'keys> {
644647
correct_key_ct.degree = degree_after_keyswitch;
645648
}
646649
// Update the noise as well
647-
correct_key_ct.set_noise_level(NoiseLevel::NOMINAL);
650+
correct_key_ct.set_noise_level(
651+
NoiseLevel::NOMINAL,
652+
self.dest_server_key.max_noise_level,
653+
);
648654
});
649655
});
650656
}
@@ -688,7 +694,10 @@ impl<'keys> KeySwitchingKeyView<'keys> {
688694
);
689695
// Update degree and noise as it's a raw PBS
690696
correct_key_ct.degree = acc.degree;
691-
correct_key_ct.set_noise_level(NoiseLevel::NOMINAL);
697+
correct_key_ct.set_noise_level(
698+
NoiseLevel::NOMINAL,
699+
self.dest_server_key.max_noise_level,
700+
);
692701
});
693702
});
694703
}
@@ -750,7 +759,10 @@ impl<'keys> KeySwitchingKeyView<'keys> {
750759
);
751760
correct_key_ct.degree = new_degree;
752761
}
753-
correct_key_ct.set_noise_level(NoiseLevel::NOMINAL);
762+
correct_key_ct.set_noise_level(
763+
NoiseLevel::NOMINAL,
764+
self.dest_server_key.max_noise_level,
765+
);
754766
});
755767
});
756768
}

‎tfhe/src/shortint/list_compression/compression.rs

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ use crate::shortint::parameters::NoiseLevel;
1111
use crate::shortint::server_key::{
1212
apply_programmable_bootstrap, generate_lookup_table, unchecked_scalar_mul_assign,
1313
};
14-
use crate::shortint::{Ciphertext, CiphertextModulus};
14+
use crate::shortint::{Ciphertext, CiphertextModulus, MaxNoiseLevel};
1515
use rayon::iter::ParallelIterator;
1616
use rayon::slice::ParallelSlice;
1717

@@ -87,8 +87,9 @@ impl CompressionKey {
8787
);
8888

8989
let mut ct = ct.clone();
90-
91-
unchecked_scalar_mul_assign(&mut ct, message_modulus.0 as u8);
90+
let max_noise_level =
91+
MaxNoiseLevel::new((ct.noise_level() * message_modulus.0).get());
92+
unchecked_scalar_mul_assign(&mut ct, message_modulus.0 as u8, max_noise_level);
9293

9394
list.extend(ct.ct.as_ref());
9495
}

‎tfhe/src/shortint/server_key/add.rs

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ use super::{CiphertextNoiseDegree, SmartCleaningOperation};
22
use crate::core_crypto::algorithms::*;
33
use crate::shortint::ciphertext::Degree;
44
use crate::shortint::server_key::CheckError;
5-
use crate::shortint::{Ciphertext, ServerKey};
5+
use crate::shortint::{Ciphertext, MaxNoiseLevel, ServerKey};
66

77
impl ServerKey {
88
/// Compute homomorphically an addition between two ciphertexts encrypting integer values.
@@ -226,7 +226,7 @@ impl ServerKey {
226226
/// assert_eq!(msg + msg, two);
227227
/// ```
228228
pub fn unchecked_add_assign(&self, ct_left: &mut Ciphertext, ct_right: &Ciphertext) {
229-
unchecked_add_assign(ct_left, ct_right);
229+
unchecked_add_assign(ct_left, ct_right, self.max_noise_level);
230230
}
231231

232232
/// Verify if ct_left and ct_right can be added together.
@@ -515,8 +515,15 @@ impl ServerKey {
515515
}
516516
}
517517

518-
pub(crate) fn unchecked_add_assign(ct_left: &mut Ciphertext, ct_right: &Ciphertext) {
518+
pub(crate) fn unchecked_add_assign(
519+
ct_left: &mut Ciphertext,
520+
ct_right: &Ciphertext,
521+
max_noise_level: MaxNoiseLevel,
522+
) {
519523
lwe_ciphertext_add_assign(&mut ct_left.ct, &ct_right.ct);
520524
ct_left.degree = Degree::new(ct_left.degree.get() + ct_right.degree.get());
521-
ct_left.set_noise_level(ct_left.noise_level() + ct_right.noise_level());
525+
ct_left.set_noise_level(
526+
ct_left.noise_level() + ct_right.noise_level(),
527+
max_noise_level,
528+
);
522529
}

‎tfhe/src/shortint/server_key/bivariate_pbs.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -158,7 +158,7 @@ impl ServerKey {
158158

159159
self.unchecked_scalar_mul_assign(ct_left, acc.ct_right_modulus.0 as u8);
160160

161-
unchecked_add_assign(ct_left, ct_right);
161+
unchecked_add_assign(ct_left, ct_right, self.max_noise_level);
162162

163163
// Compute the PBS
164164
self.apply_lookup_table_assign(ct_left, &acc.acc);

‎tfhe/src/shortint/server_key/mod.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -843,7 +843,7 @@ impl ServerKey {
843843
});
844844

845845
ct.degree = acc.degree;
846-
ct.set_noise_level(NoiseLevel::NOMINAL);
846+
ct.set_noise_level(NoiseLevel::NOMINAL, self.max_noise_level);
847847
}
848848

849849
/// Compute a keyswitch and programmable bootstrap applying several functions on an input
@@ -1172,7 +1172,7 @@ impl ServerKey {
11721172
trivially_encrypt_lwe_ciphertext(&mut ct.ct, encoded);
11731173

11741174
ct.degree = Degree::new(modular_value);
1175-
ct.set_noise_level(NoiseLevel::ZERO);
1175+
ct.set_noise_level(NoiseLevel::ZERO, self.max_noise_level);
11761176
}
11771177

11781178
pub fn bootstrapping_key_size_elements(&self) -> usize {
@@ -1326,7 +1326,7 @@ impl ServerKey {
13261326
);
13271327

13281328
output_shortint_ct.degree = *output_degree;
1329-
output_shortint_ct.set_noise_level(NoiseLevel::NOMINAL);
1329+
output_shortint_ct.set_noise_level(NoiseLevel::NOMINAL, self.max_noise_level);
13301330
outputs.push(output_shortint_ct);
13311331
}
13321332

@@ -1377,7 +1377,7 @@ impl ServerKey {
13771377
);
13781378

13791379
output_shortint_ct.degree = *output_degree;
1380-
output_shortint_ct.set_noise_level(NoiseLevel::NOMINAL);
1380+
output_shortint_ct.set_noise_level(NoiseLevel::NOMINAL, self.max_noise_level);
13811381
outputs.push(output_shortint_ct);
13821382
}
13831383

‎tfhe/src/shortint/server_key/mul.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -221,7 +221,7 @@ impl ServerKey {
221221
// ct1 + ct2
222222
let mut ct_add = ct1.clone();
223223

224-
unchecked_add_assign(&mut ct_add, ct2);
224+
unchecked_add_assign(&mut ct_add, ct2, self.max_noise_level);
225225

226226
// ct1 - ct2
227227
let (mut ct_sub, z) = self.unchecked_sub_with_correcting_term(ct1, ct2);

‎tfhe/src/shortint/server_key/scalar_mul.rs

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use crate::core_crypto::algorithms::*;
33
use crate::core_crypto::entities::*;
44
use crate::shortint::ciphertext::Degree;
55
use crate::shortint::server_key::CheckError;
6-
use crate::shortint::{Ciphertext, ServerKey};
6+
use crate::shortint::{Ciphertext, MaxNoiseLevel, ServerKey};
77

88
impl ServerKey {
99
/// Compute homomorphically a multiplication of a ciphertext by a scalar.
@@ -203,7 +203,7 @@ impl ServerKey {
203203
/// assert_eq!(3, clear);
204204
/// ```
205205
pub fn unchecked_scalar_mul_assign(&self, ct: &mut Ciphertext, scalar: u8) {
206-
unchecked_scalar_mul_assign(ct, scalar);
206+
unchecked_scalar_mul_assign(ct, scalar, self.max_noise_level);
207207
}
208208

209209
/// Multiply one ciphertext with a scalar in the case the carry space cannot fit the product
@@ -516,8 +516,12 @@ impl ServerKey {
516516
}
517517
}
518518

519-
pub(crate) fn unchecked_scalar_mul_assign(ct: &mut Ciphertext, scalar: u8) {
520-
ct.set_noise_level(ct.noise_level() * scalar as usize);
519+
pub(crate) fn unchecked_scalar_mul_assign(
520+
ct: &mut Ciphertext,
521+
scalar: u8,
522+
max_noise_level: MaxNoiseLevel,
523+
) {
524+
ct.set_noise_level(ct.noise_level() * scalar as usize, max_noise_level);
521525
ct.degree = Degree::new(ct.degree.get() * scalar as usize);
522526

523527
match scalar {

‎tfhe/src/shortint/server_key/shift.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -537,7 +537,7 @@ impl ServerKey {
537537
/// ```
538538
pub fn unchecked_scalar_left_shift_assign(&self, ct: &mut Ciphertext, shift: u8) {
539539
let scalar = 1_u8 << shift;
540-
unchecked_scalar_mul_assign(ct, scalar);
540+
unchecked_scalar_mul_assign(ct, scalar, self.max_noise_level);
541541
}
542542

543543
/// Checks if the left shift operation can be applied.

‎tfhe/src/shortint/server_key/sub.rs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -401,7 +401,10 @@ impl ServerKey {
401401

402402
lwe_ciphertext_add_assign(&mut ct_left.ct, &neg_right.ct);
403403

404-
ct_left.set_noise_level(ct_left.noise_level() + ct_right.noise_level());
404+
ct_left.set_noise_level(
405+
ct_left.noise_level() + ct_right.noise_level(),
406+
self.max_noise_level,
407+
);
405408
ct_left.degree = Degree::new(ct_left.degree.get() + z as usize);
406409

407410
z

0 commit comments

Comments
 (0)