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::{ScratchAvailable, ScratchOwnedAlloc, ScratchOwnedBorrow},
3 layouts::{Backend, Module, ScalarZnx, Scratch, ScratchOwned},
4 source::Source,
5};
6
7use crate::{
8 GGSWEncryptSk, GGSWKeyswitch, GGSWNoise, GLWESwitchingKeyEncryptSk, GLWETensorKeyEncryptSk, ScratchTakeCore,
9 encryption::SIGMA,
10 layouts::{
11 GGSW, GGSWLayout, GLWESecret, GLWESecretPreparedFactory, GLWESwitchingKey, GLWESwitchingKeyLayout,
12 GLWESwitchingKeyPreparedFactory, GLWETensorKey, GLWETensorKeyLayout, GLWETensorKeyPreparedFactory,
13 prepared::{GLWESecretPrepared, GLWESwitchingKeyPrepared, GLWETensorKeyPrepared},
14 },
15 noise::noise_ggsw_keyswitch,
16};
17
18#[allow(clippy::too_many_arguments)]
19pub fn test_ggsw_keyswitch<BE: Backend>(module: &Module<BE>)
20where
21 Module<BE>: GGSWEncryptSk<BE>
22 + GLWESwitchingKeyEncryptSk<BE>
23 + GLWETensorKeyEncryptSk<BE>
24 + GGSWKeyswitch<BE>
25 + GLWESecretPreparedFactory<BE>
26 + GLWETensorKeyPreparedFactory<BE>
27 + GLWESwitchingKeyPreparedFactory<BE>
28 + GGSWNoise<BE>,
29 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
30 Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
31{
32 let base2k: usize = 12;
33 let k_in: usize = 54;
34 let dsize: usize = k_in.div_ceil(base2k);
35 for rank in 1_usize..3 {
36 for di in 1..dsize + 1 {
37 let k_ksk: usize = k_in + base2k * di;
38 let k_tsk: usize = k_ksk;
39 let k_out: usize = k_ksk; // Better capture noise.
40
41 let n: usize = module.n();
42 let dnum: usize = k_in.div_ceil(di * base2k);
43
44 let dsize_in: usize = 1;
45
46 let ggsw_in_infos: GGSWLayout = GGSWLayout {
47 n: n.into(),
48 base2k: base2k.into(),
49 k: k_in.into(),
50 dnum: dnum.into(),
51 dsize: dsize_in.into(),
52 rank: rank.into(),
53 };
54
55 let ggsw_out_infos: GGSWLayout = GGSWLayout {
56 n: n.into(),
57 base2k: base2k.into(),
58 k: k_out.into(),
59 dnum: dnum.into(),
60 dsize: dsize_in.into(),
61 rank: rank.into(),
62 };
63
64 let tsk_infos: GLWETensorKeyLayout = GLWETensorKeyLayout {
65 n: n.into(),
66 base2k: base2k.into(),
67 k: k_tsk.into(),
68 dnum: dnum.into(),
69 dsize: di.into(),
70 rank: rank.into(),
71 };
72
73 let ksk_apply_infos: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
74 n: n.into(),
75 base2k: base2k.into(),
76 k: k_ksk.into(),
77 dnum: dnum.into(),
78 dsize: di.into(),
79 rank_in: rank.into(),
80 rank_out: rank.into(),
81 };
82
83 let mut ggsw_in: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_in_infos);
84 let mut ggsw_out: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_out_infos);
85 let mut tsk: GLWETensorKey<Vec<u8>> = GLWETensorKey::alloc_from_infos(&tsk_infos);
86 let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&ksk_apply_infos);
87 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
88
89 let mut source_xs: Source = Source::new([0u8; 32]);
90 let mut source_xe: Source = Source::new([0u8; 32]);
91 let mut source_xa: Source = Source::new([0u8; 32]);
92
93 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
94 GGSW::encrypt_sk_tmp_bytes(module, &ggsw_in_infos)
95 | GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &ksk_apply_infos)
96 | GLWETensorKey::encrypt_sk_tmp_bytes(module, &tsk_infos)
97 | GGSW::keyswitch_tmp_bytes(
98 module,
99 &ggsw_out_infos,
100 &ggsw_in_infos,
101 &ksk_apply_infos,
102 &tsk_infos,
103 ),
104 );
105
106 let var_xs: f64 = 0.5;
107
108 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
109 sk_in.fill_ternary_prob(var_xs, &mut source_xs);
110
111 let mut sk_in_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
112 sk_in_prepared.prepare(module, &sk_in);
113
114 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
115 sk_out.fill_ternary_prob(var_xs, &mut source_xs);
116
117 let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
118 sk_out_prepared.prepare(module, &sk_out);
119
120 ksk.encrypt_sk(
121 module,
122 &sk_in,
123 &sk_out,
124 &mut source_xa,
125 &mut source_xe,
126 scratch.borrow(),
127 );
128 tsk.encrypt_sk(
129 module,
130 &sk_out,
131 &mut source_xa,
132 &mut source_xe,
133 scratch.borrow(),
134 );
135
136 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
137
138 ggsw_in.encrypt_sk(
139 module,
140 &pt_scalar,
141 &sk_in_prepared,
142 &mut source_xa,
143 &mut source_xe,
144 scratch.borrow(),
145 );
146
147 let mut ksk_prepared: GLWESwitchingKeyPrepared<Vec<u8>, BE> =
148 GLWESwitchingKeyPrepared::alloc_from_infos(module, &ksk);
149 ksk_prepared.prepare(module, &ksk, scratch.borrow());
150
151 let mut tsk_prepared: GLWETensorKeyPrepared<Vec<u8>, BE> = GLWETensorKeyPrepared::alloc_from_infos(module, &tsk);
152 tsk_prepared.prepare(module, &tsk, scratch.borrow());
153
154 ggsw_out.keyswitch(
155 module,
156 &ggsw_in,
157 &ksk_prepared,
158 &tsk_prepared,
159 scratch.borrow(),
160 );
161
162 let max_noise = |col_j: usize| -> f64 {
163 noise_ggsw_keyswitch(
164 n as f64,
165 base2k * di,
166 col_j,
167 var_xs,
168 0f64,
169 SIGMA * SIGMA,
170 0f64,
171 rank as f64,
172 k_in,
173 k_ksk,
174 k_tsk,
175 ) + 0.5
176 };
177
178 ggsw_out.assert_noise(module, &sk_out_prepared, &pt_scalar, max_noise);
179 }
180 }
181}
182
183#[allow(clippy::too_many_arguments)]
184pub fn test_ggsw_keyswitch_inplace<BE: Backend>(module: &Module<BE>)
185where
186 Module<BE>: GGSWEncryptSk<BE>
187 + GLWESwitchingKeyEncryptSk<BE>
188 + GLWETensorKeyEncryptSk<BE>
189 + GGSWKeyswitch<BE>
190 + GLWESecretPreparedFactory<BE>
191 + GLWETensorKeyPreparedFactory<BE>
192 + GLWESwitchingKeyPreparedFactory<BE>
193 + GGSWNoise<BE>,
194 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
195 Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
196{
197 let base2k: usize = 12;
198 let k_out: usize = 54;
199 let dsize: usize = k_out.div_ceil(base2k);
200 for rank in 1_usize..3 {
201 for di in 1..dsize + 1 {
202 let k_ksk: usize = k_out + base2k * di;
203 let k_tsk: usize = k_ksk;
204
205 let n: usize = module.n();
206 let dnum: usize = k_out.div_ceil(di * base2k);
207
208 let dsize_in: usize = 1;
209
210 let ggsw_out_infos: GGSWLayout = GGSWLayout {
211 n: n.into(),
212 base2k: base2k.into(),
213 k: k_out.into(),
214 dnum: dnum.into(),
215 dsize: dsize_in.into(),
216 rank: rank.into(),
217 };
218
219 let tsk_infos: GLWETensorKeyLayout = GLWETensorKeyLayout {
220 n: n.into(),
221 base2k: base2k.into(),
222 k: k_tsk.into(),
223 dnum: dnum.into(),
224 dsize: di.into(),
225 rank: rank.into(),
226 };
227
228 let ksk_apply_infos: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
229 n: n.into(),
230 base2k: base2k.into(),
231 k: k_ksk.into(),
232 dnum: dnum.into(),
233 dsize: di.into(),
234 rank_in: rank.into(),
235 rank_out: rank.into(),
236 };
237
238 let mut ggsw_out: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_out_infos);
239 let mut tsk: GLWETensorKey<Vec<u8>> = GLWETensorKey::alloc_from_infos(&tsk_infos);
240 let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&ksk_apply_infos);
241 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
242
243 let mut source_xs: Source = Source::new([0u8; 32]);
244 let mut source_xe: Source = Source::new([0u8; 32]);
245 let mut source_xa: Source = Source::new([0u8; 32]);
246
247 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
248 GGSW::encrypt_sk_tmp_bytes(module, &ggsw_out_infos)
249 | GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &ksk_apply_infos)
250 | GLWETensorKey::encrypt_sk_tmp_bytes(module, &tsk_infos)
251 | GGSW::keyswitch_tmp_bytes(
252 module,
253 &ggsw_out_infos,
254 &ggsw_out_infos,
255 &ksk_apply_infos,
256 &tsk_infos,
257 ),
258 );
259
260 let var_xs: f64 = 0.5;
261
262 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
263 sk_in.fill_ternary_prob(var_xs, &mut source_xs);
264
265 let mut sk_in_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
266 sk_in_prepared.prepare(module, &sk_in);
267
268 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
269 sk_out.fill_ternary_prob(var_xs, &mut source_xs);
270
271 let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
272 sk_out_prepared.prepare(module, &sk_out);
273
274 ksk.encrypt_sk(
275 module,
276 &sk_in,
277 &sk_out,
278 &mut source_xa,
279 &mut source_xe,
280 scratch.borrow(),
281 );
282 tsk.encrypt_sk(
283 module,
284 &sk_out,
285 &mut source_xa,
286 &mut source_xe,
287 scratch.borrow(),
288 );
289
290 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
291
292 ggsw_out.encrypt_sk(
293 module,
294 &pt_scalar,
295 &sk_in_prepared,
296 &mut source_xa,
297 &mut source_xe,
298 scratch.borrow(),
299 );
300
301 let mut ksk_prepared: GLWESwitchingKeyPrepared<Vec<u8>, BE> =
302 GLWESwitchingKeyPrepared::alloc_from_infos(module, &ksk);
303 ksk_prepared.prepare(module, &ksk, scratch.borrow());
304
305 let mut tsk_prepared: GLWETensorKeyPrepared<Vec<u8>, BE> = GLWETensorKeyPrepared::alloc_from_infos(module, &tsk);
306 tsk_prepared.prepare(module, &tsk, scratch.borrow());
307
308 ggsw_out.keyswitch_inplace(module, &ksk_prepared, &tsk_prepared, scratch.borrow());
309
310 let max_noise = |col_j: usize| -> f64 {
311 noise_ggsw_keyswitch(
312 n as f64,
313 base2k * di,
314 col_j,
315 var_xs,
316 0f64,
317 SIGMA * SIGMA,
318 0f64,
319 rank as f64,
320 k_out,
321 k_ksk,
322 k_tsk,
323 ) + 0.5
324 };
325
326 ggsw_out.assert_noise(module, &sk_out_prepared, &pt_scalar, max_noise);
327 }
328 }
329}