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