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