poulpy/core/src/test_fft64/gglwe.rs
core/src/test_fft64/gglwe.rsRaw
1use base2k::{FFT64, Module, ScalarZnx, ScalarZnxAlloc, ScratchOwned, Stats, VecZnxOps, ZnxViewMut};
2use sampling::source::Source;
3
4use crate::{
5 elem::{GetRow, Infos},
6 ggsw_ciphertext::GGSWCiphertext,
7 glwe_ciphertext_fourier::GLWECiphertextFourier,
8 glwe_plaintext::GLWEPlaintext,
9 keys::{SecretKey, SecretKeyFourier},
10 keyswitch_key::GLWESwitchingKey,
11};
12
13#[test]
14fn encrypt_sk() {
15 (1..4).for_each(|rank_in| {
16 (1..4).for_each(|rank_out| {
17 println!("test encrypt_sk rank_in rank_out: {} {}", rank_in, rank_out);
18 test_encrypt_sk(11, 8, 54, 3.2, rank_in, rank_out);
19 });
20 });
21}
22
23fn test_encrypt_sk(log_n: usize, basek: usize, k_ksk: usize, sigma: f64, rank_in: usize, rank_out: usize) {
24 let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
25 let rows = (k_ksk + basek - 1) / basek;
26
27 let mut ksk: GLWESwitchingKey<Vec<u8>, FFT64> = GLWESwitchingKey::new(&module, basek, k_ksk, rows, rank_in, rank_out);
28 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ksk);
29
30 let mut source_xs: Source = Source::new([0u8; 32]);
31 let mut source_xe: Source = Source::new([0u8; 32]);
32 let mut source_xa: Source = Source::new([0u8; 32]);
33
34 let mut scratch: ScratchOwned = ScratchOwned::new(
35 GLWESwitchingKey::encrypt_sk_scratch_space(&module, rank_out, ksk.size())
36 | GLWECiphertextFourier::decrypt_scratch_space(&module, ksk.size()),
37 );
38
39 let mut sk_in: SecretKey<Vec<u8>> = SecretKey::new(&module, rank_in);
40 sk_in.fill_ternary_prob(0.5, &mut source_xs);
41
42 let mut sk_in_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank_in);
43 sk_in_dft.dft(&module, &sk_in);
44
45 let mut sk_out: SecretKey<Vec<u8>> = SecretKey::new(&module, rank_out);
46 sk_out.fill_ternary_prob(0.5, &mut source_xs);
47
48 let mut sk_out_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank_out);
49 sk_out_dft.dft(&module, &sk_out);
50
51 ksk.encrypt_sk(
52 &module,
53 &sk_in,
54 &sk_out_dft,
55 &mut source_xa,
56 &mut source_xe,
57 sigma,
58 scratch.borrow(),
59 );
60
61 let mut ct_gglwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::new(&module, basek, k_ksk, rank_out);
62
63 (0..ksk.rank_in()).for_each(|col_i| {
64 (0..ksk.rows()).for_each(|row_i| {
65 ksk.get_row(&module, row_i, 0, &mut ct_gglwe_fourier);
66 ct_gglwe_fourier.decrypt(&module, &mut pt, &sk_out_dft, scratch.borrow());
67 module.vec_znx_sub_scalar_inplace(&mut pt, 0, row_i, &sk_in, col_i);
68 let std_pt: f64 = pt.data.std(0, basek) * (k_ksk as f64).exp2();
69 assert!((sigma - std_pt).abs() <= 0.2, "{} {}", sigma, std_pt);
70 });
71 });
72}
73
74#[test]
75fn keyswitch() {
76 let module: Module<FFT64> = Module::<FFT64>::new(2048);
77 let basek: usize = 12;
78 let log_k_grlwe: usize = 60;
79 let rows: usize = (log_k_grlwe + basek - 1) / basek;
80
81 let rank: usize = 1;
82
83 let sigma: f64 = 3.2;
84
85 let mut ct_grlwe_s0s1: GLWESwitchingKey<Vec<u8>, FFT64> =
86 GLWESwitchingKey::new(&module, basek, log_k_grlwe, rows, rank, rank);
87 let mut ct_grlwe_s1s2: GLWESwitchingKey<Vec<u8>, FFT64> =
88 GLWESwitchingKey::new(&module, basek, log_k_grlwe, rows, rank, rank);
89 let mut ct_grlwe_s0s2: GLWESwitchingKey<Vec<u8>, FFT64> =
90 GLWESwitchingKey::new(&module, basek, log_k_grlwe, rows, rank, rank);
91
92 let mut source_xs: Source = Source::new([0u8; 32]);
93 let mut source_xe: Source = Source::new([0u8; 32]);
94 let mut source_xa: Source = Source::new([0u8; 32]);
95
96 let mut scratch: ScratchOwned = ScratchOwned::new(
97 GLWESwitchingKey::encrypt_sk_scratch_space(&module, rank, ct_grlwe_s0s1.size())
98 | GLWECiphertextFourier::decrypt_scratch_space(&module, ct_grlwe_s0s2.size())
99 | GLWESwitchingKey::keyswitch_scratch_space(
100 &module,
101 ct_grlwe_s0s2.size(),
102 ct_grlwe_s0s1.size(),
103 ct_grlwe_s1s2.size(),
104 ),
105 );
106
107 let mut sk0: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
108 sk0.fill_ternary_prob(0.5, &mut source_xs);
109
110 let mut sk0_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
111 sk0_dft.dft(&module, &sk0);
112
113 let mut sk1: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
114 sk1.fill_ternary_prob(0.5, &mut source_xs);
115
116 let mut sk1_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
117 sk1_dft.dft(&module, &sk1);
118
119 let mut sk2: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
120 sk2.fill_ternary_prob(0.5, &mut source_xs);
121
122 let mut sk2_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
123 sk2_dft.dft(&module, &sk2);
124
125 // GRLWE_{s1}(s0) = s0 -> s1
126 ct_grlwe_s0s1.encrypt_sk(
127 &module,
128 &sk0.data,
129 &sk1_dft,
130 &mut source_xa,
131 &mut source_xe,
132 sigma,
133 scratch.borrow(),
134 );
135
136 // GRLWE_{s2}(s1) -> s1 -> s2
137 ct_grlwe_s1s2.encrypt_sk(
138 &module,
139 &sk1.data,
140 &sk2_dft,
141 &mut source_xa,
142 &mut source_xe,
143 sigma,
144 scratch.borrow(),
145 );
146
147 // GRLWE_{s1}(s0) (x) GRLWE_{s2}(s1) = GRLWE_{s2}(s0)
148 ct_grlwe_s0s2.keyswitch(&module, &ct_grlwe_s0s1, &ct_grlwe_s1s2, scratch.borrow());
149
150 let mut ct_rlwe_dft_s0s2: GLWECiphertextFourier<Vec<u8>, FFT64> =
151 GLWECiphertextFourier::new(&module, basek, log_k_grlwe, rank);
152 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, log_k_grlwe);
153
154 (0..ct_grlwe_s0s2.rows()).for_each(|row_i| {
155 ct_grlwe_s0s2.get_row(&module, row_i, 0, &mut ct_rlwe_dft_s0s2);
156 ct_rlwe_dft_s0s2.decrypt(&module, &mut pt, &sk2_dft, scratch.borrow());
157 module.vec_znx_sub_scalar_inplace(&mut pt, 0, row_i, &sk0, 0);
158
159 let noise_have: f64 = pt.data.std(0, basek).log2();
160 let noise_want: f64 = noise_grlwe_rlwe_product(
161 module.n() as f64,
162 basek,
163 0.5,
164 0.5,
165 0f64,
166 sigma * sigma,
167 0f64,
168 log_k_grlwe,
169 log_k_grlwe,
170 );
171
172 assert!(
173 (noise_have - noise_want).abs() <= 0.1,
174 "{} {}",
175 noise_have,
176 noise_want
177 );
178 });
179}
180
181#[test]
182fn keyswitch_inplace() {
183 let module: Module<FFT64> = Module::<FFT64>::new(2048);
184 let basek: usize = 12;
185 let log_k_grlwe: usize = 60;
186 let rows: usize = (log_k_grlwe + basek - 1) / basek;
187
188 let rank: usize = 1;
189 let rank_out: usize = 1;
190
191 let sigma: f64 = 3.2;
192
193 let mut ct_grlwe_s0s1: GLWESwitchingKey<Vec<u8>, FFT64> =
194 GLWESwitchingKey::new(&module, basek, log_k_grlwe, rows, rank, rank_out);
195 let mut ct_grlwe_s1s2: GLWESwitchingKey<Vec<u8>, FFT64> =
196 GLWESwitchingKey::new(&module, basek, log_k_grlwe, rows, rank, rank_out);
197
198 let mut source_xs: Source = Source::new([0u8; 32]);
199 let mut source_xe: Source = Source::new([0u8; 32]);
200 let mut source_xa: Source = Source::new([0u8; 32]);
201
202 let mut scratch: ScratchOwned = ScratchOwned::new(
203 GLWESwitchingKey::encrypt_sk_scratch_space(&module, rank, ct_grlwe_s0s1.size())
204 | GLWECiphertextFourier::decrypt_scratch_space(&module, ct_grlwe_s0s1.size())
205 | GLWESwitchingKey::keyswitch_inplace_scratch_space(&module, ct_grlwe_s0s1.size(), ct_grlwe_s1s2.size()),
206 );
207
208 let mut sk0: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
209 sk0.fill_ternary_prob(0.5, &mut source_xs);
210
211 let mut sk0_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
212 sk0_dft.dft(&module, &sk0);
213
214 let mut sk1: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
215 sk1.fill_ternary_prob(0.5, &mut source_xs);
216
217 let mut sk1_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
218 sk1_dft.dft(&module, &sk1);
219
220 let mut sk2: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
221 sk2.fill_ternary_prob(0.5, &mut source_xs);
222
223 let mut sk2_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
224 sk2_dft.dft(&module, &sk2);
225
226 // GRLWE_{s1}(s0) = s0 -> s1
227 ct_grlwe_s0s1.encrypt_sk(
228 &module,
229 &sk0.data,
230 &sk1_dft,
231 &mut source_xa,
232 &mut source_xe,
233 sigma,
234 scratch.borrow(),
235 );
236
237 // GRLWE_{s2}(s1) -> s1 -> s2
238 ct_grlwe_s1s2.encrypt_sk(
239 &module,
240 &sk1.data,
241 &sk2_dft,
242 &mut source_xa,
243 &mut source_xe,
244 sigma,
245 scratch.borrow(),
246 );
247
248 // GRLWE_{s1}(s0) (x) GRLWE_{s2}(s1) = GRLWE_{s2}(s0)
249 ct_grlwe_s0s1.keyswitch_inplace(&module, &ct_grlwe_s1s2, scratch.borrow());
250
251 let ct_grlwe_s0s2: GLWESwitchingKey<Vec<u8>, FFT64> = ct_grlwe_s0s1;
252
253 let mut ct_rlwe_dft_s0s2: GLWECiphertextFourier<Vec<u8>, FFT64> =
254 GLWECiphertextFourier::new(&module, basek, log_k_grlwe, rank);
255 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, log_k_grlwe);
256
257 (0..ct_grlwe_s0s2.rows()).for_each(|row_i| {
258 ct_grlwe_s0s2.get_row(&module, row_i, 0, &mut ct_rlwe_dft_s0s2);
259 ct_rlwe_dft_s0s2.decrypt(&module, &mut pt, &sk2_dft, scratch.borrow());
260 module.vec_znx_sub_scalar_inplace(&mut pt, 0, row_i, &sk0, 0);
261
262 let noise_have: f64 = pt.data.std(0, basek).log2();
263 let noise_want: f64 = noise_grlwe_rlwe_product(
264 module.n() as f64,
265 basek,
266 0.5,
267 0.5,
268 0f64,
269 sigma * sigma,
270 0f64,
271 log_k_grlwe,
272 log_k_grlwe,
273 );
274
275 assert!(
276 (noise_have - noise_want).abs() <= 0.1,
277 "{} {}",
278 noise_have,
279 noise_want
280 );
281 });
282}
283
284#[test]
285fn external_product() {
286 let module: Module<FFT64> = Module::<FFT64>::new(2048);
287 let basek: usize = 12;
288 let log_k_grlwe: usize = 60;
289 let rows: usize = (log_k_grlwe + basek - 1) / basek;
290
291 let rank: usize = 1;
292 let rank_out: usize = 1;
293
294 let sigma: f64 = 3.2;
295
296 let mut ct_grlwe_in: GLWESwitchingKey<Vec<u8>, FFT64> =
297 GLWESwitchingKey::new(&module, basek, log_k_grlwe, rows, rank, rank_out);
298 let mut ct_grlwe_out: GLWESwitchingKey<Vec<u8>, FFT64> =
299 GLWESwitchingKey::new(&module, basek, log_k_grlwe, rows, rank, rank_out);
300 let mut ct_rgsw: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, log_k_grlwe, rows, rank);
301
302 let mut pt_rgsw: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
303 let mut pt_grlwe: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
304
305 let mut source_xs: Source = Source::new([0u8; 32]);
306 let mut source_xe: Source = Source::new([0u8; 32]);
307 let mut source_xa: Source = Source::new([0u8; 32]);
308
309 let mut scratch: ScratchOwned = ScratchOwned::new(
310 GLWESwitchingKey::encrypt_sk_scratch_space(&module, rank, ct_grlwe_in.size())
311 | GLWECiphertextFourier::decrypt_scratch_space(&module, ct_grlwe_out.size())
312 | GLWESwitchingKey::external_product_scratch_space(
313 &module,
314 ct_grlwe_out.size(),
315 ct_grlwe_in.size(),
316 ct_rgsw.size(),
317 )
318 | GGSWCiphertext::encrypt_sk_scratch_space(&module, rank, ct_rgsw.size()),
319 );
320
321 let k: usize = 1;
322
323 pt_rgsw.raw_mut()[k] = 1; // X^{k}
324
325 pt_grlwe.fill_ternary_prob(0, 0.5, &mut source_xs);
326
327 let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
328 sk.fill_ternary_prob(0.5, &mut source_xs);
329
330 let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
331 sk_dft.dft(&module, &sk);
332
333 // GRLWE_{s1}(s0) = s0 -> s1
334 ct_grlwe_in.encrypt_sk(
335 &module,
336 &pt_grlwe,
337 &sk_dft,
338 &mut source_xa,
339 &mut source_xe,
340 sigma,
341 scratch.borrow(),
342 );
343
344 ct_rgsw.encrypt_sk(
345 &module,
346 &pt_rgsw,
347 &sk_dft,
348 &mut source_xa,
349 &mut source_xe,
350 sigma,
351 scratch.borrow(),
352 );
353
354 // GRLWE_(m) (x) RGSW_(X^k) = GRLWE_(m * X^k)
355 ct_grlwe_out.external_product(&module, &ct_grlwe_in, &ct_rgsw, scratch.borrow());
356
357 let mut ct_rlwe_dft_s0s2: GLWECiphertextFourier<Vec<u8>, FFT64> =
358 GLWECiphertextFourier::new(&module, basek, log_k_grlwe, rank);
359 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, log_k_grlwe);
360
361 module.vec_znx_rotate_inplace(k as i64, &mut pt_grlwe, 0);
362
363 (0..ct_grlwe_out.rows()).for_each(|row_i| {
364 ct_grlwe_out.get_row(&module, row_i, 0, &mut ct_rlwe_dft_s0s2);
365 ct_rlwe_dft_s0s2.decrypt(&module, &mut pt, &sk_dft, scratch.borrow());
366 module.vec_znx_sub_scalar_inplace(&mut pt, 0, row_i, &pt_grlwe, 0);
367
368 let noise_have: f64 = pt.data.std(0, basek).log2();
369
370 let var_gct_err_lhs: f64 = sigma * sigma;
371 let var_gct_err_rhs: f64 = 0f64;
372
373 let var_msg: f64 = 1f64 / module.n() as f64; // X^{k}
374 let var_a0_err: f64 = sigma * sigma;
375 let var_a1_err: f64 = 1f64 / 12f64;
376
377 let noise_want: f64 = noise_rgsw_product(
378 module.n() as f64,
379 basek,
380 0.5,
381 var_msg,
382 var_a0_err,
383 var_a1_err,
384 var_gct_err_lhs,
385 var_gct_err_rhs,
386 log_k_grlwe,
387 log_k_grlwe,
388 );
389
390 assert!(
391 (noise_have - noise_want).abs() <= 0.1,
392 "{} {}",
393 noise_have,
394 noise_want
395 );
396 });
397}
398
399#[test]
400fn external_product_inplace() {
401 let module: Module<FFT64> = Module::<FFT64>::new(2048);
402 let basek: usize = 12;
403 let log_k_grlwe: usize = 60;
404 let rows: usize = (log_k_grlwe + basek - 1) / basek;
405
406 let rank: usize = 1;
407 let rank_out: usize = 1;
408
409 let sigma: f64 = 3.2;
410
411 let mut ct_grlwe: GLWESwitchingKey<Vec<u8>, FFT64> = GLWESwitchingKey::new(&module, basek, log_k_grlwe, rows, rank, rank_out);
412 let mut ct_rgsw: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, basek, log_k_grlwe, rows, rank);
413
414 let mut pt_rgsw: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
415 let mut pt_grlwe: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
416
417 let mut source_xs: Source = Source::new([0u8; 32]);
418 let mut source_xe: Source = Source::new([0u8; 32]);
419 let mut source_xa: Source = Source::new([0u8; 32]);
420
421 let mut scratch: ScratchOwned = ScratchOwned::new(
422 GLWESwitchingKey::encrypt_sk_scratch_space(&module, rank, ct_grlwe.size())
423 | GLWECiphertextFourier::decrypt_scratch_space(&module, ct_grlwe.size())
424 | GLWESwitchingKey::external_product_inplace_scratch_space(&module, ct_grlwe.size(), ct_rgsw.size())
425 | GGSWCiphertext::encrypt_sk_scratch_space(&module, rank, ct_rgsw.size()),
426 );
427
428 let k: usize = 1;
429
430 pt_rgsw.raw_mut()[k] = 1; // X^{k}
431
432 pt_grlwe.fill_ternary_prob(0, 0.5, &mut source_xs);
433
434 let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
435 sk.fill_ternary_prob(0.5, &mut source_xs);
436
437 let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
438 sk_dft.dft(&module, &sk);
439
440 // GRLWE_{s1}(s0) = s0 -> s1
441 ct_grlwe.encrypt_sk(
442 &module,
443 &pt_grlwe,
444 &sk_dft,
445 &mut source_xa,
446 &mut source_xe,
447 sigma,
448 scratch.borrow(),
449 );
450
451 ct_rgsw.encrypt_sk(
452 &module,
453 &pt_rgsw,
454 &sk_dft,
455 &mut source_xa,
456 &mut source_xe,
457 sigma,
458 scratch.borrow(),
459 );
460
461 // GRLWE_(m) (x) RGSW_(X^k) = GRLWE_(m * X^k)
462 ct_grlwe.external_product_inplace(&module, &ct_rgsw, scratch.borrow());
463
464 let mut ct_rlwe_dft_s0s2: GLWECiphertextFourier<Vec<u8>, FFT64> =
465 GLWECiphertextFourier::new(&module, basek, log_k_grlwe, rank);
466 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, log_k_grlwe);
467
468 module.vec_znx_rotate_inplace(k as i64, &mut pt_grlwe, 0);
469
470 (0..ct_grlwe.rows()).for_each(|row_i| {
471 ct_grlwe.get_row(&module, row_i, 0, &mut ct_rlwe_dft_s0s2);
472 ct_rlwe_dft_s0s2.decrypt(&module, &mut pt, &sk_dft, scratch.borrow());
473 module.vec_znx_sub_scalar_inplace(&mut pt, 0, row_i, &pt_grlwe, 0);
474
475 let noise_have: f64 = pt.data.std(0, basek).log2();
476
477 let var_gct_err_lhs: f64 = sigma * sigma;
478 let var_gct_err_rhs: f64 = 0f64;
479
480 let var_msg: f64 = 1f64 / module.n() as f64; // X^{k}
481 let var_a0_err: f64 = sigma * sigma;
482 let var_a1_err: f64 = 1f64 / 12f64;
483
484 let noise_want: f64 = noise_rgsw_product(
485 module.n() as f64,
486 basek,
487 0.5,
488 var_msg,
489 var_a0_err,
490 var_a1_err,
491 var_gct_err_lhs,
492 var_gct_err_rhs,
493 log_k_grlwe,
494 log_k_grlwe,
495 );
496
497 assert!(
498 (noise_have - noise_want).abs() <= 0.1,
499 "{} {}",
500 noise_have,
501 noise_want
502 );
503 });
504}
505pub(crate) fn noise_gglwe_product(
506 n: f64,
507 basek: usize,
508 var_xs: f64,
509 var_msg: f64,
510 var_a_err: f64,
511 var_gct_err_lhs: f64,
512 var_gct_err_rhs: f64,
513 rank_in: f64,
514 a_logq: usize,
515 b_logq: usize,
516) -> f64 {
517 let a_logq: usize = a_logq.min(b_logq);
518 let a_cols: usize = (a_logq + basek - 1) / basek;
519
520 let b_scale = 2.0f64.powi(b_logq as i32);
521 let a_scale: f64 = 2.0f64.powi((b_logq - a_logq) as i32);
522
523 let base: f64 = (1 << (basek)) as f64;
524 let var_base: f64 = base * base / 12f64;
525
526 // lhs = a_cols * n * (var_base * var_gct_err_lhs + var_e_a * var_msg * p^2)
527 // rhs = a_cols * n * var_base * var_gct_err_rhs * var_xs
528 let mut noise: f64 = (a_cols as f64) * n * var_base * (var_gct_err_lhs + var_xs * var_gct_err_rhs);
529 noise += var_msg * var_a_err * a_scale * a_scale * n;
530 noise *= rank_in;
531 noise = noise.sqrt();
532 noise /= b_scale;
533 noise.log2().min(-1.0) // max noise is [-2^{-1}, 2^{-1}]
534}