Skip to content

Commit

Permalink
feat(integer): add is_even/is_odd functions
Browse files Browse the repository at this point in the history
These ones are pretty simple and so are also directly done for GPU
  • Loading branch information
tmontaigu committed Aug 23, 2024
1 parent 27a4564 commit ec7ea04
Show file tree
Hide file tree
Showing 9 changed files with 523 additions and 0 deletions.
66 changes: 66 additions & 0 deletions tfhe/src/high_level_api/integers/signed/base.rs
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,72 @@ where
})
}

/// Returns a FheBool that encrypts `true` if the value is even
///
/// # Example
///
/// ```rust
/// use tfhe::prelude::*;
/// use tfhe::{generate_keys, set_server_key, ConfigBuilder, FheBool, FheInt16};
///
/// let (client_key, server_key) = generate_keys(ConfigBuilder::default());
/// set_server_key(server_key);
///
/// let a = FheInt16::encrypt(1i16, &client_key);
///
/// let result = a.is_even();
/// let decrypted = result.decrypt(&client_key);
/// assert!(!decrypted);
/// ```
pub fn is_even(&self) -> FheBool {
global_state::with_internal_keys(|key| match key {
InternalServerKey::Cpu(cpu_key) => {
let result = cpu_key
.pbs_key()
.is_even_parallelized(&*self.ciphertext.on_cpu());
FheBool::new(result)
}
#[cfg(feature = "gpu")]
InternalServerKey::Cuda(cuda_key) => with_thread_local_cuda_streams(|streams| {
let result = cuda_key.key.is_even(&*self.ciphertext.on_gpu(), streams);
FheBool::new(result)
}),
})
}

/// Returns a FheBool that encrypts `true` if the value is odd
///
/// # Example
///
/// ```rust
/// use tfhe::prelude::*;
/// use tfhe::{generate_keys, set_server_key, ConfigBuilder, FheBool, FheInt16};
///
/// let (client_key, server_key) = generate_keys(ConfigBuilder::default());
/// set_server_key(server_key);
///
/// let a = FheInt16::encrypt(1i16, &client_key);
///
/// let result = a.is_odd();
/// let decrypted = result.decrypt(&client_key);
/// assert!(decrypted);
/// ```
pub fn is_odd(&self) -> FheBool {
global_state::with_internal_keys(|key| match key {
InternalServerKey::Cpu(cpu_key) => {
let result = cpu_key
.pbs_key()
.is_odd_parallelized(&*self.ciphertext.on_cpu());
FheBool::new(result)
}
#[cfg(feature = "gpu")]
InternalServerKey::Cuda(cuda_key) => with_thread_local_cuda_streams(|streams| {
let result = cuda_key.key.is_odd(&*self.ciphertext.on_gpu(), streams);
FheBool::new(result)
}),
})
}

/// Returns the number of leading zeros in the binary representation of self.
///
/// # Example
Expand Down
66 changes: 66 additions & 0 deletions tfhe/src/high_level_api/integers/unsigned/base.rs
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,72 @@ where
self.ciphertext.move_to_device(device)
}

/// Returns a FheBool that encrypts `true` if the value is even
///
/// # Example
///
/// ```rust
/// use tfhe::prelude::*;
/// use tfhe::{generate_keys, set_server_key, ConfigBuilder, FheBool, FheUint16};
///
/// let (client_key, server_key) = generate_keys(ConfigBuilder::default());
/// set_server_key(server_key);
///
/// let a = FheUint16::encrypt(32u16, &client_key);
///
/// let result = a.is_even();
/// let decrypted = result.decrypt(&client_key);
/// assert!(decrypted);
/// ```
pub fn is_even(&self) -> FheBool {
global_state::with_internal_keys(|key| match key {
InternalServerKey::Cpu(cpu_key) => {
let result = cpu_key
.pbs_key()
.is_even_parallelized(&*self.ciphertext.on_cpu());
FheBool::new(result)
}
#[cfg(feature = "gpu")]
InternalServerKey::Cuda(cuda_key) => with_thread_local_cuda_streams(|streams| {
let result = cuda_key.key.is_even(&*self.ciphertext.on_gpu(), streams);
FheBool::new(result)
}),
})
}

