Skip to content

Commit

Permalink
chore(integer): cleanup CpuFunctionExecutor
Browse files Browse the repository at this point in the history
  • Loading branch information
@mayeul-zama
mayeul-zama committed Dec 2, 2024
1 parent 73828ba commit 483a6b0
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 166 deletions.
97 changes: 2 additions & 95 deletions tfhe/src/integer/server_key/radix_parallel/tests_signed/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,110 +24,17 @@ pub(crate) mod test_vector_comparisons;

use crate::core_crypto::prelude::SignedInteger;
use crate::integer::keycache::KEY_CACHE;
use crate::integer::server_key::radix_parallel::tests_cases_unsigned::FunctionExecutor;
use crate::integer::server_key::radix_parallel::tests_unsigned::{
nb_tests_for_params, CpuFunctionExecutor, MAX_NB_CTXT, NB_CTXT,
nb_tests_for_params, MAX_NB_CTXT, NB_CTXT,
};
use crate::integer::tests::create_parameterized_test;
use crate::integer::{
BooleanBlock, IntegerKeyKind, RadixCiphertext, RadixClientKey, ServerKey, SignedRadixCiphertext,
};
use crate::integer::IntegerKeyKind;
#[cfg(tarpaulin)]
use crate::shortint::parameters::coverage_parameters::*;
use crate::shortint::parameters::*;
use itertools::{iproduct, izip};
use rand::prelude::ThreadRng;
use rand::Rng;
use std::sync::Arc;

impl<'a, F> FunctionExecutor<&'a SignedRadixCiphertext, (SignedRadixCiphertext, BooleanBlock)>
for CpuFunctionExecutor<F>
where
F: Fn(&ServerKey, &SignedRadixCiphertext) -> (SignedRadixCiphertext, BooleanBlock),
{
fn setup(&mut self, _cks: &RadixClientKey, sks: Arc<ServerKey>) {
self.sks = Some(sks);
}

fn execute(
&mut self,
input: &'a SignedRadixCiphertext,
) -> (SignedRadixCiphertext, BooleanBlock) {
let sks = self.sks.as_ref().expect("setup was not properly called");
(self.func)(sks, input)
}
}

impl<'a, F> FunctionExecutor<&'a SignedRadixCiphertext, SignedRadixCiphertext>
for CpuFunctionExecutor<F>
where
F: Fn(&ServerKey, &SignedRadixCiphertext) -> SignedRadixCiphertext,
{
fn setup(&mut self, _cks: &RadixClientKey, sks: Arc<ServerKey>) {
self.sks = Some(sks);
}

fn execute(&mut self, input: &'a SignedRadixCiphertext) -> SignedRadixCiphertext {
let sks = self.sks.as_ref().expect("setup was not properly called");
(self.func)(sks, input)
}
}
impl<'a, F> FunctionExecutor<&'a SignedRadixCiphertext, (RadixCiphertext, BooleanBlock)>
for CpuFunctionExecutor<F>
where
F: Fn(&ServerKey, &SignedRadixCiphertext) -> (RadixCiphertext, BooleanBlock),
{
fn setup(&mut self, _cks: &RadixClientKey, sks: Arc<ServerKey>) {
self.sks = Some(sks);
}

fn execute(&mut self, input: &'a SignedRadixCiphertext) -> (RadixCiphertext, BooleanBlock) {
let sks = self.sks.as_ref().expect("setup was not properly called");
(self.func)(sks, input)
}
}

impl<'a, F> FunctionExecutor<&'a SignedRadixCiphertext, RadixCiphertext> for CpuFunctionExecutor<F>
where
F: Fn(&ServerKey, &SignedRadixCiphertext) -> RadixCiphertext,
{
fn setup(&mut self, _cks: &RadixClientKey, sks: Arc<ServerKey>) {
self.sks = Some(sks);
}

fn execute(&mut self, input: &'a SignedRadixCiphertext) -> RadixCiphertext {
let sks = self.sks.as_ref().expect("setup was not properly called");
(self.func)(sks, input)
}
}
impl<'a, F> FunctionExecutor<&'a mut SignedRadixCiphertext, ()> for CpuFunctionExecutor<F>
where
F: Fn(&ServerKey, &'a mut SignedRadixCiphertext),
{
fn setup(&mut self, _cks: &RadixClientKey, sks: Arc<ServerKey>) {
self.sks = Some(sks);
}

fn execute(&mut self, input: &'a mut SignedRadixCiphertext) {
let sks = self.sks.as_ref().expect("setup was not properly called");
(self.func)(sks, input);
}
}

impl<'a, F> FunctionExecutor<&'a mut SignedRadixCiphertext, SignedRadixCiphertext>
for CpuFunctionExecutor<F>
where
F: Fn(&ServerKey, &mut SignedRadixCiphertext) -> SignedRadixCiphertext,
{
fn setup(&mut self, _cks: &RadixClientKey, sks: Arc<ServerKey>) {
self.sks = Some(sks);
}

fn execute(&mut self, input: &'a mut SignedRadixCiphertext) -> SignedRadixCiphertext {
let sks = self.sks.as_ref().expect("setup was not properly called");
(self.func)(sks, input)
}
}

//================================================================================
// Encrypt/Decrypt Tests
Expand Down
90 changes: 19 additions & 71 deletions tfhe/src/integer/server_key/radix_parallel/tests_unsigned/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ use super::tests_cases_unsigned::*;
use crate::core_crypto::prelude::UnsignedInteger;
use crate::integer::keycache::KEY_CACHE;
use crate::integer::tests::create_parameterized_test;
use crate::integer::{BooleanBlock, IntegerKeyKind, RadixCiphertext, RadixClientKey, ServerKey};
use crate::integer::{IntegerKeyKind, RadixCiphertext, RadixClientKey, ServerKey};
use crate::shortint::ciphertext::MaxDegree;
#[cfg(tarpaulin)]
use crate::shortint::parameters::coverage_parameters::*;
Expand Down Expand Up @@ -543,83 +543,31 @@ impl<F> CpuFunctionExecutor<F> {
}
}

