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 encryption::SIGMA,
20 layouts::{
21 GGLWEAutomorphismKey, GLWEPlaintext, GLWESecret, Infos,
22 prepared::{GGLWEAutomorphismKeyPrepared, GLWESecretPrepared, Prepare, PrepareAlloc},
23 },
24 noise::log2_std_noise_gglwe_product,
25};
26
27#[allow(clippy::too_many_arguments)]
28pub fn test_gglwe_automorphism_key_automorphism<B>(
29 module: &Module<B>,
30 p0: i64,
31 p1: i64,
32 basek: usize,
33 digits: usize,
34 k_in: usize,
35 k_out: usize,
36 k_apply: usize,
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 scratch.borrow(),
110 );
111
112 // gglwe_{s2}(s1) -> s1 -> s2
113 auto_key_apply.encrypt_sk(
114 module,
115 p1,
116 &sk,
117 &mut source_xa,
118 &mut source_xe,
119 scratch.borrow(),
120 );
121
122 let mut auto_key_apply_prepared: GGLWEAutomorphismKeyPrepared<Vec<u8>, B> =
123 GGLWEAutomorphismKeyPrepared::alloc(module, n, basek, k_apply, rows_apply, digits, rank);
124
125 auto_key_apply_prepared.prepare(module, &auto_key_apply, scratch.borrow());
126
127 // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
128 auto_key_out.automorphism(
129 module,
130 &auto_key_in,
131 &auto_key_apply_prepared,
132 scratch.borrow(),
133 );
134
135 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_out);
136
137 let mut sk_auto: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
138 sk_auto.fill_zero(); // Necessary to avoid panic of unfilled sk
139 (0..rank).for_each(|i| {
140 module.vec_znx_automorphism(
141 module.galois_element_inv(p0 * p1),
142 &mut sk_auto.data.as_vec_znx_mut(),
143 i,
144 &sk.data.as_vec_znx(),
145 i,
146 );
147 });
148
149 let sk_auto_dft: GLWESecretPrepared<Vec<u8>, B> = sk_auto.prepare_alloc(module, scratch.borrow());
150
151 (0..auto_key_out.rank_in()).for_each(|col_i| {
152 (0..auto_key_out.rows()).for_each(|row_i| {
153 auto_key_out
154 .at(row_i, col_i)
155 .decrypt(module, &mut pt, &sk_auto_dft, scratch.borrow());
156
157 module.vec_znx_sub_scalar_inplace(
158 &mut pt.data,
159 0,
160 (digits_in - 1) + row_i * digits_in,
161 &sk.data,
162 col_i,
163 );
164
165 let noise_have: f64 = pt.data.std(basek, 0).log2();
166 let noise_want: f64 = log2_std_noise_gglwe_product(
167 n as f64,
168 basek * digits,
169 0.5,
170 0.5,
171 0f64,
172 SIGMA * SIGMA,
173 0f64,
174 rank as f64,
175 k_out,
176 k_apply,
177 );
178
179 assert!(
180 noise_have < noise_want + 0.5,
181 "{} {}",
182 noise_have,
183 noise_want
184 );
185 });
186 });
187}
188
189#[allow(clippy::too_many_arguments)]
190pub fn test_gglwe_automorphism_key_automorphism_inplace<B>(
191 module: &Module<B>,
192 p0: i64,
193 p1: i64,
194 basek: usize,
195 digits: usize,
196 k_in: usize,
197 k_apply: usize,
198 rank: usize,
199) where
200 Module<B>: VecZnxDftAllocBytes
201 + VecZnxBigNormalize<B>
202 + VecZnxDftFromVecZnx<B>
203 + SvpApplyInplace<B>
204 + VecZnxDftToVecZnxBigConsume<B>
205 + VecZnxNormalizeTmpBytes
206 + VecZnxFillUniform
207 + VecZnxSubABInplace
208 + VecZnxAddInplace
209 + VecZnxNormalizeInplace<B>
210 + VecZnxAddNormal
211 + VecZnxNormalize<B>
212 + VecZnxSub
213 + SvpPrepare<B>
214 + SvpPPolAllocBytes
215 + SvpPPolAlloc<B>
216 + VecZnxDftAllocBytes
217 + VmpApplyTmpBytes
218 + VecZnxBigNormalizeTmpBytes
219 + VmpApplyTmpBytes
220 + VmpApply<B>
221 + VmpApplyAdd<B>
222 + VecZnxDftFromVecZnx<B>
223 + VecZnxDftToVecZnxBigConsume<B>
224 + VecZnxBigAddSmallInplace<B>
225 + VecZnxBigNormalize<B>
226 + VecZnxAutomorphism
227 + VecZnxSwithcDegree
228 + VecZnxAddScalarInplace
229 + VecZnxAutomorphism
230 + VecZnxAutomorphismInplace
231 + VecZnxDftAllocBytes
232 + VecZnxBigAllocBytes
233 + VecZnxDftFromVecZnx<B>
234 + SvpApplyInplace<B>
235 + VecZnxDftToVecZnxBigConsume<B>
236 + VecZnxBigAddInplace<B>
237 + VecZnxBigAddSmallInplace<B>
238 + VecZnxBigNormalize<B>
239 + VecZnxNormalizeTmpBytes
240 + VecZnxSubScalarInplace
241 + VecZnxCopy
242 + VmpPMatAlloc<B>
243 + VmpPrepare<B>,
244 B: Backend
245 + ScratchOwnedAllocImpl<B>
246 + ScratchOwnedBorrowImpl<B>
247 + ScratchAvailableImpl<B>
248 + TakeScalarZnxImpl<B>
249 + TakeVecZnxDftImpl<B>
250 + TakeVecZnxImpl<B>
251 + TakeSvpPPolImpl<B>
252 + TakeVecZnxBigImpl<B>,
253{
254 let n: usize = module.n();
255 let digits_in: usize = 1;
256
257 let rows_in: usize = k_in / (basek * digits);
258 let rows_apply: usize = k_in.div_ceil(basek * digits);
259
260 let mut auto_key: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_in, rows_in, digits_in, rank);
261 let mut auto_key_apply: GGLWEAutomorphismKey<Vec<u8>> =
262 GGLWEAutomorphismKey::alloc(n, basek, k_apply, rows_apply, digits, rank);
263
264 let mut source_xs: Source = Source::new([0u8; 32]);
265 let mut source_xe: Source = Source::new([0u8; 32]);
266 let mut source_xa: Source = Source::new([0u8; 32]);
267
268 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
269 GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, n, basek, k_apply, rank)
270 | GGLWEAutomorphismKey::automorphism_inplace_scratch_space(module, n, basek, k_in, k_apply, digits, rank),
271 );
272
273 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
274 sk.fill_ternary_prob(0.5, &mut source_xs);
275
276 // gglwe_{s1}(s0) = s0 -> s1
277 auto_key.encrypt_sk(
278 module,
279 p0,
280 &sk,
281 &mut source_xa,
282 &mut source_xe,
283 scratch.borrow(),
284 );
285
286 // gglwe_{s2}(s1) -> s1 -> s2
287 auto_key_apply.encrypt_sk(
288 module,
289 p1,
290 &sk,
291 &mut source_xa,
292 &mut source_xe,
293 scratch.borrow(),
294 );
295
296 let mut auto_key_apply_prepared: GGLWEAutomorphismKeyPrepared<Vec<u8>, B> =
297 GGLWEAutomorphismKeyPrepared::alloc(module, n, basek, k_apply, rows_apply, digits, rank);
298
299 auto_key_apply_prepared.prepare(module, &auto_key_apply, scratch.borrow());
300
301 // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
302 auto_key.automorphism_inplace(module, &auto_key_apply_prepared, scratch.borrow());
303
304 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_in);
305
306 let mut sk_auto: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
307 sk_auto.fill_zero(); // Necessary to avoid panic of unfilled sk
308
309 (0..rank).for_each(|i| {
310 module.vec_znx_automorphism(
311 module.galois_element_inv(p0 * p1),
312 &mut sk_auto.data.as_vec_znx_mut(),
313 i,
314 &sk.data.as_vec_znx(),
315 i,
316 );
317 });
318
319 let sk_auto_dft: GLWESecretPrepared<Vec<u8>, B> = sk_auto.prepare_alloc(module, scratch.borrow());
320
321 (0..auto_key.rank_in()).for_each(|col_i| {
322 (0..auto_key.rows()).for_each(|row_i| {
323 auto_key
324 .at(row_i, col_i)
325 .decrypt(module, &mut pt, &sk_auto_dft, scratch.borrow());
326 module.vec_znx_sub_scalar_inplace(
327 &mut pt.data,
328 0,
329 (digits_in - 1) + row_i * digits_in,
330 &sk.data,
331 col_i,
332 );
333
334 let noise_have: f64 = pt.data.std(basek, 0).log2();
335 let noise_want: f64 = log2_std_noise_gglwe_product(
336 n as f64,
337 basek * digits,
338 0.5,
339 0.5,
340 0f64,
341 SIGMA * SIGMA,
342 0f64,
343 rank as f64,
344 k_in,
345 k_apply,
346 );
347
348 assert!(
349 noise_have < noise_want + 0.5,
350 "{} {}",
351 noise_have,
352 noise_want
353 );
354 });
355 });
356}