/// Returns a FheBool that encrypts `true` if the value is odd
///
/// # Example
///
/// ```rust
/// use tfhe::prelude::*;
/// use tfhe::{generate_keys, set_server_key, ConfigBuilder, FheBool, FheUint16};
///
/// let (client_key, server_key) = generate_keys(ConfigBuilder::default());
/// set_server_key(server_key);
///
/// let a = FheUint16::encrypt(32u16, &client_key);
///
/// let result = a.is_odd();
/// let decrypted = result.decrypt(&client_key);
/// assert!(!decrypted);
/// ```
pub fn is_odd(&self) -> FheBool {
global_state::with_internal_keys(|key| match key {
InternalServerKey::Cpu(cpu_key) => {
let result = cpu_key
.pbs_key()
.is_odd_parallelized(&*self.ciphertext.on_cpu());
FheBool::new(result)
}
#[cfg(feature = "gpu")]
InternalServerKey::Cuda(cuda_key) => with_thread_local_cuda_streams(|streams| {
let result = cuda_key.key.is_odd(&*self.ciphertext.on_gpu(), streams);
FheBool::new(result)
}),
})
}

/// Tries to decrypt a trivial ciphertext
///
/// Trivial ciphertexts are ciphertexts which are not encrypted
Expand Down
6 changes: 6 additions & 0 deletions tfhe/src/high_level_api/integers/unsigned/tests/cpu.rs
Original file line number Diff line number Diff line change
Expand Up @@ -385,6 +385,12 @@ fn test_ilog2() {
super::test_case_ilog2(&client_key);
}

#[test]
fn test_is_even_is_odd() {
let client_key = setup_default_cpu();
super::test_case_is_even_is_odd(&client_key);
}

#[test]
fn test_bitslice() {
let client_key = setup_default_cpu();
Expand Down
12 changes: 12 additions & 0 deletions tfhe/src/high_level_api/integers/unsigned/tests/gpu.rs
Original file line number Diff line number Diff line change
Expand Up @@ -109,3 +109,15 @@ fn test_sum_gpu_multibit() {
let client_key = setup_gpu(Some(PARAM_MULTI_BIT_MESSAGE_2_CARRY_2_GROUP_3_KS_PBS));
super::test_case_sum(&client_key);
}

#[test]
fn test_is_even_is_odd_gpu() {
let client_key = setup_default_gpu();
super::test_case_is_even_is_odd(&client_key);
}

#[test]
fn test_is_even_is_odd_gpu_multibit() {
let client_key = setup_gpu(Some(PARAM_MULTI_BIT_MESSAGE_2_CARRY_2_GROUP_3_KS_PBS));
super::test_case_is_even_is_odd(&client_key);
}
34 changes: 34 additions & 0 deletions tfhe/src/high_level_api/integers/unsigned/tests/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -533,3 +533,37 @@ fn test_case_sum(client_key: &ClientKey) {
assert_eq!(sum, expected_result);
}
}

fn test_case_is_even_is_odd(cks: &ClientKey) {
let mut rng = rand::thread_rng();
// This operation is cheap
for _ in 0..50 {
let clear_a = rng.gen_range(1..=u32::MAX);
let a = FheUint32::try_encrypt(clear_a, cks).unwrap();

assert_eq!(
a.is_even().decrypt(cks),
(clear_a % 2) == 0,
"Invalid is_even result for {clear_a}"
);
assert_eq!(
a.is_odd().decrypt(cks),
(clear_a % 2) == 1,
"Invalid is_odd result for {clear_a}"
);

let clear_a = rng.gen_range(i32::MIN..=i32::MAX);
let a = crate::FheInt32::try_encrypt(clear_a, cks).unwrap();
assert_eq!(
a.is_even().decrypt(cks),
(clear_a % 2) == 0,
"Invalid is_even result for {clear_a}"
);
// Use != 0 because if clear_a < 0, the returned mod is also < 0
assert_eq!(
a.is_odd().decrypt(cks),
(clear_a % 2) != 0,
"Invalid is_odd result for {clear_a}"
);
}
}
143 changes: 143 additions & 0 deletions tfhe/src/integer/gpu/server_key/radix/even_odd.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,143 @@
use crate::core_crypto::gpu::CudaStreams;
use crate::integer::gpu::ciphertext::boolean_value::CudaBooleanBlock;
use crate::integer::gpu::ciphertext::CudaIntegerRadixCiphertext;
use crate::integer::gpu::server_key::radix::{
CudaBlockInfo, CudaLweCiphertextList, CudaRadixCiphertext, CudaRadixCiphertextInfo,
LweCiphertextCount,
};
use crate::integer::gpu::server_key::CudaServerKey;
use crate::shortint::parameters::{Degree, NoiseLevel};

