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