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