impl CudaServerKey {
/// # Safety
///
/// - `stream` __must__ be synchronized to guarantee computation has finished, and inputs must
/// not be dropped until stream is synchronised
pub unsafe fn unchecked_is_even_async<T>(
&self,
ct: &T,
streams: &CudaStreams,
) -> CudaBooleanBlock
where
T: CudaIntegerRadixCiphertext,
{
let radix = ct.as_ref();
let lut = self.generate_lookup_table(|block| u64::from((block & 1) == 0));
let mut single_block = CudaRadixCiphertext {
d_blocks: CudaLweCiphertextList::new(
radix.d_blocks.0.lwe_dimension,
LweCiphertextCount(1),
radix.d_blocks.0.ciphertext_modulus,
streams,
),
info: CudaRadixCiphertextInfo {
blocks: vec![CudaBlockInfo {
degree: Degree::new(1),
message_modulus: self.message_modulus,
carry_modulus: self.carry_modulus,
pbs_order: self.pbs_order,
noise_level: NoiseLevel::NOMINAL,
}],
},
};
self.apply_lookup_table_async(&mut single_block, radix, &lut, 0..1, streams);
CudaBooleanBlock::from_cuda_radix_ciphertext(single_block)
}

pub fn unchecked_is_even<T>(&self, ct: &T, streams: &CudaStreams) -> CudaBooleanBlock
where
T: CudaIntegerRadixCiphertext,
{
let result = unsafe { self.unchecked_is_even_async(ct, streams) };
streams.synchronize();
result
}

/// # Safety
///
/// - `stream` __must__ be synchronized to guarantee computation has finished, and inputs must
/// not be dropped until stream is synchronised
pub unsafe fn is_even_async<T>(&self, ct: &T, streams: &CudaStreams) -> CudaBooleanBlock
where
T: CudaIntegerRadixCiphertext,
{
// Since the check is done on the first bit of the first block
// no need to worry about carries
self.unchecked_is_even_async(ct, streams)
}

pub fn is_even<T>(&self, ct: &T, streams: &CudaStreams) -> CudaBooleanBlock
where
T: CudaIntegerRadixCiphertext,
{
let result = unsafe { self.is_even_async(ct, streams) };
streams.synchronize();
result
}

/// # Safety
///
/// - `stream` __must__ be synchronized to guarantee computation has finished, and inputs must
/// not be dropped until stream is synchronised
pub unsafe fn unchecked_is_odd_async<T>(
&self,
ct: &T,
streams: &CudaStreams,
) -> CudaBooleanBlock
where
T: CudaIntegerRadixCiphertext,
{
let radix = ct.as_ref();
let lut = self.generate_lookup_table(|block| block & 1);
let mut single_block = CudaRadixCiphertext {
d_blocks: CudaLweCiphertextList::new(
radix.d_blocks.0.lwe_dimension,
LweCiphertextCount(1),
radix.d_blocks.0.ciphertext_modulus,
streams,
),
info: CudaRadixCiphertextInfo {
blocks: vec![CudaBlockInfo {
degree: Degree::new(1),
message_modulus: self.message_modulus,
carry_modulus: self.carry_modulus,
pbs_order: self.pbs_order,
noise_level: NoiseLevel::NOMINAL,
}],
},
};
self.apply_lookup_table_async(&mut single_block, radix, &lut, 0..1, streams);
CudaBooleanBlock::from_cuda_radix_ciphertext(single_block)
}

pub fn unchecked_is_odd<T>(&self, ct: &T, streams: &CudaStreams) -> CudaBooleanBlock
where
T: CudaIntegerRadixCiphertext,
{
let result = unsafe { self.unchecked_is_odd_async(ct, streams) };
streams.synchronize();
result
}

/// # Safety
///
/// - `stream` __must__ be synchronized to guarantee computation has finished, and inputs must
/// not be dropped until stream is synchronised
pub unsafe fn is_odd_async<T>(&self, ct: &T, streams: &CudaStreams) -> CudaBooleanBlock
where
T: CudaIntegerRadixCiphertext,
{
// Since the check is done on the first bit of the first block
// no need to worry about carries
self.unchecked_is_odd_async(ct, streams)
}

pub fn is_odd<T>(&self, ct: &T, streams: &CudaStreams) -> CudaBooleanBlock
where
T: CudaIntegerRadixCiphertext,
{
let result = unsafe { self.is_odd_async(ct, streams) };
streams.synchronize();
result
}
}
Loading

0 comments on commit ec7ea04

Please sign in to comment.