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