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