diff --git a/core/src/ggsw_ciphertext.rs b/core/src/ggsw_ciphertext.rs index 28c5961..8955b7a 100644 --- a/core/src/ggsw_ciphertext.rs +++ b/core/src/ggsw_ciphertext.rs @@ -124,6 +124,16 @@ impl GGSWCiphertext, FFT64> { res_znx + ci_dft + (ks | expand_rows | res_dft) } + pub fn keyswitch_inplace_scratch_space( + module: &Module, + out_size: usize, + ksk_size: usize, + tensor_key_size: usize, + rank: usize, + ) -> usize { + GGSWCiphertext::keyswitch_scratch_space(module, out_size, out_size, ksk_size, tensor_key_size, rank) + } + pub fn automorphism_scratch_space( module: &Module, out_size: usize, @@ -142,6 +152,23 @@ impl GGSWCiphertext, FFT64> { ) } + pub fn automorphism_inplace_scratch_space( + module: &Module, + out_size: usize, + auto_key_size: usize, + tensor_key_size: usize, + rank: usize, + ) -> usize { + GGSWCiphertext::automorphism_scratch_space( + module, + out_size, + out_size, + auto_key_size, + tensor_key_size, + rank, + ) + } + pub fn external_product_scratch_space( module: &Module, out_size: usize, @@ -389,6 +416,22 @@ where }) } + pub fn keyswitch_inplace( + &mut self, + module: &Module, + ksk: &GLWESwitchingKey, + tsk: &TensorKey, + scratch: &mut Scratch, + ) where + MatZnxDft: MatZnxDftToRef, + MatZnxDft: MatZnxDftToRef, + { + unsafe { + let self_ptr: *mut GGSWCiphertext = self as *mut GGSWCiphertext; + self.keyswitch(module, &*self_ptr, ksk, tsk, scratch); + } + } + pub fn automorphism( &mut self, module: &Module, @@ -470,6 +513,22 @@ where }) } + pub fn automorphism_inplace( + &mut self, + module: &Module, + auto_key: &AutomorphismKey, + tensor_key: &TensorKey, + scratch: &mut Scratch, + ) where + MatZnxDft: MatZnxDftToRef, + MatZnxDft: MatZnxDftToRef, + { + unsafe { + let self_ptr: *mut GGSWCiphertext = self as *mut GGSWCiphertext; + self.automorphism(module, &*self_ptr, auto_key, tensor_key, scratch); + } + } + pub fn external_product( &mut self, module: &Module, diff --git a/core/src/test_fft64/ggsw.rs b/core/src/test_fft64/ggsw.rs index cf5e3c7..f02bd87 100644 --- a/core/src/test_fft64/ggsw.rs +++ b/core/src/test_fft64/ggsw.rs @@ -25,6 +25,38 @@ fn encrypt_sk() { }); } +#[test] +fn keyswitch() { + (1..4).for_each(|rank| { + println!("test keyswitch rank: {}", rank); + test_keyswitch(12, 15, 60, rank, 3.2); + }); +} + +#[test] +fn keyswitch_inplace() { + (1..4).for_each(|rank| { + println!("test keyswitch_inplace rank: {}", rank); + test_keyswitch_inplace(12, 15, 60, rank, 3.2); + }); +} + +#[test] +fn automorphism() { + (1..4).for_each(|rank| { + println!("test automorphism rank: {}", rank); + test_automorphism(-5, 12, 15, 60, rank, 3.2); + }); +} + +#[test] +fn automorphism_inplace() { + (1..4).for_each(|rank| { + println!("test automorphism_inplace rank: {}", rank); + test_automorphism_inplace(-5, 12, 15, 60, rank, 3.2); + }); +} + #[test] fn external_product() { (1..4).for_each(|rank| { @@ -108,14 +140,6 @@ fn test_encrypt_sk(log_n: usize, basek: usize, k_ggsw: usize, sigma: f64, rank: }); } -#[test] -fn keyswitch() { - (1..4).for_each(|rank| { - println!("test keyswitch rank: {}", rank); - test_keyswitch(12, 15, 60, rank, 3.2); - }); -} - fn test_keyswitch(log_n: usize, basek: usize, k: usize, rank: usize, sigma: f64) { let module: Module = Module::::new(1 << log_n); let rows: usize = (k + basek - 1) / basek; @@ -243,6 +267,125 @@ fn test_keyswitch(log_n: usize, basek: usize, k: usize, rank: usize, sigma: f64) }); } +fn test_keyswitch_inplace(log_n: usize, basek: usize, k: usize, rank: usize, sigma: f64) { + let module: Module = Module::::new(1 << log_n); + let rows: usize = (k + basek - 1) / basek; + + let mut ct: GGSWCiphertext, FFT64> = GGSWCiphertext::new(&module, basek, k, rows, rank); + let mut tsk: TensorKey, FFT64> = TensorKey::new(&module, basek, k, rows, rank); + let mut ksk: GLWESwitchingKey, FFT64> = GLWESwitchingKey::new(&module, basek, k, rows, rank, rank); + let mut pt_have: GLWEPlaintext> = GLWEPlaintext::new(&module, basek, k); + let mut pt_want: GLWEPlaintext> = GLWEPlaintext::new(&module, basek, k); + let mut pt_scalar: ScalarZnx> = module.new_scalar_znx(1); + + let mut source_xs: Source = Source::new([0u8; 32]); + let mut source_xe: Source = Source::new([0u8; 32]); + let mut source_xa: Source = Source::new([0u8; 32]); + + let mut scratch: ScratchOwned = ScratchOwned::new( + GGSWCiphertext::encrypt_sk_scratch_space(&module, rank, ct.size()) + | GLWECiphertextFourier::decrypt_scratch_space(&module, ct.size()) + | GLWESwitchingKey::encrypt_sk_scratch_space(&module, rank, ksk.size()) + | TensorKey::encrypt_sk_scratch_space(&module, rank, ksk.size()) + | GGSWCiphertext::keyswitch_inplace_scratch_space(&module, ct.size(), ksk.size(), tsk.size(), rank), + ); + + let var_xs: f64 = 0.5; + + let mut sk_in: SecretKey> = SecretKey::new(&module, rank); + sk_in.fill_ternary_prob(var_xs, &mut source_xs); + + let mut sk_in_dft: SecretKeyFourier, FFT64> = SecretKeyFourier::new(&module, rank); + sk_in_dft.dft(&module, &sk_in); + + let mut sk_out: SecretKey> = SecretKey::new(&module, rank); + sk_out.fill_ternary_prob(var_xs, &mut source_xs); + + let mut sk_out_dft: SecretKeyFourier, FFT64> = SecretKeyFourier::new(&module, rank); + sk_out_dft.dft(&module, &sk_out); + + ksk.encrypt_sk( + &module, + &sk_in, + &sk_out_dft, + &mut source_xa, + &mut source_xe, + sigma, + scratch.borrow(), + ); + tsk.encrypt_sk( + &module, + &sk_out_dft, + &mut source_xa, + &mut source_xe, + sigma, + scratch.borrow(), + ); + + pt_scalar.fill_ternary_hw(0, module.n(), &mut source_xs); + + ct.encrypt_sk( + &module, + &pt_scalar, + &sk_in_dft, + &mut source_xa, + &mut source_xe, + sigma, + scratch.borrow(), + ); + + ct.keyswitch_inplace(&module, &ksk, &tsk, scratch.borrow()); + + let mut ct_glwe_fourier: GLWECiphertextFourier, FFT64> = GLWECiphertextFourier::new(&module, basek, k, rank); + let mut pt_dft: VecZnxDft, FFT64> = module.new_vec_znx_dft(1, ct.size()); + let mut pt_big: VecZnxBig, FFT64> = module.new_vec_znx_big(1, ct.size()); + + (0..ct.rank() + 1).for_each(|col_j| { + (0..ct.rows()).for_each(|row_i| { + module.vec_znx_add_scalar_inplace(&mut pt_want, 0, row_i, &pt_scalar, 0); + + // mul with sk[col_j-1] + if col_j > 0 { + module.vec_znx_dft(&mut pt_dft, 0, &pt_want, 0); + module.svp_apply_inplace(&mut pt_dft, 0, &sk_out_dft, col_j - 1); + module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0); + module.vec_znx_big_normalize(basek, &mut pt_want, 0, &pt_big, 0, scratch.borrow()); + } + + ct.get_row(&module, row_i, col_j, &mut ct_glwe_fourier); + + ct_glwe_fourier.decrypt(&module, &mut pt_have, &sk_out_dft, scratch.borrow()); + + module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0); + + let noise_have: f64 = pt_have.data.std(0, basek).log2(); + let noise_want: f64 = noise_ggsw_keyswitch( + module.n() as f64, + basek, + col_j, + var_xs, + 0f64, + sigma * sigma, + 0f64, + rank as f64, + k, + k, + ); + + println!("{} {}", noise_have, noise_want); + + assert!( + (noise_have - noise_want).abs() <= 0.1, + "{} {}", + noise_have, + noise_want + ); + + pt_want.data.zero(); + }); + }); +} + pub(crate) fn noise_ggsw_keyswitch( n: f64, basek: usize, @@ -292,14 +435,6 @@ pub(crate) fn noise_ggsw_keyswitch( noise.log2().min(-1.0) // max noise is [-2^{-1}, 2^{-1}] } -#[test] -fn automorphism() { - (1..4).for_each(|rank| { - println!("test automorphism rank: {}", rank); - test_automorphism(-5, 12, 15, 60, rank, 3.2); - }); -} - fn test_automorphism(p: i64, log_n: usize, basek: usize, k: usize, rank: usize, sigma: f64) { let module: Module = Module::::new(1 << log_n); let rows: usize = (k + basek - 1) / basek; @@ -421,6 +556,119 @@ fn test_automorphism(p: i64, log_n: usize, basek: usize, k: usize, rank: usize, }); } +fn test_automorphism_inplace(p: i64, log_n: usize, basek: usize, k: usize, rank: usize, sigma: f64) { + let module: Module = Module::::new(1 << log_n); + let rows: usize = (k + basek - 1) / basek; + + let mut ct: GGSWCiphertext, FFT64> = GGSWCiphertext::new(&module, basek, k, rows, rank); + let mut tensor_key: TensorKey, FFT64> = TensorKey::new(&module, basek, k, rows, rank); + let mut auto_key: AutomorphismKey, FFT64> = AutomorphismKey::new(&module, basek, k, rows, rank); + let mut pt_have: GLWEPlaintext> = GLWEPlaintext::new(&module, basek, k); + let mut pt_want: GLWEPlaintext> = GLWEPlaintext::new(&module, basek, k); + let mut pt_scalar: ScalarZnx> = module.new_scalar_znx(1); + + let mut source_xs: Source = Source::new([0u8; 32]); + let mut source_xe: Source = Source::new([0u8; 32]); + let mut source_xa: Source = Source::new([0u8; 32]); + + let mut scratch: ScratchOwned = ScratchOwned::new( + GGSWCiphertext::encrypt_sk_scratch_space(&module, rank, ct.size()) + | GLWECiphertextFourier::decrypt_scratch_space(&module, ct.size()) + | AutomorphismKey::encrypt_sk_scratch_space(&module, rank, auto_key.size()) + | TensorKey::encrypt_sk_scratch_space(&module, rank, tensor_key.size()) + | GGSWCiphertext::automorphism_inplace_scratch_space(&module, ct.size(), auto_key.size(), tensor_key.size(), rank), + ); + + let var_xs: f64 = 0.5; + + let mut sk: SecretKey> = SecretKey::new(&module, rank); + sk.fill_ternary_prob(var_xs, &mut source_xs); + + let mut sk_dft: SecretKeyFourier, FFT64> = SecretKeyFourier::new(&module, rank); + sk_dft.dft(&module, &sk); + + auto_key.encrypt_sk( + &module, + p, + &sk, + &mut source_xa, + &mut source_xe, + sigma, + scratch.borrow(), + ); + tensor_key.encrypt_sk( + &module, + &sk_dft, + &mut source_xa, + &mut source_xe, + sigma, + scratch.borrow(), + ); + + pt_scalar.fill_ternary_hw(0, module.n(), &mut source_xs); + + ct.encrypt_sk( + &module, + &pt_scalar, + &sk_dft, + &mut source_xa, + &mut source_xe, + sigma, + scratch.borrow(), + ); + + ct.automorphism_inplace(&module, &auto_key, &tensor_key, scratch.borrow()); + + module.scalar_znx_automorphism_inplace(p, &mut pt_scalar, 0); + + let mut ct_glwe_fourier: GLWECiphertextFourier, FFT64> = GLWECiphertextFourier::new(&module, basek, k, rank); + let mut pt_dft: VecZnxDft, FFT64> = module.new_vec_znx_dft(1, ct.size()); + let mut pt_big: VecZnxBig, FFT64> = module.new_vec_znx_big(1, ct.size()); + + (0..ct.rank() + 1).for_each(|col_j| { + (0..ct.rows()).for_each(|row_i| { + module.vec_znx_add_scalar_inplace(&mut pt_want, 0, row_i, &pt_scalar, 0); + + // mul with sk[col_j-1] + if col_j > 0 { + module.vec_znx_dft(&mut pt_dft, 0, &pt_want, 0); + module.svp_apply_inplace(&mut pt_dft, 0, &sk_dft, col_j - 1); + module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0); + module.vec_znx_big_normalize(basek, &mut pt_want, 0, &pt_big, 0, scratch.borrow()); + } + + ct.get_row(&module, row_i, col_j, &mut ct_glwe_fourier); + + ct_glwe_fourier.decrypt(&module, &mut pt_have, &sk_dft, scratch.borrow()); + + module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0); + + let noise_have: f64 = pt_have.data.std(0, basek).log2(); + let noise_want: f64 = noise_ggsw_keyswitch( + module.n() as f64, + basek, + col_j, + var_xs, + 0f64, + sigma * sigma, + 0f64, + rank as f64, + k, + k, + ); + + assert!( + (noise_have - noise_want).abs() <= 0.1, + "{} {}", + noise_have, + noise_want + ); + + pt_want.data.zero(); + }); + }); +} + fn test_external_product(log_n: usize, basek: usize, k_ggsw: usize, rank: usize, sigma: f64) { let module: Module = Module::::new(1 << log_n);