poulpy/core/src/test_fft64/ggsw.rs
core/src/test_fft64/ggsw.rsRaw
1use base2k::{
2 FFT64, Module, ScalarZnx, ScalarZnxAlloc, ScalarZnxDftOps, ScratchOwned, Stats, VecZnxBig, VecZnxBigAlloc, VecZnxBigOps,
3 VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, VecZnxOps, VecZnxToMut, ZnxViewMut, ZnxZero,
4};
5use sampling::source::Source;
6
7use crate::{
8 elem::{GetRow, Infos},
9 ggsw_ciphertext::GGSWCiphertext,
10 glwe_ciphertext_fourier::GLWECiphertextFourier,
11 glwe_plaintext::GLWEPlaintext,
12 keys::{SecretKey, SecretKeyFourier},
13 keyswitch_key::GLWESwitchingKey,
14};
15
16#[test]
17fn encrypt_sk() {
18 (1..4).for_each(|rank| {
19 println!("test encrypt_sk rank: {}", rank);
20 test_encrypt_sk(11, 8, 54, 3.2, rank);
21 });
22}
23
24#[test]
25fn external_product() {
26 (1..4).for_each(|rank| {
27 println!("test external_product rank: {}", rank);
28 test_external_product(12, 12, 60, rank, 3.2);
29 });
30}
31
32#[test]
33fn external_product_inplace() {
34 (1..4).for_each(|rank| {
35 println!("test external_product rank: {}", rank);
36 test_external_product_inplace(12, 15, 60, rank, 3.2);
37 });
38}
39
40fn test_encrypt_sk(log_n: usize, basek: usize, k_ggsw: usize, sigma: f64, rank: usize) {
41 let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
42
43 let rows: usize = (k_ggsw + basek - 1) / basek;
44
45 let mut ct: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, k_ggsw, rows, rank);
46 let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ggsw);
47 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ggsw);
48 let mut pt_scalar: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
49
50 let mut source_xs: Source = Source::new([0u8; 32]);
51 let mut source_xe: Source = Source::new([0u8; 32]);
52 let mut source_xa: Source = Source::new([0u8; 32]);
53
54 pt_scalar.fill_ternary_hw(0, module.n(), &mut source_xs);
55
56 let mut scratch: ScratchOwned = ScratchOwned::new(
57 GGSWCiphertext::encrypt_sk_scratch_space(&module, rank, ct.size())
58 | GLWECiphertextFourier::decrypt_scratch_space(&module, ct.size()),
59 );
60
61 let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
62 sk.fill_ternary_prob(0.5, &mut source_xs);
63
64 let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
65 sk_dft.dft(&module, &sk);
66
67 ct.encrypt_sk(
68 &module,
69 &pt_scalar,
70 &sk_dft,
71 &mut source_xa,
72 &mut source_xe,
73 sigma,
74 scratch.borrow(),
75 );
76
77 let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::new(&module, basek, k_ggsw, rank);
78 let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct.size());
79 let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct.size());
80
81 (0..ct.rank() + 1).for_each(|col_j| {
82 (0..ct.rows()).for_each(|row_i| {
83 module.vec_znx_add_scalar_inplace(&mut pt_want, 0, row_i, &pt_scalar, 0);
84
85 // mul with sk[col_j-1]
86 if col_j > 0 {
87 module.vec_znx_dft(&mut pt_dft, 0, &pt_want, 0);
88 module.svp_apply_inplace(&mut pt_dft, 0, &sk_dft, col_j - 1);
89 module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
90 module.vec_znx_big_normalize(basek, &mut pt_want, 0, &pt_big, 0, scratch.borrow());
91 }
92
93 ct.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
94
95 ct_glwe_fourier.decrypt(&module, &mut pt_have, &sk_dft, scratch.borrow());
96
97 module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0);
98
99 let std_pt: f64 = pt_have.data.std(0, basek) * (k_ggsw as f64).exp2();
100 assert!((sigma - std_pt).abs() <= 0.2, "{} {}", sigma, std_pt);
101
102 pt_want.data.zero();
103 });
104 });
105}
106
107fn test_external_product(log_n: usize, basek: usize, k_ggsw: usize, rank: usize, sigma: f64) {
108 let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
109
110 let rows: usize = (k_ggsw + basek - 1) / basek;
111
112 let mut ct_ggsw_rhs: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, k_ggsw, rows, rank);
113 let mut ct_ggsw_lhs_in: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, k_ggsw, rows, rank);
114 let mut ct_ggsw_lhs_out: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, k_ggsw, rows, rank);
115 let mut pt_ggsw_lhs: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
116 let mut pt_ggsw_rhs: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
117
118 let mut source_xs: Source = Source::new([0u8; 32]);
119 let mut source_xe: Source = Source::new([0u8; 32]);
120 let mut source_xa: Source = Source::new([0u8; 32]);
121
122 pt_ggsw_lhs.fill_ternary_prob(0, 0.5, &mut source_xs);
123
124 let k: usize = 1;
125
126 pt_ggsw_rhs.to_mut().raw_mut()[k] = 1; //X^{k}
127
128 let mut scratch: ScratchOwned = ScratchOwned::new(
129 GLWESwitchingKey::encrypt_sk_scratch_space(&module, rank, ct_ggsw_rhs.size())
130 | GLWECiphertextFourier::decrypt_scratch_space(&module, ct_ggsw_lhs_out.size())
131 | GGSWCiphertext::encrypt_sk_scratch_space(&module, rank, ct_ggsw_lhs_in.size())
132 | GGSWCiphertext::external_product_scratch_space(
133 &module,
134 ct_ggsw_lhs_out.size(),
135 ct_ggsw_lhs_in.size(),
136 ct_ggsw_rhs.size(),
137 rank,
138 ),
139 );
140
141 let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
142 sk.fill_ternary_prob(0.5, &mut source_xs);
143
144 let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
145 sk_dft.dft(&module, &sk);
146
147 ct_ggsw_rhs.encrypt_sk(
148 &module,
149 &pt_ggsw_rhs,
150 &sk_dft,
151 &mut source_xa,
152 &mut source_xe,
153 sigma,
154 scratch.borrow(),
155 );
156
157 ct_ggsw_lhs_in.encrypt_sk(
158 &module,
159 &pt_ggsw_lhs,
160 &sk_dft,
161 &mut source_xa,
162 &mut source_xe,
163 sigma,
164 scratch.borrow(),
165 );
166
167 ct_ggsw_lhs_out.external_product(&module, &ct_ggsw_lhs_in, &ct_ggsw_rhs, scratch.borrow());
168
169 let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::new(&module, basek, k_ggsw, rank);
170 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ggsw);
171 let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct_ggsw_lhs_out.size());
172 let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct_ggsw_lhs_out.size());
173 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ggsw);
174
175 module.vec_znx_rotate_inplace(k as i64, &mut pt_ggsw_lhs, 0);
176
177 (0..ct_ggsw_lhs_out.rank() + 1).for_each(|col_j| {
178 (0..ct_ggsw_lhs_out.rows()).for_each(|row_i| {
179 module.vec_znx_add_scalar_inplace(&mut pt_want, 0, row_i, &pt_ggsw_lhs, 0);
180
181 if col_j > 0 {
182 module.vec_znx_dft(&mut pt_dft, 0, &pt_want, 0);
183 module.svp_apply_inplace(&mut pt_dft, 0, &sk_dft, col_j - 1);
184 module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
185 module.vec_znx_big_normalize(basek, &mut pt_want, 0, &pt_big, 0, scratch.borrow());
186 }
187
188 ct_ggsw_lhs_out.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
189 ct_glwe_fourier.decrypt(&module, &mut pt, &sk_dft, scratch.borrow());
190
191 module.vec_znx_sub_ab_inplace(&mut pt, 0, &pt_want, 0);
192
193 let noise_have: f64 = pt.data.std(0, basek).log2();
194
195 let var_gct_err_lhs: f64 = sigma * sigma;
196 let var_gct_err_rhs: f64 = 0f64;
197
198 let var_msg: f64 = 1f64 / module.n() as f64; // X^{k}
199 let var_a0_err: f64 = sigma * sigma;
200 let var_a1_err: f64 = 1f64 / 12f64;
201
202 let noise_want: f64 = noise_ggsw_product(
203 module.n() as f64,
204 basek,
205 0.5,
206 var_msg,
207 var_a0_err,
208 var_a1_err,
209 var_gct_err_lhs,
210 var_gct_err_rhs,
211 rank as f64,
212 k_ggsw,
213 k_ggsw,
214 );
215
216 assert!(
217 (noise_have - noise_want).abs() <= 0.1,
218 "have: {} want: {}",
219 noise_have,
220 noise_want
221 );
222
223 pt_want.data.zero();
224 });
225 });
226}
227
228fn test_external_product_inplace(log_n: usize, basek: usize, k_ggsw: usize, rank: usize, sigma: f64) {
229 let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
230 let rows: usize = (k_ggsw + basek - 1) / basek;
231
232 let mut ct_ggsw_rhs: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, k_ggsw, rows, rank);
233 let mut ct_ggsw_lhs: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, k_ggsw, rows, rank);
234 let mut pt_ggsw_lhs: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
235 let mut pt_ggsw_rhs: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
236
237 let mut source_xs: Source = Source::new([0u8; 32]);
238 let mut source_xe: Source = Source::new([0u8; 32]);
239 let mut source_xa: Source = Source::new([0u8; 32]);
240
241 pt_ggsw_lhs.fill_ternary_prob(0, 0.5, &mut source_xs);
242
243 let k: usize = 1;
244
245 pt_ggsw_rhs.to_mut().raw_mut()[k] = 1; //X^{k}
246
247 let mut scratch: ScratchOwned = ScratchOwned::new(
248 GLWESwitchingKey::encrypt_sk_scratch_space(&module, rank, ct_ggsw_rhs.size())
249 | GLWECiphertextFourier::decrypt_scratch_space(&module, ct_ggsw_lhs.size())
250 | GGSWCiphertext::encrypt_sk_scratch_space(&module, rank, ct_ggsw_lhs.size())
251 | GGSWCiphertext::external_product_inplace_scratch_space(&module, ct_ggsw_lhs.size(), ct_ggsw_rhs.size(), rank),
252 );
253
254 let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
255 sk.fill_ternary_prob(0.5, &mut source_xs);
256
257 let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
258 sk_dft.dft(&module, &sk);
259
260 ct_ggsw_rhs.encrypt_sk(
261 &module,
262 &pt_ggsw_rhs,
263 &sk_dft,
264 &mut source_xa,
265 &mut source_xe,
266 sigma,
267 scratch.borrow(),
268 );
269
270 ct_ggsw_lhs.encrypt_sk(
271 &module,
272 &pt_ggsw_lhs,
273 &sk_dft,
274 &mut source_xa,
275 &mut source_xe,
276 sigma,
277 scratch.borrow(),
278 );
279
280 ct_ggsw_lhs.external_product_inplace(&module, &ct_ggsw_rhs, scratch.borrow());
281
282 let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::new(&module, basek, k_ggsw, rank);
283 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ggsw);
284 let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct_ggsw_lhs.size());
285 let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct_ggsw_lhs.size());
286 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ggsw);
287
288 module.vec_znx_rotate_inplace(k as i64, &mut pt_ggsw_lhs, 0);
289
290 (0..ct_ggsw_lhs.rank() + 1).for_each(|col_j| {
291 (0..ct_ggsw_lhs.rows()).for_each(|row_i| {
292 module.vec_znx_add_scalar_inplace(&mut pt_want, 0, row_i, &pt_ggsw_lhs, 0);
293
294 if col_j > 0 {
295 module.vec_znx_dft(&mut pt_dft, 0, &pt_want, 0);
296 module.svp_apply_inplace(&mut pt_dft, 0, &sk_dft, col_j - 1);
297 module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
298 module.vec_znx_big_normalize(basek, &mut pt_want, 0, &pt_big, 0, scratch.borrow());
299 }
300
301 ct_ggsw_lhs.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
302 ct_glwe_fourier.decrypt(&module, &mut pt, &sk_dft, scratch.borrow());
303
304 module.vec_znx_sub_ab_inplace(&mut pt, 0, &pt_want, 0);
305
306 let noise_have: f64 = pt.data.std(0, basek).log2();
307
308 let var_gct_err_lhs: f64 = sigma * sigma;
309 let var_gct_err_rhs: f64 = 0f64;
310
311 let var_msg: f64 = 1f64 / module.n() as f64; // X^{k}
312 let var_a0_err: f64 = sigma * sigma;
313 let var_a1_err: f64 = 1f64 / 12f64;
314
315 let noise_want: f64 = noise_ggsw_product(
316 module.n() as f64,
317 basek,
318 0.5,
319 var_msg,
320 var_a0_err,
321 var_a1_err,
322 var_gct_err_lhs,
323 var_gct_err_rhs,
324 rank as f64,
325 k_ggsw,
326 k_ggsw,
327 );
328
329 assert!(
330 (noise_have - noise_want).abs() <= 0.1,
331 "have: {} want: {}",
332 noise_have,
333 noise_want
334 );
335
336 pt_want.data.zero();
337 });
338 });
339}
340pub(crate) fn noise_ggsw_product(
341 n: f64,
342 basek: usize,
343 var_xs: f64,
344 var_msg: f64,
345 var_a0_err: f64,
346 var_a1_err: f64,
347 var_gct_err_lhs: f64,
348 var_gct_err_rhs: f64,
349 rank: f64,
350 a_logq: usize,
351 b_logq: usize,
352) -> f64 {
353 let a_logq: usize = a_logq.min(b_logq);
354 let a_cols: usize = (a_logq + basek - 1) / basek;
355
356 let b_scale = 2.0f64.powi(b_logq as i32);
357 let a_scale: f64 = 2.0f64.powi((b_logq - a_logq) as i32);
358
359 let base: f64 = (1 << (basek)) as f64;
360 let var_base: f64 = base * base / 12f64;
361
362 // lhs = a_cols * n * (var_base * var_gct_err_lhs + var_e_a * var_msg * p^2)
363 // rhs = a_cols * n * var_base * var_gct_err_rhs * var_xs
364 let mut noise: f64 = (rank + 1.0) * (a_cols as f64) * n * var_base * (var_gct_err_lhs + var_xs * var_gct_err_rhs);
365 noise += var_msg * var_a0_err * a_scale * a_scale * n;
366 noise += var_msg * var_a1_err * a_scale * a_scale * n * var_xs * rank;
367 noise = noise.sqrt();
368 noise /= b_scale;
369 noise.log2().min(-1.0) // max noise is [-2^{-1}, 2^{-1}]
370}