poulpy/core/src/gglwe/tests/generics_automorphism_key.rs
core/src/gglwe/tests/generics_automorphism_key.rsRaw
1use backend::hal::{
2 api::{
3 ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxAutomorphismInplace, VecZnxCopy,
4 VecZnxStd, VecZnxSubScalarInplace, VecZnxSwithcDegree,
5 },
6 layouts::{Backend, Module, ScratchOwned},
7 oep::{
8 ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
9 TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl,
10 },
11};
12use sampling::source::Source;
13
14use crate::{
15 AutomorphismKey, AutomorphismKeyCompressed, AutomorphismKeyEncryptSkFamily, AutomorphismKeyExec, GGLWEExecLayoutFamily,
16 GLWEDecryptFamily, GLWEKeyswitchFamily, GLWEPlaintext, GLWESecret, GLWESecretExec, Infos,
17 noise::log2_std_noise_gglwe_product,
18};
19
20pub(crate) trait AutomorphismTestModuleFamily<B: Backend> = AutomorphismKeyEncryptSkFamily<B>
21 + GLWEKeyswitchFamily<B>
22 + VecZnxAutomorphism
23 + GGLWEExecLayoutFamily<B>
24 + VecZnxSwithcDegree
25 + VecZnxAddScalarInplace
26 + VecZnxAutomorphism
27 + VecZnxAutomorphismInplace
28 + GLWEDecryptFamily<B>
29 + VecZnxSubScalarInplace
30 + VecZnxStd
31 + VecZnxCopy;
32pub(crate) trait AutomorphismTestScratchFamily<B: Backend> = ScratchOwnedAllocImpl<B>
33 + ScratchOwnedBorrowImpl<B>
34 + ScratchAvailableImpl<B>
35 + TakeScalarZnxImpl<B>
36 + TakeVecZnxDftImpl<B>
37 + TakeVecZnxImpl<B>
38 + TakeSvpPPolImpl<B>
39 + TakeVecZnxBigImpl<B>;
40
41pub(crate) fn test_automorphisk_key_encrypt_sk<B: Backend>(
42 module: &Module<B>,
43 basek: usize,
44 k_ksk: usize,
45 digits: usize,
46 rank: usize,
47 sigma: f64,
48) where
49 Module<B>: AutomorphismTestModuleFamily<B>,
50 B: AutomorphismTestScratchFamily<B>,
51{
52 let n: usize = module.n();
53 let rows: usize = (k_ksk - digits * basek) / (digits * basek);
54
55 let mut atk: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(n, basek, k_ksk, rows, digits, rank);
56
57 let mut source_xs: Source = Source::new([0u8; 32]);
58 let mut source_xe: Source = Source::new([0u8; 32]);
59 let mut source_xa: Source = Source::new([0u8; 32]);
60
61 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(AutomorphismKey::encrypt_sk_scratch_space(
62 module, n, basek, k_ksk, rank,
63 ));
64
65 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
66 sk.fill_ternary_prob(0.5, &mut source_xs);
67
68 let p = -5;
69
70 atk.encrypt_sk(
71 module,
72 p,
73 &sk,
74 &mut source_xa,
75 &mut source_xe,
76 sigma,
77 scratch.borrow(),
78 );
79
80 let mut sk_out: GLWESecret<Vec<u8>> = sk.clone();
81 (0..atk.rank()).for_each(|i| {
82 module.vec_znx_automorphism(
83 module.galois_element_inv(p),
84 &mut sk_out.data.as_vec_znx_mut(),
85 i,
86 &sk.data.as_vec_znx(),
87 i,
88 );
89 });
90 let sk_out_exec = GLWESecretExec::from(module, &sk_out);
91
92 atk.key
93 .key
94 .assert_noise(module, &sk_out_exec, &sk.data, sigma);
95}
96
97pub(crate) fn test_automorphisk_key_encrypt_sk_compressed<B: Backend>(
98 module: &Module<B>,
99 basek: usize,
100 k_ksk: usize,
101 digits: usize,
102 rank: usize,
103 sigma: f64,
104) where
105 Module<B>: AutomorphismTestModuleFamily<B>,
106 B: AutomorphismTestScratchFamily<B>,
107{
108 let n: usize = module.n();
109 let rows: usize = (k_ksk - digits * basek) / (digits * basek);
110
111 let mut atk_compressed: AutomorphismKeyCompressed<Vec<u8>> =
112 AutomorphismKeyCompressed::alloc(n, basek, k_ksk, rows, digits, rank);
113
114 let mut source_xs: Source = Source::new([0u8; 32]);
115 let mut source_xe: Source = Source::new([0u8; 32]);
116
117 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(AutomorphismKey::encrypt_sk_scratch_space(
118 module, n, basek, k_ksk, rank,
119 ));
120
121 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
122 sk.fill_ternary_prob(0.5, &mut source_xs);
123
124 let p = -5;
125
126 let seed_xa: [u8; 32] = [1u8; 32];
127
128 atk_compressed.encrypt_sk(
129 module,
130 p,
131 &sk,
132 seed_xa,
133 &mut source_xe,
134 sigma,
135 scratch.borrow(),
136 );
137
138 let mut sk_out: GLWESecret<Vec<u8>> = sk.clone();
139 (0..atk_compressed.rank()).for_each(|i| {
140 module.vec_znx_automorphism(
141 module.galois_element_inv(p),
142 &mut sk_out.data.as_vec_znx_mut(),
143 i,
144 &sk.data.as_vec_znx(),
145 i,
146 );
147 });
148 let sk_out_exec = GLWESecretExec::from(module, &sk_out);
149
150 let mut atk: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(n, basek, k_ksk, rows, digits, rank);
151 atk.decompress(module, &atk_compressed);
152
153 atk.key
154 .key
155 .assert_noise(module, &sk_out_exec, &sk.data, sigma);
156}
157
158pub(crate) fn test_gglwe_automorphism<B: Backend>(
159 module: &Module<B>,
160 p0: i64,
161 p1: i64,
162 basek: usize,
163 digits: usize,
164 k_in: usize,
165 k_out: usize,
166 k_apply: usize,
167 sigma: f64,
168 rank: usize,
169) where
170 Module<B>: AutomorphismTestModuleFamily<B>,
171 B: AutomorphismTestScratchFamily<B>,
172{
173 let n: usize = module.n();
174 let digits_in: usize = 1;
175
176 let rows_in: usize = k_in / (basek * digits);
177 let rows_apply: usize = k_in.div_ceil(basek * digits);
178
179 let mut auto_key_in: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(n, basek, k_in, rows_in, digits_in, rank);
180 let mut auto_key_out: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(n, basek, k_out, rows_in, digits_in, rank);
181 let mut auto_key_apply: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(n, basek, k_apply, rows_apply, digits, rank);
182
183 let mut source_xs: Source = Source::new([0u8; 32]);
184 let mut source_xe: Source = Source::new([0u8; 32]);
185 let mut source_xa: Source = Source::new([0u8; 32]);
186
187 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
188 AutomorphismKey::encrypt_sk_scratch_space(module, n, basek, k_apply, rank)
189 | AutomorphismKey::automorphism_scratch_space(module, n, basek, k_out, k_in, k_apply, digits, rank),
190 );
191
192 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
193 sk.fill_ternary_prob(0.5, &mut source_xs);
194
195 // gglwe_{s1}(s0) = s0 -> s1
196 auto_key_in.encrypt_sk(
197 module,
198 p0,
199 &sk,
200 &mut source_xa,
201 &mut source_xe,
202 sigma,
203 scratch.borrow(),
204 );
205
206 // gglwe_{s2}(s1) -> s1 -> s2
207 auto_key_apply.encrypt_sk(
208 module,
209 p1,
210 &sk,
211 &mut source_xa,
212 &mut source_xe,
213 sigma,
214 scratch.borrow(),
215 );
216
217 let mut auto_key_apply_exec: AutomorphismKeyExec<Vec<u8>, B> =
218 AutomorphismKeyExec::alloc(module, n, basek, k_apply, rows_apply, digits, rank);
219
220 auto_key_apply_exec.prepare(module, &auto_key_apply, scratch.borrow());
221
222 // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
223 auto_key_out.automorphism(module, &auto_key_in, &auto_key_apply_exec, scratch.borrow());
224
225 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_out);
226
227 let mut sk_auto: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
228 sk_auto.fill_zero(); // Necessary to avoid panic of unfilled sk
229 (0..rank).for_each(|i| {
230 module.vec_znx_automorphism(
231 module.galois_element_inv(p0 * p1),
232 &mut sk_auto.data.as_vec_znx_mut(),
233 i,
234 &sk.data.as_vec_znx(),
235 i,
236 );
237 });
238
239 let sk_auto_dft: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_auto);
240
241 (0..auto_key_out.rank_in()).for_each(|col_i| {
242 (0..auto_key_out.rows()).for_each(|row_i| {
243 auto_key_out
244 .at(row_i, col_i)
245 .decrypt(module, &mut pt, &sk_auto_dft, scratch.borrow());
246
247 module.vec_znx_sub_scalar_inplace(
248 &mut pt.data,
249 0,
250 (digits_in - 1) + row_i * digits_in,
251 &sk.data,
252 col_i,
253 );
254
255 let noise_have: f64 = module.vec_znx_std(basek, &pt.data, 0).log2();
256 let noise_want: f64 = log2_std_noise_gglwe_product(
257 n as f64,
258 basek * digits,
259 0.5,
260 0.5,
261 0f64,
262 sigma * sigma,
263 0f64,
264 rank as f64,
265 k_out,
266 k_apply,
267 );
268
269 assert!(
270 noise_have < noise_want + 0.5,
271 "{} {}",
272 noise_have,
273 noise_want
274 );
275 });
276 });
277}
278
279pub(crate) fn test_gglwe_automorphism_inplace<B: Backend>(
280 module: &Module<B>,
281 p0: i64,
282 p1: i64,
283 basek: usize,
284 digits: usize,
285 k_in: usize,
286 k_apply: usize,
287 sigma: f64,
288 rank: usize,
289) where
290 Module<B>: AutomorphismTestModuleFamily<B>,
291 B: AutomorphismTestScratchFamily<B>,
292{
293 let n: usize = module.n();
294 let digits_in: usize = 1;
295
296 let rows_in: usize = k_in / (basek * digits);
297 let rows_apply: usize = k_in.div_ceil(basek * digits);
298
299 let mut auto_key: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(n, basek, k_in, rows_in, digits_in, rank);
300 let mut auto_key_apply: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(n, basek, k_apply, rows_apply, digits, rank);
301
302 let mut source_xs: Source = Source::new([0u8; 32]);
303 let mut source_xe: Source = Source::new([0u8; 32]);
304 let mut source_xa: Source = Source::new([0u8; 32]);
305
306 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
307 AutomorphismKey::encrypt_sk_scratch_space(module, n, basek, k_apply, rank)
308 | AutomorphismKey::automorphism_inplace_scratch_space(module, n, basek, k_in, k_apply, digits, rank),
309 );
310
311 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
312 sk.fill_ternary_prob(0.5, &mut source_xs);
313
314 // gglwe_{s1}(s0) = s0 -> s1
315 auto_key.encrypt_sk(
316 module,
317 p0,
318 &sk,
319 &mut source_xa,
320 &mut source_xe,
321 sigma,
322 scratch.borrow(),
323 );
324
325 // gglwe_{s2}(s1) -> s1 -> s2
326 auto_key_apply.encrypt_sk(
327 module,
328 p1,
329 &sk,
330 &mut source_xa,
331 &mut source_xe,
332 sigma,
333 scratch.borrow(),
334 );
335
336 let mut auto_key_apply_exec: AutomorphismKeyExec<Vec<u8>, B> =
337 AutomorphismKeyExec::alloc(module, n, basek, k_apply, rows_apply, digits, rank);
338
339 auto_key_apply_exec.prepare(module, &auto_key_apply, scratch.borrow());
340
341 // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
342 auto_key.automorphism_inplace(module, &auto_key_apply_exec, scratch.borrow());
343
344 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_in);
345
346 let mut sk_auto: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
347 sk_auto.fill_zero(); // Necessary to avoid panic of unfilled sk
348
349 (0..rank).for_each(|i| {
350 module.vec_znx_automorphism(
351 module.galois_element_inv(p0 * p1),
352 &mut sk_auto.data.as_vec_znx_mut(),
353 i,
354 &sk.data.as_vec_znx(),
355 i,
356 );
357 });
358
359 let sk_auto_dft: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_auto);
360
361 (0..auto_key.rank_in()).for_each(|col_i| {
362 (0..auto_key.rows()).for_each(|row_i| {
363 auto_key
364 .at(row_i, col_i)
365 .decrypt(module, &mut pt, &sk_auto_dft, scratch.borrow());
366 module.vec_znx_sub_scalar_inplace(
367 &mut pt.data,
368 0,
369 (digits_in - 1) + row_i * digits_in,
370 &sk.data,
371 col_i,
372 );
373
374 let noise_have: f64 = module.vec_znx_std(basek, &pt.data, 0).log2();
375 let noise_want: f64 = log2_std_noise_gglwe_product(
376 n as f64,
377 basek * digits,
378 0.5,
379 0.5,
380 0f64,
381 sigma * sigma,
382 0f64,
383 rank as f64,
384 k_in,
385 k_apply,
386 );
387
388 assert!(
389 noise_have < noise_want + 0.5,
390 "{} {}",
391 noise_have,
392 noise_want
393 );
394 });
395 });
396}