chore(core): add utils to test noise distribution for power of 2 q

tfhe/src/core_crypto/algorithms/misc.rs

@@ -64,6 +64,50 @@ where
     y.wrapping_rem(Scalar::ONE.shl(log2_modulo))
 }
 
+/// Compute the smallest signed difference between two torus elements
+pub fn torus_modular_diff<T: UnsignedInteger>(
+    first: T,
+    other: T,
+    modulus: CiphertextModulus<T>,
+) -> f64 {
+    if modulus.is_native_modulus() {
+        let bits = T::BITS as i32;
+        // Using the [0; 1[ torus to reason
+        // Example with first = 0.1 and other = 0.9
+        // d0 = first - other = -0.8 = 0.2 mod 1
+        // d1 = other - first = 0.8
+        // d0 < d1 return 0.2
+        // if other and first are inverted we get
+        // d0 = 0.8
+        // d1 = 0.2
+        // d1 <= d0 return -0.2, the minus here can be seen as taking first as a reference
+        // In the first example adding 0.2 to other (0.9 + 0.2 mod 1 = 0.1) gets us to first
+        // In the second example adding -0.2 to other (0.1 - 0.2 mod 1 = 0.9) gets us to first
+        let d0 = first.wrapping_sub(other);
+        let d1 = other.wrapping_sub(first);
+        if d0 < d1 {
+            let d: f64 = d0.cast_into();
+            d / 2_f64.powi(bits)
+        } else {
+            let d: f64 = d1.cast_into();
+            -d / 2_f64.powi(bits)
+        }
+    } else {
+        let custom_modulus = T::cast_from(modulus.get_custom_modulus());
+        let d0 = first.wrapping_sub_custom_mod(other, custom_modulus);
+        let d1 = other.wrapping_sub_custom_mod(first, custom_modulus);
+        if d0 < d1 {
+            let d: f64 = d0.cast_into();
+            let cm_f: f64 = custom_modulus.cast_into();
+            d / cm_f
+        } else {
+            let d: f64 = d1.cast_into();
+            let cm_f: f64 = custom_modulus.cast_into();
+            -d / cm_f
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;

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

@@ -1,5 +1,5 @@
 use super::*;
-use crate::core_crypto::commons::test_tools::{torus_modular_distance, variance};
+use crate::core_crypto::commons::test_tools::{torus_modular_diff, variance};
 
 // This is 1 / 16 which is exactly representable in an f64 (even an f32)
 // 1 / 32 is too strict and fails the tests
@@ -57,21 +57,27 @@ fn lwe_encrypt_decrypt_noise_distribution_custom_mod<Scalar: UnsignedTorus + Cas
 
             assert_eq!(msg, decoded);
 
-            let torus_distance =
-                torus_modular_distance(plaintext.0, decrypted.0, ciphertext_modulus);
+            let torus_distance = torus_modular_diff(plaintext.0, decrypted.0, ciphertext_modulus);
             noise_samples.push(torus_distance);
         }
     }
 
     let measured_variance = variance(&noise_samples);
+
+    let var_abs_diff = (expected_variance.0 - measured_variance.0).abs();
+    let tolerance_threshold = RELATIVE_TOLERANCE * expected_variance.0;
     assert!(
-        (expected_variance.0 - measured_variance.0).abs()
-            < RELATIVE_TOLERANCE * expected_variance.0
+        var_abs_diff < tolerance_threshold,
+        "Absolute difference for variance: {var_abs_diff}, \
+        tolerance threshold: {tolerance_threshold}, \
+        got variance: {measured_variance:?}, \
+        expected variance: {expected_variance:?}"
     );
 }
 
 create_parametrized_test!(lwe_encrypt_decrypt_noise_distribution_custom_mod {
-    TEST_PARAMS_4_BITS_NATIVE_U64
+    TEST_PARAMS_4_BITS_NATIVE_U64,
+    TEST_PARAMS_3_BITS_63_U64
 });
 
 fn lwe_compact_public_key_encryption_expected_variance(
@@ -158,16 +164,20 @@ fn lwe_compact_public_encrypt_noise_distribution_custom_mod<
 
             assert_eq!(msg, decoded);
 
-            let torus_distance =
-                torus_modular_distance(plaintext.0, decrypted.0, ciphertext_modulus);
+            let torus_distance = torus_modular_diff(plaintext.0, decrypted.0, ciphertext_modulus);
             noise_samples.push(torus_distance);
         }
     }
 
     let measured_variance = variance(&noise_samples);
