poulpy/poulpy-core/src/tests/test_suite/keyswitch/ggsw_ct.rs
poulpy-core/src/tests/test_suite/keyswitch/ggsw_ct.rsRaw
1use poulpy_hal::{
2 api::{
3 ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyDftToDft, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolAllocBytes,
4 SvpPrepare, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace,
5 VecZnxBigAlloc, VecZnxBigAllocBytes, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftAddInplace, VecZnxDftAlloc,
6 VecZnxDftAllocBytes, VecZnxDftApply, VecZnxDftCopy, VecZnxFillUniform, VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA,
7 VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubInplace, VecZnxSwitchRing,
8 VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes, VmpPMatAlloc, VmpPrepare,
9 },
10 layouts::{Backend, Module, ScalarZnx, 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 GGLWESwitchingKey, GGLWESwitchingKeyLayout, GGLWETensorKey, GGLWETensorKeyLayout, GGSWCiphertext, GGSWCiphertextLayout,
22 GLWESecret,
23 prepared::{GGLWESwitchingKeyPrepared, GGLWETensorKeyPrepared, GLWESecretPrepared, PrepareAlloc},
24 },
25 noise::noise_ggsw_keyswitch,
26};
27
28#[allow(clippy::too_many_arguments)]
29pub fn test_ggsw_keyswitch<B>(module: &Module<B>)
30where
31 Module<B>: VecZnxDftAllocBytes
32 + VecZnxBigNormalize<B>
33 + VecZnxDftApply<B>
34 + SvpApplyDftToDftInplace<B>
35 + VecZnxIdftApplyConsume<B>
36 + VecZnxFillUniform
37 + VecZnxSubInplace
38 + VecZnxAddInplace
39 + VecZnxNormalizeInplace<B>
40 + VecZnxAddNormal
41 + VecZnxNormalize<B>
42 + VecZnxSub
43 + SvpPrepare<B>
44 + SvpPPolAllocBytes
45 + SvpPPolAlloc<B>
46 + VecZnxBigAllocBytes
47 + VecZnxBigAddInplace<B>
48 + VecZnxBigAddSmallInplace<B>
49 + VecZnxNormalizeTmpBytes
50 + VecZnxAddScalarInplace
51 + VmpPMatAlloc<B>
52 + VmpPrepare<B>
53 + VmpApplyDftToDftTmpBytes
54 + VmpApplyDftToDft<B>
55 + VmpApplyDftToDftAdd<B>
56 + VecZnxBigNormalizeTmpBytes
57 + VecZnxSwitchRing
58 + SvpApplyDftToDft<B>
59 + VecZnxIdftApplyTmpA<B>
60 + VecZnxDftCopy<B>
61 + VecZnxDftAddInplace<B>
62 + VecZnxBigAlloc<B>
63 + VecZnxDftAlloc<B>,
64 B: Backend
65 + TakeVecZnxDftImpl<B>
66 + TakeVecZnxBigImpl<B>
67 + TakeSvpPPolImpl<B>
68 + ScratchOwnedAllocImpl<B>
69 + ScratchOwnedBorrowImpl<B>
70 + ScratchAvailableImpl<B>
71 + TakeScalarZnxImpl<B>
72 + TakeVecZnxImpl<B>,
73{
74 let base2k: usize = 12;
75 let k_in: usize = 54;
76 let digits: usize = k_in.div_ceil(base2k);
77 for rank in 1_usize..3 {
78 for di in 1..digits + 1 {
79 let k_ksk: usize = k_in + base2k * di;
80 let k_tsk: usize = k_ksk;
81 let k_out: usize = k_ksk; // Better capture noise.
82
83 let n: usize = module.n();
84 let rows: usize = k_in.div_ceil(di * base2k);
85
86 let digits_in: usize = 1;
87
88 let ggsw_in_infos: GGSWCiphertextLayout = GGSWCiphertextLayout {
89 n: n.into(),
90 base2k: base2k.into(),
91 k: k_in.into(),
92 rows: rows.into(),
93 digits: digits_in.into(),
94 rank: rank.into(),
95 };
96
97 let ggsw_out_infos: GGSWCiphertextLayout = GGSWCiphertextLayout {
98 n: n.into(),
99 base2k: base2k.into(),
100 k: k_out.into(),
101 rows: rows.into(),
102 digits: digits_in.into(),
103 rank: rank.into(),
104 };
105
106 let tsk_infos: GGLWETensorKeyLayout = GGLWETensorKeyLayout {
107 n: n.into(),
108 base2k: base2k.into(),
109 k: k_tsk.into(),
110 rows: rows.into(),
111 digits: di.into(),
112 rank: rank.into(),
113 };
114
115 let ksk_apply_infos: GGLWESwitchingKeyLayout = GGLWESwitchingKeyLayout {
116 n: n.into(),
117 base2k: base2k.into(),
118 k: k_ksk.into(),
119 rows: rows.into(),
120 digits: di.into(),
121 rank_in: rank.into(),
122 rank_out: rank.into(),
123 };
124
125 let mut ggsw_in: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(&ggsw_in_infos);
126 let mut ggsw_out: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(&ggsw_out_infos);
127 let mut tsk: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(&tsk_infos);
128 let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(&ksk_apply_infos);
129 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
130
131 let mut source_xs: Source = Source::new([0u8; 32]);
132 let mut source_xe: Source = Source::new([0u8; 32]);
133 let mut source_xa: Source = Source::new([0u8; 32]);
134
135 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
136 GGSWCiphertext::encrypt_sk_scratch_space(module, &ggsw_in_infos)
137 | GGLWESwitchingKey::encrypt_sk_scratch_space(module, &ksk_apply_infos)
138 | GGLWETensorKey::encrypt_sk_scratch_space(module, &tsk_infos)
139 | GGSWCiphertext::keyswitch_scratch_space(
140 module,
141 &ggsw_out_infos,
142 &ggsw_in_infos,
143 &ksk_apply_infos,
144 &tsk_infos,
145 ),
146 );
147
148 let var_xs: f64 = 0.5;
149
150 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc_with(n.into(), rank.into());
151 sk_in.fill_ternary_prob(var_xs, &mut source_xs);
152 let sk_in_dft: GLWESecretPrepared<Vec<u8>, B> = sk_in.prepare_alloc(module, scratch.borrow());
153
154 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc_with(n.into(), rank.into());
155 sk_out.fill_ternary_prob(var_xs, &mut source_xs);
156 let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
157
158 ksk.encrypt_sk(
159 module,
160 &sk_in,
161 &sk_out,
162 &mut source_xa,
163 &mut source_xe,
164 scratch.borrow(),
165 );
166 tsk.encrypt_sk(
167 module,
168 &sk_out,
169 &mut source_xa,
170 &mut source_xe,
171 scratch.borrow(),
172 );
173
174 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
175
176 ggsw_in.encrypt_sk(
177 module,
178 &pt_scalar,
179 &sk_in_dft,
180 &mut source_xa,
181 &mut source_xe,
182 scratch.borrow(),
183 );
184
185 let ksk_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
186 let tsk_prepared: GGLWETensorKeyPrepared<Vec<u8>, B> = tsk.prepare_alloc(module, scratch.borrow());
187
188 ggsw_out.keyswitch(
189 module,
190 &ggsw_in,
191 &ksk_prepared,
192 &tsk_prepared,
193 scratch.borrow(),
194 );
195
196 let max_noise = |col_j: usize| -> f64 {
197 noise_ggsw_keyswitch(
198 n as f64,
199 base2k * di,
200 col_j,
201 var_xs,
202 0f64,
203 SIGMA * SIGMA,
204 0f64,
205 rank as f64,
206 k_in,
207 k_ksk,
208 k_tsk,
209 ) + 0.5
210 };
211
212 ggsw_out.assert_noise(module, &sk_out_prepared, &pt_scalar, max_noise);
213 }
214 }
215}
216
217#[allow(clippy::too_many_arguments)]
218pub fn test_ggsw_keyswitch_inplace<B>(module: &Module<B>)
219where
220 Module<B>: VecZnxDftAllocBytes
221 + VecZnxBigNormalize<B>
222 + VecZnxDftApply<B>
223 + SvpApplyDftToDftInplace<B>
224 + VecZnxIdftApplyConsume<B>
225 + VecZnxFillUniform
226 + VecZnxSubInplace
227 + VecZnxAddInplace
228 + VecZnxNormalizeInplace<B>
229 + VecZnxAddNormal
230 + VecZnxNormalize<B>
231 + VecZnxSub
232 + SvpPrepare<B>
233 + SvpPPolAllocBytes
234 + SvpPPolAlloc<B>
235 + VecZnxBigAllocBytes
236 + VecZnxBigAddInplace<B>
237 + VecZnxBigAddSmallInplace<B>
238 + VecZnxNormalizeTmpBytes
239 + VecZnxAddScalarInplace
240 + VmpPMatAlloc<B>
241 + VmpPrepare<B>
242 + VmpApplyDftToDftTmpBytes
243 + VmpApplyDftToDft<B>
244 + VmpApplyDftToDftAdd<B>
245 + VecZnxBigNormalizeTmpBytes
246 + VecZnxSwitchRing
247 + SvpApplyDftToDft<B>
248 + VecZnxIdftApplyTmpA<B>
249 + VecZnxDftCopy<B>
250 + VecZnxDftAddInplace<B>
251 + VecZnxBigAlloc<B>
252 + VecZnxDftAlloc<B>,
253 B: Backend
254 + TakeVecZnxDftImpl<B>
255 + TakeVecZnxBigImpl<B>
256 + TakeSvpPPolImpl<B>
257 + ScratchOwnedAllocImpl<B>
258 + ScratchOwnedBorrowImpl<B>
259 + ScratchAvailableImpl<B>
260 + TakeScalarZnxImpl<B>
261 + TakeVecZnxImpl<B>,
262{
263 let base2k: usize = 12;
264 let k_out: usize = 54;
265 let digits: usize = k_out.div_ceil(base2k);
266 for rank in 1_usize..3 {
267 for di in 1..digits + 1 {
268 let k_ksk: usize = k_out + base2k * di;
269 let k_tsk: usize = k_ksk;
270
271 let n: usize = module.n();
272 let rows: usize = k_out.div_ceil(di * base2k);
273
274 let digits_in: usize = 1;
275
276 let ggsw_out_infos: GGSWCiphertextLayout = GGSWCiphertextLayout {
277 n: n.into(),
278 base2k: base2k.into(),
279 k: k_out.into(),
280 rows: rows.into(),
281 digits: digits_in.into(),
282 rank: rank.into(),
283 };
284
285 let tsk_infos: GGLWETensorKeyLayout = GGLWETensorKeyLayout {
286 n: n.into(),
287 base2k: base2k.into(),
288 k: k_tsk.into(),
289 rows: rows.into(),
290 digits: di.into(),
291 rank: rank.into(),
292 };
293
294 let ksk_apply_infos: GGLWESwitchingKeyLayout = GGLWESwitchingKeyLayout {
295 n: n.into(),
296 base2k: base2k.into(),
297 k: k_ksk.into(),
298 rows: rows.into(),
299 digits: di.into(),
300 rank_in: rank.into(),
301 rank_out: rank.into(),
302 };
303
304 let mut ggsw_out: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(&ggsw_out_infos);
305 let mut tsk: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(&tsk_infos);
306 let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(&ksk_apply_infos);
307 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
308
309 let mut source_xs: Source = Source::new([0u8; 32]);
310 let mut source_xe: Source = Source::new([0u8; 32]);
311 let mut source_xa: Source = Source::new([0u8; 32]);
312
313 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
314 GGSWCiphertext::encrypt_sk_scratch_space(module, &ggsw_out_infos)
315 | GGLWESwitchingKey::encrypt_sk_scratch_space(module, &ksk_apply_infos)
316 | GGLWETensorKey::encrypt_sk_scratch_space(module, &tsk_infos)
317 | GGSWCiphertext::keyswitch_inplace_scratch_space(module, &ggsw_out_infos, &ksk_apply_infos, &tsk_infos),
318 );
319
320 let var_xs: f64 = 0.5;
321
322 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc_with(n.into(), rank.into());
323 sk_in.fill_ternary_prob(var_xs, &mut source_xs);
324 let sk_in_dft: GLWESecretPrepared<Vec<u8>, B> = sk_in.prepare_alloc(module, scratch.borrow());
325
326 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc_with(n.into(), rank.into());
327 sk_out.fill_ternary_prob(var_xs, &mut source_xs);
328 let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
329
330 ksk.encrypt_sk(
331 module,
332 &sk_in,
333 &sk_out,
334 &mut source_xa,
335 &mut source_xe,
336 scratch.borrow(),
337 );
338 tsk.encrypt_sk(
339 module,
340 &sk_out,
341 &mut source_xa,
342 &mut source_xe,
343 scratch.borrow(),
344 );
345
346 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
347
348 ggsw_out.encrypt_sk(
349 module,
350 &pt_scalar,
351 &sk_in_dft,
352 &mut source_xa,
353 &mut source_xe,
354 scratch.borrow(),
355 );
356
357 let ksk_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
358 let tsk_prepared: GGLWETensorKeyPrepared<Vec<u8>, B> = tsk.prepare_alloc(module, scratch.borrow());
359
360 ggsw_out.keyswitch_inplace(module, &ksk_prepared, &tsk_prepared, scratch.borrow());
361
362 let max_noise = |col_j: usize| -> f64 {
363 noise_ggsw_keyswitch(
364 n as f64,
365 base2k * di,
366 col_j,
367 var_xs,
368 0f64,
369 SIGMA * SIGMA,
370 0f64,
371 rank as f64,
372 k_out,
373 k_ksk,
374 k_tsk,
375 ) + 0.5
376 };
377
378 ggsw_out.assert_noise(module, &sk_out_prepared, &pt_scalar, max_noise);
379 }
380 }
381}