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