poulpy/core/src/tests/generics/automorphism/gglwe_atk.rs
core/src/tests/generics/automorphism/gglwe_atk.rsRaw
1use backend::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};
16use sampling::source::Source;
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
26pub fn test_gglwe_automorphism_key_automorphism<B: Backend>(
27 module: &Module<B>,
28 p0: i64,
29 p1: i64,
30 basek: usize,
31 digits: usize,
32 k_in: usize,
33 k_out: usize,
34 k_apply: usize,
35 sigma: f64,
36 rank: usize,
37) where
38 Module<B>: VecZnxDftAllocBytes
39 + VmpApplyTmpBytes
40 + VecZnxBigNormalizeTmpBytes
41 + VmpApply<B>
42 + VmpApplyAdd<B>
43 + VecZnxDftFromVecZnx<B>
44 + VecZnxDftToVecZnxBigConsume<B>
45 + VecZnxBigAddSmallInplace<B>
46 + VecZnxBigNormalize<B>
47 + VecZnxAutomorphism
48 + VecZnxAutomorphismInplace
49 + SvpPPolAllocBytes
50 + VecZnxDftAllocBytes
51 + VecZnxNormalizeTmpBytes
52 + VmpPMatAlloc<B>
53 + VmpPrepare<B>
54 + SvpPrepare<B>
55 + SvpApplyInplace<B>
56 + VecZnxAddScalarInplace
57 + VecZnxFillUniform
58 + VecZnxSubABInplace
59 + VecZnxAddInplace
60 + VecZnxNormalizeInplace<B>
61 + VecZnxAddNormal
62 + VecZnxNormalize<B>
63 + VecZnxSub
64 + VecZnxSwithcDegree
65 + SvpPPolAlloc<B>
66 + VecZnxBigAddInplace<B>
67 + VecZnxSubScalarInplace,
68 B: ScratchOwnedAllocImpl<B>
69 + ScratchOwnedBorrowImpl<B>
70 + ScratchAvailableImpl<B>
71 + TakeScalarZnxImpl<B>
72 + TakeVecZnxDftImpl<B>
73 + TakeVecZnxImpl<B>
74 + TakeSvpPPolImpl<B>
75 + TakeVecZnxBigImpl<B>,
76{
77 let n: usize = module.n();
78 let digits_in: usize = 1;
79
80 let rows_in: usize = k_in / (basek * digits);
81 let rows_apply: usize = k_in.div_ceil(basek * digits);
82
83 let mut auto_key_in: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_in, rows_in, digits_in, rank);
84 let mut auto_key_out: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_out, rows_in, digits_in, rank);
85 let mut auto_key_apply: GGLWEAutomorphismKey<Vec<u8>> =
86 GGLWEAutomorphismKey::alloc(n, basek, k_apply, rows_apply, digits, rank);
87
88 let mut source_xs: Source = Source::new([0u8; 32]);
89 let mut source_xe: Source = Source::new([0u8; 32]);
90 let mut source_xa: Source = Source::new([0u8; 32]);
91
92 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
93 GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, n, basek, k_apply, rank)
94 | GGLWEAutomorphismKey::automorphism_scratch_space(module, n, basek, k_out, k_in, k_apply, digits, rank),
95 );
96
97 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
98 sk.fill_ternary_prob(0.5, &mut source_xs);
99
100 // gglwe_{s1}(s0) = s0 -> s1
101 auto_key_in.encrypt_sk(
102 module,
103 p0,
104 &sk,
105 &mut source_xa,
106 &mut source_xe,
107 sigma,
108 scratch.borrow(),
109 );
110
111 // gglwe_{s2}(s1) -> s1 -> s2
112 auto_key_apply.encrypt_sk(
113 module,
114 p1,
115 &sk,
116 &mut source_xa,
117 &mut source_xe,
118 sigma,
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
189pub fn test_gglwe_automorphism_key_automorphism_inplace<B: Backend>(
190 module: &Module<B>,
191 p0: i64,
192 p1: i64,
193 basek: usize,
194 digits: usize,
195 k_in: usize,
196 k_apply: usize,
197 sigma: f64,
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: ScratchOwnedAllocImpl<B>
245 + ScratchOwnedBorrowImpl<B>
246 + ScratchAvailableImpl<B>
247 + TakeScalarZnxImpl<B>
248 + TakeVecZnxDftImpl<B>
249 + TakeVecZnxImpl<B>
250 + TakeSvpPPolImpl<B>
251 + TakeVecZnxBigImpl<B>,
252{
253 let n: usize = module.n();
254 let digits_in: usize = 1;
255
256 let rows_in: usize = k_in / (basek * digits);
257 let rows_apply: usize = k_in.div_ceil(basek * digits);
258
259 let mut auto_key: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_in, rows_in, digits_in, rank);
260 let mut auto_key_apply: GGLWEAutomorphismKey<Vec<u8>> =
261 GGLWEAutomorphismKey::alloc(n, basek, k_apply, rows_apply, digits, rank);
262
263 let mut source_xs: Source = Source::new([0u8; 32]);
264 let mut source_xe: Source = Source::new([0u8; 32]);
265 let mut source_xa: Source = Source::new([0u8; 32]);
266
267 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
268 GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, n, basek, k_apply, rank)
269 | GGLWEAutomorphismKey::automorphism_inplace_scratch_space(module, n, basek, k_in, k_apply, digits, rank),
270 );
271
272 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
273 sk.fill_ternary_prob(0.5, &mut source_xs);
274
275 // gglwe_{s1}(s0) = s0 -> s1
276 auto_key.encrypt_sk(
277 module,
278 p0,
279 &sk,
280 &mut source_xa,
281 &mut source_xe,
282 sigma,
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 sigma,
294 scratch.borrow(),
295 );
296
297 let mut auto_key_apply_prepared: GGLWEAutomorphismKeyPrepared<Vec<u8>, B> =
298 GGLWEAutomorphismKeyPrepared::alloc(module, n, basek, k_apply, rows_apply, digits, rank);
299
300 auto_key_apply_prepared.prepare(module, &auto_key_apply, scratch.borrow());
301
302 // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
303 auto_key.automorphism_inplace(module, &auto_key_apply_prepared, scratch.borrow());
304
305 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_in);
306
307 let mut sk_auto: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
308 sk_auto.fill_zero(); // Necessary to avoid panic of unfilled sk
309
310 (0..rank).for_each(|i| {
311 module.vec_znx_automorphism(
312 module.galois_element_inv(p0 * p1),
313 &mut sk_auto.data.as_vec_znx_mut(),
314 i,
315 &sk.data.as_vec_znx(),
316 i,
317 );
318 });
319
320 let sk_auto_dft: GLWESecretPrepared<Vec<u8>, B> = sk_auto.prepare_alloc(module, scratch.borrow());
321
322 (0..auto_key.rank_in()).for_each(|col_i| {
323 (0..auto_key.rows()).for_each(|row_i| {
324 auto_key
325 .at(row_i, col_i)
326 .decrypt(module, &mut pt, &sk_auto_dft, scratch.borrow());
327 module.vec_znx_sub_scalar_inplace(
328 &mut pt.data,
329 0,
330 (digits_in - 1) + row_i * digits_in,
331 &sk.data,
332 col_i,
333 );
334
335 let noise_have: f64 = pt.data.std(basek, 0).log2();
336 let noise_want: f64 = log2_std_noise_gglwe_product(
337 n as f64,
338 basek * digits,
339 0.5,
340 0.5,
341 0f64,
342 sigma * sigma,
343 0f64,
344 rank as f64,
345 k_in,
346 k_apply,
347 );
348
349 assert!(
350 noise_have < noise_want + 0.5,
351 "{} {}",
352 noise_have,
353 noise_want
354 );
355 });
356 });
357}