+    let var_abs_diff = (expected_variance.0 - measured_variance.0).abs();
+    let tolerance_threshold = RELATIVE_TOLERANCE * expected_variance.0;
     assert!(
-        (expected_variance.0 - measured_variance.0).abs()
-            < RELATIVE_TOLERANCE * expected_variance.0
+        var_abs_diff < tolerance_threshold,
+        "Absolute difference for variance: {var_abs_diff}, \
+        tolerance threshold: {tolerance_threshold}, \
+        got variance: {measured_variance:?}, \
+        expected variance: {expected_variance:?}"
     );
 }
 

tfhe/src/core_crypto/commons/mod.rs

@@ -56,7 +56,7 @@ pub mod traits;
 pub mod test_tools {
     use rand::Rng;
 
-    use crate::core_crypto::commons::ciphertext_modulus::CiphertextModulus;
+    pub use crate::core_crypto::algorithms::misc::torus_modular_diff;
     use crate::core_crypto::commons::dispersion::{DispersionParameter, Variance};
     use crate::core_crypto::commons::generators::{
         EncryptionRandomGenerator, SecretRandomGenerator,
@@ -77,32 +77,6 @@ pub mod test_tools {
         d0.min(d1)
     }
 
-    pub fn torus_modular_distance<T: UnsignedInteger>(
-        first: T,
-        other: T,
-        modulus: CiphertextModulus<T>,
-    ) -> f64 {
-        if modulus.is_compatible_with_native_modulus() {
-            let bits = if modulus.is_native_modulus() {
-                T::BITS as i32
-            } else {
-                modulus.get_custom_modulus().ilog2() as i32
-            };
-
-            let d0 = first.wrapping_sub(other);
-            let d1 = other.wrapping_sub(first);
-            if d0 < d1 {
-                let d: f64 = d0.cast_into();
-                d / 2_f64.powi(bits)
-            } else {
-                let d: f64 = d1.cast_into();
-                -d / 2_f64.powi(bits)
-            }
-        } else {
-            todo!("Currently unimplemented for non power of 2 moduli")
-        }
-    }
-
     pub fn variance(samples: &[f64]) -> Variance {
         let num_samples = samples.len();
         let mean = samples.iter().sum::<f64>() / (num_samples as f64);

tfhe/src/core_crypto/commons/numeric/unsigned.rs

@@ -46,6 +46,12 @@ pub trait UnsignedInteger:
     /// Compute a subtraction, modulo the max of the type.
     #[must_use]
     fn wrapping_sub(self, other: Self) -> Self;
+    /// Compute an addition, modulo a custom modulus.
+    #[must_use]
+    fn wrapping_add_custom_mod(self, other: Self, custom_modulus: Self) -> Self;
+    /// Compute a subtraction, modulo a custom modulus.
+    #[must_use]
+    fn wrapping_sub_custom_mod(self, other: Self, custom_modulus: Self) -> Self;
     /// Compute a division, modulo the max of the type.
     #[must_use]
     fn wrapping_div(self, other: Self) -> Self;
@@ -128,6 +134,35 @@ macro_rules! implement {
                 self.wrapping_sub(other)
             }
             #[inline]
+            fn wrapping_add_custom_mod(self, other: Self, custom_modulus: Self) -> Self {
+                if Self::BITS < 128 {
+                    let self_u128: u128 = self.cast_into();
+                    let other_u128: u128 = other.cast_into();
+                    let custom_modulus_u128: u128 = custom_modulus.cast_into();
+                    self_u128
+                        .wrapping_add(other_u128)
+                        .wrapping_rem(custom_modulus_u128)
+                        .cast_into()
+                } else {
+                    todo!("wrapping_add_custom_mod is not yet implemented for types wider than u64")
+                }
+            }
+            #[inline]
+            fn wrapping_sub_custom_mod(self, other: Self, custom_modulus: Self) -> Self {
+                if Self::BITS < 128 {
+                    let self_u128: u128 = self.cast_into();
+                    let other_u128: u128 = other.cast_into();
+                    let custom_modulus_u128: u128 = custom_modulus.cast_into();
+                    self_u128
+                        .wrapping_add(custom_modulus_u128)
+                        .wrapping_sub(other_u128)
+                        .wrapping_rem(custom_modulus_u128)
+                        .cast_into()
+                } else {
+                    todo!("wrapping_sub_custom_mod is not yet implemented for types wider than u64")
+                }
+            }
+            #[inline]
             fn wrapping_div(self, other: Self) -> Self {
                 self.wrapping_div(other)
             }
@@ -228,4 +263,77 @@ mod test {
                 .to_string()
         );
     }
+
+    #[test]
+    fn test_wrapping_add_custom_mod() {
+        let custom_modulus_u128 = (1u128 << 64) - (1 << 32) + 1;
+        let custom_modulus = custom_modulus_u128 as u64;
+        let a = u64::MAX % custom_modulus;
+        let b = u64::MAX % custom_modulus;
+
+        let a_u128: u128 = a.into();
+        let b_u128: u128 = b.into();
+
+        let expected_res = ((a_u128 + b_u128) % custom_modulus_u128) as u64;
+
+        let res = a.wrapping_add_custom_mod(b, custom_modulus);
+        assert_eq!(expected_res, res);
+
+        const NB_REPS: usize = 100_000_000;
+
+        use rand::Rng;
+        let mut thread_rng = rand::thread_rng();
+        for _ in 0..NB_REPS {
+            let a = thread_rng.gen::<u64>() % custom_modulus;
+            let b = thread_rng.gen::<u64>() % custom_modulus;
+
+            let a_u128: u128 = a.into();
+            let b_u128: u128 = b.into();
+
+            let expected_res = ((a_u128 + b_u128) % custom_modulus_u128) as u64;
+
+            let res = a.wrapping_add_custom_mod(b, custom_modulus);
+            assert_eq!(expected_res, res, "a: {a}, b: {b}");
+        }
+    }
+
+    #[test]
+    fn test_wrapping_sub_custom_mod() {
+        let custom_modulus_u128 = (1u128 << 64) - (1 << 32) + 1;
+        let custom_modulus = custom_modulus_u128 as u64;
+
+        let a = 0u64;
+        let b = u64::MAX % custom_modulus;
+
+        let a_u128: u128 = a.into();
+        let b_u128: u128 = b.into();
+
+        let expected_res = ((a_u128
+            .wrapping_add(custom_modulus_u128)
+            .wrapping_sub(b_u128))
+            % custom_modulus_u128) as u64;
+
+        let res = a.wrapping_sub_custom_mod(b, custom_modulus);
+        assert_eq!(expected_res, res);
+
+        const NB_REPS: usize = 100_000_000;
+
+        use rand::Rng;
+        let mut thread_rng = rand::thread_rng();
+        for _ in 0..NB_REPS {
+            let a = thread_rng.gen::<u64>() % custom_modulus;
+            let b = thread_rng.gen::<u64>() % custom_modulus;
+
+            let a_u128: u128 = a.into();
+            let b_u128: u128 = b.into();
+
+            let expected_res = ((a_u128
+                .wrapping_add(custom_modulus_u128)
+                .wrapping_sub(b_u128))
+                % custom_modulus_u128) as u64;
+
+            let res = a.wrapping_sub_custom_mod(b, custom_modulus);
+            assert_eq!(expected_res, res, "a: {a}, b: {b}");
+        }
+    }
 }