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