Add prepare multi thread

poulpy-core/src/scratch.rs

@@ -1,5 +1,5 @@
 use poulpy_hal::{
-    api::{ModuleN, ScratchAvailable, ScratchTakeBasic, SvpPPolBytesOf, VecZnxDftBytesOf, VmpPMatBytesOf},
+    api::{ModuleN, ScratchAvailable, ScratchFromBytes, ScratchTakeBasic, SvpPPolBytesOf, VecZnxDftBytesOf, VmpPMatBytesOf},
     layouts::{Backend, Scratch},
 };
 
@@ -7,7 +7,7 @@ use crate::{
     dist::Distribution,
     layouts::{
         Degree, GGLWE, GGLWEInfos, GGLWELayout, GGSW, GGSWInfos, GLWE, GLWEAutomorphismKey, GLWEInfos, GLWEPlaintext,
-        GLWEPrepared, GLWEPublicKey, GLWESecret, GLWESecretTensor, GLWESwitchingKey, GLWETensorKey, Rank,
+        GLWEPrepared, GLWEPublicKey, GLWESecret, GLWESecretTensor, GLWESwitchingKey, GLWETensorKey, LWE, LWEInfos, Rank,
         prepared::{
             GGLWEPrepared, GGSWPrepared, GLWEAutomorphismKeyPrepared, GLWEPublicKeyPrepared, GLWESecretPrepared,
             GLWESwitchingKeyPrepared, GLWETensorKeyPrepared,
@@ -17,8 +17,23 @@ use crate::{
 
 pub trait ScratchTakeCore<B: Backend>
 where
-    Self: ScratchTakeBasic + ScratchAvailable,
+    Self: ScratchTakeBasic + ScratchAvailable + ScratchFromBytes<B>,
 {
+    fn take_lwe<A>(&mut self, infos: &A) -> (LWE<&mut [u8]>, &mut Self)
+    where
+        A: LWEInfos,
+    {
+        let (data, scratch) = self.take_zn(infos.n().into(), 1, infos.size());
+        (
+            LWE {
+                k: infos.k(),
+                base2k: infos.base2k(),
+                data,
+            },
+            scratch,
+        )
+    }
+
     fn take_glwe<A>(&mut self, infos: &A) -> (GLWE<&mut [u8]>, &mut Self)
     where
         A: GLWEInfos,
@@ -367,4 +382,4 @@ where
     }
 }
 
-impl<B: Backend> ScratchTakeCore<B> for Scratch<B> where Self: ScratchTakeBasic + ScratchAvailable {}
+impl<B: Backend> ScratchTakeCore<B> for Scratch<B> where Self: ScratchTakeBasic + ScratchAvailable + ScratchFromBytes<B> {}

poulpy-hal/src/api/scratch.rs

@@ -1,6 +1,6 @@
 use crate::{
     api::{ModuleN, SvpPPolBytesOf, VecZnxBigBytesOf, VecZnxDftBytesOf, VmpPMatBytesOf},
-    layouts::{Backend, MatZnx, ScalarZnx, Scratch, SvpPPol, VecZnx, VecZnxBig, VecZnxDft, VmpPMat},
+    layouts::{Backend, MatZnx, ScalarZnx, Scratch, SvpPPol, VecZnx, VecZnxBig, VecZnxDft, VmpPMat, Zn},
 };
 
 /// Allocates a new [crate::layouts::ScratchOwned] of `size` aligned bytes.
@@ -28,7 +28,17 @@ pub trait TakeSlice {
     fn take_slice<T>(&mut self, len: usize) -> (&mut [T], &mut Self);
 }
 
-impl<B: Backend> ScratchTakeBasic for Scratch<B> where Self: TakeSlice {}
+impl<BE: Backend> Scratch<BE>
+where
+    Self: TakeSlice + ScratchFromBytes<BE>,
+{
+    pub fn split_at_mut(&mut self, len: usize) -> (&mut Scratch<BE>, &mut Self) {
+        let (take_slice, rem_slice) = self.take_slice(len);
+        (Self::from_bytes(take_slice), rem_slice)
+    }
+}
+
+impl<B: Backend> ScratchTakeBasic for Scratch<B> where Self: TakeSlice + ScratchFromBytes<B> {}
 
 pub trait ScratchTakeBasic
 where
@@ -47,6 +57,11 @@ where
         (SvpPPol::from_data(take_slice, module.n(), cols), rem_slice)
     }
 
+    fn take_zn(&mut self, n: usize, cols: usize, size: usize) -> (Zn<&mut [u8]>, &mut Self) {
+        let (take_slice, rem_slice) = self.take_slice(Zn::bytes_of(n, cols, size));
+        (Zn::from_data(take_slice, n, cols, size), rem_slice)
+    }
+
     fn take_vec_znx(&mut self, n: usize, cols: usize, size: usize) -> (VecZnx<&mut [u8]>, &mut Self) {
         let (take_slice, rem_slice) = self.take_slice(VecZnx::bytes_of(n, cols, size));
         (VecZnx::from_data(take_slice, n, cols, size), rem_slice)

poulpy-hal/src/layouts/mod.rs

@@ -28,8 +28,8 @@ pub use zn::*;
 pub use znx_base::*;
 
 pub trait Data = PartialEq + Eq + Sized + Default;
-pub trait DataRef = Data + AsRef<[u8]>;
+pub trait DataRef = Data + AsRef<[u8]> + Sync;
-pub trait DataMut = DataRef + AsMut<[u8]>;
+pub trait DataMut = DataRef + AsMut<[u8]> + Send;
 
 pub trait ToOwnedDeep {
     type Owned;

poulpy-hal/src/layouts/module.rs

@@ -13,7 +13,7 @@ use crate::{
 };
 
 #[allow(clippy::missing_safety_doc)]
-pub trait Backend: Sized {
+pub trait Backend: Sized + Sync + Send {
     type ScalarBig: Copy + Zero + Display + Debug + Pod;
     type ScalarPrep: Copy + Zero + Display + Debug + Pod;
     type Handle: 'static;

poulpy-schemes/Cargo.toml

@@ -20,5 +20,5 @@ rand = "0.9.2"
 
 
 [[bench]]
-name = "circuit_bootstrapping"
+name = "fhe_uint_prepare"
 harness = false

poulpy-schemes/benches/fhe_uint_prepare.rs

@@ -0,0 +1,126 @@
+use std::hint::black_box;
+
+use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
+use poulpy_backend::{FFT64Avx, FFT64Ref};
+use poulpy_core::{
+    GGSWNoise, GLWEDecrypt, GLWEEncryptSk, GLWENoise, ScratchTakeCore,
+    layouts::{GGSWLayout, GLWELayout, GLWESecretPreparedFactory, prepared::GLWESecretPrepared},
+};
+use poulpy_hal::{
+    api::{ModuleNew, ScratchOwnedAlloc, ScratchOwnedBorrow},
+    layouts::{Backend, Module, Scratch, ScratchOwned},
+    source::Source,
+};
+use rand::RngCore;
+
+use poulpy_schemes::tfhe::{
+    bdd_arithmetic::{
+        BDDKeyEncryptSk, BDDKeyPrepared, BDDKeyPreparedFactory, ExecuteBDDCircuit2WTo1W, FheUint, FheUintPrepare,
+        FheUintPrepareDebug, FheUintPrepared, FheUintPreparedEncryptSk, FheUintPreparedFactory,
+        tests::test_suite::TestContext,
+    },
+    blind_rotation::{BlindRotationAlgo, BlindRotationKey, BlindRotationKeyFactory, CGGI},
+};
+
+pub fn benc_bdd_prepare<BRA: BlindRotationAlgo, BE: Backend>(
+    c: &mut Criterion,
+    label: &str,
+    test_context: &TestContext<BRA, BE>,
+) where
+    Module<BE>: ModuleNew<BE>
+        + GLWESecretPreparedFactory<BE>
+        + GLWEDecrypt<BE>
+        + GLWENoise<BE>
+        + FheUintPreparedFactory<u32, BE>
+        + FheUintPreparedEncryptSk<u32, BE>
+        + FheUintPrepareDebug<BRA, u32, BE>
+        + BDDKeyEncryptSk<BRA, BE>
+        + BDDKeyPreparedFactory<BRA, BE>
+        + GGSWNoise<BE>
+        + FheUintPrepare<BRA, BE>
+        + ExecuteBDDCircuit2WTo1W<u32, BE>
+        + GLWEEncryptSk<BE>,
+    BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,
+    ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
+    Scratch<BE>: ScratchTakeCore<BE>,
+{
+    let group_name: String = format!("bdd_prepare::{label}");
+
+    let mut group = c.benchmark_group(group_name);
+
+    fn runner<BE: Backend, BRA: BlindRotationAlgo>(test_context: &TestContext<BRA, BE>) -> impl FnMut()
+    where
+        Module<BE>: ModuleNew<BE>
+            + GLWESecretPreparedFactory<BE>
+            + GLWEDecrypt<BE>
+            + GLWENoise<BE>
+            + FheUintPreparedFactory<u32, BE>
+            + FheUintPreparedEncryptSk<u32, BE>
+            + FheUintPrepareDebug<BRA, u32, BE>
+            + BDDKeyEncryptSk<BRA, BE>
+            + BDDKeyPreparedFactory<BRA, BE>
+            + GGSWNoise<BE>
+            + FheUintPrepare<BRA, BE>
+            + ExecuteBDDCircuit2WTo1W<u32, BE>
+            + GLWEEncryptSk<BE>,
+        BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,
+        ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
+        Scratch<BE>: ScratchTakeCore<BE>,
+    {
+        let glwe_infos: GLWELayout = test_context.glwe_infos();
+        let ggsw_infos: GGSWLayout = test_context.ggsw_infos();
+
+        let module: &Module<BE> = &test_context.module;
+        let sk_glwe_prep: &GLWESecretPrepared<Vec<u8>, BE> = &test_context.sk_glwe;
+        let bdd_key_prepared: &BDDKeyPrepared<Vec<u8>, BRA, BE> = &test_context.bdd_key;
+
+        let mut source: Source = Source::new([6u8; 32]);
+
+        let mut source_xa: Source = Source::new([2u8; 32]);
+        let mut source_xe: Source = Source::new([3u8; 32]);
+
+        let threads = 1;
+
+        let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc((1 << 22) * threads);
+
+        // GLWE(value)
+        let mut c_enc: FheUint<Vec<u8>, u32> = FheUint::alloc_from_infos(&glwe_infos);
+        let value: u32 = source.next_u32();
+        c_enc.encrypt_sk(
+            module,
+            value,
+            sk_glwe_prep,
+            &mut source_xa,
+            &mut source_xe,
+            scratch.borrow(),
+        );
+
+        // GGSW(0)
+        let mut c_enc_prep: FheUintPrepared<Vec<u8>, u32, BE> =
+            FheUintPrepared::<Vec<u8>, u32, BE>::alloc_from_infos(module, &ggsw_infos);
+
+        // GGSW(value)
+        move || {
+            c_enc_prep.prepare_custom_multi_thread(threads, module, &c_enc, 0, 32, bdd_key_prepared, scratch.borrow());
+            black_box(());
+        }
+    }
+
+    let id: BenchmarkId = BenchmarkId::from_parameter(format!("n_glwe: {}", test_context.module.n()));
+    let mut runner = runner::<BE, BRA>(test_context);
+    group.bench_with_input(id, &(), |b, _| b.iter(&mut runner));
+
+    group.finish();
+}
+
+fn bench_bdd_prepare_cpu_ref_fft64(c: &mut Criterion) {
+    benc_bdd_prepare::<CGGI, FFT64Avx>(
+        c,
+        "bdd_prepare_fft64_ref",
+        &TestContext::<CGGI, FFT64Avx>::new(),
+    );
+}
+
+criterion_group!(benches, bench_bdd_prepare_cpu_ref_fft64,);
+
+criterion_main!(benches);

poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint_prepared.rs

@@ -1,15 +1,17 @@
 use std::marker::PhantomData;
+use std::thread;
 
 use poulpy_core::layouts::{
     Base2K, Dnum, Dsize, GGSWInfos, GGSWPreparedFactory, GLWEInfos, LWEInfos, Rank, TorusPrecision, prepared::GGSWPrepared,
 };
 use poulpy_core::layouts::{
-    GGLWEInfos, GGLWEPreparedToRef, GGSWPreparedToMut, GGSWPreparedToRef, GLWEAutomorphismKeyHelper, GetGaloisElement, LWE,
+    GGLWEInfos, GGLWEPreparedToRef, GGSW, GGSWLayout, GGSWPreparedToMut, GGSWPreparedToRef, GLWEAutomorphismKeyHelper,
+    GetGaloisElement, LWE,
 };
 use poulpy_core::{GLWECopy, GLWEDecrypt, GLWEPacking, LWEFromGLWE};
 
 use poulpy_core::{GGSWEncryptSk, ScratchTakeCore, layouts::GLWESecretPreparedToRef};
-use poulpy_hal::api::ModuleLogN;
+use poulpy_hal::api::{ModuleLogN, ScratchAvailable, ScratchFromBytes};
 use poulpy_hal::layouts::{Backend, Data, DataRef, Module};
 
 use poulpy_hal::{
@@ -21,7 +23,7 @@ use poulpy_hal::{
 use crate::tfhe::bdd_arithmetic::{BDDKey, BDDKeyHelper, BDDKeyInfos, BDDKeyPrepared, BDDKeyPreparedFactory, FheUint, ToBits};
 use crate::tfhe::bdd_arithmetic::{Cmux, FromBits, ScratchTakeBDD, UnsignedInteger};
 use crate::tfhe::blind_rotation::BlindRotationAlgo;
-use crate::tfhe::circuit_bootstrapping::CirtuitBootstrappingExecute;
+use crate::tfhe::circuit_bootstrapping::{CircuitBootstrappingKeyInfos, CirtuitBootstrappingExecute};
 
 /// A prepared FHE ciphertext encrypting the bits of an [UnsignedInteger].
 pub struct FheUintPrepared<D: Data, T: UnsignedInteger, B: Backend> {
@@ -219,12 +221,13 @@ impl<D: DataMut, BRA: BlindRotationAlgo, BE: Backend> BDDKeyPrepared<D, BRA, BE>
     }
 }
 
-pub trait FheUintPrepare<BRA: BlindRotationAlgo, T: UnsignedInteger, BE: Backend> {
+pub trait FheUintPrepare<BRA: BlindRotationAlgo, BE: Backend> {
-    fn fhe_uint_prepare_tmp_bytes<R, A>(&self, block_size: usize, extension_factor: usize, res_infos: &R, infos: &A) -> usize
+    fn fhe_uint_prepare_tmp_bytes<R, A, B>(&self, block_size: usize, extension_factor: usize, res_infos: &R, bits_infos: &A, bdd_infos: &B) -> usize
     where
         R: GGSWInfos,
-        A: BDDKeyInfos;
+        A: GLWEInfos,
-    fn fhe_uint_prepare<DM, DB, DK, K>(
+        B: BDDKeyInfos;
+    fn fhe_uint_prepare<DM, DB, DK, K, T: UnsignedInteger>(
         &self,
         res: &mut FheUintPrepared<DM, T, BE>,
         bits: &FheUint<DB, T>,
@@ -234,79 +237,119 @@ pub trait FheUintPrepare<BRA: BlindRotationAlgo, T: UnsignedInteger, BE: Backend
         DM: DataMut,
         DB: DataRef,
         DK: DataRef,
-        K: BDDKeyHelper<DK, BRA, BE>;
+        K: BDDKeyHelper<DK, BRA, BE> + BDDKeyInfos,
-    fn fhe_uint_prepare_custom<DM, DB, DK, K>(
+        Scratch<BE>: ScratchFromBytes<BE>,
+    {
+        self.fhe_uint_prepare_custom(res, bits, 0, T::BITS as usize, key, scratch);
+    }
+    fn fhe_uint_prepare_custom<DM, DB, DK, K, T: UnsignedInteger>(
         &self,
         res: &mut FheUintPrepared<DM, T, BE>,
         bits: &FheUint<DB, T>,
         bit_start: usize,
-        bit_end: usize,
+        bit_count: usize,
         key: &K,
         scratch: &mut Scratch<BE>,
     ) where
         DM: DataMut,
         DB: DataRef,
         DK: DataRef,
-        K: BDDKeyHelper<DK, BRA, BE>;
+        K: BDDKeyHelper<DK, BRA, BE> + BDDKeyInfos,
+    {
+        self.fhe_uint_prepare_custom_multi_thread(1, res, bits, bit_start, bit_count, key, scratch)
+    }
+    fn fhe_uint_prepare_custom_multi_thread<DM, DB, DK, K, T: UnsignedInteger>(
+        &self,
+        threads: usize,
+        res: &mut FheUintPrepared<DM, T, BE>,
+        bits: &FheUint<DB, T>,
+        bit_start: usize,
+        bit_count: usize,
+        key: &K,
+        scratch: &mut Scratch<BE>,
+    ) where
+        DM: DataMut,
+        DB: DataRef,
+        DK: DataRef,
+        K: BDDKeyHelper<DK, BRA, BE> + BDDKeyInfos;
 }
 
-impl<BRA: BlindRotationAlgo, BE: Backend, T: UnsignedInteger> FheUintPrepare<BRA, T, BE> for Module<BE>
+impl<BRA: BlindRotationAlgo, BE: Backend> FheUintPrepare<BRA, BE> for Module<BE>
 where
     Self: LWEFromGLWE<BE> + CirtuitBootstrappingExecute<BRA, BE> + GGSWPreparedFactory<BE>,
     Scratch<BE>: ScratchTakeCore<BE>,
 {
-    fn fhe_uint_prepare_tmp_bytes<R, A>(&self, block_size: usize, extension_factor: usize, res_infos: &R, bdd_infos: &A) -> usize
+    fn fhe_uint_prepare_tmp_bytes<R, A, B>(&self, block_size: usize, extension_factor: usize, res_infos: &R, bits_infos: &A, bdd_infos: &B) -> usize
     where
         R: GGSWInfos,
-        A: BDDKeyInfos,
+        A: GLWEInfos,
+        B: BDDKeyInfos,
     {
         self.circuit_bootstrapping_execute_tmp_bytes(
             block_size,
             extension_factor,
             res_infos,
             &bdd_infos.cbt_infos(),
-        )
+        ) + GGSW::bytes_of_from_infos(res_infos)
+            + LWE::bytes_of_from_infos(bits_infos)
     }
 
-    fn fhe_uint_prepare<DM, DB, DK, K>(
+    fn fhe_uint_prepare_custom_multi_thread<DM, DB, DK, K, T: UnsignedInteger>(
-        &self,
-        res: &mut FheUintPrepared<DM, T, BE>,
-        bits: &FheUint<DB, T>,
-        key: &K,
-        scratch: &mut Scratch<BE>,
-    ) where
-        DM: DataMut,
-        DB: DataRef,
-        DK: DataRef,
-        K: BDDKeyHelper<DK, BRA, BE>,
-    {
-        self.fhe_uint_prepare_custom(res, bits, 0, T::BITS as usize, key, scratch);
-    }
-
-    fn fhe_uint_prepare_custom<DM, DB, DK, K>(
         &self,
+        threads: usize,
         res: &mut FheUintPrepared<DM, T, BE>,
         bits: &FheUint<DB, T>,
         bit_start: usize,
-        bit_end: usize,
+        bit_count: usize,
         key: &K,
-        scratch: &mut Scratch<BE>,
+        mut scratch: &mut Scratch<BE>,
     ) where
         DM: DataMut,
         DB: DataRef,
         DK: DataRef,
-        K: BDDKeyHelper<DK, BRA, BE>,
+        K: BDDKeyHelper<DK, BRA, BE> + BDDKeyInfos,
     {
+        let bit_end = bit_start + bit_count;
         let (cbt, ks) = key.get_cbt_key();
 
-        let mut lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(bits); //TODO: add TakeLWE
+        assert!(bit_end <= T::BITS as usize);
-        let (mut tmp_ggsw, scratch_1) = scratch.take_ggsw(res);
+
-        for (bit, dst) in res.bits[bit_start..bit_end].iter_mut().enumerate() {
+        let scratch_thread_size = self.fhe_uint_prepare_tmp_bytes(cbt.block_size(), 1, res, bits, key);
-            // TODO: set the rest of the bits to a prepared zero GGSW
+
-            bits.get_bit_lwe(self, bit, &mut lwe, ks, scratch_1);
+        assert!(scratch.available() >= threads * scratch_thread_size);
-            cbt.execute_to_constant(self, &mut tmp_ggsw, &lwe, 1, 1, scratch_1);
+
-            dst.prepare(self, &tmp_ggsw, scratch_1);
+        // How many bits we need to process
+        let chunk_size: usize = bit_count.div_ceil(threads); // ceil division
+
+        let mut scratches = Vec::new();
+        for _ in 0..(threads - 1) {
+            let (tmp, scratch_new) = scratch.split_at_mut(scratch_thread_size);
+            scratch = scratch_new;
+            scratches.push(tmp);
         }
+        scratches.push(scratch);
+
+        let ggsw_infos: &GGSWLayout = &res.ggsw_layout();
+
+        thread::scope(|scope| {
+            for (thread_index, (scratch_thread, res_bits_chunk)) in scratches
+                .iter_mut()
+                .zip(res.bits[bit_start..bit_end].chunks_mut(chunk_size))
+                .enumerate()
+            {
+                let start: usize = bit_start + thread_index * chunk_size;
+
+                scope.spawn(move || {
+                    let (mut tmp_ggsw, scratch_1) = scratch_thread.take_ggsw(ggsw_infos);
+                    let (mut tmp_lwe, scratch_2) = scratch_1.take_lwe(bits);
+                    for (local_bit, dst) in res_bits_chunk.iter_mut().enumerate() {
+                        bits.get_bit_lwe(self, start + local_bit, &mut tmp_lwe, ks, scratch_2);
+                        cbt.execute_to_constant(self, &mut tmp_ggsw, &tmp_lwe, 1, 1, scratch_2);
+                        dst.prepare(self, &tmp_ggsw, scratch_2);
+                    }
+                });
+            }
+        });
 
         for i in 0..bit_start {
             res.bits[i].zero(self);
@@ -324,8 +367,8 @@ impl<D: DataMut, T: UnsignedInteger, BE: Backend> FheUintPrepared<D, T, BE> {
         BRA: BlindRotationAlgo,
         O: DataRef,
         DK: DataRef,
-        K: BDDKeyHelper<DK, BRA, BE>,
+        K: BDDKeyHelper<DK, BRA, BE> + BDDKeyInfos,
-        M: FheUintPrepare<BRA, T, BE>,
+        M: FheUintPrepare<BRA, BE>,
         Scratch<BE>: ScratchTakeCore<BE>,
     {
         module.fhe_uint_prepare(self, other, key, scratch);
@@ -342,10 +385,30 @@ impl<D: DataMut, T: UnsignedInteger, BE: Backend> FheUintPrepared<D, T, BE> {
         BRA: BlindRotationAlgo,
         O: DataRef,
         DK: DataRef,
-        K: BDDKeyHelper<DK, BRA, BE>,
+        K: BDDKeyHelper<DK, BRA, BE> + BDDKeyInfos,
-        M: FheUintPrepare<BRA, T, BE>,
+        M: FheUintPrepare<BRA, BE>,
         Scratch<BE>: ScratchTakeCore<BE>,
     {
         module.fhe_uint_prepare_custom(self, other, bit_start, bit_end, key, scratch);
     }
+
+    pub fn prepare_custom_multi_thread<BRA, M, O, K, DK>(
+        &mut self,
+        threads: usize,
+        module: &M,
+        other: &FheUint<O, T>,
+        bit_start: usize,
+        bit_end: usize,
+        key: &K,
+        scratch: &mut Scratch<BE>,
+    ) where
+        BRA: BlindRotationAlgo,
+        O: DataRef,
+        DK: DataRef,
+        K: BDDKeyHelper<DK, BRA, BE> + BDDKeyInfos,
+        M: FheUintPrepare<BRA, BE>,
+        Scratch<BE>: ScratchTakeCore<BE>,
+    {
+        module.fhe_uint_prepare_custom_multi_thread(threads, self, other, bit_start, bit_end, key, scratch);
+    }
 }

poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint_prepared_debug.rs

@@ -10,7 +10,7 @@ use poulpy_core::layouts::{Base2K, Dnum, Dsize, Rank, TorusPrecision};
 use poulpy_core::layouts::{GGSW, GLWESecretPreparedToRef};
 use poulpy_core::{
     LWEFromGLWE, ScratchTakeCore,
-    layouts::{GGSWInfos, GGSWPreparedFactory, GLWEInfos, LWE, LWEInfos},
+    layouts::{GGSWInfos, GGSWPreparedFactory, GLWEInfos, LWEInfos},
 };
 
 use poulpy_hal::api::ModuleN;
@@ -125,10 +125,12 @@ where
         DR0: DataRef,
         DR1: DataRef,
     {   
-        let mut lwe: LWE<Vec<u8>> = LWE::alloc_from_infos(bits); //TODO: add TakeLWE
+
+        let (_, scratch_1) = scratch.take_ggsw(res);
+        let (mut tmp_lwe, scratch_2) = scratch_1.take_lwe(bits);
         for (bit, dst) in res.bits.iter_mut().enumerate() {
-            bits.get_bit_lwe(self, bit, &mut lwe, &key.ks, scratch);
+            bits.get_bit_lwe(self, bit, &mut tmp_lwe, &key.ks, scratch_2);
-            key.cbt.execute_to_constant(self, dst, &lwe, 1, 1, scratch);
+            key.cbt.execute_to_constant(self, dst, &tmp_lwe, 1, 1, scratch_2);
         }
     }
 }

poulpy-schemes/src/tfhe/bdd_arithmetic/key.rs

@@ -1,4 +1,5 @@
 use crate::tfhe::bdd_arithmetic::FheUintPreparedDebug;
+use crate::tfhe::circuit_bootstrapping::CircuitBootstrappingKeyInfos;
 use crate::tfhe::{
     bdd_arithmetic::{FheUint, UnsignedInteger},
     blind_rotation::{BlindRotationAlgo, BlindRotationKey, BlindRotationKeyFactory},
@@ -8,7 +9,7 @@ use crate::tfhe::{
     },
 };
 
-use poulpy_core::layouts::{GLWEAutomorphismKeyHelper, GLWEAutomorphismKeyPrepared};
+use poulpy_core::layouts::{GGLWEInfos, GLWEAutomorphismKeyHelper, GLWEAutomorphismKeyPrepared};
 use poulpy_core::{
     GLWEToLWESwitchingKeyEncryptSk, GetDistribution, ScratchTakeCore,
     layouts::{
@@ -135,6 +136,21 @@ where
     pub(crate) ks: GLWEToLWEKeyPrepared<D, BE>,
 }
 
+impl<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> BDDKeyInfos for BDDKeyPrepared<D, BRA, BE>{
+    fn cbt_infos(&self) -> CircuitBootstrappingKeyLayout {
+        CircuitBootstrappingKeyLayout { layout_brk: self.cbt.brk_infos(), layout_atk: self.cbt.atk_infos(), layout_tsk: self.cbt.tsk_infos() }
+    }
+    fn ks_infos(&self) -> GLWEToLWEKeyLayout {
+        GLWEToLWEKeyLayout{
+            n: self.ks.n(),
+            base2k: self.ks.base2k(),
+            k: self.ks.k(),
+            rank_in: self.ks.rank_in(),
+            dnum: self.ks.dnum()
+        }
+    }
+}
+
 impl<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> GLWEAutomorphismKeyHelper<GLWEAutomorphismKeyPrepared<D, BE>, BE>
     for BDDKeyPrepared<D, BRA, BE>
 {

poulpy-schemes/src/tfhe/bdd_arithmetic/mod.rs

@@ -16,7 +16,7 @@ pub use key::*;
 
 pub mod tests;
 
-pub trait UnsignedInteger: Copy + 'static {
+pub trait UnsignedInteger: Copy + Sync + Send + 'static {
     const BITS: u32;
     const LOG_BITS: u32;
     const LOG_BYTES: u32;

poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/add.rs

@@ -30,7 +30,7 @@ where
         + BDDKeyEncryptSk<BRA, BE>
         + BDDKeyPreparedFactory<BRA, BE>
         + GGSWNoise<BE>
-        + FheUintPrepare<BRA, u32, BE>
+        + FheUintPrepare<BRA, BE>
         + ExecuteBDDCircuit2WTo1W<u32, BE>
         + GLWEEncryptSk<BE>,
     BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,

poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/and.rs

@@ -30,7 +30,7 @@ where
         + BDDKeyEncryptSk<BRA, BE>
         + BDDKeyPreparedFactory<BRA, BE>
         + GGSWNoise<BE>
-        + FheUintPrepare<BRA, u32, BE>
+        + FheUintPrepare<BRA, BE>
         + ExecuteBDDCircuit2WTo1W<u32, BE>
         + GLWEEncryptSk<BE>,
     BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,

poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/mod.rs

@@ -75,6 +75,14 @@ where
 }
 
 impl<BRA: BlindRotationAlgo, BE: Backend> TestContext<BRA, BE> {
+    pub fn glwe_infos(&self) -> GLWELayout {
+        TEST_GLWE_INFOS
+    }
+
+    pub fn ggsw_infos(&self) -> GGSWLayout {
+        TEST_GGSW_INFOS
+    }
+
     pub fn new() -> Self
     where
         Module<BE>: ModuleNew<BE>
@@ -125,8 +133,8 @@ impl<BRA: BlindRotationAlgo, BE: Backend> TestContext<BRA, BE> {
     }
 }
 
-pub(crate) const TEST_N_GLWE: u32 = 256;
+pub(crate) const TEST_N_GLWE: u32 = 1024;
-pub(crate) const TEST_N_LWE: u32 = 77;
+pub(crate) const TEST_N_LWE: u32 = 574;
 pub(crate) const TEST_BASE2K: u32 = 13;
 pub(crate) const TEST_K_GLWE: u32 = 26;
 pub(crate) const TEST_K_GGSW: u32 = 39;

poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/or.rs

@@ -30,7 +30,7 @@ where
         + BDDKeyEncryptSk<BRA, BE>
         + BDDKeyPreparedFactory<BRA, BE>
         + GGSWNoise<BE>
-        + FheUintPrepare<BRA, u32, BE>
+        + FheUintPrepare<BRA, BE>
         + ExecuteBDDCircuit2WTo1W<u32, BE>
         + GLWEEncryptSk<BE>,
     BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,

poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/prepare.rs

@@ -30,7 +30,7 @@ where
         + BDDKeyEncryptSk<BRA, BE>
         + BDDKeyPreparedFactory<BRA, BE>
         + GGSWNoise<BE>
-        + FheUintPrepare<BRA, u32, BE>
+        + FheUintPrepare<BRA, BE>
         + ExecuteBDDCircuit2WTo1W<u32, BE>
         + GLWEEncryptSk<BE>,
     BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,
@@ -67,8 +67,10 @@ where
     let mut c_enc_prep_debug: FheUintPreparedDebug<Vec<u8>, u32> =
         FheUintPreparedDebug::<Vec<u8>, u32>::alloc_from_infos(module, &ggsw_infos);
 
+    let mut scratch_2 = ScratchOwned::alloc(module.fhe_uint_prepare_tmp_bytes(7, 1, &c_enc_prep_debug, &c_enc, bdd_key_prepared));
+
     // GGSW(value)
-    c_enc_prep_debug.prepare(module, &c_enc, bdd_key_prepared, scratch.borrow());
+    c_enc_prep_debug.prepare(module, &c_enc, bdd_key_prepared, scratch_2.borrow());
 
     let max_noise = |col_i: usize| {
         let mut noise: f64 = -(ggsw_infos.size() as f64 * TEST_BASE2K as f64) + SIGMA.log2() + 2.0;

poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/sll.rs

@@ -30,7 +30,7 @@ where
         + BDDKeyEncryptSk<BRA, BE>
         + BDDKeyPreparedFactory<BRA, BE>
         + GGSWNoise<BE>
-        + FheUintPrepare<BRA, u32, BE>
+        + FheUintPrepare<BRA, BE>
         + ExecuteBDDCircuit2WTo1W<u32, BE>
         + GLWEEncryptSk<BE>,
     BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,

poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/slt.rs

@@ -30,7 +30,7 @@ where
         + BDDKeyEncryptSk<BRA, BE>
         + BDDKeyPreparedFactory<BRA, BE>
         + GGSWNoise<BE>
-        + FheUintPrepare<BRA, u32, BE>
+        + FheUintPrepare<BRA, BE>
         + ExecuteBDDCircuit2WTo1W<u32, BE>
         + GLWEEncryptSk<BE>,
     BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,

poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/sltu.rs

@@ -30,7 +30,7 @@ where
         + BDDKeyEncryptSk<BRA, BE>
         + BDDKeyPreparedFactory<BRA, BE>
         + GGSWNoise<BE>
-        + FheUintPrepare<BRA, u32, BE>
+        + FheUintPrepare<BRA, BE>
         + ExecuteBDDCircuit2WTo1W<u32, BE>
         + GLWEEncryptSk<BE>,
     BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,

poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/sra.rs

@@ -30,7 +30,7 @@ where
         + BDDKeyEncryptSk<BRA, BE>
         + BDDKeyPreparedFactory<BRA, BE>
         + GGSWNoise<BE>
-        + FheUintPrepare<BRA, u32, BE>
+        + FheUintPrepare<BRA, BE>
         + ExecuteBDDCircuit2WTo1W<u32, BE>
         + GLWEEncryptSk<BE>,
     BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,

poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/srl.rs

@@ -30,7 +30,7 @@ where
         + BDDKeyEncryptSk<BRA, BE>
         + BDDKeyPreparedFactory<BRA, BE>
         + GGSWNoise<BE>
-        + FheUintPrepare<BRA, u32, BE>
+        + FheUintPrepare<BRA, BE>
         + ExecuteBDDCircuit2WTo1W<u32, BE>
         + GLWEEncryptSk<BE>,
     BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,

poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/sub.rs

@@ -30,7 +30,7 @@ where
         + BDDKeyEncryptSk<BRA, BE>
         + BDDKeyPreparedFactory<BRA, BE>
         + GGSWNoise<BE>
-        + FheUintPrepare<BRA, u32, BE>
+        + FheUintPrepare<BRA, BE>
         + ExecuteBDDCircuit2WTo1W<u32, BE>
         + GLWEEncryptSk<BE>,
     BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,

poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/xor.rs

@@ -30,7 +30,7 @@ where
         + BDDKeyEncryptSk<BRA, BE>
         + BDDKeyPreparedFactory<BRA, BE>
         + GGSWNoise<BE>
-        + FheUintPrepare<BRA, u32, BE>
+        + FheUintPrepare<BRA, BE>
         + ExecuteBDDCircuit2WTo1W<u32, BE>
         + GLWEEncryptSk<BE>,
     BlindRotationKey<Vec<u8>, BRA>: BlindRotationKeyFactory<BRA>,

poulpy-schemes/src/tfhe/blind_rotation/algorithms/mod.rs

@@ -11,7 +11,7 @@ use poulpy_hal::layouts::{Backend, DataMut, DataRef, Scratch, ZnxView};
 
 use crate::tfhe::blind_rotation::{BlindRotationKeyInfos, BlindRotationKeyPrepared, LookUpTableRotationDirection, LookupTable};
 
-pub trait BlindRotationAlgo {}
+pub trait BlindRotationAlgo: Sync {}
 
 pub trait BlindRotationExecute<BRA: BlindRotationAlgo, BE: Backend> {
     fn blind_rotation_execute_tmp_bytes<G, B>(

poulpy-schemes/src/tfhe/blind_rotation/layouts/key.rs

@@ -188,8 +188,7 @@ impl<D: DataRef, BRT: BlindRotationAlgo> BlindRotationKeyInfos for BlindRotation
 }
 
 impl<D: DataRef, BRT: BlindRotationAlgo> BlindRotationKey<D, BRT> {
-    #[allow(dead_code)]
+    pub fn block_size(&self) -> usize {
-    fn block_size(&self) -> usize {
         match self.dist {
             Distribution::BinaryBlock(value) => value,
             _ => 1,

poulpy-schemes/src/tfhe/circuit_bootstrapping/circuit.rs

@@ -1,15 +1,14 @@
 use std::collections::HashMap;
 
 use poulpy_hal::{
-    api::{ModuleLogN, ModuleN, ScratchOwnedAlloc, ScratchOwnedBorrow},
+    api::{ModuleLogN, ModuleN, ScratchAvailable, ScratchOwnedAlloc, ScratchOwnedBorrow},
     layouts::{Backend, DataRef, Module, Scratch, ScratchOwned, ToOwnedDeep},
 };
 
 use poulpy_core::{
     GGSWFromGGLWE, GLWEDecrypt, GLWEPacking, GLWERotate, GLWETrace, ScratchTakeCore,
     layouts::{
-        Dsize, GGLWEInfos, GGLWELayout, GGLWEPreparedToRef, GGSWInfos, GGSWToMut, GLWEAutomorphismKeyHelper, GLWEInfos,
+        Dsize, GGLWE, GGLWEInfos, GGLWELayout, GGLWEPreparedToRef, GGSWInfos, GGSWToMut, GLWEAutomorphismKeyHelper, GLWEInfos, GLWESecretPreparedFactory, GLWEToMut, GLWEToRef, GetGaloisElement, LWEInfos, LWEToRef, Rank
-        GLWESecretPreparedFactory, GLWEToMut, GLWEToRef, GetGaloisElement, LWEInfos, LWEToRef,
     },
 };
 
@@ -132,6 +131,17 @@ where
         R: GGSWInfos,
         A: CircuitBootstrappingKeyInfos,
     {
+
+        let gglwe_infos: GGLWELayout = GGLWELayout {
+            n: res_infos.n(),
+            base2k: res_infos.base2k(),
+            k: res_infos.k(),
+            dnum: res_infos.dnum(),
+            dsize: Dsize(1),
+            rank_in: res_infos.rank().max(Rank(1)).into(),
+            rank_out: res_infos.rank(),
+        };
+
         self.blind_rotation_execute_tmp_bytes(
             block_size,
             extension_factor,
@@ -139,7 +149,7 @@ where
             &cbt_infos.brk_infos(),
         )
         .max(self.glwe_trace_tmp_bytes(res_infos, res_infos, &cbt_infos.atk_infos()))
-        .max(self.ggsw_from_gglwe_tmp_bytes(res_infos, &cbt_infos.tsk_infos()))
+        .max(self.ggsw_from_gglwe_tmp_bytes(res_infos, &cbt_infos.tsk_infos())) + GLWE::bytes_of_from_infos(res_infos) + GGLWE::bytes_of_from_infos(&gglwe_infos)
     }
 
     fn circuit_bootstrapping_execute_to_constant<R, L, D>(
@@ -155,6 +165,9 @@ where
         L: LWEToRef + LWEInfos,
         D: DataRef,
     {   
+
+        assert!(scratch.available() >= self.circuit_bootstrapping_execute_tmp_bytes(key.block_size(), extension_factor, res, key));
+
         circuit_bootstrap_core(
             false,
             self,
@@ -182,6 +195,9 @@ where
         L: LWEToRef + LWEInfos,
         D: DataRef,
     {   
+
+        assert!(scratch.available() >= self.circuit_bootstrapping_execute_tmp_bytes(key.block_size(), extension_factor, res, key));
+
         circuit_bootstrap_core(
             true,
             self,

poulpy-schemes/src/tfhe/circuit_bootstrapping/key.rs

@@ -19,6 +19,7 @@ use crate::tfhe::blind_rotation::{
 };
 
 pub trait CircuitBootstrappingKeyInfos {
+    fn block_size(&self) -> usize;
     fn brk_infos(&self) -> BlindRotationKeyLayout;
     fn atk_infos(&self) -> GLWEAutomorphismKeyLayout;
     fn tsk_infos(&self) -> GGLWEToGGSWKeyLayout;
@@ -32,6 +33,10 @@ pub struct CircuitBootstrappingKeyLayout {
 }
 
 impl CircuitBootstrappingKeyInfos for CircuitBootstrappingKeyLayout {
+    fn block_size(&self) -> usize {
+        unimplemented!("unimplemented for CircuitBootstrappingKeyLayout")
+    }
+
     fn atk_infos(&self) -> GLWEAutomorphismKeyLayout {
         self.layout_atk
     }
@@ -164,6 +169,10 @@ where
 }
 
 impl<D: DataRef, BRA: BlindRotationAlgo> CircuitBootstrappingKeyInfos for CircuitBootstrappingKey<D, BRA> {
+    fn block_size(&self) -> usize {
+        self.brk.block_size()
+    }
+
     fn atk_infos(&self) -> GLWEAutomorphismKeyLayout {
         let (_, atk) = self.atk.iter().next().expect("atk is empty");
         GLWEAutomorphismKeyLayout {

poulpy-schemes/src/tfhe/circuit_bootstrapping/key_prepared.rs

@@ -122,6 +122,10 @@ impl<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> GLWEAutomorphismKeyHelper<
 }
 
 impl<D: DataRef, BRA: BlindRotationAlgo, B: Backend> CircuitBootstrappingKeyInfos for CircuitBootstrappingKeyPrepared<D, BRA, B> {
+    fn block_size(&self) -> usize {
+        self.brk.block_size()
+    }
+
     fn atk_infos(&self) -> GLWEAutomorphismKeyLayout {
         let (_, atk) = self.atk.iter().next().expect("atk is empty");
         GLWEAutomorphismKeyLayout {