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, VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxAutomorphismInplace, VecZnxCopy,
4 VecZnxSubScalarInplace, VecZnxSwithcDegree, VmpPMatAlloc, VmpPMatPrepare,
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 layouts::{
16 GGLWEAutomorphismKey, GLWEPlaintext, GLWESecret, Infos,
17 prepared::{GGLWEAutomorphismKeyPrepared, GLWESecretPrepared, Prepare, PrepareAlloc},
18 },
19 noise::log2_std_noise_gglwe_product,
20 trait_families::{GLWEDecryptFamily, GLWEKeyswitchFamily},
21};
22
23use crate::trait_families::{GGLWEAutomorphismKeyEncryptSkFamily, GLWESecretPreparedModuleFamily};
24
25pub fn test_gglwe_automorphism_key_automorphism<B: Backend>(
26 module: &Module<B>,
27 p0: i64,
28 p1: i64,
29 basek: usize,
30 digits: usize,
31 k_in: usize,
32 k_out: usize,
33 k_apply: usize,
34 sigma: f64,
35 rank: usize,
36) where
37 Module<B>: GGLWEAutomorphismKeyEncryptSkFamily<B>
38 + GLWESecretPreparedModuleFamily<B>
39 + GLWEKeyswitchFamily<B>
40 + VecZnxAutomorphism
41 + VecZnxSwithcDegree
42 + VecZnxAddScalarInplace
43 + VecZnxAutomorphism
44 + VecZnxAutomorphismInplace
45 + GLWEDecryptFamily<B>
46 + VecZnxSubScalarInplace
47 + VecZnxCopy
48 + VmpPMatAlloc<B>
49 + VmpPMatPrepare<B>,
50 B: ScratchOwnedAllocImpl<B>
51 + ScratchOwnedBorrowImpl<B>
52 + ScratchAvailableImpl<B>
53 + TakeScalarZnxImpl<B>
54 + TakeVecZnxDftImpl<B>
55 + TakeVecZnxImpl<B>
56 + TakeSvpPPolImpl<B>
57 + TakeVecZnxBigImpl<B>,
58{
59 let n: usize = module.n();
60 let digits_in: usize = 1;
61
62 let rows_in: usize = k_in / (basek * digits);
63 let rows_apply: usize = k_in.div_ceil(basek * digits);
64
65 let mut auto_key_in: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_in, rows_in, digits_in, rank);
66 let mut auto_key_out: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_out, rows_in, digits_in, rank);
67 let mut auto_key_apply: GGLWEAutomorphismKey<Vec<u8>> =
68 GGLWEAutomorphismKey::alloc(n, basek, k_apply, rows_apply, digits, rank);
69
70 let mut source_xs: Source = Source::new([0u8; 32]);
71 let mut source_xe: Source = Source::new([0u8; 32]);
72 let mut source_xa: Source = Source::new([0u8; 32]);
73
74 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
75 GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, n, basek, k_apply, rank)
76 | GGLWEAutomorphismKey::automorphism_scratch_space(module, n, basek, k_out, k_in, k_apply, digits, rank),
77 );
78
79 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
80 sk.fill_ternary_prob(0.5, &mut source_xs);
81
82 // gglwe_{s1}(s0) = s0 -> s1
83 auto_key_in.encrypt_sk(
84 module,
85 p0,
86 &sk,
87 &mut source_xa,
88 &mut source_xe,
89 sigma,
90 scratch.borrow(),
91 );
92
93 // gglwe_{s2}(s1) -> s1 -> s2
94 auto_key_apply.encrypt_sk(
95 module,
96 p1,
97 &sk,
98 &mut source_xa,
99 &mut source_xe,
100 sigma,
101 scratch.borrow(),
102 );
103
104 let mut auto_key_apply_prepared: GGLWEAutomorphismKeyPrepared<Vec<u8>, B> =
105 GGLWEAutomorphismKeyPrepared::alloc(module, n, basek, k_apply, rows_apply, digits, rank);
106
107 auto_key_apply_prepared.prepare(module, &auto_key_apply, scratch.borrow());
108
109 // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
110 auto_key_out.automorphism(
111 module,
112 &auto_key_in,
113 &auto_key_apply_prepared,
114 scratch.borrow(),
115 );
116
117 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_out);
118
119 let mut sk_auto: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
120 sk_auto.fill_zero(); // Necessary to avoid panic of unfilled sk
121 (0..rank).for_each(|i| {
122 module.vec_znx_automorphism(
123 module.galois_element_inv(p0 * p1),
124 &mut sk_auto.data.as_vec_znx_mut(),
125 i,
126 &sk.data.as_vec_znx(),
127 i,
128 );
129 });
130
131 let sk_auto_dft: GLWESecretPrepared<Vec<u8>, B> = sk_auto.prepare_alloc(module, scratch.borrow());
132
133 (0..auto_key_out.rank_in()).for_each(|col_i| {
134 (0..auto_key_out.rows()).for_each(|row_i| {
135 auto_key_out
136 .at(row_i, col_i)
137 .decrypt(module, &mut pt, &sk_auto_dft, scratch.borrow());
138
139 module.vec_znx_sub_scalar_inplace(
140 &mut pt.data,
141 0,
142 (digits_in - 1) + row_i * digits_in,
143 &sk.data,
144 col_i,
145 );
146
147 let noise_have: f64 = pt.data.std(basek, 0).log2();
148 let noise_want: f64 = log2_std_noise_gglwe_product(
149 n as f64,
150 basek * digits,
151 0.5,
152 0.5,
153 0f64,
154 sigma * sigma,
155 0f64,
156 rank as f64,
157 k_out,
158 k_apply,
159 );
160
161 assert!(
162 noise_have < noise_want + 0.5,
163 "{} {}",
164 noise_have,
165 noise_want
166 );
167 });
168 });
169}
170
171pub fn test_gglwe_automorphism_key_automorphism_inplace<B: Backend>(
172 module: &Module<B>,
173 p0: i64,
174 p1: i64,
175 basek: usize,
176 digits: usize,
177 k_in: usize,
178 k_apply: usize,
179 sigma: f64,
180 rank: usize,
181) where
182 Module<B>: GGLWEAutomorphismKeyEncryptSkFamily<B>
183 + GLWESecretPreparedModuleFamily<B>
184 + GLWEKeyswitchFamily<B>
185 + VecZnxAutomorphism
186 + VecZnxSwithcDegree
187 + VecZnxAddScalarInplace
188 + VecZnxAutomorphism
189 + VecZnxAutomorphismInplace
190 + GLWEDecryptFamily<B>
191 + VecZnxSubScalarInplace
192 + VecZnxCopy
193 + VmpPMatAlloc<B>
194 + VmpPMatPrepare<B>,
195 B: ScratchOwnedAllocImpl<B>
196 + ScratchOwnedBorrowImpl<B>
197 + ScratchAvailableImpl<B>
198 + TakeScalarZnxImpl<B>
199 + TakeVecZnxDftImpl<B>
200 + TakeVecZnxImpl<B>
201 + TakeSvpPPolImpl<B>
202 + TakeVecZnxBigImpl<B>,
203{
204 let n: usize = module.n();
205 let digits_in: usize = 1;
206
207 let rows_in: usize = k_in / (basek * digits);
208 let rows_apply: usize = k_in.div_ceil(basek * digits);
209
210 let mut auto_key: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_in, rows_in, digits_in, rank);
211 let mut auto_key_apply: GGLWEAutomorphismKey<Vec<u8>> =
212 GGLWEAutomorphismKey::alloc(n, basek, k_apply, rows_apply, digits, rank);
213
214 let mut source_xs: Source = Source::new([0u8; 32]);
215 let mut source_xe: Source = Source::new([0u8; 32]);
216 let mut source_xa: Source = Source::new([0u8; 32]);
217
218 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
219 GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, n, basek, k_apply, rank)
220 | GGLWEAutomorphismKey::automorphism_inplace_scratch_space(module, n, basek, k_in, k_apply, digits, rank),
221 );
222
223 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
224 sk.fill_ternary_prob(0.5, &mut source_xs);
225
226 // gglwe_{s1}(s0) = s0 -> s1
227 auto_key.encrypt_sk(
228 module,
229 p0,
230 &sk,
231 &mut source_xa,
232 &mut source_xe,
233 sigma,
234 scratch.borrow(),
235 );
236
237 // gglwe_{s2}(s1) -> s1 -> s2
238 auto_key_apply.encrypt_sk(
239 module,
240 p1,
241 &sk,
242 &mut source_xa,
243 &mut source_xe,
244 sigma,
245 scratch.borrow(),
246 );
247
248 let mut auto_key_apply_prepared: GGLWEAutomorphismKeyPrepared<Vec<u8>, B> =
249 GGLWEAutomorphismKeyPrepared::alloc(module, n, basek, k_apply, rows_apply, digits, rank);
250
251 auto_key_apply_prepared.prepare(module, &auto_key_apply, scratch.borrow());
252
253 // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
254 auto_key.automorphism_inplace(module, &auto_key_apply_prepared, scratch.borrow());
255
256 let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_in);
257
258 let mut sk_auto: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
259 sk_auto.fill_zero(); // Necessary to avoid panic of unfilled sk
260
261 (0..rank).for_each(|i| {
262 module.vec_znx_automorphism(
263 module.galois_element_inv(p0 * p1),
264 &mut sk_auto.data.as_vec_znx_mut(),
265 i,
266 &sk.data.as_vec_znx(),
267 i,
268 );
269 });
270
271 let sk_auto_dft: GLWESecretPrepared<Vec<u8>, B> = sk_auto.prepare_alloc(module, scratch.borrow());
272
273 (0..auto_key.rank_in()).for_each(|col_i| {
274 (0..auto_key.rows()).for_each(|row_i| {
275 auto_key
276 .at(row_i, col_i)
277 .decrypt(module, &mut pt, &sk_auto_dft, scratch.borrow());
278 module.vec_znx_sub_scalar_inplace(
279 &mut pt.data,
280 0,
281 (digits_in - 1) + row_i * digits_in,
282 &sk.data,
283 col_i,
284 );
285
286 let noise_have: f64 = pt.data.std(basek, 0).log2();
287 let noise_want: f64 = log2_std_noise_gglwe_product(
288 n as f64,
289 basek * digits,
290 0.5,
291 0.5,
292 0f64,
293 sigma * sigma,
294 0f64,
295 rank as f64,
296 k_in,
297 k_apply,
298 );
299
300 assert!(
301 noise_have < noise_want + 0.5,
302 "{} {}",
303 noise_have,
304 noise_want
305 );
306 });
307 });
308}