Improve WoPBS API to avoid mistakes when using ciphertexts with varying degrees, consider having a carry invariant radix LUT generation

[cgouert]

Describe the bug
After a smart_mul, the WoPBS does not yield the expected answer.

To Reproduce

use std::{collections::HashMap};
use tfhe::{
    integer::{
        gen_keys_radix, wopbs::*,
    },
    shortint::parameters::{
        parameters_wopbs_message_carry::WOPBS_PARAM_MESSAGE_2_CARRY_2_KS_PBS,
        PARAM_MESSAGE_2_CARRY_2_KS_PBS,
    },
};

fn foo(x: u64, lut_entries: &HashMap<u64, u64>) -> u64 {
    lut_entries[&x]
}

fn main() {
    let nb_blocks: usize = 4;

    // Generate radix keys
    let (client_key, server_key) = gen_keys_radix(PARAM_MESSAGE_2_CARRY_2_KS_PBS, nb_blocks.into());

    // Generate key for PBS (without padding)
    let wopbs_key = WopbsKey::new_wopbs_key(
        &client_key,
        &server_key,
        &WOPBS_PARAM_MESSAGE_2_CARRY_2_KS_PBS,
    );

    // Create ciphertexts 
    let mut ct = client_key.encrypt(2_u64);
    let mut ct_2 = client_key.encrypt(4_u64);

    // Generate LUTs for WoPBS
    let mut lut_1_map : HashMap<u64, u64> = HashMap::new();
    let mut lut_2_map : HashMap<u64, u64> = HashMap::new();
    for i in 0..256 {
      lut_1_map.insert(i, 2*i % 256);
      lut_2_map.insert(i, 3*i % 256);
    }
    let lut_1 = wopbs_key.generate_lut_radix(&ct, |x: u64| foo(x, &lut_1_map));
    let lut_2 = wopbs_key.generate_lut_radix(&ct, |x: u64| foo(x, &lut_2_map));
    
    // Multiply input ciphertexts: 2 * 4 = 8
    ct = server_key.smart_mul(&mut ct, &mut ct_2);

    // Apply LUT #1
    ct = wopbs_key.keyswitch_to_wopbs_params(&server_key, &ct);
    let lut_1_res = wopbs_key.wopbs(&ct, &lut_1);
    let lut_1_res = wopbs_key.keyswitch_to_pbs_params(&lut_1_res);
    let test: u64 = client_key.decrypt(&lut_1_res);
    println!("Lut #1 result: {:?}", &test);
    println!("Expected result: {:?}", &lut_1_map[&8]);

    // Apply LUT #2
    let lut_2_res = wopbs_key.wopbs(&ct, &lut_2);
    let lut_2_res = wopbs_key.keyswitch_to_pbs_params(&lut_2_res);
    let test: u64 = client_key.decrypt(&lut_2_res);
    println!("Lut #2 result: {:?}", &test);
    println!("Expected result: {:?}", &lut_2_map[&8]);

}

Expected behaviour
For the above code, the decrypted results do not match the expected results.

Evidence

Configuration(please complete the following information):

• OS: Ubuntu 22.04

cc: @jimouris

[IceTDrinker]

try using decrypt_without_padding ?

https://docs.rs/tfhe/latest/tfhe/integer/client_key/struct.RadixClientKey.html#method.decrypt_without_padding

[IceTDrinker]

I see some mention of without padding in your code example ?

[IceTDrinker]

ah no my bad I mixed things up 😵‍💫

[IceTDrinker]

can confirm it reproes on latest main

[IceTDrinker]

mul_parallelized does not have the issue, looks like a bad carry management somewhere, could be keyswitching or the lut generation which does not handle this properly

[IceTDrinker]

This is not a bug on our end but a hard to use feature @cgouert

see updated code below, the main thing is to move the lut generation (which is adapted to the ciphertext degree) right before applying the wopbs

use std::collections::HashMap;
use tfhe::integer::gen_keys_radix;
use tfhe::integer::wopbs::*;
use tfhe::shortint::parameters::parameters_wopbs_message_carry::WOPBS_PARAM_MESSAGE_2_CARRY_2_KS_PBS;
use tfhe::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS_PBS;

fn foo(x: u64, lut_entries: &HashMap<u64, u64>) -> u64 {
    lut_entries[&x]
}

