poulpy/core/src/test_fft64/ggsw.rs
core/src/test_fft64/ggsw.rsRaw
1use backend::{
2 FFT64, Module, ScalarZnx, ScalarZnxAlloc, ScalarZnxDftOps, ScalarZnxOps, ScratchOwned, Stats, VecZnxBig, VecZnxBigAlloc,
3 VecZnxBigOps, VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, VecZnxOps, VecZnxToMut, ZnxViewMut, ZnxZero,
4};
5use sampling::source::Source;
6
7use crate::{
8 GGSWCiphertext, GLWECiphertextFourier, GLWEPlaintext, GLWESecret, GLWESwitchingKey, GetRow, Infos, TensorKey,
9 automorphism::AutomorphismKey,
10 div_ceil,
11 test_fft64::{noise_ggsw_keyswitch, noise_ggsw_product},
12};
13
14#[test]
15fn encrypt_sk() {
16 let log_n: usize = 8;
17 let basek: usize = 12;
18 let k_ct: usize = 54;
19 let digits: usize = k_ct / basek;
20 (1..4).for_each(|rank| {
21 (1..digits + 1).for_each(|di| {
22 println!("test encrypt_sk digits: {} rank: {}", di, rank);
23 test_encrypt_sk(log_n, basek, k_ct, di, rank, 3.2);
24 });
25 });
26}
27
28#[test]
29fn keyswitch() {
30 let log_n: usize = 8;
31 let basek: usize = 12;
32 let k_in: usize = 54;
33 let digits: usize = div_ceil(k_in, basek);
34 (1..4).for_each(|rank| {
35 (1..digits + 1).for_each(|di| {
36 let k_ksk: usize = k_in + basek * di;
37 let k_tsk: usize = k_ksk;
38 println!("test keyswitch digits: {} rank: {}", di, rank);
39 let k_out: usize = k_ksk; // Better capture noise.
40 test_keyswitch(log_n, basek, k_out, k_in, k_ksk, k_tsk, di, rank, 3.2);
41 });
42 });
43}
44
45#[test]
46fn keyswitch_inplace() {
47 let log_n: usize = 8;
48 let basek: usize = 12;
49 let k_ct: usize = 54;
50 let digits: usize = div_ceil(k_ct, basek);
51 (1..4).for_each(|rank| {
52 (1..digits + 1).for_each(|di| {
53 let k_ksk: usize = k_ct + basek * di;
54 let k_tsk: usize = k_ksk;
55 println!("test keyswitch_inplace digits: {} rank: {}", di, rank);
56 test_keyswitch_inplace(log_n, basek, k_ct, k_ksk, k_tsk, di, rank, 3.2);
57 });
58 });
59}
60
61#[test]
62fn automorphism() {
63 let log_n: usize = 8;
64 let basek: usize = 12;
65 let k_in: usize = 54;
66 let digits: usize = div_ceil(k_in, basek);
67 (1..4).for_each(|rank| {
68 (1..digits + 1).for_each(|di| {
69 let k_ksk: usize = k_in + basek * di;
70 let k_tsk: usize = k_ksk;
71 println!("test automorphism rank: {}", rank);
72 let k_out: usize = k_ksk; // Better capture noise.
73 test_automorphism(-5, log_n, basek, k_out, k_in, k_ksk, k_tsk, di, rank, 3.2);
74 });
75 });
76}
77
78#[test]
79fn automorphism_inplace() {
80 let log_n: usize = 8;
81 let basek: usize = 12;
82 let k_ct: usize = 54;
83 let digits: usize = div_ceil(k_ct, basek);
84 (1..4).for_each(|rank| {
85 (1..digits + 1).for_each(|di| {
86 let k_ksk: usize = k_ct + basek * di;
87 let k_tsk: usize = k_ksk;
88 println!("test automorphism_inplace rank: {}", rank);
89 test_automorphism_inplace(-5, log_n, basek, k_ct, k_ksk, k_tsk, di, rank, 3.2);
90 });
91 });
92}
93
94#[test]
95fn external_product() {
96 let log_n: usize = 8;
97 let basek: usize = 12;
98 let k_in: usize = 60;
99 let digits: usize = div_ceil(k_in, basek);
100 (1..4).for_each(|rank| {
101 (1..digits + 1).for_each(|di| {
102 let k_ggsw: usize = k_in + basek * di;
103 println!("test external_product digits: {} ranks: {}", di, rank);
104 let k_out: usize = k_in; // Better capture noise.
105 test_external_product(log_n, basek, k_in, k_out, k_ggsw, di, rank, 3.2);
106 });
107 });
108}
109
110#[test]
111fn external_product_inplace() {
112 let log_n: usize = 5;
113 let basek: usize = 12;
114 let k_ct: usize = 60;
115 let digits: usize = div_ceil(k_ct, basek);
116 (1..4).for_each(|rank| {
117 (1..digits).for_each(|di| {
118 let k_ggsw: usize = k_ct + basek * di;
119 println!("test external_product digits: {} rank: {}", di, rank);
120 test_external_product_inplace(log_n, basek, k_ct, k_ggsw, di, rank, 3.2);
121 });
122 });
123}
124
125fn test_encrypt_sk(log_n: usize, basek: usize, k: usize, digits: usize, rank: usize, sigma: f64) {
126 let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
127
128 let rows: usize = (k - digits * basek) / (digits * basek);
129
130 let mut ct: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k, rows, digits, rank);
131 let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k);
132 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k);
133 let mut pt_scalar: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
134
135 let mut source_xs: Source = Source::new([0u8; 32]);
136 let mut source_xe: Source = Source::new([0u8; 32]);
137 let mut source_xa: Source = Source::new([0u8; 32]);
138
139 pt_scalar.fill_ternary_hw(0, module.n(), &mut source_xs);
140
141 let mut scratch: ScratchOwned = ScratchOwned::new(
142 GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, k, rank)
143 | GLWECiphertextFourier::decrypt_scratch_space(&module, basek, k),
144 );
145
146 let mut sk: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
147 sk.fill_ternary_prob(&module, 0.5, &mut source_xs);
148
149 ct.encrypt_sk(
150 &module,
151 &pt_scalar,
152 &sk,
153 &mut source_xa,
154 &mut source_xe,
155 sigma,
156 scratch.borrow(),
157 );
158
159 let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k, rank);
160 let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct.size());
161 let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct.size());
162
163 (0..ct.rank() + 1).for_each(|col_j| {
164 (0..ct.rows()).for_each(|row_i| {
165 module.vec_znx_add_scalar_inplace(
166 &mut pt_want.data,
167 0,
168 (digits - 1) + row_i * digits,
169 &pt_scalar,
170 0,
171 );
172
173 // mul with sk[col_j-1]
174 if col_j > 0 {
175 module.vec_znx_dft(1, 0, &mut pt_dft, 0, &pt_want.data, 0);
176 module.svp_apply_inplace(&mut pt_dft, 0, &sk.data_fourier, col_j - 1);
177 module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
178 module.vec_znx_big_normalize(basek, &mut pt_want.data, 0, &pt_big, 0, scratch.borrow());
179 }
180
181 ct.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
182
183 ct_glwe_fourier.decrypt(&module, &mut pt_have, &sk, scratch.borrow());
184
185 module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt_want.data, 0);
186
187 let std_pt: f64 = pt_have.data.std(0, basek) * (k as f64).exp2();
188 assert!((sigma - std_pt).abs() <= 0.5, "{} {}", sigma, std_pt);
189
190 pt_want.data.zero();
191 });
192 });
193}
194
195fn test_keyswitch(
196 log_n: usize,
197 basek: usize,
198 k_out: usize,
199 k_in: usize,
200 k_ksk: usize,
201 k_tsk: usize,
202 digits: usize,
203 rank: usize,
204 sigma: f64,
205) {
206 let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
207 let rows: usize = div_ceil(k_in, digits * basek);
208
209 let digits_in: usize = 1;
210
211 let mut ct_in: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k_in, rows, digits_in, rank);
212 let mut ct_out: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k_out, rows, digits_in, rank);
213 let mut tsk: TensorKey<Vec<u8>, FFT64> = TensorKey::alloc(&module, basek, k_ksk, rows, digits, rank);
214 let mut ksk: GLWESwitchingKey<Vec<u8>, FFT64> = GLWESwitchingKey::alloc(&module, basek, k_ksk, rows, digits, rank, rank);
215 let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_out);
216 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_out);
217 let mut pt_scalar: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
218
219 let mut source_xs: Source = Source::new([0u8; 32]);
220 let mut source_xe: Source = Source::new([0u8; 32]);
221 let mut source_xa: Source = Source::new([0u8; 32]);
222
223 let mut scratch: ScratchOwned = ScratchOwned::new(
224 GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, k_in, rank)
225 | GLWECiphertextFourier::decrypt_scratch_space(&module, basek, k_out)
226 | GLWESwitchingKey::encrypt_sk_scratch_space(&module, basek, k_ksk, rank)
227 | TensorKey::generate_from_sk_scratch_space(&module, basek, k_tsk, rank)
228 | GGSWCiphertext::keyswitch_scratch_space(
229 &module, basek, k_out, k_in, k_ksk, digits, k_tsk, digits, rank,
230 ),
231 );
232
233 let var_xs: f64 = 0.5;
234
235 let mut sk_in: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
236 sk_in.fill_ternary_prob(&module, var_xs, &mut source_xs);
237
238 let mut sk_out: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
239 sk_out.fill_ternary_prob(&module, var_xs, &mut source_xs);
240
241 ksk.generate_from_sk(
242 &module,
243 &sk_in,
244 &sk_out,
245 &mut source_xa,
246 &mut source_xe,
247 sigma,
248 scratch.borrow(),
249 );
250 tsk.generate_from_sk(
251 &module,
252 &sk_out,
253 &mut source_xa,
254 &mut source_xe,
255 sigma,
256 scratch.borrow(),
257 );
258
259 pt_scalar.fill_ternary_hw(0, module.n(), &mut source_xs);
260
261 ct_in.encrypt_sk(
262 &module,
263 &pt_scalar,
264 &sk_in,
265 &mut source_xa,
266 &mut source_xe,
267 sigma,
268 scratch.borrow(),
269 );
270
271 ct_out.keyswitch(&module, &ct_in, &ksk, &tsk, scratch.borrow());
272
273 let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k_out, rank);
274 let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct_out.size());
275 let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct_out.size());
276
277 (0..ct_out.rank() + 1).for_each(|col_j| {
278 (0..ct_out.rows()).for_each(|row_i| {
279 module.vec_znx_add_scalar_inplace(&mut pt_want.data, 0, row_i, &pt_scalar, 0);
280
281 // mul with sk[col_j-1]
282 if col_j > 0 {
283 module.vec_znx_dft(1, 0, &mut pt_dft, 0, &pt_want.data, 0);
284 module.svp_apply_inplace(&mut pt_dft, 0, &sk_out.data_fourier, col_j - 1);
285 module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
286 module.vec_znx_big_normalize(basek, &mut pt_want.data, 0, &pt_big, 0, scratch.borrow());
287 }
288
289 ct_out.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
290
291 ct_glwe_fourier.decrypt(&module, &mut pt_have, &sk_out, scratch.borrow());
292
293 module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt_want.data, 0);
294
295 let noise_have: f64 = pt_have.data.std(0, basek).log2();
296 let noise_want: f64 = noise_ggsw_keyswitch(
297 module.n() as f64,
298 basek * digits,
299 col_j,
300 var_xs,
301 0f64,
302 sigma * sigma,
303 0f64,
304 rank as f64,
305 k_in,
306 k_ksk,
307 k_tsk,
308 );
309
310 println!("{} {}", noise_have, noise_want);
311
312 assert!(
313 noise_have < noise_want + 0.5,
314 "{} {}",
315 noise_have,
316 noise_want
317 );
318
319 pt_want.data.zero();
320 });
321 });
322}
323
324fn test_keyswitch_inplace(
325 log_n: usize,
326 basek: usize,
327 k_ct: usize,
328 k_ksk: usize,
329 k_tsk: usize,
330 digits: usize,
331 rank: usize,
332 sigma: f64,
333) {
334 let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
335 let rows: usize = div_ceil(k_ct, digits * basek);
336
337 let digits_in: usize = 1;
338
339 let mut ct: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k_ct, rows, digits_in, rank);
340 let mut tsk: TensorKey<Vec<u8>, FFT64> = TensorKey::alloc(&module, basek, k_tsk, rows, digits, rank);
341 let mut ksk: GLWESwitchingKey<Vec<u8>, FFT64> = GLWESwitchingKey::alloc(&module, basek, k_ksk, rows, digits, rank, rank);
342 let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
343 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
344 let mut pt_scalar: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
345
346 let mut source_xs: Source = Source::new([0u8; 32]);
347 let mut source_xe: Source = Source::new([0u8; 32]);
348 let mut source_xa: Source = Source::new([0u8; 32]);
349
350 let mut scratch: ScratchOwned = ScratchOwned::new(
351 GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, k_ct, rank)
352 | GLWECiphertextFourier::decrypt_scratch_space(&module, basek, k_ct)
353 | GLWESwitchingKey::encrypt_sk_scratch_space(&module, basek, k_ksk, rank)
354 | TensorKey::generate_from_sk_scratch_space(&module, basek, k_tsk, rank)
355 | GGSWCiphertext::keyswitch_inplace_scratch_space(&module, basek, k_ct, k_ksk, digits, k_tsk, digits, rank),
356 );
357
358 let var_xs: f64 = 0.5;
359
360 let mut sk_in: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
361 sk_in.fill_ternary_prob(&module, var_xs, &mut source_xs);
362
363 let mut sk_out: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
364 sk_out.fill_ternary_prob(&module, var_xs, &mut source_xs);
365
366 ksk.generate_from_sk(
367 &module,
368 &sk_in,
369 &sk_out,
370 &mut source_xa,
371 &mut source_xe,
372 sigma,
373 scratch.borrow(),
374 );
375 tsk.generate_from_sk(
376 &module,
377 &sk_out,
378 &mut source_xa,
379 &mut source_xe,
380 sigma,
381 scratch.borrow(),
382 );
383
384 pt_scalar.fill_ternary_hw(0, module.n(), &mut source_xs);
385
386 ct.encrypt_sk(
387 &module,
388 &pt_scalar,
389 &sk_in,
390 &mut source_xa,
391 &mut source_xe,
392 sigma,
393 scratch.borrow(),
394 );
395
396 ct.keyswitch_inplace(&module, &ksk, &tsk, scratch.borrow());
397
398 let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k_ct, rank);
399 let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct.size());
400 let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct.size());
401
402 (0..ct.rank() + 1).for_each(|col_j| {
403 (0..ct.rows()).for_each(|row_i| {
404 module.vec_znx_add_scalar_inplace(
405 &mut pt_want.data,
406 0,
407 (digits_in - 1) + row_i * digits_in,
408 &pt_scalar,
409 0,
410 );
411
412 // mul with sk[col_j-1]
413 if col_j > 0 {
414 module.vec_znx_dft(1, 0, &mut pt_dft, 0, &pt_want.data, 0);
415 module.svp_apply_inplace(&mut pt_dft, 0, &sk_out.data_fourier, col_j - 1);
416 module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
417 module.vec_znx_big_normalize(basek, &mut pt_want.data, 0, &pt_big, 0, scratch.borrow());
418 }
419
420 ct.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
421
422 ct_glwe_fourier.decrypt(&module, &mut pt_have, &sk_out, scratch.borrow());
423
424 module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt_want.data, 0);
425
426 let noise_have: f64 = pt_have.data.std(0, basek).log2();
427 let noise_want: f64 = noise_ggsw_keyswitch(
428 module.n() as f64,
429 basek * digits,
430 col_j,
431 var_xs,
432 0f64,
433 sigma * sigma,
434 0f64,
435 rank as f64,
436 k_ct,
437 k_ksk,
438 k_tsk,
439 );
440
441 println!("{} {}", noise_have, noise_want);
442
443 assert!(
444 noise_have < noise_want + 0.5,
445 "{} {}",
446 noise_have,
447 noise_want
448 );
449
450 pt_want.data.zero();
451 });
452 });
453}
454
455fn test_automorphism(
456 p: i64,
457 log_n: usize,
458 basek: usize,
459 k_out: usize,
460 k_in: usize,
461 k_ksk: usize,
462 k_tsk: usize,
463 digits: usize,
464 rank: usize,
465 sigma: f64,
466) {
467 let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
468 let rows: usize = div_ceil(k_in, basek * digits);
469 let rows_in: usize = k_in / (basek * digits);
470
471 let digits_in: usize = 1;
472
473 let mut ct_in: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k_in, rows_in, digits_in, rank);
474 let mut ct_out: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k_out, rows_in, digits_in, rank);
475 let mut tensor_key: TensorKey<Vec<u8>, FFT64> = TensorKey::alloc(&module, basek, k_tsk, rows, digits, rank);
476 let mut auto_key: AutomorphismKey<Vec<u8>, FFT64> = AutomorphismKey::alloc(&module, basek, k_ksk, rows, digits, rank);
477 let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_out);
478 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_out);
479 let mut pt_scalar: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
480
481 let mut source_xs: Source = Source::new([0u8; 32]);
482 let mut source_xe: Source = Source::new([0u8; 32]);
483 let mut source_xa: Source = Source::new([0u8; 32]);
484
485 let mut scratch: ScratchOwned = ScratchOwned::new(
486 GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, k_in, rank)
487 | GLWECiphertextFourier::decrypt_scratch_space(&module, basek, k_out)
488 | AutomorphismKey::generate_from_sk_scratch_space(&module, basek, k_ksk, rank)
489 | TensorKey::generate_from_sk_scratch_space(&module, basek, k_tsk, rank)
490 | GGSWCiphertext::automorphism_scratch_space(
491 &module, basek, k_out, k_in, k_ksk, digits, k_tsk, digits, rank,
492 ),
493 );
494
495 let var_xs: f64 = 0.5;
496
497 let mut sk: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
498 sk.fill_ternary_prob(&module, var_xs, &mut source_xs);
499
500 auto_key.generate_from_sk(
501 &module,
502 p,
503 &sk,
504 &mut source_xa,
505 &mut source_xe,
506 sigma,
507 scratch.borrow(),
508 );
509 tensor_key.generate_from_sk(
510 &module,
511 &sk,
512 &mut source_xa,
513 &mut source_xe,
514 sigma,
515 scratch.borrow(),
516 );
517
518 pt_scalar.fill_ternary_hw(0, module.n(), &mut source_xs);
519
520 ct_in.encrypt_sk(
521 &module,
522 &pt_scalar,
523 &sk,
524 &mut source_xa,
525 &mut source_xe,
526 sigma,
527 scratch.borrow(),
528 );
529
530 ct_out.automorphism(&module, &ct_in, &auto_key, &tensor_key, scratch.borrow());
531
532 module.scalar_znx_automorphism_inplace(p, &mut pt_scalar, 0);
533
534 let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k_out, rank);
535 let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct_out.size());
536 let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct_out.size());
537
538 (0..ct_out.rank() + 1).for_each(|col_j| {
539 (0..ct_out.rows()).for_each(|row_i| {
540 module.vec_znx_add_scalar_inplace(&mut pt_want.data, 0, row_i, &pt_scalar, 0);
541
542 // mul with sk[col_j-1]
543 if col_j > 0 {
544 module.vec_znx_dft(1, 0, &mut pt_dft, 0, &pt_want.data, 0);
545 module.svp_apply_inplace(&mut pt_dft, 0, &sk.data_fourier, col_j - 1);
546 module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
547 module.vec_znx_big_normalize(basek, &mut pt_want.data, 0, &pt_big, 0, scratch.borrow());
548 }
549
550 ct_out.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
551
552 ct_glwe_fourier.decrypt(&module, &mut pt_have, &sk, scratch.borrow());
553
554 module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt_want.data, 0);
555
556 let noise_have: f64 = pt_have.data.std(0, basek).log2();
557 let noise_want: f64 = noise_ggsw_keyswitch(
558 module.n() as f64,
559 basek * digits,
560 col_j,
561 var_xs,
562 0f64,
563 sigma * sigma,
564 0f64,
565 rank as f64,
566 k_in,
567 k_ksk,
568 k_tsk,
569 );
570
571 assert!(
572 noise_have < noise_want + 0.5,
573 "{} {}",
574 noise_have,
575 noise_want
576 );
577
578 pt_want.data.zero();
579 });
580 });
581}
582
583fn test_automorphism_inplace(
584 p: i64,
585 log_n: usize,
586 basek: usize,
587 k_ct: usize,
588 k_ksk: usize,
589 k_tsk: usize,
590 digits: usize,
591 rank: usize,
592 sigma: f64,
593) {
594 let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
595 let rows: usize = div_ceil(k_ct, digits * basek);
596 let rows_in: usize = k_ct / (basek * digits);
597 let digits_in: usize = 1;
598
599 let mut ct: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k_ct, rows_in, digits_in, rank);
600 let mut tensor_key: TensorKey<Vec<u8>, FFT64> = TensorKey::alloc(&module, basek, k_tsk, rows, digits, rank);
601 let mut auto_key: AutomorphismKey<Vec<u8>, FFT64> = AutomorphismKey::alloc(&module, basek, k_ksk, rows, digits, rank);
602 let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
603 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
604 let mut pt_scalar: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
605
606 let mut source_xs: Source = Source::new([0u8; 32]);
607 let mut source_xe: Source = Source::new([0u8; 32]);
608 let mut source_xa: Source = Source::new([0u8; 32]);
609
610 let mut scratch: ScratchOwned = ScratchOwned::new(
611 GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, k_ct, rank)
612 | GLWECiphertextFourier::decrypt_scratch_space(&module, basek, k_ct)
613 | AutomorphismKey::generate_from_sk_scratch_space(&module, basek, k_ksk, rank)
614 | TensorKey::generate_from_sk_scratch_space(&module, basek, k_tsk, rank)
615 | GGSWCiphertext::automorphism_inplace_scratch_space(&module, basek, k_ct, k_ksk, digits, k_tsk, digits, rank),
616 );
617
618 let var_xs: f64 = 0.5;
619
620 let mut sk: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
621 sk.fill_ternary_prob(&module, var_xs, &mut source_xs);
622
623 auto_key.generate_from_sk(
624 &module,
625 p,
626 &sk,
627 &mut source_xa,
628 &mut source_xe,
629 sigma,
630 scratch.borrow(),
631 );
632 tensor_key.generate_from_sk(
633 &module,
634 &sk,
635 &mut source_xa,
636 &mut source_xe,
637 sigma,
638 scratch.borrow(),
639 );
640
641 pt_scalar.fill_ternary_hw(0, module.n(), &mut source_xs);
642
643 ct.encrypt_sk(
644 &module,
645 &pt_scalar,
646 &sk,
647 &mut source_xa,
648 &mut source_xe,
649 sigma,
650 scratch.borrow(),
651 );
652
653 ct.automorphism_inplace(&module, &auto_key, &tensor_key, scratch.borrow());
654
655 module.scalar_znx_automorphism_inplace(p, &mut pt_scalar, 0);
656
657 let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k_ct, rank);
658 let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct.size());
659 let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct.size());
660
661 (0..ct.rank() + 1).for_each(|col_j| {
662 (0..ct.rows()).for_each(|row_i| {
663 module.vec_znx_add_scalar_inplace(&mut pt_want.data, 0, row_i, &pt_scalar, 0);
664
665 // mul with sk[col_j-1]
666 if col_j > 0 {
667 module.vec_znx_dft(1, 0, &mut pt_dft, 0, &pt_want.data, 0);
668 module.svp_apply_inplace(&mut pt_dft, 0, &sk.data_fourier, col_j - 1);
669 module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
670 module.vec_znx_big_normalize(basek, &mut pt_want.data, 0, &pt_big, 0, scratch.borrow());
671 }
672
673 ct.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
674
675 ct_glwe_fourier.decrypt(&module, &mut pt_have, &sk, scratch.borrow());
676
677 module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt_want.data, 0);
678
679 let noise_have: f64 = pt_have.data.std(0, basek).log2();
680 let noise_want: f64 = noise_ggsw_keyswitch(
681 module.n() as f64,
682 basek * digits,
683 col_j,
684 var_xs,
685 0f64,
686 sigma * sigma,
687 0f64,
688 rank as f64,
689 k_ct,
690 k_ksk,
691 k_tsk,
692 );
693
694 assert!(
695 noise_have <= noise_want + 0.5,
696 "{} {}",
697 noise_have,
698 noise_want
699 );
700
701 pt_want.data.zero();
702 });
703 });
704}
705
706fn test_external_product(
707 log_n: usize,
708 basek: usize,
709 k_in: usize,
710 k_out: usize,
711 k_ggsw: usize,
712 digits: usize,
713 rank: usize,
714 sigma: f64,
715) {
716 let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
717
718 let rows: usize = div_ceil(k_in, basek * digits);
719 let rows_in: usize = k_in / (basek * digits);
720 let digits_in: usize = 1;
721
722 let mut ct_ggsw_lhs_in: GGSWCiphertext<Vec<u8>, FFT64> =
723 GGSWCiphertext::alloc(&module, basek, k_in, rows_in, digits_in, rank);
724 let mut ct_ggsw_lhs_out: GGSWCiphertext<Vec<u8>, FFT64> =
725 GGSWCiphertext::alloc(&module, basek, k_out, rows_in, digits_in, rank);
726 let mut ct_ggsw_rhs: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k_ggsw, rows, digits, rank);
727 let mut pt_ggsw_lhs: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
728 let mut pt_ggsw_rhs: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
729
730 let mut source_xs: Source = Source::new([0u8; 32]);
731 let mut source_xe: Source = Source::new([0u8; 32]);
732 let mut source_xa: Source = Source::new([0u8; 32]);
733
734 pt_ggsw_lhs.fill_ternary_prob(0, 0.5, &mut source_xs);
735
736 let k: usize = 1;
737
738 pt_ggsw_rhs.to_mut().raw_mut()[k] = 1; //X^{k}
739
740 let mut scratch: ScratchOwned = ScratchOwned::new(
741 GLWECiphertextFourier::decrypt_scratch_space(&module, basek, k_out)
742 | GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, k_ggsw, rank)
743 | GGSWCiphertext::external_product_scratch_space(&module, basek, k_out, k_in, k_ggsw, digits, rank),
744 );
745
746 let mut sk: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
747 sk.fill_ternary_prob(&module, 0.5, &mut source_xs);
748
749 ct_ggsw_rhs.encrypt_sk(
750 &module,
751 &pt_ggsw_rhs,
752 &sk,
753 &mut source_xa,
754 &mut source_xe,
755 sigma,
756 scratch.borrow(),
757 );
758
759 ct_ggsw_lhs_in.encrypt_sk(
760 &module,
761 &pt_ggsw_lhs,
762 &sk,
763 &mut source_xa,
764 &mut source_xe,
765 sigma,
766 scratch.borrow(),
767 );
768
769 ct_ggsw_lhs_out.external_product(&module, &ct_ggsw_lhs_in, &ct_ggsw_rhs, scratch.borrow());
770
771 let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k_out, rank);
772 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_out);
773 let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct_ggsw_lhs_out.size());
774 let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct_ggsw_lhs_out.size());
775 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_out);
776
777 module.vec_znx_rotate_inplace(k as i64, &mut pt_ggsw_lhs, 0);
778
779 (0..ct_ggsw_lhs_out.rank() + 1).for_each(|col_j| {
780 (0..ct_ggsw_lhs_out.rows()).for_each(|row_i| {
781 module.vec_znx_add_scalar_inplace(
782 &mut pt_want.data,
783 0,
784 (digits_in - 1) + row_i * digits_in,
785 &pt_ggsw_lhs,
786 0,
787 );
788
789 if col_j > 0 {
790 module.vec_znx_dft(1, 0, &mut pt_dft, 0, &pt_want.data, 0);
791 module.svp_apply_inplace(&mut pt_dft, 0, &sk.data_fourier, col_j - 1);
792 module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
793 module.vec_znx_big_normalize(basek, &mut pt_want.data, 0, &pt_big, 0, scratch.borrow());
794 }
795
796 ct_ggsw_lhs_out.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
797 ct_glwe_fourier.decrypt(&module, &mut pt, &sk, scratch.borrow());
798
799 module.vec_znx_sub_ab_inplace(&mut pt.data, 0, &pt_want.data, 0);
800
801 let noise_have: f64 = pt.data.std(0, basek).log2();
802
803 let var_gct_err_lhs: f64 = sigma * sigma;
804 let var_gct_err_rhs: f64 = 0f64;
805
806 let var_msg: f64 = 1f64 / module.n() as f64; // X^{k}
807 let var_a0_err: f64 = sigma * sigma;
808 let var_a1_err: f64 = 1f64 / 12f64;
809
810 let noise_want: f64 = noise_ggsw_product(
811 module.n() as f64,
812 basek * digits,
813 0.5,
814 var_msg,
815 var_a0_err,
816 var_a1_err,
817 var_gct_err_lhs,
818 var_gct_err_rhs,
819 rank as f64,
820 k_in,
821 k_ggsw,
822 );
823
824 assert!(
825 noise_have <= noise_want + 0.5,
826 "have: {} want: {}",
827 noise_have,
828 noise_want
829 );
830
831 println!("{} {}", noise_have, noise_want);
832
833 pt_want.data.zero();
834 });
835 });
836}
837
838fn test_external_product_inplace(log_n: usize, basek: usize, k_ct: usize, k_ggsw: usize, digits: usize, rank: usize, sigma: f64) {
839 let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
840 let rows: usize = div_ceil(k_ct, digits * basek);
841 let rows_in: usize = k_ct / (basek * digits);
842 let digits_in: usize = 1;
843
844 let mut ct_ggsw_lhs: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k_ct, rows_in, digits_in, rank);
845 let mut ct_ggsw_rhs: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k_ggsw, rows, digits, rank);
846
847 let mut pt_ggsw_lhs: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
848 let mut pt_ggsw_rhs: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
849
850 let mut source_xs: Source = Source::new([0u8; 32]);
851 let mut source_xe: Source = Source::new([0u8; 32]);
852 let mut source_xa: Source = Source::new([0u8; 32]);
853
854 pt_ggsw_lhs.fill_ternary_prob(0, 0.5, &mut source_xs);
855
856 let k: usize = 1;
857
858 pt_ggsw_rhs.to_mut().raw_mut()[k] = 1; //X^{k}
859
860 let mut scratch: ScratchOwned = ScratchOwned::new(
861 GLWECiphertextFourier::decrypt_scratch_space(&module, basek, k_ct)
862 | GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, k_ggsw, rank)
863 | GGSWCiphertext::external_product_inplace_scratch_space(&module, basek, k_ct, k_ggsw, digits, rank),
864 );
865
866 let mut sk: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
867 sk.fill_ternary_prob(&module, 0.5, &mut source_xs);
868
869 ct_ggsw_rhs.encrypt_sk(
870 &module,
871 &pt_ggsw_rhs,
872 &sk,
873 &mut source_xa,
874 &mut source_xe,
875 sigma,
876 scratch.borrow(),
877 );
878
879 ct_ggsw_lhs.encrypt_sk(
880 &module,
881 &pt_ggsw_lhs,
882 &sk,
883 &mut source_xa,
884 &mut source_xe,
885 sigma,
886 scratch.borrow(),
887 );
888
889 ct_ggsw_lhs.external_product_inplace(&module, &ct_ggsw_rhs, scratch.borrow());
890
891 let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k_ct, rank);
892 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
893 let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct_ggsw_lhs.size());
894 let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct_ggsw_lhs.size());
895 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
896
897 module.vec_znx_rotate_inplace(k as i64, &mut pt_ggsw_lhs, 0);
898
899 (0..ct_ggsw_lhs.rank() + 1).for_each(|col_j| {
900 (0..ct_ggsw_lhs.rows()).for_each(|row_i| {
901 module.vec_znx_add_scalar_inplace(
902 &mut pt_want.data,
903 0,
904 (digits_in - 1) + row_i * digits_in,
905 &pt_ggsw_lhs,
906 0,
907 );
908
909 if col_j > 0 {
910 module.vec_znx_dft(1, 0, &mut pt_dft, 0, &pt_want.data, 0);
911 module.svp_apply_inplace(&mut pt_dft, 0, &sk.data_fourier, col_j - 1);
912 module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
913 module.vec_znx_big_normalize(basek, &mut pt_want.data, 0, &pt_big, 0, scratch.borrow());
914 }
915
916 ct_ggsw_lhs.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
917 ct_glwe_fourier.decrypt(&module, &mut pt, &sk, scratch.borrow());
918
919 module.vec_znx_sub_ab_inplace(&mut pt.data, 0, &pt_want.data, 0);
920
921 let noise_have: f64 = pt.data.std(0, basek).log2();
922
923 let var_gct_err_lhs: f64 = sigma * sigma;
924 let var_gct_err_rhs: f64 = 0f64;
925
926 let var_msg: f64 = 1f64 / module.n() as f64; // X^{k}
927 let var_a0_err: f64 = sigma * sigma;
928 let var_a1_err: f64 = 1f64 / 12f64;
929
930 let noise_want: f64 = noise_ggsw_product(
931 module.n() as f64,
932 basek * digits,
933 0.5,
934 var_msg,
935 var_a0_err,
936 var_a1_err,
937 var_gct_err_lhs,
938 var_gct_err_rhs,
939 rank as f64,
940 k_ct,
941 k_ggsw,
942 );
943
944 assert!(
945 noise_have <= noise_want + 0.5,
946 "have: {} want: {}",
947 noise_have,
948 noise_want
949 );
950
951 pt_want.data.zero();
952 });
953 });
954}