fix(hpu): Correctly select adder configuration in ERC_20/ERC_20_SIMD

Add knobs to select ripple or kogge adder in ERC_20/ERC_20_SIMD.
Previously, it was hardcoded to ripple carry and thus degraded latency
performance of ERC_20.

backends/tfhe-hpu-backend/src/fw/fw_impl/llt/mod.rs

@@ -113,6 +113,7 @@ pub fn iop_add_simd(prog: &mut Program) {
         prog,
         crate::asm::iop::SIMD_N,
         fw_impl::llt::iop_add_ripple_rtl,
+        None,
     );
 }
 
@@ -227,14 +228,18 @@ pub fn iop_muls(prog: &mut Program) {
 pub fn iop_erc_20(prog: &mut Program) {
     // Add Comment header
     prog.push_comment("ERC_20 (new_from, new_to) <- (from, to, amount)".to_string());
-    iop_erc_20_rtl(prog, 0).add_to_prog(prog);
+    // TODO: Make sweep of kogge_blk_w
+    // All these little parameters would be very handy to write an
+    // exploration/compilation program which would try to minimize latency by
+    // playing with these.
+    iop_erc_20_rtl(prog, 0, Some(10)).add_to_prog(prog);
 }
 
 #[instrument(level = "trace", skip(prog))]
 pub fn iop_erc_20_simd(prog: &mut Program) {
     // Add Comment header
     prog.push_comment("ERC_20_SIMD (new_from, new_to) <- (from, to, amount)".to_string());
-    simd(prog, crate::asm::iop::SIMD_N, fw_impl::llt::iop_erc_20_rtl);
+    simd(prog, crate::asm::iop::SIMD_N, fw_impl::llt::iop_erc_20_rtl, None);
 }
 
 #[instrument(level = "trace", skip(prog))]
@@ -381,7 +386,7 @@ pub fn iop_rotate_scalar_left(prog: &mut Program) {
 ///     (dst_from[0], dst_to[0], ..., dst_from[N-1], dst_to[N-1])
 /// Where N is the batch size
 #[instrument(level = "trace", skip(prog))]
-pub fn iop_erc_20_rtl(prog: &mut Program, batch_index: u8) -> Rtl {
+pub fn iop_erc_20_rtl(prog: &mut Program, batch_index: u8, kogge_blk_w: Option<usize>) -> Rtl {
     // Allocate metavariables:
     // Dest -> Operand
     let dst_from = prog.iop_template_var(OperandKind::Dst, 2 * batch_index);
@@ -392,13 +397,6 @@ pub fn iop_erc_20_rtl(prog: &mut Program, batch_index: u8) -> Rtl {
     // Src Amount -> Operand
     let src_amount = prog.iop_template_var(OperandKind::Src, 3 * batch_index + 2);
 
-    // TODO: Make this a parameter or sweep this
-    // All these little parameters would be very handy to write an
-    // exploration/compilation program which would try to minimize latency by
-    // playing with these.
-    let kogge_blk_w = 10;
-    let ripple = true;
-
     {
         let props = prog.params();
         let tfhe_params: asm::DigitParameters = props.clone().into();
@@ -429,19 +427,20 @@ pub fn iop_erc_20_rtl(prog: &mut Program, batch_index: u8) -> Rtl {
             })
             .collect::<Vec<_>>();
 
-        if ripple {
+        if let Some(blk_w) = kogge_blk_w {
+            kogge::add(prog, dst_to, src_to, src_amount.clone(), None, blk_w)
+                + kogge::sub(prog, dst_from, src_from, src_amount, blk_w)
+        } else { // Default to ripple carry
             kogge::ripple_add(dst_to, src_to, src_amount.clone(), None)
                 + kogge::ripple_sub(prog, dst_from, src_from, src_amount)
-        } else {
-            kogge::add(prog, dst_to, src_to, src_amount.clone(), None, kogge_blk_w)
-                + kogge::sub(prog, dst_from, src_from, src_amount, kogge_blk_w)
         }
     }
 }
 
 /// A SIMD implementation of add for maximum throughput
+/// NB: No use of kogge_blk_w here, impl force to use ripple carry
 #[instrument(level = "trace", skip(prog))]
-pub fn iop_add_ripple_rtl(prog: &mut Program, i: u8) -> Rtl {
+pub fn iop_add_ripple_rtl(prog: &mut Program, i: u8, _kogge_blk_w: Option<usize>) -> Rtl {
     // Allocate metavariables:
     let dst = prog.iop_template_var(OperandKind::Dst, i);
     let src_a = prog.iop_template_var(OperandKind::Src, 2 * i);
@@ -899,13 +898,13 @@ fn bw_inv(prog: &mut Program, b: Vec<VarCell>) -> Vec<VarCell> {
 /// Maybe this should go into a SIMD firmware implementation... At some point we
 /// would need a mechanism to choose between implementations on the fly to make
 /// real good use of all of this.
-fn simd<F>(prog: &mut Program, batch_size: usize, rtl_closure: F)
+fn simd<F>(prog: &mut Program, batch_size: usize, rtl_closure: F, kogge_blk_w: Option<usize>)
 where
-    F: Fn(&mut Program, u8) -> Rtl,
+    F: Fn(&mut Program, u8, Option<usize>) -> Rtl,
 {
     (0..batch_size)
         .map(|i| i as u8)
-        .map(|i| rtl_closure(prog, i))
+        .map(|i| rtl_closure(prog, i, kogge_blk_w))
         .sum::<Rtl>()
         .add_to_prog(prog);
 }

tfhe-benchmark/benches/high_level_api/erc20.rs

@@ -28,13 +28,12 @@ pub fn transfer_whitepaper<FheType>(
     amount: &FheType,
 ) -> (FheType, FheType)
 where
-    FheType: Add<Output = FheType> + for<'a> FheOrd<&'a FheType> + FheTrivialEncrypt<u64>,
-    FheBool: IfThenElse<FheType>,
+    FheType: Add<Output = FheType> + for<'a> FheOrd<&'a FheType>,
+    FheBool: IfThenZero<FheType>,
     for<'a> &'a FheType: Add<Output = FheType> + Sub<Output = FheType>,
 {
     let has_enough_funds = (from_amount).ge(amount);
-    let zero_amount = FheType::encrypt_trivial(0u64);
-    let amount_to_transfer = has_enough_funds.select(amount, &zero_amount);
+    let amount_to_transfer = has_enough_funds.if_then_zero(amount);
 
     let new_to_amount = to_amount + &amount_to_transfer;
     let new_from_amount = from_amount - &amount_to_transfer;
@@ -51,12 +50,13 @@ pub fn par_transfer_whitepaper<FheType>(
 where
     FheType:
         Add<Output = FheType> + for<'a> FheOrd<&'a FheType> + Send + Sync + FheTrivialEncrypt<u64>,
-    FheBool: IfThenElse<FheType>,
+    FheBool: IfThenZero<FheType>,
     for<'a> &'a FheType: Add<Output = FheType> + Sub<Output = FheType>,
 {
     let has_enough_funds = (from_amount).ge(amount);
-    let zero_amount = FheType::encrypt_trivial(0u64);
-    let amount_to_transfer = has_enough_funds.select(amount, &zero_amount);
+    //let zero_amount = FheType::encrypt_trivial(0u64);
+    //let amount_to_transfer = has_enough_funds.select(amount, &zero_amount);
+    let amount_to_transfer = has_enough_funds.if_then_zero(amount);
 
     let (new_to_amount, new_from_amount) = rayon::join(
         || to_amount + &amount_to_transfer,

tfhe/src/high_level_api/booleans/base.rs

@@ -7,7 +7,7 @@ use crate::high_level_api::integers::{FheInt, FheIntId, FheUint, FheUintId};
 use crate::high_level_api::keys::InternalServerKey;
 use crate::high_level_api::re_randomization::ReRandomizationMetadata;
 use crate::high_level_api::traits::{
-    FheEq, Flip, IfThenElse, ReRandomize, ScalarIfThenElse, Tagged,
+    FheEq, Flip, IfThenElse, IfThenZero, ReRandomize, ScalarIfThenElse, Tagged,
 };
 use crate::high_level_api::{global_state, CompactPublicKey};
 use crate::integer::block_decomposition::DecomposableInto;
@@ -552,6 +552,66 @@ where
     }
 }
 
+impl<Id> IfThenZero<FheUint<Id>> for FheBool
+where
+    Id: FheUintId,
+{
+    /// Conditional selection.
+    ///
+    /// The output value returned depends on the value of `self`.
+    ///
+    /// - if `self` is true, the output will have the value of `ct_then`
+    /// - if `self` is false, the output will be an encryption of 0
+    fn if_then_zero(&self, ct_then: &FheUint<Id>) -> FheUint<Id> {
+        global_state::with_internal_keys(|sks| match sks {
+            InternalServerKey::Cpu(cpu_sks) => {
+                let ct_condition = self;
+                let mut ct_out = ct_then.ciphertext.on_cpu().clone();
+                cpu_sks.pbs_key().zero_out_if_condition_is_false(
+                    &mut ct_out,
+                    &ct_condition.ciphertext.on_cpu().0,
+                );
+                FheUint::new(
+                    ct_out,
+                    cpu_sks.tag.clone(),
+                    ReRandomizationMetadata::default(),
+                )
+            }
+            #[cfg(feature = "gpu")]
+            InternalServerKey::Cuda(_) => {
+                panic!("Cuda does not support if_then_zero")
+            }
+            #[cfg(feature = "hpu")]
+            InternalServerKey::Hpu(device) => {
+                let hpu_then = ct_then.ciphertext.on_hpu(device);
+                let hpu_cond = self.ciphertext.on_hpu(device);
+
+                let (opcode, proto) = {
+                    let asm_iop = &hpu_asm::iop::IOP_IF_THEN_ZERO;
+                    (
+                        asm_iop.opcode(),
+                        &asm_iop.format().expect("Unspecified IOP format").proto,
+                    )
+                };
+                // These clones are cheap are they are just Arc
+                let hpu_result = HpuRadixCiphertext::exec(
+                    proto,
+                    opcode,
+                    &[hpu_then.clone(), hpu_cond.clone()],
+                    &[],
+                )
+                .pop()
+                .unwrap();
+                FheUint::new(
+                    hpu_result,
+                    device.tag.clone(),
+                    ReRandomizationMetadata::default(),
+                )
+            }
+        })
+    }
+}
+
 impl<Id: FheIntId> IfThenElse<FheInt<Id>> for FheBool {
     /// Conditional selection.
     ///

tfhe/src/high_level_api/prelude.rs

@@ -11,6 +11,7 @@ pub use crate::high_level_api::traits::{
     FheOrd, FheTrivialEncrypt, FheTryEncrypt, FheTryTrivialEncrypt, FheWait, Flip, IfThenElse,
     OverflowingAdd, OverflowingMul, OverflowingNeg, OverflowingSub, ReRandomize, RotateLeft,
     RotateLeftAssign, RotateRight, RotateRightAssign, ScalarIfThenElse, SquashNoise, Tagged,
+    IfThenZero,
 };
 #[cfg(feature = "hpu")]
 pub use crate::high_level_api::traits::{FheHpu, HpuHandle};

tfhe/src/high_level_api/traits.rs

@@ -149,6 +149,10 @@ pub trait IfThenElse<Ciphertext> {
     }
 }
 
+pub trait IfThenZero<Ciphertext> {
+    fn if_then_zero(&self, ct_then: &Ciphertext) -> Ciphertext;
+}
+
 pub trait ScalarIfThenElse<Lhs, Rhs> {
     type Output;