/// For unary function
///
/// Note, we don't do
/// impl<F, I, O> TestExecutor<I, O> for CpuTestExecutor<F>
/// where F: Fn(&ServerKey, I) -> O {}
/// As it would conflict with other impls.
///
/// impl<F, I1, O> TestExecutor<(I,), O> for CpuTestExecutor<F>
/// would be possible tho.
impl<'a, F> FunctionExecutor<&'a RadixCiphertext, (RadixCiphertext, BooleanBlock)>
for CpuFunctionExecutor<F>
where
F: Fn(&ServerKey, &RadixCiphertext) -> (RadixCiphertext, BooleanBlock),
{
fn setup(&mut self, _cks: &RadixClientKey, sks: Arc<ServerKey>) {
self.sks = Some(sks);
}
pub(crate) trait NotTuple {}

fn execute(&mut self, input: &'a RadixCiphertext) -> (RadixCiphertext, BooleanBlock) {
let sks = self.sks.as_ref().expect("setup was not properly called");
(self.func)(sks, input)
}
}
impl<T> NotTuple for &crate::integer::ciphertext::BaseRadixCiphertext<T> {}

impl<'a, F> FunctionExecutor<&'a RadixCiphertext, RadixCiphertext> for CpuFunctionExecutor<F>
where
F: Fn(&ServerKey, &RadixCiphertext) -> RadixCiphertext,
{
fn setup(&mut self, _cks: &RadixClientKey, sks: Arc<ServerKey>) {
self.sks = Some(sks);
}
impl<T> NotTuple for &crate::integer::ciphertext::BaseSignedRadixCiphertext<T> {}

fn execute(&mut self, input: &'a RadixCiphertext) -> RadixCiphertext {
let sks = self.sks.as_ref().expect("setup was not properly called");
(self.func)(sks, input)
}
}
impl<T> NotTuple for &mut crate::integer::ciphertext::BaseRadixCiphertext<T> {}

impl<'a, F> FunctionExecutor<&'a Vec<RadixCiphertext>, Option<RadixCiphertext>>
for CpuFunctionExecutor<F>
where
F: Fn(&ServerKey, &Vec<RadixCiphertext>) -> Option<RadixCiphertext>,
{
fn setup(&mut self, _cks: &RadixClientKey, sks: Arc<ServerKey>) {
self.sks = Some(sks);
}

fn execute(&mut self, input: &'a Vec<RadixCiphertext>) -> Option<RadixCiphertext> {
let sks = self.sks.as_ref().expect("setup was not properly called");
(self.func)(sks, input)
}
}

/// Unary assign fn
impl<'a, F> FunctionExecutor<&'a mut RadixCiphertext, ()> for CpuFunctionExecutor<F>
where
F: Fn(&ServerKey, &'a mut RadixCiphertext),
{
fn setup(&mut self, _cks: &RadixClientKey, sks: Arc<ServerKey>) {
self.sks = Some(sks);
}
impl<T> NotTuple for &mut crate::integer::ciphertext::BaseSignedRadixCiphertext<T> {}

fn execute(&mut self, input: &'a mut RadixCiphertext) {
let sks = self.sks.as_ref().expect("setup was not properly called");
(self.func)(sks, input);
}
}
impl<T> NotTuple for &Vec<T> {}

impl<'a, F> FunctionExecutor<&'a mut RadixCiphertext, RadixCiphertext> for CpuFunctionExecutor<F>
/// For unary operations
///
/// Note, we need to `NotTuple` constraint to avoid conflicts with binary or ternary operations
impl<F, I1, O> FunctionExecutor<I1, O> for CpuFunctionExecutor<F>
where
F: Fn(&ServerKey, &mut RadixCiphertext) -> RadixCiphertext,
F: Fn(&ServerKey, I1) -> O,
I1: NotTuple,
{
fn setup(&mut self, _cks: &RadixClientKey, sks: Arc<ServerKey>) {
self.sks = Some(sks);
}

fn execute(&mut self, input: &'a mut RadixCiphertext) -> RadixCiphertext {
fn execute(&mut self, input: I1) -> O {
let sks = self.sks.as_ref().expect("setup was not properly called");
(self.func)(sks, input)
}
Expand Down Expand Up @@ -655,18 +603,18 @@ where
}
}

// Bitnot
impl<'a, F> FunctionExecutor<&'a BooleanBlock, BooleanBlock> for CpuFunctionExecutor<F>
/// For 4-ary operations
impl<F, I1, I2, I3, I4, O> FunctionExecutor<(I1, I2, I3, I4), O> for CpuFunctionExecutor<F>
where
F: Fn(&ServerKey, &BooleanBlock) -> BooleanBlock,
F: Fn(&ServerKey, I1, I2, I3, I4) -> O,
{
fn setup(&mut self, _cks: &RadixClientKey, sks: Arc<ServerKey>) {
self.sks = Some(sks);
}

fn execute(&mut self, input: &'a BooleanBlock) -> BooleanBlock {
fn execute(&mut self, input: (I1, I2, I3, I4)) -> O {
let sks = self.sks.as_ref().expect("setup was not properly called");
(self.func)(sks, input)
(self.func)(sks, input.0, input.1, input.2, input.3)
}
}

Expand Down

0 comments on commit 483a6b0

Please sign in to comment.