poulpy/poulpy-core/src/tests/test_suite/glwe_tensor.rs
poulpy-core/src/tests/test_suite/glwe_tensor.rsRaw
1use poulpy_hal::{
2 api::{
3 ScratchAvailable, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxCopy, VecZnxFillUniform, VecZnxNormalize,
4 VecZnxNormalizeInplace,
5 },
6 layouts::{Backend, FillUniform, Module, Scratch, ScratchOwned, VecZnx, ZnxViewMut},
7 source::Source,
8 test_suite::convolution::bivariate_convolution_naive,
9};
10use rand::RngCore;
11use std::f64::consts::SQRT_2;
12
13use crate::{
14 GLWEDecrypt, GLWEEncryptSk, GLWEMulConst, GLWEMulPlain, GLWESub, GLWETensorKeyEncryptSk, GLWETensoring, ScratchTakeCore,
15 layouts::{
16 Dsize, GLWE, GLWELayout, GLWEPlaintext, GLWESecret, GLWESecretPreparedFactory, GLWESecretTensor, GLWESecretTensorFactory,
17 GLWESecretTensorPrepared, GLWETensor, GLWETensorKey, GLWETensorKeyLayout, GLWETensorKeyPrepared,
18 GLWETensorKeyPreparedFactory, LWEInfos, TorusPrecision, prepared::GLWESecretPrepared,
19 },
20};
21
22pub fn test_glwe_tensoring<BE: Backend>(module: &Module<BE>)
23where
24 Module<BE>: GLWETensoring<BE>
25 + GLWEEncryptSk<BE>
26 + GLWEDecrypt<BE>
27 + VecZnxFillUniform
28 + GLWESecretPreparedFactory<BE>
29 + GLWESub
30 + VecZnxNormalizeInplace<BE>
31 + GLWESecretTensorFactory<BE>
32 + VecZnxCopy
33 + VecZnxNormalize<BE>
34 + GLWETensorKeyEncryptSk<BE>
35 + GLWETensorKeyPreparedFactory<BE>,
36 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
37 Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
38{
39 let in_base2k: usize = 16;
40 let out_base2k: usize = 13;
41 let tsk_base2k: usize = 15;
42 let k: usize = 128;
43
44 for rank in 1_usize..=3 {
45 let n: usize = module.n();
46
47 let glwe_in_infos: GLWELayout = GLWELayout {
48 n: n.into(),
49 base2k: in_base2k.into(),
50 k: k.into(),
51 rank: rank.into(),
52 };
53
54 let glwe_out_infos: GLWELayout = GLWELayout {
55 n: n.into(),
56 base2k: out_base2k.into(),
57 k: k.into(),
58 rank: rank.into(),
59 };
60
61 let tsk_infos: GLWETensorKeyLayout = GLWETensorKeyLayout {
62 n: n.into(),
63 base2k: tsk_base2k.into(),
64 k: (k + tsk_base2k).into(),
65 rank: rank.into(),
66 dnum: k.div_ceil(tsk_base2k).into(),
67 dsize: Dsize(1),
68 };
69
70 let mut a: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_in_infos);
71 let mut b: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_in_infos);
72 let mut res_tensor: GLWETensor<Vec<u8>> = GLWETensor::alloc_from_infos(&glwe_out_infos);
73 let mut res_relin: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_out_infos);
74 let mut pt_in: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_in_infos);
75 let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_out_infos);
76 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_out_infos);
77 let mut pt_tmp: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_out_infos);
78
79 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
80 GLWE::encrypt_sk_tmp_bytes(module, &glwe_in_infos)
81 .max(GLWE::decrypt_tmp_bytes(module, &glwe_out_infos))
82 .max(module.glwe_tensor_apply_tmp_bytes(&res_tensor, 0, &a, &b))
83 .max(module.glwe_secret_tensor_prepare_tmp_bytes(rank.into()))
84 .max(module.glwe_tensor_relinearize_tmp_bytes(&res_relin, &res_tensor, &tsk_infos)),
85 );
86
87 let mut source_xs: Source = Source::new([0u8; 32]);
88 let mut source_xe: Source = Source::new([0u8; 32]);
89 let mut source_xa: Source = Source::new([0u8; 32]);
90
91 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(module.n().into(), rank.into());
92 sk.fill_ternary_prob(0.5, &mut source_xs);
93
94 let mut sk_dft: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc_from_infos(module, &sk);
95 sk_dft.prepare(module, &sk);
96
97 let mut sk_tensor: GLWESecretTensor<Vec<u8>> = GLWESecretTensor::alloc(module.n().into(), rank.into());
98 sk_tensor.prepare(module, &sk, scratch.borrow());
99
100 let mut sk_tensor_prep: GLWESecretTensorPrepared<Vec<u8>, BE> = GLWESecretTensorPrepared::alloc(module, rank.into());
101 sk_tensor_prep.prepare(module, &sk_tensor);
102
103 let mut tsk: GLWETensorKey<Vec<u8>> = GLWETensorKey::alloc_from_infos(&tsk_infos);
104 tsk.encrypt_sk(module, &sk, &mut source_xa, &mut source_xe, scratch.borrow());
105
106 let mut tsk_prep: GLWETensorKeyPrepared<Vec<u8>, BE> = GLWETensorKeyPrepared::alloc_from_infos(module, &tsk_infos);
107 tsk_prep.prepare(module, &tsk, scratch.borrow());
108
109 let scale: usize = 2 * in_base2k;
110
111 let mut data = vec![0i64; n];
112 for i in data.iter_mut() {
113 *i = (source_xa.next_i64() & 7) - 4;
114 }
115
116 pt_in.encode_vec_i64(&data, TorusPrecision(scale as u32));
117
118 let mut pt_want_base2k_in = VecZnx::alloc(n, 1, pt_in.size());
119 bivariate_convolution_naive(
120 module,
121 in_base2k,
122 2,
123 &mut pt_want_base2k_in,
124 0,
125 pt_in.data(),
126 0,
127 pt_in.data(),
128 0,
129 scratch.borrow(),
130 );
131
132 a.encrypt_sk(module, &pt_in, &sk_dft, &mut source_xa, &mut source_xe, scratch.borrow());
133 b.encrypt_sk(module, &pt_in, &sk_dft, &mut source_xa, &mut source_xe, scratch.borrow());
134
135 for res_offset in 0..scale {
136 module.glwe_tensor_apply(&mut res_tensor, scale + res_offset, &a, &b, scratch.borrow());
137
138 res_tensor.decrypt(module, &mut pt_have, &sk_dft, &sk_tensor_prep, scratch.borrow());
139 module.vec_znx_normalize(
140 pt_want.data_mut(),
141 out_base2k,
142 res_offset as i64,
143 0,
144 &pt_want_base2k_in,
145 in_base2k,
146 0,
147 scratch.borrow(),
148 );
149
150 module.glwe_sub(&mut pt_tmp, &pt_have, &pt_want);
151 module.vec_znx_normalize_inplace(pt_tmp.base2k().as_usize(), &mut pt_tmp.data, 0, scratch.borrow());
152
153 let noise_have: f64 = pt_tmp.stats().std().log2();
154 let noise_want = -((k - scale - res_offset - module.log_n()) as f64 - ((rank - 1) as f64) / SQRT_2);
155
156 assert!(noise_have - noise_want <= 0.5, "{} > {}", noise_have, noise_want);
157
158 module.glwe_tensor_relinearize(&mut res_relin, &res_tensor, &tsk_prep, tsk_prep.size(), scratch.borrow());
159 res_relin.decrypt(module, &mut pt_have, &sk_dft, scratch.borrow());
160
161 module.glwe_sub(&mut pt_tmp, &pt_have, &pt_want);
162 module.vec_znx_normalize_inplace(pt_tmp.base2k().as_usize(), &mut pt_tmp.data, 0, scratch.borrow());
163
164 // We can reuse the same noise bound because the relinearization noise (which is additive)
165 // is much smaller than the tensoring noise (which is multiplicative)
166 let noise_have: f64 = pt_tmp.stats().std().log2();
167 assert!(noise_have - noise_want <= 0.5, "{} > {}", noise_have, noise_want);
168 }
169 }
170}
171
172pub fn test_glwe_mul_plain<BE: Backend>(module: &Module<BE>)
173where
174 Module<BE>: GLWEEncryptSk<BE>
175 + GLWEDecrypt<BE>
176 + VecZnxFillUniform
177 + GLWESecretPreparedFactory<BE>
178 + GLWESub
179 + VecZnxNormalizeInplace<BE>
180 + VecZnxCopy
181 + VecZnxNormalize<BE>
182 + GLWEMulPlain<BE>,
183 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
184 Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
185{
186 let in_base2k: usize = 16;
187 let out_base2k: usize = 13;
188 let k: usize = 128;
189
190 for rank in 1_usize..=3 {
191 let n: usize = module.n();
192
193 let glwe_in_infos: GLWELayout = GLWELayout {
194 n: n.into(),
195 base2k: in_base2k.into(),
196 k: k.into(),
197 rank: rank.into(),
198 };
199
200 let glwe_out_infos: GLWELayout = GLWELayout {
201 n: n.into(),
202 base2k: out_base2k.into(),
203 k: k.into(),
204 rank: rank.into(),
205 };
206
207 let mut a: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_in_infos);
208 let mut res: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_out_infos);
209 let mut pt_a: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_in_infos);
210 let mut pt_b: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module.n().into(), in_base2k.into(), (2 * in_base2k).into());
211 let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_out_infos);
212 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_out_infos);
213 let mut pt_tmp: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_out_infos);
214
215 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
216 GLWE::encrypt_sk_tmp_bytes(module, &glwe_in_infos)
217 .max(GLWE::decrypt_tmp_bytes(module, &glwe_out_infos))
218 .max(module.glwe_mul_plain_tmp_bytes(&res, 0, &a, &pt_b)),
219 );
220
221 let mut source_xs: Source = Source::new([0u8; 32]);
222 let mut source_xe: Source = Source::new([0u8; 32]);
223 let mut source_xa: Source = Source::new([0u8; 32]);
224
225 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(module.n().into(), rank.into());
226 sk.fill_ternary_prob(0.5, &mut source_xs);
227
228 let mut sk_dft: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc_from_infos(module, &sk);
229 sk_dft.prepare(module, &sk);
230
231 let scale: usize = 2 * in_base2k;
232
233 pt_b.data_mut().fill_uniform(in_base2k, &mut source_xa);
234 pt_a.data_mut().fill_uniform(in_base2k, &mut source_xa);
235
236 let mut pt_want_base2k_in = VecZnx::alloc(n, 1, pt_a.size() + pt_b.size());
237 bivariate_convolution_naive(
238 module,
239 in_base2k,
240 2,
241 &mut pt_want_base2k_in,
242 0,
243 pt_a.data(),
244 0,
245 pt_b.data(),
246 0,
247 scratch.borrow(),
248 );
249
250 a.encrypt_sk(module, &pt_a, &sk_dft, &mut source_xa, &mut source_xe, scratch.borrow());
251
252 for res_offset in 0..scale {
253 module.glwe_mul_plain(&mut res, scale + res_offset, &a, &pt_b, scratch.borrow());
254
255 res.decrypt(module, &mut pt_have, &sk_dft, scratch.borrow());
256 module.vec_znx_normalize(
257 pt_want.data_mut(),
258 out_base2k,
259 res_offset as i64,
260 0,
261 &pt_want_base2k_in,
262 in_base2k,
263 0,
264 scratch.borrow(),
265 );
266
267 module.glwe_sub(&mut pt_tmp, &pt_have, &pt_want);
268 module.vec_znx_normalize_inplace(pt_tmp.base2k().as_usize(), &mut pt_tmp.data, 0, scratch.borrow());
269
270 let noise_have: f64 = pt_tmp.stats().std().log2();
271 let noise_want = -((k - scale - res_offset - module.log_n()) as f64 - ((rank - 1) as f64) / SQRT_2);
272
273 assert!(noise_have - noise_want <= 0.5, "{} > {}", noise_have, noise_want);
274 }
275 }
276}
277
278pub fn test_glwe_mul_const<BE: Backend>(module: &Module<BE>)
279where
280 Module<BE>: GLWEEncryptSk<BE>
281 + GLWEDecrypt<BE>
282 + VecZnxFillUniform
283 + GLWESecretPreparedFactory<BE>
284 + GLWESub
285 + VecZnxNormalizeInplace<BE>
286 + VecZnxCopy
287 + VecZnxNormalize<BE>
288 + GLWEMulConst<BE>,
289 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
290 Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
291{
292 let in_base2k: usize = 16;
293 let out_base2k: usize = 13;
294 let k: usize = 128;
295 let b_size = 3;
296
297 for rank in 1_usize..=3 {
298 let n: usize = module.n();
299
300 let glwe_in_infos: GLWELayout = GLWELayout {
301 n: n.into(),
302 base2k: in_base2k.into(),
303 k: k.into(),
304 rank: rank.into(),
305 };
306
307 let glwe_out_infos: GLWELayout = GLWELayout {
308 n: n.into(),
309 base2k: out_base2k.into(),
310 k: k.into(),
311 rank: rank.into(),
312 };
313
314 let mut a: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_in_infos);
315 let mut res: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_out_infos);
316 let mut pt_a: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_in_infos);
317 let mut pt_b: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(module.n().into(), in_base2k.into(), (2 * in_base2k).into());
318 let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_out_infos);
319 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_out_infos);
320 let mut pt_tmp: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_out_infos);
321
322 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
323 GLWE::encrypt_sk_tmp_bytes(module, &glwe_in_infos)
324 .max(GLWE::decrypt_tmp_bytes(module, &glwe_out_infos))
325 .max(module.glwe_mul_const_tmp_bytes(&res, 0, &a, b_size)),
326 );
327
328 let mut source_xs: Source = Source::new([0u8; 32]);
329 let mut source_xe: Source = Source::new([0u8; 32]);
330 let mut source_xa: Source = Source::new([0u8; 32]);
331
332 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(module.n().into(), rank.into());
333 sk.fill_ternary_prob(0.5, &mut source_xs);
334
335 let mut sk_dft: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc_from_infos(module, &sk);
336 sk_dft.prepare(module, &sk);
337
338 let scale: usize = 2 * in_base2k;
339
340 pt_a.data_mut().fill_uniform(in_base2k, &mut source_xa);
341
342 let mut b_const = vec![0i64; b_size];
343 let mask = (1 << in_base2k) - 1;
344 for (j, x) in b_const[..1].iter_mut().enumerate() {
345 let r = source_xa.next_u64() & mask;
346 *x = ((r << (64 - in_base2k)) as i64) >> (64 - in_base2k);
347 pt_b.data_mut().at_mut(0, j)[0] = *x
348 }
349
350 let mut pt_want_base2k_in = VecZnx::alloc(n, 1, pt_a.size() + pt_b.size());
351 bivariate_convolution_naive(
352 module,
353 in_base2k,
354 2,
355 &mut pt_want_base2k_in,
356 0,
357 pt_a.data(),
358 0,
359 pt_b.data(),
360 0,
361 scratch.borrow(),
362 );
363
364 a.encrypt_sk(module, &pt_a, &sk_dft, &mut source_xa, &mut source_xe, scratch.borrow());
365
366 for res_offset in 0..scale {
367 module.glwe_mul_const(&mut res, scale + res_offset, &a, &b_const, scratch.borrow());
368
369 res.decrypt(module, &mut pt_have, &sk_dft, scratch.borrow());
370 module.vec_znx_normalize(
371 pt_want.data_mut(),
372 out_base2k,
373 res_offset as i64,
374 0,
375 &pt_want_base2k_in,
376 in_base2k,
377 0,
378 scratch.borrow(),
379 );
380
381 module.glwe_sub(&mut pt_tmp, &pt_have, &pt_want);
382 module.vec_znx_normalize_inplace(pt_tmp.base2k().as_usize(), &mut pt_tmp.data, 0, scratch.borrow());
383
384 let noise_have: f64 = pt_tmp.stats().std().log2();
385 let noise_want = -((k - scale - res_offset - module.log_n()) as f64 - ((rank - 1) as f64) / SQRT_2);
386
387 assert!(noise_have - noise_want <= 0.5, "{} > {}", noise_have, noise_want);
388 }
389 }
390}