diff --git a/tfhe/Cargo.toml b/tfhe/Cargo.toml index bf049c980..911edf781 100644 --- a/tfhe/Cargo.toml +++ b/tfhe/Cargo.toml @@ -27,6 +27,7 @@ rand_distr = "0.4.3" criterion = "0.5.1" doc-comment = "0.3.3" serde_json = "1.0.94" +num-bigint = "0.4.6" # clap has to be pinned as its minimum supported rust version # changes often between minor releases, which breaks our CI clap = { version = "=4.5.30", features = ["derive"] } diff --git a/tfhe/docs/fhe-computation/advanced-features/encrypted-prf.md b/tfhe/docs/fhe-computation/advanced-features/encrypted-prf.md index 1e4abd43c..cd11a59d7 100644 --- a/tfhe/docs/fhe-computation/advanced-features/encrypted-prf.md +++ b/tfhe/docs/fhe-computation/advanced-features/encrypted-prf.md @@ -2,14 +2,30 @@ This document explains the mechanism and steps to generate an oblivious encrypted random value using only server keys. -The goal is to give to the server the possibility to generate a random value, which will be obtained in an encrypted format and will remain unknown to the server. The implementation is based on [this article](https://eprint.iacr.org/2024/665). +The goal is to give to the server the possibility to generate a random value, which will be obtained in an encrypted format and will remain unknown to the server. -This is possible through two methods on `FheUint` and `FheInt`: +The main method for this is `FheUint::generate_oblivious_pseudo_random_custom_range` which returns an integer in the given range. +Currently the range can only be in the form `[0, excluded_upper_bound[` with any `excluded_upper_bound` in `[1, 2^64[` +It follows a distribution close to the uniform. + +This function guarantees the norm-1 distance (defined as ∆(P,Q) := 1/2 Sum[ω∈Ω] |P(ω) − Q(ω)|) +between the actual distribution and the target uniform distribution will be below the `max_distance` argument (which must be in ]0, 1[). +The higher the distance, the more dissimilar the actual distribution is from the target uniform distribution. + +The default value for `max_distance` is `2^-128` if `None` is provided. + +Higher values allow better performance but must be considered carefully in the context of their target application as it may have serious unintended consequences. + +If the range is a power of 2, the distribution is uniform (for any `max_distance`) and the cost is smaller. + + +For powers of 2 specifically there are two methods on `FheUint` and `FheInt` (based on [this article](https://eprint.iacr.org/2024/665)): - `generate_oblivious_pseudo_random` which return an integer taken uniformly in the full integer range (`[0; 2^N[` for a `FheUintN` and `[-2^(N-1); 2^(N-1)[` for a `FheIntN`). - `generate_oblivious_pseudo_random_bounded` which return an integer taken uniformly in `[0; 2^random_bits_count[`. For a `FheUintN`, we must have `random_bits_count <= N`. For a `FheIntN`, we must have `random_bits_count <= N - 1`. -Both methods functions take a seed `Seed` as input, which could be any `u128` value. -They both rely on the use of the usual server key. + +These method functions take a seed `Seed` as input, which could be any `u128` value. +They rely on the use of the usual server key. The output is reproducible, i.e., the function is deterministic from the inputs: assuming the same hardware, seed and server key, this function outputs the same random encrypted value. @@ -18,7 +34,8 @@ Here is an example of the usage: ```rust use tfhe::prelude::FheDecrypt; -use tfhe::{generate_keys, set_server_key, ConfigBuilder, FheUint8, FheInt8, Seed}; +use tfhe::{generate_keys, set_server_key, ConfigBuilder, FheUint8, FheInt8, RangeForRandom, Seed}; +use std::num::NonZeroU64; pub fn main() { let config = ConfigBuilder::default().build(); @@ -26,23 +43,30 @@ pub fn main() { set_server_key(server_key); - let random_bits_count = 3; - - let ct_res = FheUint8::generate_oblivious_pseudo_random(Seed(0)); + let excluded_upper_bound = NonZeroU64::new(3).unwrap(); + let range = RangeForRandom::new_from_excluded_upper_bound(excluded_upper_bound); + // in [0, excluded_upper_bound[ = {0, 1, 2} + let ct_res = FheUint8::generate_oblivious_pseudo_random_custom_range(Seed(0), &range, None); let dec_result: u8 = ct_res.decrypt(&client_key); - let ct_res = FheUint8::generate_oblivious_pseudo_random_bounded(Seed(0), random_bits_count); + let random_bits_count = 3; + // in [0, 2^8[ + let ct_res = FheUint8::generate_oblivious_pseudo_random(Seed(0)); + let dec_result: u8 = ct_res.decrypt(&client_key); + + // in [0, 2^random_bits_count[ = [0, 8[ + let ct_res = FheUint8::generate_oblivious_pseudo_random_bounded(Seed(0), random_bits_count); let dec_result: u8 = ct_res.decrypt(&client_key); assert!(dec_result < (1 << random_bits_count)); + // in [-2^7, 2^7[ let ct_res = FheInt8::generate_oblivious_pseudo_random(Seed(0)); - let dec_result: i8 = ct_res.decrypt(&client_key); + // in [0, 2^random_bits_count[ = [0, 8[ let ct_res = FheInt8::generate_oblivious_pseudo_random_bounded(Seed(0), random_bits_count); - let dec_result: i8 = ct_res.decrypt(&client_key); assert!(dec_result < (1 << random_bits_count)); } diff --git a/tfhe/src/high_level_api/integers/oprf.rs b/tfhe/src/high_level_api/integers/oprf.rs index fe0d16815..5fc5a74f1 100644 --- a/tfhe/src/high_level_api/integers/oprf.rs +++ b/tfhe/src/high_level_api/integers/oprf.rs @@ -4,7 +4,9 @@ use crate::high_level_api::keys::InternalServerKey; use crate::high_level_api::re_randomization::ReRandomizationMetadata; #[cfg(feature = "gpu")] use crate::integer::gpu::ciphertext::{CudaSignedRadixCiphertext, CudaUnsignedRadixCiphertext}; +use crate::shortint::MessageModulus; use crate::{FheInt, Seed}; +use std::num::NonZeroU64; impl FheUint { /// Generates an encrypted unsigned integer @@ -150,6 +152,103 @@ impl FheUint { } }) } + + /// Generates an encrypted unsigned integer + /// taken almost uniformly in the given range using the given seed. + /// Currently the range can only be in the form `[0, excluded_upper_bound[` + /// with any `excluded_upper_bound` in `[1, 2^64[`. + /// + /// The encrypted value is oblivious to the server. + /// It can be useful to make server random generation deterministic. + /// + /// This function guarantees the the norm-1 distance + /// (defined as ∆(P,Q) := 1/2 Sum[ω∈Ω] |P(ω) − Q(ω)|) + /// between the actual distribution and the target uniform distribution + /// will be below the `max_distance` argument (which must be in ]0, 1[). + /// The higher the distance, the more dissimilar the actual distribution is + /// from the target uniform distribution. + /// + /// The default value for `max_distance` is `2^-128` if `None` is provided. + /// + /// Higher values allow better performance but must be considered carefully in the context of + /// their target application as it may have serious unintended consequences. + /// + /// If the range is a power of 2, the distribution is uniform (for any `max_distance`) and + /// the cost is smaller. + /// + /// ```rust + /// use std::num::NonZeroU64; + /// use tfhe::prelude::FheDecrypt; + /// use tfhe::{generate_keys, set_server_key, ConfigBuilder, FheUint8, RangeForRandom, Seed}; + /// + /// let config = ConfigBuilder::default().build(); + /// let (client_key, server_key) = generate_keys(config); + /// + /// set_server_key(server_key); + /// + /// let excluded_upper_bound = NonZeroU64::new(3).unwrap(); + /// + /// let range = RangeForRandom::new_from_excluded_upper_bound(excluded_upper_bound); + /// + /// let ct_res = FheUint8::generate_oblivious_pseudo_random_custom_range(Seed(0), &range, None); + /// + /// let dec_result: u16 = ct_res.decrypt(&client_key); + /// assert!(dec_result < excluded_upper_bound.get() as u16); + /// ``` + pub fn generate_oblivious_pseudo_random_custom_range( + seed: Seed, + range: &RangeForRandom, + max_distance: Option, + ) -> Self { + let excluded_upper_bound = range.excluded_upper_bound; + + if excluded_upper_bound.is_power_of_two() { + let random_bits_count = excluded_upper_bound.ilog2() as u64; + + Self::generate_oblivious_pseudo_random_bounded(seed, random_bits_count) + } else { + let max_distance = max_distance.unwrap_or_else(|| 2_f64.powi(-128)); + + assert!( + 0_f64 < max_distance && max_distance < 1_f64, + "max_distance (={max_distance}) should be in ]0, 1[" + ); + + global_state::with_internal_keys(|key| match key { + InternalServerKey::Cpu(key) => { + let message_modulus = key.message_modulus(); + + let num_input_random_bits = num_input_random_bits_for_max_distance( + excluded_upper_bound, + max_distance, + message_modulus, + ); + + let num_blocks_output = Id::num_blocks(key.message_modulus()) as u64; + + let ct = key + .pbs_key() + .par_generate_oblivious_pseudo_random_unsigned_custom_range( + seed, + num_input_random_bits, + excluded_upper_bound, + num_blocks_output, + ); + + Self::new(ct, key.tag.clone(), ReRandomizationMetadata::default()) + } + #[cfg(feature = "gpu")] + InternalServerKey::Cuda(_cuda_key) => { + panic!("Gpu does not support this operation yet.") + } + #[cfg(feature = "hpu")] + InternalServerKey::Hpu(_device) => { + panic!("Hpu does not support this operation yet.") + } + }) + } + } + #[cfg(feature = "gpu")] /// Returns the amount of memory required to execute generate_oblivious_pseudo_random_bounded /// @@ -367,10 +466,350 @@ impl FheInt { } } +pub struct RangeForRandom { + excluded_upper_bound: NonZeroU64, +} + +impl RangeForRandom { + pub fn new_from_excluded_upper_bound(excluded_upper_bound: NonZeroU64) -> Self { + Self { + excluded_upper_bound, + } + } +} + +fn num_input_random_bits_for_max_distance( + excluded_upper_bound: NonZeroU64, + max_distance: f64, + message_modulus: MessageModulus, +) -> u64 { + assert!(message_modulus.0.is_power_of_two()); + let log_message_modulus = message_modulus.0.ilog2() as u64; + + let mut random_block_count = 1; + + let random_block_count = loop { + let random_bit_count = random_block_count * log_message_modulus; + + let distance = distance(excluded_upper_bound.get(), random_bit_count); + + if distance < max_distance { + break random_block_count; + } + + random_block_count += 1; + }; + + random_block_count * log_message_modulus +} + +fn distance(excluded_upper_bound: u64, random_bit_count: u64) -> f64 { + let remainder = mod_pow_2(random_bit_count, excluded_upper_bound); + + remainder as f64 * (excluded_upper_bound - remainder) as f64 + / (2_f64.powi(random_bit_count as i32) * excluded_upper_bound as f64) +} + +// Computes 2^exponent % modulus +fn mod_pow_2(exponent: u64, modulus: u64) -> u64 { + assert_ne!(modulus, 0); + + if modulus == 1 { + return 0; + } + + let mut result: u128 = 1; + let mut base: u128 = 2; // We are calculating 2^i + + // We cast exponent to u128 to match the loop, though u64 is fine + let mut exp = exponent; + let mod_val = modulus as u128; + + while exp > 0 { + // If exponent is odd, multiply result with base + if exp % 2 == 1 { + result = (result * base) % mod_val; + } + + // Square the base + base = (base * base) % mod_val; + + // Divide exponent by 2 + exp /= 2; + } + + result as u64 +} + +#[cfg(test)] +mod test { + + use super::*; + use crate::integer::server_key::radix_parallel::tests_unsigned::test_oprf::{ + oprf_density_function, p_value_upper_bound_oprf_almost_uniformity_from_values, + probability_density_function_from_density, + }; + use crate::prelude::FheDecrypt; + use crate::shortint::oprf::test::test_uniformity; + use crate::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS32_PBS_TUNIFORM_2M128; + use crate::{generate_keys, set_server_key, ClientKey, ConfigBuilder, FheUint8, Seed}; + use num_bigint::BigUint; + use rand::{thread_rng, Rng}; + use rayon::iter::{IntoParallelIterator, ParallelIterator}; + + // Helper: The "Oracle" implementation using BigInt + // This is slow but mathematically guaranteed to be correct. + fn oracle_mod_pow_2(exponent: u64, modulus: u64) -> u64 { + assert_ne!(modulus, 0); + + if modulus == 1 { + return 0; + } + + let base = BigUint::from(2u32); + let exp = BigUint::from(exponent); + let modu = BigUint::from(modulus); + + let res = base.modpow(&exp, &modu); + res.iter_u64_digits().next().unwrap_or(0) + } + + #[test] + fn test_edge_cases() { + // 2^0 % 10 = 1 + assert_eq!(mod_pow_2(0, 10), 1, "Failed exponent 0"); + + // 2^10 % 1 = 0 + assert_eq!(mod_pow_2(10, 1), 0, "Failed modulus 1"); + + // 2^1 % 10 = 2 + assert_eq!(mod_pow_2(1, 10), 2, "Failed exponent 1"); + + // 2^3 % 5 = 8 % 5 = 3 + assert_eq!(mod_pow_2(3, 5), 3, "Failed small calc"); + } + + #[test] + fn test_boundaries_and_overflow() { + assert_eq!(mod_pow_2(2, u64::MAX), 4); + + assert_eq!(mod_pow_2(u64::MAX, 3), 2); + + assert_eq!(mod_pow_2(5, 32), 0); + } + + #[test] + fn test_against_oracle() { + let mut rng = thread_rng(); + for _ in 0..1_000_000 { + let exp: u64 = rng.gen(); + let mod_val: u64 = rng.gen(); + + let mod_val = if mod_val == 0 { 1 } else { mod_val }; + + let expected = oracle_mod_pow_2(exp, mod_val); + let actual = mod_pow_2(exp, mod_val); + + assert_eq!( + actual, expected, + "Mismatch! 2^{exp} % {mod_val} => Ours: {actual}, Oracle: {expected}", + ); + } + } + + #[test] + fn test_distance_with_uniform() { + for excluded_upper_bound in 1..20 { + for num_input_random_bits in 0..20 { + let density = oprf_density_function(excluded_upper_bound, num_input_random_bits); + + let theoretical_pdf = probability_density_function_from_density(&density); + + let p_uniform = 1. / excluded_upper_bound as f64; + + let actual_distance: f64 = 1. / 2. + * theoretical_pdf + .iter() + .map(|p| (*p - p_uniform).abs()) + .sum::(); + + let theoretical_distance = distance(excluded_upper_bound, num_input_random_bits); + + assert!( + (theoretical_distance - actual_distance).abs() + <= theoretical_distance / 1_000_000., + "{theoretical_distance} != {actual_distance}" + ); + } + } + } + + #[test] + fn test_uniformity_scalar_mul_shift() { + let max_distance = 2_f64.powi(-20); + + let message_modulus = MessageModulus(4); + + let excluded_upper_bound = 3; + + let num_input_random_bits = num_input_random_bits_for_max_distance( + NonZeroU64::new(excluded_upper_bound).unwrap(), + max_distance, + message_modulus, + ); + + let sample_count: usize = 10_000_000; + + let p_value_limit: f64 = 0.001; + + // The distribution is not exactly uniform + // This check ensures than with the given low max_distance, + // the distribution is indistinguishable from the uniform with at the given sample count + test_uniformity(sample_count, p_value_limit, excluded_upper_bound, |_seed| { + oprf_clear_equivalent(excluded_upper_bound, num_input_random_bits) + }); + } + + fn oprf_clear_equivalent(excluded_upper_bound: u64, num_input_random_bits: u64) -> u64 { + let random_input_upper_bound = 1 << num_input_random_bits; + + let random_input = thread_rng().gen_range(0..random_input_upper_bound); + + (random_input * excluded_upper_bound) >> num_input_random_bits + } + + #[test] + fn test_uniformity_generate_oblivious_pseudo_random_custom_range() { + let base_sample_count: usize = 10_000; + + let p_value_limit: f64 = 0.001; + + let params = PARAM_MESSAGE_2_CARRY_2_KS32_PBS_TUNIFORM_2M128; + let config = ConfigBuilder::with_custom_parameters(params).build(); + + let (cks, sks) = generate_keys(config); + rayon::broadcast(|_| set_server_key(sks.clone())); + + let message_modulus = params.message_modulus; + + // [0.7, 0.1] for `max_distance` chosen to have `num_input_random_bits` be [2, 4] + // for any of the listed `excluded_upper_bound` + for (expected_num_input_random_bits, max_distance, excluded_upper_bounds) in + [(2, 0.7, [3, 5, 6, 7]), (4, 0.1, [3, 5, 6, 7])] + { + for excluded_upper_bound in excluded_upper_bounds { + let sample_count = base_sample_count * excluded_upper_bound as usize; + + let excluded_upper_bound = NonZeroU64::new(excluded_upper_bound).unwrap(); + + let num_input_random_bits = num_input_random_bits_for_max_distance( + excluded_upper_bound, + max_distance, + message_modulus, + ); + + assert_eq!(num_input_random_bits, expected_num_input_random_bits); + + test_uniformity_generate_oblivious_pseudo_random_custom_range2( + sample_count, + p_value_limit, + message_modulus, + &cks, + excluded_upper_bound, + max_distance, + ); + } + } + } + + fn test_uniformity_generate_oblivious_pseudo_random_custom_range2( + sample_count: usize, + p_value_limit: f64, + message_modulus: MessageModulus, + cks: &ClientKey, + excluded_upper_bound: NonZeroU64, + max_distance: f64, + ) { + let num_input_random_bits = num_input_random_bits_for_max_distance( + excluded_upper_bound, + max_distance, + message_modulus, + ); + + let range = RangeForRandom::new_from_excluded_upper_bound(excluded_upper_bound); + + let real_values: Vec = (0..sample_count) + .into_par_iter() + .map(|_| { + let img = FheUint8::generate_oblivious_pseudo_random_custom_range( + Seed(rand::thread_rng().gen::()), + &range, + Some(max_distance), + ); + + img.decrypt(cks) + }) + .collect(); + + let excluded_upper_bound = excluded_upper_bound.get(); + + let uniform_values: Vec = (0..sample_count) + .into_par_iter() + .map(|_| thread_rng().gen_range(0..excluded_upper_bound)) + .collect(); + + let clear_oprf_value_lower_num_input_random_bits = (0..sample_count) + .into_par_iter() + .map(|_| oprf_clear_equivalent(excluded_upper_bound, num_input_random_bits - 1)) + .collect(); + + let clear_oprf_value_same_num_input_random_bits = (0..sample_count) + .into_par_iter() + .map(|_| oprf_clear_equivalent(excluded_upper_bound, num_input_random_bits)) + .collect(); + + let clear_oprf_value_higher_num_input_random_bits = (0..sample_count) + .into_par_iter() + .map(|_| oprf_clear_equivalent(excluded_upper_bound, num_input_random_bits + 1)) + .collect(); + + for (values, should_have_low_p_value) in [ + (&real_values, false), + // to test that the same distribution passes + (&clear_oprf_value_same_num_input_random_bits, false), + // to test that other distribution don't pass + // (makes sure the test is statistically powerful) + (&uniform_values, true), + (&clear_oprf_value_lower_num_input_random_bits, true), + (&clear_oprf_value_higher_num_input_random_bits, true), + ] { + let p_value_upper_bound = p_value_upper_bound_oprf_almost_uniformity_from_values( + values, + num_input_random_bits, + excluded_upper_bound, + ); + + println!("p_value_upper_bound: {p_value_upper_bound}"); + + if should_have_low_p_value { + assert!( + p_value_upper_bound < p_value_limit, + "p_value_upper_bound (={p_value_upper_bound}) expected to be smaller than {p_value_limit}" + ); + } else { + assert!( + p_value_limit < p_value_upper_bound , + "p_value_upper_bound (={p_value_upper_bound}) expected to be bigger than {p_value_limit}" + ); + } + } + } +} + #[cfg(test)] #[cfg(feature = "gpu")] #[allow(unused_imports)] -mod test { +mod test_gpu { use crate::prelude::*; use crate::{ generate_keys, set_server_key, ConfigBuilder, FheInt128, FheUint32, FheUint64, GpuIndex, diff --git a/tfhe/src/high_level_api/mod.rs b/tfhe/src/high_level_api/mod.rs index f3ca30059..b294e3a73 100644 --- a/tfhe/src/high_level_api/mod.rs +++ b/tfhe/src/high_level_api/mod.rs @@ -48,6 +48,7 @@ macro_rules! export_concrete_array_types { } pub use crate::core_crypto::commons::math::random::{Seed, XofSeed}; +pub use crate::high_level_api::integers::oprf::RangeForRandom; pub use crate::integer::server_key::MatchValues; use crate::{error, Error, Versionize}; use backward_compatibility::compressed_ciphertext_list::SquashedNoiseCiphertextStateVersions;