fn main() {
    let nb_blocks: usize = 4;

    // Generate radix keys
    let (client_key, server_key) = gen_keys_radix(PARAM_MESSAGE_2_CARRY_2_KS_PBS, nb_blocks.into());

    // Generate key for PBS (without padding)
    let wopbs_key = WopbsKey::new_wopbs_key(
        &client_key,
        &server_key,
        &WOPBS_PARAM_MESSAGE_2_CARRY_2_KS_PBS,
    );

    let clear_1 = 2_u64;
    let clear_2 = 4_u64;

    // Create ciphertexts
    let mut ct = client_key.encrypt(clear_1);
    let mut ct_2 = client_key.encrypt(clear_2);

    // Generate LUTs for WoPBS
    let mut lut_1_map: HashMap<u64, u64> = HashMap::new();
    let mut lut_2_map: HashMap<u64, u64> = HashMap::new();
    for i in 0..256 {
        lut_1_map.insert(i, 2 * i % 256);
        lut_2_map.insert(i, 3 * i % 256);
    }

    let f1 = |x: u64| foo(x, &lut_1_map);
    let f2 = |x: u64| foo(x, &lut_2_map);

    // Multiply input ciphertexts: 2 * 4 = 8
    ct = server_key.smart_mul(&mut ct, &mut ct_2);

    let sanity_dec: u64 = client_key.decrypt(&ct);
    let clear_prod = clear_1 * clear_2;
    assert_eq!(sanity_dec, clear_prod);

    // Lut generation just in time
    let lut_1 = wopbs_key.generate_lut_radix(&ct, f1);
    let lut_2 = wopbs_key.generate_lut_radix(&ct, f2);

    // Apply LUT #1
    ct = wopbs_key.keyswitch_to_wopbs_params(&server_key, &ct);
    let lut_1_res = wopbs_key.wopbs(&ct, &lut_1);
    let lut_1_res = wopbs_key.keyswitch_to_pbs_params(&lut_1_res);
    let test: u64 = client_key.decrypt(&lut_1_res);
    println!("Lut #1 result: {:?}", &test);
    println!("Expected result: {:?}", f1(clear_prod));

    // Apply LUT #2
    let lut_2_res = wopbs_key.wopbs(&ct, &lut_2);
    let lut_2_res = wopbs_key.keyswitch_to_pbs_params(&lut_2_res);
    let test: u64 = client_key.decrypt(&lut_2_res);
    println!("Lut #2 result: {:?}", &test);
    println!("Expected result: {:?}", f2(clear_prod));
}

RUSTFLAGS="-C target-cpu=native" cargo run --profile devo --features=x86_64-unix,integer,internal-keycache --example wop_smart_mul -p tfhe
   Compiling tfhe v0.6.0 (/home/***/Documents/zama/code/tfhe-rs/tfhe)
    Finished devo [optimized + debuginfo] target(s) in 8.01s
     Running `target/devo/examples/wop_smart_mul`
Lut #1 result: 16
Expected result: 16
Lut #2 result: 24
Expected result: 24

[IceTDrinker]

we'll consider adding an "apply_function" that just takes the function and manages the JIT LUT generation

[cgouert]

Thanks for your help with this @IceTDrinker! Just to clarify, are you saying the smart_mul changes the degree of the ciphertexts? If this is the case, the LUT doesn't work as expected because it was generated based on the original degree of a fresh encryption?

[jimouris]

@IceTDrinker Thanks for your responses!

It seems that just putting the smart_mul before the generate_lut_radix solves the issue but what if we want to do something like:

smart_mul(..)
let lut_1_res = wopbs_key.wopbs(&ct, &lut_1);
let lut_1_res = wopbs_key.keyswitch_to_pbs_params(&lut_1_res);

smart_mul(..)
let lut_1_res = wopbs_key.wopbs(&ct, &lut_1);
let lut_1_res = wopbs_key.keyswitch_to_pbs_params(&lut_1_res);

smart_mul(..)
let lut_1_res = wopbs_key.wopbs(&ct, &lut_1);
let lut_1_res = wopbs_key.keyswitch_to_pbs_params(&lut_1_res);

In this case, we'd have to run generate_lut_radix three times although the LUTs will be the same?

[IceTDrinker]

Thanks for your help with this @IceTDrinker! Just to clarify, are you saying the smart_mul changes the degree of the ciphertexts? If this is the case, the LUT doesn't work as expected because it was generated based on the original degree of a fresh encryption?

exactly and yes it's expected that multiplying changes the degrees of the underlying blocks :)

[IceTDrinker]

@IceTDrinker Thanks for your responses!

It seems that just putting the smart_mul before the generate_lut_radix solves the issue but what if we want to do something like:

smart_mul(..)
let lut_1_res = wopbs_key.wopbs(&ct, &lut_1);
let lut_1_res = wopbs_key.keyswitch_to_pbs_params(&lut_1_res);

smart_mul(..)
let lut_1_res = wopbs_key.wopbs(&ct, &lut_1);
let lut_1_res = wopbs_key.keyswitch_to_pbs_params(&lut_1_res);

smart_mul(..)
let lut_1_res = wopbs_key.wopbs(&ct, &lut_1);
let lut_1_res = wopbs_key.keyswitch_to_pbs_params(&lut_1_res);

In this case, we'd have to run generate_lut_radix three times although the LUTs will be the same?

my best guess here is that it does a lazy lut evaluation filling as little coefficients as it can, as having more 0s in the LUT IIRC will result in less noise in the output (could be wrong on that, but if you take a trivial 0 lut you have a noiseless encryption of 0 so there may be some of that)

I don't think the LUT generation would show up on a performance measurement when compared to the runtime of a wopbs so yes you likely need a LUT generation each time, you could do a small wrapper to generate the lut and apply it right afterwards to never have issues with degrees