Refactor of GGSW key-switch to enable easier implementation of GGSW automorphism

This commit is contained in:
Jean-Philippe Bossuat
2025-05-20 17:42:43 +02:00
parent a803127424
commit 640ff9ea61
9 changed files with 404 additions and 288 deletions

View File

@@ -13,7 +13,6 @@ use crate::{
keys::{SecretKey, SecretKeyFourier},
keyswitch_key::GLWESwitchingKey,
tensor_key::TensorKey,
test_fft64::gglwe::log2_std_noise_gglwe_product,
};
use super::gglwe::var_noise_gglwe_product;
@@ -230,6 +229,8 @@ fn test_keyswitch(log_n: usize, basek: usize, k: usize, rank: usize, sigma: f64)
k,
);
println!("{} {}", noise_have, noise_want);
assert!(
(noise_have - noise_want).abs() <= 0.1,
"{} {}",
@@ -291,122 +292,130 @@ pub(crate) fn noise_ggsw_keyswitch(
noise.log2().min(-1.0) // max noise is [-2^{-1}, 2^{-1}]
}
// fn test_automorphism(p: i64, log_n: usize, basek: usize, k: usize, rank: usize, sigma: f64) {
// let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
// let rows: usize = (k_ggsw + basek - 1) / basek;
//
// let mut ct_ggsw_in: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, k, rows, rank);
// let mut ct_ggsw_out: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, k, rows, rank);
// let mut auto_key: AutomorphismKey<Vec<u8>, FFT64> = AutomorphismKey::new(&module, basek, k, rows, rank);
//
// let mut pt_ggsw_in: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
// let mut pt_ggsw_out: ScalarZnx<Vec<u8>> = 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]);
//
// pt_ggsw_in.fill_ternary_prob(0, 0.5, &mut source_xs);
//
// let mut scratch: ScratchOwned = ScratchOwned::new(
// AutomorphismKey::encrypt_sk_scratch_space(&module, rank, auto_key.size())
// | GLWECiphertextFourier::decrypt_scratch_space(&module, ct_ggsw_out.size())
// | GGSWCiphertext::encrypt_sk_scratch_space(&module, rank, ct_ggsw_in.size())
// | GGSWCiphertext::automorphism_scratch_space(
// &module,
// ct_ggsw_out.size(),
// ct_ggsw_in.size(),
// auto_key.size(),
// rank,
// ),
// );
//
// let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
// sk.fill_ternary_prob(0.5, &mut source_xs);
//
// let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
// sk_dft.dft(&module, &sk);
//
// ct_ggsw_in.encrypt_sk(
// &module,
// &pt_ggsw_in,
// &sk_dft,
// &mut source_xa,
// &mut source_xe,
// sigma,
// scratch.borrow(),
// );
//
// auto_key.encrypt_sk(
// &module,
// p,
// &sk,
// &mut source_xa,
// &mut source_xe,
// sigma,
// scratch.borrow(),
// );
//
// ct_ggsw_out.automorphism(&module, &ct_ggsw_in, &auto_key, scratch.borrow());
//
// let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::new(&module, basek, k_ggsw, rank);
// let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ggsw);
// let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct_ggsw_lhs_out.size());
// let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct_ggsw_lhs_out.size());
// let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ggsw);
//
// module.vec_znx_rotate_inplace(k as i64, &mut pt_ggsw_lhs, 0);
//
// (0..ct_ggsw_lhs_out.rank() + 1).for_each(|col_j| {
// (0..ct_ggsw_lhs_out.rows()).for_each(|row_i| {
// module.vec_znx_add_scalar_inplace(&mut pt_want, 0, row_i, &pt_ggsw_lhs, 0);
//
// 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_ggsw_lhs_out.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
// ct_glwe_fourier.decrypt(&module, &mut pt, &sk_dft, scratch.borrow());
//
// module.vec_znx_sub_ab_inplace(&mut pt, 0, &pt_want, 0);
//
// let noise_have: f64 = pt.data.std(0, basek).log2();
//
// let var_gct_err_lhs: f64 = sigma * sigma;
// let var_gct_err_rhs: f64 = 0f64;
//
// let var_msg: f64 = 1f64 / module.n() as f64; // X^{k}
// let var_a0_err: f64 = sigma * sigma;
// let var_a1_err: f64 = 1f64 / 12f64;
//
// let noise_want: f64 = noise_ggsw_product(
// module.n() as f64,
// basek,
// 0.5,
// var_msg,
// var_a0_err,
// var_a1_err,
// var_gct_err_lhs,
// var_gct_err_rhs,
// rank as f64,
// k_ggsw,
// k_ggsw,
// );
//
// assert!(
// (noise_have - noise_want).abs() <= 0.1,
// "have: {} want: {}",
// noise_have,
// noise_want
// );
//
// pt_want.data.zero();
// });
// });
// }
fn test_automorphism(p: i64, log_n: usize, basek: usize, k: usize, rank: usize, sigma: f64) {
let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
let rows: usize = (k + basek - 1) / basek;
let mut ct_in: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, k, rows, rank);
let mut ct_out: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, k, rows, rank);
let mut tsk: TensorKey<Vec<u8>, FFT64> = TensorKey::new(&module, basek, k, rows, rank);
let mut ksk: GLWESwitchingKey<Vec<u8>, FFT64> = GLWESwitchingKey::new(&module, basek, k, rows, rank, rank);
let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k);
let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k);
let mut pt_scalar: ScalarZnx<Vec<u8>> = 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_in.size())
| GLWECiphertextFourier::decrypt_scratch_space(&module, ct_out.size())
| GLWESwitchingKey::encrypt_sk_scratch_space(&module, rank, ksk.size())
| TensorKey::encrypt_sk_scratch_space(&module, rank, ksk.size())
| GGSWCiphertext::keyswitch_scratch_space(
&module,
ct_out.size(),
ct_in.size(),
ksk.size(),
tsk.size(),
rank,
),
);
let var_xs: f64 = 0.5;
let mut sk_in: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
sk_in.fill_ternary_prob(var_xs, &mut source_xs);
let mut sk_in_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
sk_in_dft.dft(&module, &sk_in);
let mut sk_out: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
sk_out.fill_ternary_prob(var_xs, &mut source_xs);
let mut sk_out_dft: SecretKeyFourier<Vec<u8>, 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_in.encrypt_sk(
&module,
&pt_scalar,
&sk_in_dft,
&mut source_xa,
&mut source_xe,
sigma,
scratch.borrow(),
);
ct_out.keyswitch(&module, &ct_in, &ksk, &tsk, scratch.borrow());
let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::new(&module, basek, k, rank);
let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct_out.size());
let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct_out.size());
(0..ct_out.rank() + 1).for_each(|col_j| {
(0..ct_out.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_out.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,
);
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<FFT64> = Module::<FFT64>::new(1 << log_n);