three_operand_HS.patch

diff --git a/tfhe-rs/tfhe/src/integer/server_key/radix_parallel/add.rs b/tfhe-rs/tfhe/src/integer/server_key/radix_parallel/add.rs
index 4776825..788981a 100644
--- a/tfhe-rs/tfhe/src/integer/server_key/radix_parallel/add.rs
+++ b/tfhe-rs/tfhe/src/integer/server_key/radix_parallel/add.rs
@@ -1,9 +1,11 @@
 use crate::core_crypto::commons::numeric::UnsignedInteger;
+use crate::core_crypto::prelude::{lwe_ciphertext_plaintext_add_assign, lwe_ciphertext_plaintext_sub_assign, GlweCiphertext, Plaintext};
 use crate::integer::ciphertext::IntegerRadixCiphertext;
 use crate::integer::{BooleanBlock, RadixCiphertext, ServerKey, SignedRadixCiphertext};
 use crate::shortint::ciphertext::Degree;
 use crate::shortint::Ciphertext;
 use rayon::prelude::*;
+use crate::shortint::server_key::LookupTableOwned;
 
 #[derive(Copy, Clone, PartialEq, Eq, Debug)]
 pub(crate) enum CarryPropagationAlgorithm {
@@ -75,7 +77,13 @@ fn should_parallel_propagation_be_faster(
     let hillis_steel_depth = if num_carry_to_resolve == 0 {
         0
     } else {
-        num_carry_to_resolve.ceil_ilog2()
+        let mut k = 0;
+        let mut power = 1u64;
+        while power < num_carry_to_resolve as u64 {
+            power *= 3;
+            k += 1;
+        }
+        k
     };
 
     let parallel_algo_uses_sequential_to_resolve_grouping_carries =
@@ -85,13 +93,18 @@ fn should_parallel_propagation_be_faster(
         parallel_expected_latency += sequential_depth as usize
             * compute_latency_of_one_layer(grouping_size as usize, num_threads);
     } else {
-        let max_depth = num_blocks.ceil_ilog2();
+        let mut max_depth = 0;
+        let mut power = 1u64;
+        while power < num_blocks as u64 {
+            power *= 3;
+            max_depth += 1;
+        }
         let mut space = 1;
         for _ in 0..max_depth {
             let num_block_at_iter = num_blocks - space;
             let iter_latency = compute_latency_of_one_layer(num_block_at_iter, num_threads);
             parallel_expected_latency += iter_latency;
-            space *= 2;
+            space *= 3;
         }
     }
 
@@ -1061,7 +1074,7 @@ impl ServerKey {
 
         // This stores the LUTs that given a cum sum block in the first grouping
         // tells if a carry is generated or not
-        let first_grouping_inner_propagation_luts = (0..grouping_size - 1)
+        let first_grouping_inner_propagation_luts = (1..grouping_size - 1)
             .map(|index| {
                 self.key.generate_lookup_table(|propa_cum_sum_block| {
                     let carry = (propa_cum_sum_block >> index) & 1;
@@ -1076,7 +1089,7 @@ impl ServerKey {
 
         // This stores the LUTs that given a cum sum in non first grouping
         // tells if a carry is generated or propagated or neither of these
-        let other_groupings_inner_propagation_luts = (0..grouping_size)
+        let other_groupings_inner_propagation_luts = (1..grouping_size)
             .map(|index| {
                 self.key.generate_lookup_table(|propa_cum_sum_block| {
                     let mask = (2 << index) - 1;
@@ -1091,12 +1104,6 @@ impl ServerKey {
             })
             .collect::<Vec<_>>();
 
-        // This stores the LUT that outputs the propagation result of the first grouping
-        let first_grouping_outer_propagation_lut = self.key.generate_lookup_table(|block| {
-            // Check if the last bit of the block is set
-            (block >> (num_bits_in_block - 1)) & 1
-        });
-
         let num_groupings = num_blocks.div_ceil(grouping_size);
         let num_carry_to_resolve = num_groupings - 1;
 
@@ -1105,12 +1112,33 @@ impl ServerKey {
         let hillis_steel_depth = if num_carry_to_resolve == 0 {
             0
         } else {
-            num_carry_to_resolve.ceil_ilog2()
+            let mut k = 0;
+            let mut power = 1u64;
+            while power < num_carry_to_resolve as u64 {
+                power *= 3;
+                k += 1;
+            }
+            k
         };
 
         let use_sequential_algorithm_to_resolved_grouping_carries =
             sequential_depth <= hillis_steel_depth;
 
+        // This stores the LUT that outputs the propagation result of the first grouping
+        let first_grouping_outer_propagation_lut = if use_sequential_algorithm_to_resolved_grouping_carries{
+            self.key.generate_lookup_table(|block| {
+                // Check if the last bit of the block is set
+                (block >> (num_bits_in_block - 1)) & 1
+            })
+        } else {
+            self.key.generate_lookup_table(|block| {
+                // Check if the last bit of the block is set
+                ((block >> (num_bits_in_block - 1)) & 1) + 1
+            })
+        };
+
+        
+
         // This stores the LUTs that output the propagation result of the other groupings
         let grouping_chunk_pgn_luts = if use_sequential_algorithm_to_resolved_grouping_carries {
             // When using the sequential algorithm for the propagation of one grouping to the
@@ -1143,7 +1171,7 @@ impl ServerKey {
             vec![self.key.generate_lookup_table(|block| {
                 if block == (block_modulus - 1) {
                     // All bits set to 1 (e.g. 0b1111), means propagate
-                    2
+                    0
                 } else {
                     // u64::MAX is -1 in two's complement
                     // We apply the modulus including the padding bit
@@ -1174,24 +1202,33 @@ impl ServerKey {
                 let grouping_index = i / grouping_size;
                 let is_in_first_grouping = grouping_index == 0;
                 let index_in_grouping = i % grouping_size;
+                let is_first_index_in_grouping = index_in_grouping == 0;
 
-                let lut = if is_in_first_grouping {
+                let lut = if is_first_index_in_grouping {
+                    None
+                } else if is_in_first_grouping {
                     if index_in_grouping == grouping_size - 1 {
-                        &first_grouping_outer_propagation_lut
+                        Some(&first_grouping_outer_propagation_lut)
                     } else {
-                        &first_grouping_inner_propagation_luts[index_in_grouping]
+                        Some(&first_grouping_inner_propagation_luts[index_in_grouping-1])
                     }
                 } else if index_in_grouping == grouping_size - 1 {
                     if use_sequential_algorithm_to_resolved_grouping_carries {
-                        &grouping_chunk_pgn_luts[(grouping_index - 1) % (grouping_size - 1)]
+                        Some(&grouping_chunk_pgn_luts[(grouping_index - 1) % (grouping_size - 1)])
                     } else {
-                        &grouping_chunk_pgn_luts[0]
+                        Some(&grouping_chunk_pgn_luts[0])
                     }
                 } else {
-                    &other_groupings_inner_propagation_luts[index_in_grouping]
+                    Some(&other_groupings_inner_propagation_luts[index_in_grouping-1])
                 };
 
-                self.key.apply_lookup_table_assign(cum_sum_block, lut);
+                if is_first_index_in_grouping {
+                    if is_in_first_grouping {
+                        self.key.unchecked_scalar_mul_assign(cum_sum_block, 2);
+                    }
+                } else {
+                    self.key.apply_lookup_table_assign(cum_sum_block, lut.unwrap());
+                }
 
                 let may_have_its_padding_bit_set =
                     !is_in_first_grouping && index_in_grouping == grouping_size - 1;
@@ -1231,12 +1268,150 @@ impl ServerKey {
         } else if use_sequential_algorithm_to_resolved_grouping_carries {
             self.resolve_carries_of_groups_sequentially(groupings_pgns, grouping_size)
         } else {
-            self.resolve_carries_of_groups_using_hillis_steele(groupings_pgns)
+            self.resolve_carries_of_groups_using_log3_hillis_steele(groupings_pgns, block_modulus)
         };
 
         (propagation_simulators, resolved_carries)
     }
 
+    /// This resolves the carries using a log3 Hillis-Steele algorithm
+    ///
+    /// Blocks must have a value in 
+    /// - 2 for generate
+    /// - 1 for propagate
+    /// - 0 for no carry
+    ///
+    /// The returned Vec of blocks encrypting 1 if a carry is generated, 0 if not
+    pub(crate) fn resolve_carries_of_groups_using_log3_hillis_steele(
+        &self,
+        mut groupings_pgns: Vec<Ciphertext>,
+        block_modulus: u64,
+    ) -> Vec<Ciphertext> {
+        let num_bits_in_block = block_modulus.ilog2();
+        
+        let mut acc0 = GlweCiphertext::new(0, self.key.bootstrapping_key.glwe_size(), self.key.bootstrapping_key.polynomial_size(), self.key.ciphertext_modulus);
+        self.generate_resolved_lut(&mut acc0, block_modulus as usize, num_bits_in_block);
+        let to_resolved_lut = LookupTableOwned {
+            acc: acc0,
+            degree: Degree::new(2),
+        };
+        let mut acc1 = GlweCiphertext::new(0, self.key.bootstrapping_key.glwe_size(), self.key.bootstrapping_key.polynomial_size(), self.key.ciphertext_modulus);
+        self.generate_unresolved_lut(&mut acc1, block_modulus as usize, num_bits_in_block);
+        let to_unresolved_lut = LookupTableOwned {
+            acc: acc1,
+            degree: Degree::new(2),
+        };
+
+        let num_blocks = groupings_pgns.len();
+        let mut num_steps: usize = 0;
+        let mut power = 1u64;
+        while power < num_blocks as u64 {
+            power *= 3;
+            num_steps += 1;
+        }
+
+        let mut space = 1;
+        let mut step_output = groupings_pgns.clone();
+        for _ in 0..num_steps {
+            step_output[space..num_blocks]
+                .par_iter_mut()
+                .enumerate()
+                .for_each(|(i, block)| {
+                    if i<space {
+                        self.key.unchecked_scalar_mul_assign(block, 3);
+                        self.key.unchecked_add_assign(block, &groupings_pgns[i]);
+                        self.key.unchecked_scalar_mul_assign(block, 3);
+                        self.key.apply_lookup_table_assign(block, &to_resolved_lut);
+                        lwe_ciphertext_plaintext_add_assign(&mut block.ct, Plaintext(3*2u64.pow(62-num_bits_in_block)));
+                    } else if i<2*space {
+                        self.key.unchecked_scalar_mul_assign(block, 3);
+                        self.key.unchecked_add_assign(block, &groupings_pgns[i]);
+                        self.key.unchecked_scalar_mul_assign(block, 3);
+                        self.key.unchecked_add_assign(block, &groupings_pgns[i-space]);
+                        self.key.apply_lookup_table_assign(block, &to_resolved_lut);
+                        lwe_ciphertext_plaintext_add_assign(&mut block.ct, Plaintext(3*2u64.pow(62-num_bits_in_block)));
+                    } else {
+                        self.key.unchecked_scalar_mul_assign(block, 3);
+                        self.key.unchecked_add_assign(block, &groupings_pgns[i]);
+                        self.key.unchecked_scalar_mul_assign(block, 3);
+                        self.key.unchecked_add_assign(block, &groupings_pgns[i-space]);
+                        self.key.apply_lookup_table_assign(block, &to_unresolved_lut);
+                        lwe_ciphertext_plaintext_add_assign(&mut block.ct, Plaintext(2u64.pow(63-num_bits_in_block)));
+                    }
+                });
+            for i in space..num_blocks {
+                groupings_pgns[i].clone_from(&step_output[i]);
+            }
+
+            space *= 3;
+        }
+
+        groupings_pgns.par_iter_mut().for_each(|block| {
+            lwe_ciphertext_plaintext_sub_assign(&mut block.ct, Plaintext(2u64.pow(63-num_bits_in_block)));
+        });
+
+        groupings_pgns.insert(0, self.key.create_trivial(0));
+        groupings_pgns
+    }
+
+    fn generate_resolved_lut(&self, acc : &mut  GlweCiphertext<Vec<u64>>, block_modulus: usize, num_bits_in_block: u32) 
+    {
+        // generate special LUT
+        let mut accumulator_view = acc.as_mut_view();
+
+        // N/(p/2) = size of each block
+        let box_size = self.key.bootstrapping_key.polynomial_size().0 / block_modulus;
+        let mut body = accumulator_view.get_mut_body();
+        let accumulator_u64 = body.as_mut();
+        for i in 0..block_modulus {
+            let index = i * box_size;
+            let f_eval = if i <= 13 {
+                (2u64.pow(num_bits_in_block+2)-1) * 2u64.pow(62-num_bits_in_block)
+            } else {
+                2u64.pow(62-num_bits_in_block)
+            };
+            accumulator_u64[index..index + box_size].fill(f_eval);
+        }
+        let half_box_size = box_size / 2;
+        // Negate the first half_box_size coefficients
+        for a_i in accumulator_u64[0..half_box_size].iter_mut() {
+            *a_i = (*a_i).wrapping_neg();
+        }
+
+        // Rotate the accumulator
+        accumulator_u64.rotate_left(half_box_size);
+    }
+
+    fn generate_unresolved_lut(&self, acc : &mut  GlweCiphertext<Vec<u64>>, block_modulus: usize, num_bits_in_block: u32) 
+    {
+        // generate special LUT
+        let mut accumulator_view = acc.as_mut_view();
+
+        // N/(p/2) = size of each block
+        let box_size = self.key.bootstrapping_key.polynomial_size().0 / block_modulus;
+        let mut body = accumulator_view.get_mut_body();
+        let accumulator_u64 = body.as_mut();
+        for i in 0..block_modulus {
+            let index = i * box_size;
+            let f_eval = if i<13 {
+                (2u64.pow(num_bits_in_block+1)-1) * 2u64.pow(63-num_bits_in_block)
+            } else if i==13 {
+                0
+            } else {
+                2u64.pow(63-num_bits_in_block)
+            };
+            accumulator_u64[index..index + box_size].fill(f_eval);
+        }
+        let half_box_size = box_size / 2;
+        // Negate the first half_box_size coefficients
+        for a_i in accumulator_u64[0..half_box_size].iter_mut() {
+            *a_i = (*a_i).wrapping_neg();
+        }
+
+        // Rotate the accumulator
+        accumulator_u64.rotate_left(half_box_size);
+    }
+    
     /// This resolves the carries using a Hillis-Steele algorithm
     ///
     /// Blocks must have a value in
diff --git a/tfhe-rs/tfhe/src/integer/server_key/radix_parallel/mod.rs b/tfhe-rs/tfhe/src/integer/server_key/radix_parallel/mod.rs
index 399759a..a6ab908 100644
--- a/tfhe-rs/tfhe/src/integer/server_key/radix_parallel/mod.rs
+++ b/tfhe-rs/tfhe/src/integer/server_key/radix_parallel/mod.rs
@@ -134,13 +134,15 @@ impl ServerKey {
         };
 
         if self.is_eligible_for_parallel_single_carry_propagation(blocks.len()) {
+            //if at least 4 message/carry bits and if parallel is faster
             let highest_degree = blocks
                 .iter()
                 .max_by(|block_a, block_b| block_a.degree.get().cmp(&block_b.degree.get()))
                 .map(|block| block.degree.get())
                 .unwrap(); // We checked for emptiness earlier
 
-            if highest_degree >= (self.key.message_modulus.0 - 1) * 2 {
+            if highest_degree > (self.key.message_modulus.0 - 1) * 2 {
+                // if true, we cannot do the HS procedure yet, but first have to go to a result with lower degree
                 // At least one of the blocks has more than one carry,
                 // we need to extract message and carries, then add + propagate
                 let (mut message_blocks, carry_blocks) = extract_message_and_carry_blocks(blocks);
diff --git a/tfhe-rs/tfhe/src/shortint/parameters/v1_1/classic/gaussian/p_fail_2_minus_128/ks_pbs.rs b/tfhe-rs/tfhe/src/shortint/parameters/v1_1/classic/gaussian/p_fail_2_minus_128/ks_pbs.rs
index 7920830..ead5047 100644
--- a/tfhe-rs/tfhe/src/shortint/parameters/v1_1/classic/gaussian/p_fail_2_minus_128/ks_pbs.rs
+++ b/tfhe-rs/tfhe/src/shortint/parameters/v1_1/classic/gaussian/p_fail_2_minus_128/ks_pbs.rs
@@ -316,35 +316,77 @@ pub const V1_1_PARAM_MESSAGE_2_CARRY_1_KS_PBS_GAUSSIAN_2M128: ClassicPBSParamete
         }),
     };
 
-// p-fail = 2^-128.377, algorithmic cost ~ 110, 2-norm = 5
-// Average number of encryptions of 0s ~ 17, peak noise ~ Variance(0.00000141649065433221)
-pub const V1_1_PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M128: ClassicPBSParameters =
+// // p-fail = 2^-128.377, algorithmic cost ~ 110, 2-norm = 5
+// // Average number of encryptions of 0s ~ 17, peak noise ~ Variance(0.00000141649065433221)
+// pub const V1_1_PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M128: ClassicPBSParameters =
+//     ClassicPBSParameters {
+//         lwe_dimension: LweDimension(866),
+//         glwe_dimension: GlweDimension(1),
+//         polynomial_size: PolynomialSize(2048),
+//         lwe_noise_distribution: DynamicDistribution::new_gaussian_from_std_dev(StandardDev(
+//             2.046151696979124e-06,
+//         )),
+//         glwe_noise_distribution: DynamicDistribution::new_gaussian_from_std_dev(StandardDev(
+//             2.845267479601915e-15,
+//         )),
+//         pbs_base_log: DecompositionBaseLog(23),
+//         pbs_level: DecompositionLevelCount(1),
+//         ks_base_log: DecompositionBaseLog(3),
+//         ks_level: DecompositionLevelCount(5),
+//         message_modulus: MessageModulus(4),
+//         carry_modulus: CarryModulus(4),
+//         max_noise_level: MaxNoiseLevel::new(5),
+//         log2_p_fail: -128.377,
+//         ciphertext_modulus: CiphertextModulus::new_native(),
+//         encryption_key_choice: EncryptionKeyChoice::Big,
+//         modulus_switch_noise_reduction_params: Some(ModulusSwitchNoiseReductionParams {
+//             modulus_switch_zeros_count: LweCiphertextCount(1446),
+//             ms_bound: NoiseEstimationMeasureBound(288230376151711744f64),
+//             ms_r_sigma_factor: RSigmaFactor(13.128441378136914f64),
+//             ms_input_variance: Variance(3.38639994643900E-7f64),
+//         }),
+//     };
+
+// // Custom Parameter set for baseline
+// pub const V1_1_PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M128: ClassicPBSParameters = 
+//     ClassicPBSParameters {
+//         lwe_dimension: LweDimension(831),
+//         glwe_dimension: GlweDimension(1),
+//         polynomial_size: PolynomialSize(2048),
+//         lwe_noise_distribution: DynamicDistribution::new_gaussian_from_std_dev(StandardDev(2.325084026726674e-06)),
+//         glwe_noise_distribution: DynamicDistribution::new_gaussian_from_std_dev(StandardDev(7.28604160123004e-16)),
+//         pbs_base_log: DecompositionBaseLog(24),
+//         pbs_level: DecompositionLevelCount(1),
+//         ks_base_log: DecompositionBaseLog(3),
+//         ks_level: DecompositionLevelCount(6),
+//         message_modulus: MessageModulus(4),
+//         carry_modulus: CarryModulus(4),
+//         max_noise_level: MaxNoiseLevel::new(135),
+//         log2_p_fail: -128.0,
+//         ciphertext_modulus: CiphertextModulus::new_native(),
+//         encryption_key_choice: EncryptionKeyChoice::Big,
+//         modulus_switch_noise_reduction_params: None,
+//     };
+
+// Custom Parameter set for altered LUT
+pub const V1_1_PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M128: ClassicPBSParameters = 
     ClassicPBSParameters {
-        lwe_dimension: LweDimension(866),
+        lwe_dimension: LweDimension(837),
         glwe_dimension: GlweDimension(1),
         polynomial_size: PolynomialSize(2048),
-        lwe_noise_distribution: DynamicDistribution::new_gaussian_from_std_dev(StandardDev(
-            2.046151696979124e-06,
-        )),
-        glwe_noise_distribution: DynamicDistribution::new_gaussian_from_std_dev(StandardDev(
-            2.845267479601915e-15,
-        )),
-        pbs_base_log: DecompositionBaseLog(23),
+        lwe_noise_distribution: DynamicDistribution::new_gaussian_from_std_dev(StandardDev(2.1058836244083634e-06)),
+        glwe_noise_distribution: DynamicDistribution::new_gaussian_from_std_dev(StandardDev(7.28604160123004e-16)),
+        pbs_base_log: DecompositionBaseLog(24),
         pbs_level: DecompositionLevelCount(1),
         ks_base_log: DecompositionBaseLog(3),
         ks_level: DecompositionLevelCount(5),
         message_modulus: MessageModulus(4),
         carry_modulus: CarryModulus(4),
-        max_noise_level: MaxNoiseLevel::new(5),
-        log2_p_fail: -128.377,
+        max_noise_level: MaxNoiseLevel::new(300),
+        log2_p_fail: -128.0,
         ciphertext_modulus: CiphertextModulus::new_native(),
         encryption_key_choice: EncryptionKeyChoice::Big,
-        modulus_switch_noise_reduction_params: Some(ModulusSwitchNoiseReductionParams {
-            modulus_switch_zeros_count: LweCiphertextCount(1446),
-            ms_bound: NoiseEstimationMeasureBound(288230376151711744f64),
-            ms_r_sigma_factor: RSigmaFactor(13.128441378136914f64),
-            ms_input_variance: Variance(3.38639994643900E-7f64),
-        }),
+        modulus_switch_noise_reduction_params: None,
     };
 
 // p-fail = 2^-128.419, algorithmic cost ~ 373, 2-norm = 10