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