diff --git a/poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint.rs b/poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint.rs index 697bcc4..f4cdd83 100644 --- a/poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint.rs +++ b/poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint.rs @@ -1,5 +1,5 @@ use poulpy_core::{ - GLWECopy, GLWEDecrypt, GLWEEncryptSk, GLWEPacking, GLWERotate, LWEFromGLWE, ScratchTakeCore, + GLWEAdd, GLWECopy, GLWEDecrypt, GLWEEncryptSk, GLWEPacking, GLWERotate, GLWESub, GLWETrace, LWEFromGLWE, ScratchTakeCore, layouts::{ Base2K, Degree, GGLWEInfos, GGLWEPreparedToRef, GLWE, GLWEInfos, GLWEPlaintextLayout, GLWESecretPreparedToRef, GLWEToMut, GLWEToRef, LWEInfos, LWEToMut, Rank, TorusPrecision, @@ -170,54 +170,109 @@ impl FheUint { module.glwe_copy(&mut self.bits, cts.remove(&0).unwrap()); } - // pub fn copy_byte( - // &mut self, - // module: &M, - // byte_self: usize, - // byte_a: usize, - // a: &FheUint, - // keys: &BDDKeyPrepared, - // scratch: &mut Scratch, - // ) where - // D0: DataRef, - // D1: DataRef, - // BRA: BlindRotationAlgo, - // M:ModuleLogN + GLWERotate + GLWETrace + GLWESub + GLWEAdd, - // Scratch: ScratchTakeBDD, - // { - // let (mut tmp_fhe_uint_byte, scratch_1) = scratch.take_fhe_uint(a); - // - // - // let log_gap: usize = module.log_n() - T::LOG_BITS as usize; - // module.glwe_rotate(-((T::bit_index(byte_a << 3) << log_gap) as i64), tmp_fhe_uint_byte, self); - // module.glwe_trace_inplace(&mut tmp_fhe_uint_byte, module.log_n() - 3, module.log_n(),&keys.cbt.atk, scratch); - // - // let log_gap: usize = module.log_n() - T::LOG_BITS as usize; - // let rot: i64 = (T::bit_index(byte_self << 3) << log_gap) as i64; - // - // Move starting byte index to first coefficient - // module.glwe_rotate_inplace(-rot, &mut self.bits, scratch); - // - // Stores this byte (everything else zeroed) into tmp_trace - // let (mut tmp_trace, scratch_1) = scratch.take_glwe(self); - // module.glwe_trace( - // &mut tmp_trace, - // module.log_n() - 3, - // module.log_n(), - // self, - // &keys.cbt.atk, - // scratch_1, - // ); - // - // Subtracts the byte - // module.glwe_sub_inplace(&mut self.bits, &tmp_trace); - // - // module.glwe_add_inplace(&mut self.bits, &tmp_fhe_uint_byte); - // - // Moves back into the original position - // module.glwe_rotate_inplace(-rot, &mut self.bits, scratch); - // - // } + #[allow(clippy::too_many_arguments)] + pub fn splice_u16( + &mut self, + module: &M, + dst: usize, + src: usize, + a: &A, + b: &B, + keys: &BDDKeyPrepared, + scratch: &mut Scratch, + ) where + D0: DataRef, + A: GLWEToRef + GLWEInfos, + B: GLWEToRef + GLWEInfos, + BRA: BlindRotationAlgo, + M: ModuleLogN + GLWERotate + GLWETrace + GLWESub + GLWEAdd + GLWECopy, + Scratch: ScratchTakeBDD, + { + assert!(dst < (T::BITS >> 4) as usize); + assert!(src < (T::BITS >> 4) as usize); + + let (mut tmp, scratch_1) = scratch.take_fhe_uint(self); + tmp.splice_u8(module, dst << 1, src << 1, a, b, keys, scratch_1); + self.splice_u8( + module, + (dst << 1) + 1, + (src << 1) + 1, + &tmp, + b, + keys, + scratch_1, + ); + } + + #[allow(clippy::too_many_arguments)] + // Store on the receiver a where the byte_a-th byte of a has been replaced by byte_src2 of src2. + pub fn splice_u8( + &mut self, + module: &M, + dst: usize, + src: usize, + a: &A, + b: &B, + keys: &BDDKeyPrepared, + scratch: &mut Scratch, + ) where + D0: DataRef, + A: GLWEToRef + GLWEInfos, + B: GLWEToRef + GLWEInfos, + BRA: BlindRotationAlgo, + M: ModuleLogN + GLWERotate + GLWETrace + GLWESub + GLWEAdd + GLWECopy, + Scratch: ScratchTakeBDD, + { + assert!(dst < (T::BITS >> 3) as usize); + assert!(src < (T::BITS >> 3) as usize); + + // 1) Zero the byte receiver + let log_gap: usize = module.log_n() - T::LOG_BITS as usize; + let trace_start = (T::LOG_BITS - T::LOG_BYTES) as usize; + let rot: i64 = (T::bit_index(dst << 3) << log_gap) as i64; + + // Move a to self and align byte + module.glwe_rotate(-rot, &mut self.bits, a); + + // Stores this byte (everything else zeroed) into tmp_trace + let (mut tmp_trace, scratch_1) = scratch.take_glwe(a); + module.glwe_trace( + &mut tmp_trace, + trace_start, + module.log_n(), + self, + &keys.cbt.atk, + scratch_1, + ); + + // Subtracts to self to zero it + module.glwe_sub_inplace(&mut self.bits, &tmp_trace); + + // Isolate the byte to transfer from a + let (mut tmp_fhe_uint_byte, scratch_1) = scratch.take_fhe_uint(b); + + // Move a[byte_a] into a[0] + module.glwe_rotate( + -((T::bit_index(src << 3) << log_gap) as i64), + &mut tmp_fhe_uint_byte, + b, + ); + + // Zeroes all other bytes + module.glwe_trace_inplace( + &mut tmp_fhe_uint_byte, + trace_start, + module.log_n(), + &keys.cbt.atk, + scratch_1, + ); + + // Add self[0] += a[0] + module.glwe_add_inplace(&mut self.bits, &tmp_fhe_uint_byte); + + // Moves back self[0] to self[byte_tg] + module.glwe_rotate_inplace(rot, &mut self.bits, scratch); + } } impl GLWEToMut for FheUint { diff --git a/poulpy-schemes/src/tfhe/bdd_arithmetic/mod.rs b/poulpy-schemes/src/tfhe/bdd_arithmetic/mod.rs index 4898b98..66b0699 100644 --- a/poulpy-schemes/src/tfhe/bdd_arithmetic/mod.rs +++ b/poulpy-schemes/src/tfhe/bdd_arithmetic/mod.rs @@ -22,7 +22,7 @@ pub trait UnsignedInteger: Copy + 'static { #[inline(always)] fn bit_index(i: usize) -> usize { - ((i & Self::LOG_BYTES_MASK) << 3) | (i >> Self::LOG_BYTES) + ((i & 7) << Self::LOG_BYTES) | (i >> 3) } } diff --git a/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/fft64_ref.rs b/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/fft64_ref.rs index bf60d25..e4a6de4 100644 --- a/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/fft64_ref.rs +++ b/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/fft64_ref.rs @@ -5,8 +5,8 @@ use poulpy_backend::FFT64Ref; use crate::tfhe::{ bdd_arithmetic::tests::test_suite::{ TestContext, test_bdd_add, test_bdd_and, test_bdd_or, test_bdd_prepare, test_bdd_sll, test_bdd_slt, test_bdd_sltu, - test_bdd_sra, test_bdd_srl, test_bdd_sub, test_bdd_xor, test_glwe_to_glwe_blind_rotation, - test_scalar_to_ggsw_blind_rotation, + test_bdd_sra, test_bdd_srl, test_bdd_sub, test_bdd_xor, test_fhe_uint_splice_u8, test_fhe_uint_splice_u16, + test_glwe_to_glwe_blind_rotation, test_scalar_to_ggsw_blind_rotation, }, blind_rotation::CGGI, }; @@ -14,6 +14,16 @@ use crate::tfhe::{ static TEST_CONTEXT_CGGI_FFT64_REF: LazyLock> = LazyLock::new(|| TestContext::::new()); +#[test] +fn test_fhe_uint_splice_u8_fft64_ref() { + test_fhe_uint_splice_u8(&TEST_CONTEXT_CGGI_FFT64_REF) +} + +#[test] +fn test_fhe_uint_splice_u16_fft64_ref() { + test_fhe_uint_splice_u16(&TEST_CONTEXT_CGGI_FFT64_REF) +} + #[test] fn test_glwe_to_glwe_blind_rotation_fft64_ref() { test_glwe_to_glwe_blind_rotation(&TEST_CONTEXT_CGGI_FFT64_REF) diff --git a/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/fheuint.rs b/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/fheuint.rs new file mode 100644 index 0000000..0c82084 --- /dev/null +++ b/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/fheuint.rs @@ -0,0 +1,128 @@ +use poulpy_core::{ + GLWEAdd, GLWEDecrypt, GLWEEncryptSk, GLWERotate, GLWESub, GLWETrace, + layouts::{GLWELayout, GLWESecretPrepared}, +}; +use poulpy_hal::{ + api::{ScratchOwnedAlloc, ScratchOwnedBorrow}, + layouts::{Backend, Module, Scratch, ScratchOwned}, + source::Source, +}; + +use crate::tfhe::{ + bdd_arithmetic::{ + BDDKeyPrepared, FheUint, ScratchTakeBDD, + tests::test_suite::{TEST_GLWE_INFOS, TestContext}, + }, + blind_rotation::BlindRotationAlgo, +}; + +pub fn test_fhe_uint_splice_u8(test_context: &TestContext) +where + Module: GLWEEncryptSk + GLWERotate + GLWETrace + GLWESub + GLWEAdd + GLWEDecrypt, + ScratchOwned: ScratchOwnedAlloc + ScratchOwnedBorrow, + Scratch: ScratchTakeBDD, +{ + let glwe_infos: GLWELayout = TEST_GLWE_INFOS; + + let module: &Module = &test_context.module; + let sk: &GLWESecretPrepared, BE> = &test_context.sk_glwe; + let keys: &BDDKeyPrepared, BRA, BE> = &test_context.bdd_key; + + let mut source_xa: Source = Source::new([2u8; 32]); + let mut source_xe: Source = Source::new([3u8; 32]); + + let mut scratch: ScratchOwned = ScratchOwned::alloc(1 << 22); + + let mut a_enc: FheUint, u32> = FheUint::, u32>::alloc_from_infos(&glwe_infos); + let mut b_enc: FheUint, u32> = FheUint::, u32>::alloc_from_infos(&glwe_infos); + let mut c_enc: FheUint, u32> = FheUint::, u32>::alloc_from_infos(&glwe_infos); + + let a: u32 = 0xFFFFFFFF; + let b: u32 = 0xAABBCCDD; + + b_enc.encrypt_sk( + module, + b, + sk, + &mut source_xa, + &mut source_xe, + scratch.borrow(), + ); + a_enc.encrypt_sk( + module, + a, + sk, + &mut source_xa, + &mut source_xe, + scratch.borrow(), + ); + + for dst in 0..4 { + for src in 0..4 { + c_enc.splice_u8(module, dst, src, &a_enc, &b_enc, keys, scratch.borrow()); + + let rj: u32 = (dst << 3) as u32; + let ri: u32 = (src << 3) as u32; + let a_r: u32 = a.rotate_right(rj); + let b_r: u32 = b.rotate_right(ri); + + let c_want: u32 = ((a_r & 0xFFFF_FF00) | (b_r & 0x0000_00FF)).rotate_left(rj); + + assert_eq!(c_want, c_enc.decrypt(module, sk, scratch.borrow())); + } + } +} + +pub fn test_fhe_uint_splice_u16(test_context: &TestContext) +where + Module: GLWEEncryptSk + GLWERotate + GLWETrace + GLWESub + GLWEAdd + GLWEDecrypt, + ScratchOwned: ScratchOwnedAlloc + ScratchOwnedBorrow, + Scratch: ScratchTakeBDD, +{ + let glwe_infos: GLWELayout = TEST_GLWE_INFOS; + + let module: &Module = &test_context.module; + let sk: &GLWESecretPrepared, BE> = &test_context.sk_glwe; + let keys: &BDDKeyPrepared, BRA, BE> = &test_context.bdd_key; + + let mut source_xa: Source = Source::new([2u8; 32]); + let mut source_xe: Source = Source::new([3u8; 32]); + + let mut scratch: ScratchOwned = ScratchOwned::alloc(1 << 22); + + let mut a_enc: FheUint, u32> = FheUint::, u32>::alloc_from_infos(&glwe_infos); + let mut b_enc: FheUint, u32> = FheUint::, u32>::alloc_from_infos(&glwe_infos); + let mut c_enc: FheUint, u32> = FheUint::, u32>::alloc_from_infos(&glwe_infos); + + let a: u32 = 0xFFFFFFFF; + let b: u32 = 0xAABBCCDD; + + b_enc.encrypt_sk( + module, + b, + sk, + &mut source_xa, + &mut source_xe, + scratch.borrow(), + ); + a_enc.encrypt_sk( + module, + a, + sk, + &mut source_xa, + &mut source_xe, + scratch.borrow(), + ); + + for dst in 0..2 { + for src in 0..2 { + c_enc.splice_u16(module, dst, src, &a_enc, &b_enc, keys, scratch.borrow()); + let rj: u32 = (dst << 4) as u32; + let ri: u32 = (src << 4) as u32; + let a_r: u32 = a.rotate_right(rj); + let b_r: u32 = b.rotate_right(ri); + let c_want: u32 = ((a_r & 0xFFFF_0000) | (b_r & 0x0000_FFFF)).rotate_left(rj); + assert_eq!(c_want, c_enc.decrypt(module, sk, scratch.borrow())); + } + } +} diff --git a/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/mod.rs b/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/mod.rs index 766ae73..53a02ae 100644 --- a/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/mod.rs +++ b/poulpy-schemes/src/tfhe/bdd_arithmetic/tests/test_suite/mod.rs @@ -1,5 +1,6 @@ mod add; mod and; +mod fheuint; mod ggsw_blind_rotations; mod glwe_blind_rotation; mod or; @@ -14,6 +15,7 @@ mod xor; pub use add::*; pub use and::*; +pub use fheuint::*; pub use ggsw_blind_rotations::*; pub use glwe_blind_rotation::*; pub use or::*;