poulpy/poulpy-core/src/tests/test_suite/keyswitch/glwe_ct.rs
poulpy-core/src/tests/test_suite/keyswitch/glwe_ct.rsRaw
1use poulpy_hal::{
2 api::{
3 ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyDftToDftInplace, SvpPPolAlloc, SvpPPolBytesOf, SvpPrepare,
4 VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace,
5 VecZnxBigBytesOf, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftApply, VecZnxDftBytesOf, VecZnxFillUniform,
6 VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubInplace,
7 VecZnxSwitchRing, VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes, VmpPMatAlloc, VmpPrepare,
8 },
9 layouts::{Backend, Module, ScratchOwned},
10 oep::{
11 ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
12 TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl,
13 },
14 source::Source,
15};
16
17use crate::{
18 encryption::SIGMA,
19 layouts::{
20 GLWE, GLWELayout, GLWEPlaintext, GLWESecret, GLWESwitchingKey, GLWESwitchingKeyLayout,
21 prepared::{GLWESecretPrepared, GLWESwitchingKeyPrepared, PrepareAlloc},
22 },
23 noise::log2_std_noise_gglwe_product,
24};
25
26#[allow(clippy::too_many_arguments)]
27pub fn test_glwe_keyswitch<B>(module: &Module<B>)
28where
29 Module<B>: VecZnxDftBytesOf
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 + SvpPPolBytesOf
43 + SvpPPolAlloc<B>
44 + VecZnxBigBytesOf
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 B: Backend
57 + TakeVecZnxDftImpl<B>
58 + TakeVecZnxBigImpl<B>
59 + TakeSvpPPolImpl<B>
60 + ScratchOwnedAllocImpl<B>
61 + ScratchOwnedBorrowImpl<B>
62 + ScratchAvailableImpl<B>
63 + TakeScalarZnxImpl<B>
64 + TakeVecZnxImpl<B>,
65{
66 let base2k: usize = 12;
67 let k_in: usize = 45;
68 let dsize: usize = k_in.div_ceil(base2k);
69
70 for rank_in in 1_usize..3 {
71 for rank_out in 1_usize..3 {
72 for di in 1_usize..dsize + 1 {
73 let k_ksk: usize = k_in + base2k * di;
74 let k_out: usize = k_ksk; // better capture noise
75
76 let n: usize = module.n();
77 let dnum: usize = k_in.div_ceil(base2k * dsize);
78
79 let glwe_in_infos: GLWELayout = GLWELayout {
80 n: n.into(),
81 base2k: base2k.into(),
82 k: k_in.into(),
83 rank: rank_in.into(),
84 };
85
86 let glwe_out_infos: GLWELayout = GLWELayout {
87 n: n.into(),
88 base2k: base2k.into(),
89 k: k_out.into(),
90 rank: rank_out.into(),
91 };
92
93 let key_apply: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
94 n: n.into(),
95 base2k: base2k.into(),
96 k: k_ksk.into(),
97 dnum: dnum.into(),
98 dsize: di.into(),
99 rank_in: rank_in.into(),
100 rank_out: rank_out.into(),
101 };
102
103 let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&key_apply);
104 let mut glwe_in: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_in_infos);
105 let mut glwe_out: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_out_infos);
106 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_in_infos);
107
108 let mut source_xs: Source = Source::new([0u8; 32]);
109 let mut source_xe: Source = Source::new([0u8; 32]);
110 let mut source_xa: Source = Source::new([0u8; 32]);
111
112 module.vec_znx_fill_uniform(base2k, &mut pt_want.data, 0, &mut source_xa);
113
114 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
115 GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &key_apply)
116 | GLWE::encrypt_sk_tmp_bytes(module, &glwe_in_infos)
117 | GLWE::keyswitch_tmp_bytes(module, &glwe_out_infos, &glwe_in_infos, &key_apply),
118 );
119
120 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_in.into());
121 sk_in.fill_ternary_prob(0.5, &mut source_xs);
122 let sk_in_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_in.prepare_alloc(module, scratch.borrow());
123
124 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_out.into());
125 sk_out.fill_ternary_prob(0.5, &mut source_xs);
126 let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
127
128 ksk.encrypt_sk(
129 module,
130 &sk_in,
131 &sk_out,
132 &mut source_xa,
133 &mut source_xe,
134 scratch.borrow(),
135 );
136
137 glwe_in.encrypt_sk(
138 module,
139 &pt_want,
140 &sk_in_prepared,
141 &mut source_xa,
142 &mut source_xe,
143 scratch.borrow(),
144 );
145
146 let ksk_prepared: GLWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
147
148 glwe_out.keyswitch(module, &glwe_in, &ksk_prepared, scratch.borrow());
149
150 let max_noise: f64 = log2_std_noise_gglwe_product(
151 module.n() as f64,
152 base2k * dsize,
153 0.5,
154 0.5,
155 0f64,
156 SIGMA * SIGMA,
157 0f64,
158 rank_in as f64,
159 k_in,
160 k_ksk,
161 );
162
163 glwe_out.assert_noise(module, &sk_out_prepared, &pt_want, max_noise + 0.5);
164 }
165 }
166 }
167}
168
169pub fn test_glwe_keyswitch_inplace<B>(module: &Module<B>)
170where
171 Module<B>: VecZnxDftBytesOf
172 + VecZnxBigNormalize<B>
173 + VecZnxDftApply<B>
174 + SvpApplyDftToDftInplace<B>
175 + VecZnxIdftApplyConsume<B>
176 + VecZnxFillUniform
177 + VecZnxSubInplace
178 + VecZnxAddInplace
179 + VecZnxNormalizeInplace<B>
180 + VecZnxAddNormal
181 + VecZnxNormalize<B>
182 + VecZnxSub
183 + SvpPrepare<B>
184 + SvpPPolBytesOf
185 + SvpPPolAlloc<B>
186 + VecZnxBigBytesOf
187 + VecZnxBigAddInplace<B>
188 + VecZnxBigAddSmallInplace<B>
189 + VecZnxNormalizeTmpBytes
190 + VecZnxAddScalarInplace
191 + VmpPMatAlloc<B>
192 + VmpPrepare<B>
193 + VmpApplyDftToDftTmpBytes
194 + VmpApplyDftToDft<B>
195 + VmpApplyDftToDftAdd<B>
196 + VecZnxBigNormalizeTmpBytes
197 + VecZnxSwitchRing,
198 B: Backend
199 + TakeVecZnxDftImpl<B>
200 + TakeVecZnxBigImpl<B>
201 + TakeSvpPPolImpl<B>
202 + ScratchOwnedAllocImpl<B>
203 + ScratchOwnedBorrowImpl<B>
204 + ScratchAvailableImpl<B>
205 + TakeScalarZnxImpl<B>
206 + TakeVecZnxImpl<B>,
207{
208 let base2k: usize = 12;
209 let k_out: usize = 45;
210 let dsize: usize = k_out.div_ceil(base2k);
211
212 for rank in 1_usize..3 {
213 for di in 1..dsize + 1 {
214 let k_ksk: usize = k_out + base2k * di;
215
216 let n: usize = module.n();
217 let dnum: usize = k_out.div_ceil(base2k * dsize);
218
219 let glwe_out_infos: GLWELayout = GLWELayout {
220 n: n.into(),
221 base2k: base2k.into(),
222 k: k_out.into(),
223 rank: rank.into(),
224 };
225
226 let key_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 key_apply: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&key_apply_infos);
237 let mut glwe_out: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_out_infos);
238 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_out_infos);
239
240 let mut source_xs: Source = Source::new([0u8; 32]);
241 let mut source_xe: Source = Source::new([0u8; 32]);
242 let mut source_xa: Source = Source::new([0u8; 32]);
243
244 module.vec_znx_fill_uniform(base2k, &mut pt_want.data, 0, &mut source_xa);
245
246 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
247 GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &key_apply_infos)
248 | GLWE::encrypt_sk_tmp_bytes(module, &glwe_out_infos)
249 | GLWE::keyswitch_inplace_tmp_bytes(module, &glwe_out_infos, &key_apply_infos),
250 );
251
252 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
253 sk_in.fill_ternary_prob(0.5, &mut source_xs);
254 let sk_in_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_in.prepare_alloc(module, scratch.borrow());
255
256 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
257 sk_out.fill_ternary_prob(0.5, &mut source_xs);
258 let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
259
260 key_apply.encrypt_sk(
261 module,
262 &sk_in,
263 &sk_out,
264 &mut source_xa,
265 &mut source_xe,
266 scratch.borrow(),
267 );
268
269 glwe_out.encrypt_sk(
270 module,
271 &pt_want,
272 &sk_in_prepared,
273 &mut source_xa,
274 &mut source_xe,
275 scratch.borrow(),
276 );
277
278 let ksk_prepared: GLWESwitchingKeyPrepared<Vec<u8>, B> = key_apply.prepare_alloc(module, scratch.borrow());
279
280 glwe_out.keyswitch_inplace(module, &ksk_prepared, scratch.borrow());
281
282 let max_noise: f64 = log2_std_noise_gglwe_product(
283 module.n() as f64,
284 base2k * dsize,
285 0.5,
286 0.5,
287 0f64,
288 SIGMA * SIGMA,
289 0f64,
290 rank as f64,
291 k_out,
292 k_ksk,
293 );
294
295 glwe_out.assert_noise(module, &sk_out_prepared, &pt_want, max_noise + 0.5);
296 }
297 }
298}