poulpy/core/src/test_fft64/glwe_fourier.rs
core/src/test_fft64/glwe_fourier.rsRaw
1use crate::{
2 GGSWCiphertext, GLWECiphertext, GLWECiphertextFourier, GLWEOps, GLWEPlaintext, GLWESecret, GLWESwitchingKey, Infos, div_ceil,
3 test_fft64::{log2_std_noise_gglwe_product, noise_ggsw_product},
4};
5use backend::{FFT64, FillUniform, Module, ScalarZnx, ScalarZnxAlloc, ScratchOwned, Stats, VecZnxOps, ZnxViewMut};
6use sampling::source::Source;
7
8#[test]
9fn keyswitch() {
10 let log_n: usize = 8;
11 let basek: usize = 12;
12 let k_in: usize = 45;
13 let digits: usize = div_ceil(k_in, basek);
14 (1..4).for_each(|rank_in| {
15 (1..4).for_each(|rank_out| {
16 (1..digits + 1).for_each(|di| {
17 let k_ksk: usize = k_in + basek * di;
18 println!(
19 "test keyswitch digits: {} rank_in: {} rank_out: {}",
20 di, rank_in, rank_out
21 );
22 let k_out: usize = k_ksk; // Better capture noise.
23 test_keyswitch(log_n, basek, k_in, k_out, k_ksk, di, rank_in, rank_out, 3.2);
24 })
25 });
26 });
27}
28
29#[test]
30fn keyswitch_inplace() {
31 let log_n: usize = 8;
32 let basek: usize = 12;
33 let k_ct: usize = 45;
34 let digits: usize = div_ceil(k_ct, basek);
35 (1..4).for_each(|rank| {
36 (1..digits + 1).for_each(|di| {
37 let k_ksk: usize = k_ct + basek * di;
38 println!("test keyswitch_inplace digits: {} rank: {}", di, rank);
39 test_keyswitch_inplace(log_n, basek, k_ct, k_ksk, di, rank, 3.2);
40 });
41 });
42}
43
44#[test]
45fn external_product() {
46 let log_n: usize = 8;
47 let basek: usize = 12;
48 let k_in: usize = 45;
49 let digits: usize = div_ceil(k_in, basek);
50 (1..4).for_each(|rank| {
51 (1..digits + 1).for_each(|di| {
52 let k_ggsw: usize = k_in + basek * di;
53 println!("test external_product digits: {} rank: {}", di, rank);
54 let k_out: usize = k_ggsw; // Better capture noise.
55 test_external_product(log_n, basek, k_out, k_in, k_ggsw, di, rank, 3.2);
56 });
57 });
58}
59
60#[test]
61fn external_product_inplace() {
62 let log_n: usize = 8;
63 let basek: usize = 12;
64 let k_ct: usize = 60;
65 let digits: usize = div_ceil(k_ct, basek);
66 (1..4).for_each(|rank| {
67 (1..digits + 1).for_each(|di| {
68 let k_ggsw: usize = k_ct + basek * di;
69 println!("test external_product digits: {} rank: {}", di, rank);
70 test_external_product_inplace(log_n, basek, k_ct, k_ggsw, di, rank, 3.2);
71 });
72 });
73}
74
75fn test_keyswitch(
76 log_n: usize,
77 basek: usize,
78 k_in: usize,
79 k_out: usize,
80 k_ksk: usize,
81 digits: usize,
82 rank_in: usize,
83 rank_out: usize,
84 sigma: f64,
85) {
86 let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
87
88 let rows: usize = div_ceil(k_in, basek * digits);
89
90 let mut ksk: GLWESwitchingKey<Vec<u8>, FFT64> =
91 GLWESwitchingKey::alloc(&module, basek, k_ksk, rows, digits, rank_in, rank_out);
92 let mut ct_glwe_in: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_in, rank_in);
93 let mut ct_glwe_dft_in: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k_in, rank_in);
94 let mut ct_glwe_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_out, rank_out);
95 let mut ct_glwe_dft_out: GLWECiphertextFourier<Vec<u8>, FFT64> =
96 GLWECiphertextFourier::alloc(&module, basek, k_out, rank_out);
97 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_in);
98 let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_out);
99
100 let mut source_xs: Source = Source::new([0u8; 32]);
101 let mut source_xe: Source = Source::new([0u8; 32]);
102 let mut source_xa: Source = Source::new([0u8; 32]);
103
104 // Random input plaintext
105 pt_want
106 .data
107 .fill_uniform(basek, 0, pt_want.size(), &mut source_xa);
108
109 let mut scratch: ScratchOwned = ScratchOwned::new(
110 GLWESwitchingKey::encrypt_sk_scratch_space(&module, basek, k_ksk, rank_out)
111 | GLWECiphertext::decrypt_scratch_space(&module, basek, k_out)
112 | GLWECiphertext::encrypt_sk_scratch_space(&module, basek, k_in)
113 | GLWECiphertextFourier::keyswitch_scratch_space(
114 &module,
115 basek,
116 ct_glwe_out.k(),
117 ksk.k(),
118 ct_glwe_in.k(),
119 digits,
120 rank_in,
121 rank_out,
122 ),
123 );
124
125 let mut sk_in: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank_in);
126 sk_in.fill_ternary_prob(&module, 0.5, &mut source_xs);
127
128 let mut sk_out: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank_out);
129 sk_out.fill_ternary_prob(&module, 0.5, &mut source_xs);
130
131 ksk.generate_from_sk(
132 &module,
133 &sk_in,
134 &sk_out,
135 &mut source_xa,
136 &mut source_xe,
137 sigma,
138 scratch.borrow(),
139 );
140
141 ct_glwe_in.encrypt_sk(
142 &module,
143 &pt_want,
144 &sk_in,
145 &mut source_xa,
146 &mut source_xe,
147 sigma,
148 scratch.borrow(),
149 );
150
151 ct_glwe_in.dft(&module, &mut ct_glwe_dft_in);
152 ct_glwe_dft_out.keyswitch(&module, &ct_glwe_dft_in, &ksk, scratch.borrow());
153 ct_glwe_dft_out.idft(&module, &mut ct_glwe_out, scratch.borrow());
154
155 ct_glwe_out.decrypt(&module, &mut pt_have, &sk_out, scratch.borrow());
156
157 module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt_want.data, 0);
158
159 let noise_have: f64 = pt_have.data.std(0, basek).log2();
160 let noise_want: f64 = log2_std_noise_gglwe_product(
161 module.n() as f64,
162 basek * digits,
163 0.5,
164 0.5,
165 0f64,
166 sigma * sigma,
167 0f64,
168 rank_in as f64,
169 k_in,
170 k_ksk,
171 );
172
173 assert!(
174 (noise_have - noise_want).abs() <= 0.5,
175 "{} {}",
176 noise_have,
177 noise_want
178 );
179}
180
181fn test_keyswitch_inplace(log_n: usize, basek: usize, k_ct: usize, k_ksk: usize, digits: usize, rank: usize, sigma: f64) {
182 let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
183
184 let rows: usize = div_ceil(k_ct, basek * digits);
185
186 let mut ksk: GLWESwitchingKey<Vec<u8>, FFT64> = GLWESwitchingKey::alloc(&module, basek, k_ksk, rows, digits, rank, rank);
187 let mut ct_glwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_ct, rank);
188 let mut ct_rlwe_dft: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k_ct, rank);
189 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
190 let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
191
192 let mut source_xs: Source = Source::new([0u8; 32]);
193 let mut source_xe: Source = Source::new([0u8; 32]);
194 let mut source_xa: Source = Source::new([0u8; 32]);
195
196 // Random input plaintext
197 pt_want
198 .data
199 .fill_uniform(basek, 0, pt_want.size(), &mut source_xa);
200
201 let mut scratch: ScratchOwned = ScratchOwned::new(
202 GLWESwitchingKey::encrypt_sk_scratch_space(&module, basek, ksk.k(), rank)
203 | GLWECiphertext::decrypt_scratch_space(&module, basek, ct_glwe.k())
204 | GLWECiphertext::encrypt_sk_scratch_space(&module, basek, ct_glwe.k())
205 | GLWECiphertextFourier::keyswitch_inplace_scratch_space(&module, basek, ct_rlwe_dft.k(), ksk.k(), digits, rank),
206 );
207
208 let mut sk_in: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
209 sk_in.fill_ternary_prob(&module, 0.5, &mut source_xs);
210
211 let mut sk_out: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
212 sk_out.fill_ternary_prob(&module, 0.5, &mut source_xs);
213
214 ksk.generate_from_sk(
215 &module,
216 &sk_in,
217 &sk_out,
218 &mut source_xa,
219 &mut source_xe,
220 sigma,
221 scratch.borrow(),
222 );
223
224 ct_glwe.encrypt_sk(
225 &module,
226 &pt_want,
227 &sk_in,
228 &mut source_xa,
229 &mut source_xe,
230 sigma,
231 scratch.borrow(),
232 );
233
234 ct_glwe.dft(&module, &mut ct_rlwe_dft);
235 ct_rlwe_dft.keyswitch_inplace(&module, &ksk, scratch.borrow());
236 ct_rlwe_dft.idft(&module, &mut ct_glwe, scratch.borrow());
237
238 ct_glwe.decrypt(&module, &mut pt_have, &sk_out, scratch.borrow());
239
240 module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt_want.data, 0);
241
242 let noise_have: f64 = pt_have.data.std(0, basek).log2();
243 let noise_want: f64 = log2_std_noise_gglwe_product(
244 module.n() as f64,
245 basek * digits,
246 0.5,
247 0.5,
248 0f64,
249 sigma * sigma,
250 0f64,
251 rank as f64,
252 k_ct,
253 k_ksk,
254 );
255
256 assert!(
257 (noise_have - noise_want).abs() <= 0.5,
258 "{} {}",
259 noise_have,
260 noise_want
261 );
262}
263
264fn test_external_product(
265 log_n: usize,
266 basek: usize,
267 k_out: usize,
268 k_in: usize,
269 k_ggsw: usize,
270 digits: usize,
271 rank: usize,
272 sigma: f64,
273) {
274 let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
275
276 let rows: usize = div_ceil(k_in, digits * basek);
277
278 let mut ct_ggsw: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k_ggsw, rows, digits, rank);
279 let mut ct_in: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_in, rank);
280 let mut ct_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_out, rank);
281 let mut ct_in_dft: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k_in, rank);
282 let mut ct_out_dft: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k_out, rank);
283 let mut pt_rgsw: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
284 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_in);
285 let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_out);
286
287 let mut source_xs: Source = Source::new([0u8; 32]);
288 let mut source_xe: Source = Source::new([0u8; 32]);
289 let mut source_xa: Source = Source::new([0u8; 32]);
290
291 // Random input plaintext
292 pt_want
293 .data
294 .fill_uniform(basek, 0, pt_want.size(), &mut source_xa);
295
296 pt_want.data.at_mut(0, 0)[1] = 1;
297
298 let k: i64 = 1;
299
300 pt_rgsw.raw_mut()[0] = 1; // X^{0}
301 module.vec_znx_rotate_inplace(k, &mut pt_rgsw, 0); // X^{k}
302
303 let mut scratch: ScratchOwned = ScratchOwned::new(
304 GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, ct_ggsw.k(), rank)
305 | GLWECiphertext::decrypt_scratch_space(&module, basek, ct_out.k())
306 | GLWECiphertext::encrypt_sk_scratch_space(&module, basek, ct_in.k())
307 | GLWECiphertextFourier::external_product_scratch_space(
308 &module,
309 basek,
310 ct_out.k(),
311 ct_in.k(),
312 ct_ggsw.k(),
313 digits,
314 rank,
315 ),
316 );
317
318 let mut sk: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
319 sk.fill_ternary_prob(&module, 0.5, &mut source_xs);
320
321 ct_ggsw.encrypt_sk(
322 &module,
323 &pt_rgsw,
324 &sk,
325 &mut source_xa,
326 &mut source_xe,
327 sigma,
328 scratch.borrow(),
329 );
330
331 ct_in.encrypt_sk(
332 &module,
333 &pt_want,
334 &sk,
335 &mut source_xa,
336 &mut source_xe,
337 sigma,
338 scratch.borrow(),
339 );
340
341 ct_in.dft(&module, &mut ct_in_dft);
342 ct_out_dft.external_product(&module, &ct_in_dft, &ct_ggsw, scratch.borrow());
343 ct_out_dft.idft(&module, &mut ct_out, scratch.borrow());
344
345 ct_out.decrypt(&module, &mut pt_have, &sk, scratch.borrow());
346
347 pt_want.rotate_inplace(&module, k);
348 pt_have.sub_inplace_ab(&module, &pt_want);
349
350 let noise_have: f64 = pt_have.data.std(0, basek).log2();
351
352 let var_gct_err_lhs: f64 = sigma * sigma;
353 let var_gct_err_rhs: f64 = 0f64;
354
355 let var_msg: f64 = 1f64 / module.n() as f64; // X^{k}
356 let var_a0_err: f64 = sigma * sigma;
357 let var_a1_err: f64 = 1f64 / 12f64;
358
359 let noise_want: f64 = noise_ggsw_product(
360 module.n() as f64,
361 basek * digits,
362 0.5,
363 var_msg,
364 var_a0_err,
365 var_a1_err,
366 var_gct_err_lhs,
367 var_gct_err_rhs,
368 rank as f64,
369 k_in,
370 k_ggsw,
371 );
372
373 assert!(
374 (noise_have - noise_want).abs() <= 0.5,
375 "{} {}",
376 noise_have,
377 noise_want
378 );
379}
380
381fn test_external_product_inplace(log_n: usize, basek: usize, k_ct: usize, k_ggsw: usize, digits: usize, rank: usize, sigma: f64) {
382 let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
383 let rows: usize = div_ceil(k_ct, digits * basek);
384
385 let mut ct_ggsw: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k_ggsw, rows, digits, rank);
386 let mut ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k_ct, rank);
387 let mut ct_rlwe_dft: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k_ct, rank);
388 let mut pt_rgsw: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
389 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
390 let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ct);
391
392 let mut source_xs: Source = Source::new([0u8; 32]);
393 let mut source_xe: Source = Source::new([0u8; 32]);
394 let mut source_xa: Source = Source::new([0u8; 32]);
395
396 // Random input plaintext
397 pt_want
398 .data
399 .fill_uniform(basek, 0, pt_want.size(), &mut source_xa);
400
401 pt_want.data.at_mut(0, 0)[1] = 1;
402
403 let k: i64 = 1;
404
405 pt_rgsw.raw_mut()[0] = 1; // X^{0}
406 module.vec_znx_rotate_inplace(k, &mut pt_rgsw, 0); // X^{k}
407
408 let mut scratch: ScratchOwned = ScratchOwned::new(
409 GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, ct_ggsw.k(), rank)
410 | GLWECiphertext::decrypt_scratch_space(&module, basek, ct.k())
411 | GLWECiphertext::encrypt_sk_scratch_space(&module, basek, ct.k())
412 | GLWECiphertextFourier::external_product_inplace_scratch_space(&module, basek, ct.k(), ct_ggsw.k(), digits, rank),
413 );
414
415 let mut sk: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
416 sk.fill_ternary_prob(&module, 0.5, &mut source_xs);
417
418 ct_ggsw.encrypt_sk(
419 &module,
420 &pt_rgsw,
421 &sk,
422 &mut source_xa,
423 &mut source_xe,
424 sigma,
425 scratch.borrow(),
426 );
427
428 ct.encrypt_sk(
429 &module,
430 &pt_want,
431 &sk,
432 &mut source_xa,
433 &mut source_xe,
434 sigma,
435 scratch.borrow(),
436 );
437
438 ct.dft(&module, &mut ct_rlwe_dft);
439 ct_rlwe_dft.external_product_inplace(&module, &ct_ggsw, scratch.borrow());
440 ct_rlwe_dft.idft(&module, &mut ct, scratch.borrow());
441
442 ct.decrypt(&module, &mut pt_have, &sk, scratch.borrow());
443
444 pt_want.rotate_inplace(&module, k);
445 pt_have.sub_inplace_ab(&module, &pt_want);
446
447 let noise_have: f64 = pt_have.data.std(0, basek).log2();
448
449 let var_gct_err_lhs: f64 = sigma * sigma;
450 let var_gct_err_rhs: f64 = 0f64;
451
452 let var_msg: f64 = 1f64 / module.n() as f64; // X^{k}
453 let var_a0_err: f64 = sigma * sigma;
454 let var_a1_err: f64 = 1f64 / 12f64;
455
456 let noise_want: f64 = noise_ggsw_product(
457 module.n() as f64,
458 basek * digits,
459 0.5,
460 var_msg,
461 var_a0_err,
462 var_a1_err,
463 var_gct_err_lhs,
464 var_gct_err_rhs,
465 rank as f64,
466 k_ct,
467 k_ggsw,
468 );
469
470 assert!(
471 (noise_have - noise_want).abs() <= 0.5,
472 "{} {}",
473 noise_have,
474 noise_want
475 );
476
477 println!("{} {}", noise_have, noise_want);
478}