WIP

zama-ai · Sep 18, 2024 · 649673b · 649673b
1 parent 47848fd
commit 649673b
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 13 deletions.
diff --git a/tfhe/src/high_level_api/integers/signed/base.rs b/tfhe/src/high_level_api/integers/signed/base.rs
@@ -462,7 +462,7 @@ where
                     result,
                     crate::FheUint32Id::num_blocks(cpu_key.pbs_key().message_modulus()),
                 );
-                crate::FheUint32::new(result)
+                crate::FheUint32::new(result, cpu_key.tag.clone())
             }
             #[cfg(feature = "gpu")]
             InternalServerKey::Cuda(_) => {
@@ -499,7 +499,7 @@ where
                     result,
                     crate::FheUint32Id::num_blocks(cpu_key.pbs_key().message_modulus()),
                 );
-                crate::FheUint32::new(result)
+                crate::FheUint32::new(result, cpu_key.tag.clone())
             }
             #[cfg(feature = "gpu")]
             InternalServerKey::Cuda(_) => {

diff --git a/tfhe/src/high_level_api/integers/unsigned/base.rs b/tfhe/src/high_level_api/integers/unsigned/base.rs
@@ -561,7 +561,7 @@ where
                     result,
                     super::FheUint32Id::num_blocks(cpu_key.pbs_key().message_modulus()),
                 );
-                super::FheUint32::new(result)
+                super::FheUint32::new(result, cpu_key.tag.clone())
             }
             #[cfg(feature = "gpu")]
             InternalServerKey::Cuda(_) => {
@@ -598,7 +598,7 @@ where
                     result,
                     super::FheUint32Id::num_blocks(cpu_key.pbs_key().message_modulus()),
                 );
-                super::FheUint32::new(result)
+                super::FheUint32::new(result, cpu_key.tag.clone())
             }
             #[cfg(feature = "gpu")]
             InternalServerKey::Cuda(_) => {

diff --git a/tfhe/src/integer/server_key/mod.rs b/tfhe/src/integer/server_key/mod.rs
@@ -33,7 +33,7 @@ use super::backward_compatibility::server_key::{CompressedServerKeyVersions, Ser
 #[derive(Serialize, Deserialize, Clone, Versionize)]
 #[versionize(ServerKeyVersions)]
 pub struct ServerKey {
-    pub key: crate::shortint::ServerKey,
+    pub(crate) key: crate::shortint::ServerKey,
 }
 
 impl From<ServerKey> for crate::shortint::ServerKey {

diff --git a/tfhe/src/integer/server_key/radix_parallel/count_zeros_ones.rs b/tfhe/src/integer/server_key/radix_parallel/count_zeros_ones.rs
@@ -1,5 +1,6 @@
 use super::ServerKey;
 use crate::integer::{IntegerRadixCiphertext, RadixCiphertext, SignedRadixCiphertext};
+use crate::shortint::ciphertext::Degree;
 
 use rayon::prelude::*;
 
@@ -47,7 +48,7 @@ impl ServerKey {
     where
         T: IntegerRadixCiphertext,
     {
-        if self.message_modulus().0 == 4 && self.carry_modulus().0 >= 4 {
+        if self.message_modulus().0 == 4 && self.carry_modulus().0 == 4 {
             self.count_bits_2_2(ct, kind)
         } else {
             self.count_bits_naive(ct, kind)
@@ -151,7 +152,7 @@ impl ServerKey {
 
         // We can pack the block if the carry space allow it, but more importantly,
         // if the number of bits in 2 blocks does not exceed the message modulus
-        // e.g. 1_1 -> 2 bits in one block -> 2 blocks = 2 bits -> 2 >= 2**1 (2)-> cant pack
+        // e.g. 1_1 -> 1 bits in one block -> 2 blocks = 2 bits -> 2 >= 2**1 (2)-> cant pack
         //      3_3 -> 3 bits in one block -> 2 blocks = 6 bits  -> 6 < 2**3 (8) -> can pack
         let can_pack = self.carry_modulus().0 >= self.message_modulus().0
             && (num_bits_in_block * 2) < (self.message_modulus().0 as u32);
@@ -325,7 +326,7 @@ impl ServerKey {
         // Now, to go further, with 3 blocks, which is 6 bits, we can do 2 bivariate PBS, to split
         // the count in two blocks with value in 0…=3
         // [b0,b1] [b2, b3] [b4, b5]
-        // PBS 1 -> [b0, b2, b2, b3] -> count(b0, b2,b2)
+        // PBS 1 -> [b0, b1, b2, b3] -> count(b0, b1, b2)
         // PBS 2 -> [b2, b3, b4, b5] -> count(b3, b4, b5)
         // This also mean 2 PBS for 3 blocks, instead of 3 for 3
         //
@@ -363,7 +364,7 @@ impl ServerKey {
         num_full_chunks += num_single_completer_blocks;
         num_non_chunked -= num_single_completer_blocks;
 
-        // No we go a bit beyond again
+        // Now, we go a bit beyond again
         //
         // A non-full chunk has 3 packed blocks, so 6 ciphertexts
         // 3 ciphertexts can be split into 2 blocks, to complete 2 non-full chunks
@@ -468,7 +469,7 @@ impl ServerKey {
                     .par_iter()
                     .enumerate()
                     .map(|(i, packed_block)| {
-                        self.key.apply_lookup_table(packed_block, &luts[(i) % 2])
+                        self.key.apply_lookup_table(packed_block, &luts[i % 2])
                     })
                     .chain(single_completer_blocks.par_iter().map(|block| {
                         self.key
@@ -478,8 +479,8 @@ impl ServerKey {
             },
         );
 
-        let num_sum = ((self.message_modulus().0 * self.carry_modulus().0) - 1)
-            / (self.message_modulus().0 * self.carry_modulus().0).ilog2() as usize;
+        // Since each block encrypts 4 bits, max count is 4
+        let num_sum = self.max_sum_size(Degree::new(4));
 
         let mut pre_count = bit_count_of_packed_blocks
             .chunks_exact(num_sum)

diff --git a/tfhe/src/integer/server_key/radix_parallel/tests_signed/test_count_zeros_ones.rs b/tfhe/src/integer/server_key/radix_parallel/tests_signed/test_count_zeros_ones.rs
@@ -67,9 +67,12 @@ pub(crate) fn signed_default_count_zeros_ones_test<P, E1, E2>(
         else {
             break;
         };
-        if modulus > i128::MAX as u128 || modulus <= 2 {
+        if modulus > i128::MAX as u128 {
             break;
         }
+        if modulus <= 2 {
+            continue;
+        }
         let half_modulus = modulus / 2;
 
         for _ in 0..nb_tests {