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::{ScratchAvailable, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxFillUniform},
3 layouts::{Backend, Module, Scratch, ScratchOwned},
4 source::Source,
5};
6
7use crate::{
8 GLWEEncryptSk, GLWEKeyswitch, GLWENoise, GLWENormalize, GLWESwitchingKeyEncryptSk, ScratchTakeCore,
9 encryption::SIGMA,
10 layouts::{
11 GLWE, GLWELayout, GLWEPlaintext, GLWESecret, GLWESecretPreparedFactory, GLWESwitchingKey, GLWESwitchingKeyLayout,
12 GLWESwitchingKeyPreparedFactory, LWEInfos,
13 prepared::{GLWESecretPrepared, GLWESwitchingKeyPrepared},
14 },
15 var_noise_gglwe_product_v2,
16};
17
18#[allow(clippy::too_many_arguments)]
19pub fn test_glwe_keyswitch<BE: Backend>(module: &Module<BE>)
20where
21 Module<BE>: VecZnxFillUniform
22 + GLWESwitchingKeyEncryptSk<BE>
23 + GLWEEncryptSk<BE>
24 + GLWEKeyswitch<BE>
25 + GLWESecretPreparedFactory<BE>
26 + GLWESwitchingKeyPreparedFactory<BE>
27 + GLWENoise<BE>
28 + GLWENormalize<BE>,
29 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
30 Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
31{
32 let base2k_in: usize = 17;
33 let base2k_key: usize = 13;
34 let base2k_out: usize = 15;
35 let k_in: usize = 102;
36 let max_dsize: usize = k_in.div_ceil(base2k_key);
37
38 for rank_in in 1_usize..3 {
39 for rank_out in 1_usize..3 {
40 for dsize in 1_usize..max_dsize + 1 {
41 let k_ksk: usize = k_in + base2k_key * dsize;
42 let k_out: usize = k_ksk; // better capture noise
43
44 let n: usize = module.n();
45 let dnum: usize = k_in.div_ceil(base2k_key * dsize);
46
47 let glwe_in_infos: GLWELayout = GLWELayout {
48 n: n.into(),
49 base2k: base2k_in.into(),
50 k: k_in.into(),
51 rank: rank_in.into(),
52 };
53
54 let glwe_out_infos: GLWELayout = GLWELayout {
55 n: n.into(),
56 base2k: base2k_out.into(),
57 k: k_out.into(),
58 rank: rank_out.into(),
59 };
60
61 let ksk: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
62 n: n.into(),
63 base2k: base2k_key.into(),
64 k: k_ksk.into(),
65 dnum: dnum.into(),
66 dsize: dsize.into(),
67 rank_in: rank_in.into(),
68 rank_out: rank_out.into(),
69 };
70
71 let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&ksk);
72 let mut glwe_in: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_in_infos);
73 let mut glwe_out: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_out_infos);
74 let mut pt_in: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_in_infos);
75 let mut pt_out: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_out_infos);
76
77 let mut source_xs: Source = Source::new([0u8; 32]);
78 let mut source_xe: Source = Source::new([0u8; 32]);
79 let mut source_xa: Source = Source::new([0u8; 32]);
80
81 module.vec_znx_fill_uniform(pt_in.base2k().into(), &mut pt_in.data, 0, &mut source_xa);
82
83 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
84 GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &ksk)
85 | GLWE::encrypt_sk_tmp_bytes(module, &glwe_in_infos)
86 | GLWE::keyswitch_tmp_bytes(module, &glwe_out_infos, &glwe_in_infos, &ksk),
87 );
88
89 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_in.into());
90 sk_in.fill_ternary_prob(0.5, &mut source_xs);
91
92 let mut sk_in_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank_in.into());
93 sk_in_prepared.prepare(module, &sk_in);
94
95 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_out.into());
96 sk_out.fill_ternary_prob(0.5, &mut source_xs);
97
98 let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank_out.into());
99 sk_out_prepared.prepare(module, &sk_out);
100
101 ksk.encrypt_sk(
102 module,
103 &sk_in,
104 &sk_out,
105 &mut source_xa,
106 &mut source_xe,
107 scratch.borrow(),
108 );
109
110 glwe_in.encrypt_sk(
111 module,
112 &pt_in,
113 &sk_in_prepared,
114 &mut source_xa,
115 &mut source_xe,
116 scratch.borrow(),
117 );
118
119 let mut ksk_prepared: GLWESwitchingKeyPrepared<Vec<u8>, BE> =
120 GLWESwitchingKeyPrepared::alloc_from_infos(module, &ksk);
121 ksk_prepared.prepare(module, &ksk, scratch.borrow());
122
123 glwe_out.keyswitch(module, &glwe_in, &ksk_prepared, scratch.borrow());
124
125 let max_noise: f64 = var_noise_gglwe_product_v2(
126 module.n() as f64,
127 k_ksk,
128 dnum,
129 dsize,
130 base2k_key,
131 0.5,
132 0.5,
133 0f64,
134 SIGMA * SIGMA,
135 0f64,
136 rank_in as f64,
137 )
138 .sqrt()
139 .log2();
140
141 module.glwe_normalize(&mut pt_out, &pt_in, scratch.borrow());
142
143 assert!(
144 glwe_out
145 .noise(module, &pt_out, &sk_out_prepared, scratch.borrow())
146 .std()
147 .log2()
148 <= max_noise + 1.0
149 )
150 }
151 }
152 }
153}
154
155pub fn test_glwe_keyswitch_inplace<BE: Backend>(module: &Module<BE>)
156where
157 Module<BE>: VecZnxFillUniform
158 + GLWESwitchingKeyEncryptSk<BE>
159 + GLWEEncryptSk<BE>
160 + GLWEKeyswitch<BE>
161 + GLWESecretPreparedFactory<BE>
162 + GLWESwitchingKeyPreparedFactory<BE>
163 + GLWENoise<BE>,
164 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
165 Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
166{
167 let base2k_out: usize = 17;
168 let base2k_key: usize = 13;
169
170 let k_out: usize = 102;
171 let max_dsize: usize = k_out.div_ceil(base2k_key);
172
173 for rank in 1_usize..3 {
174 for dsize in 1..max_dsize + 1 {
175 let k_ksk: usize = k_out + base2k_key * dsize;
176
177 let n: usize = module.n();
178 let dnum: usize = k_out.div_ceil(base2k_key * dsize);
179 let glwe_out_infos: GLWELayout = GLWELayout {
180 n: n.into(),
181 base2k: base2k_out.into(),
182 k: k_out.into(),
183 rank: rank.into(),
184 };
185
186 let ksk_infos: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
187 n: n.into(),
188 base2k: base2k_key.into(),
189 k: k_ksk.into(),
190 dnum: dnum.into(),
191 dsize: dsize.into(),
192 rank_in: rank.into(),
193 rank_out: rank.into(),
194 };
195
196 let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&ksk_infos);
197 let mut glwe_out: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_out_infos);
198 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_out_infos);
199
200 let mut source_xs: Source = Source::new([0u8; 32]);
201 let mut source_xe: Source = Source::new([0u8; 32]);
202 let mut source_xa: Source = Source::new([0u8; 32]);
203
204 module.vec_znx_fill_uniform(
205 pt_want.base2k().into(),
206 &mut pt_want.data,
207 0,
208 &mut source_xa,
209 );
210
211 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
212 GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &ksk_infos)
213 | GLWE::encrypt_sk_tmp_bytes(module, &glwe_out_infos)
214 | GLWE::keyswitch_tmp_bytes(module, &glwe_out_infos, &glwe_out_infos, &ksk_infos),
215 );
216
217 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
218 sk_in.fill_ternary_prob(0.5, &mut source_xs);
219
220 let mut sk_in_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
221 sk_in_prepared.prepare(module, &sk_in);
222
223 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
224 sk_out.fill_ternary_prob(0.5, &mut source_xs);
225
226 let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
227 sk_out_prepared.prepare(module, &sk_out);
228
229 ksk.encrypt_sk(
230 module,
231 &sk_in,
232 &sk_out,
233 &mut source_xa,
234 &mut source_xe,
235 scratch.borrow(),
236 );
237
238 glwe_out.encrypt_sk(
239 module,
240 &pt_want,
241 &sk_in_prepared,
242 &mut source_xa,
243 &mut source_xe,
244 scratch.borrow(),
245 );
246
247 let mut ksk_prepared: GLWESwitchingKeyPrepared<Vec<u8>, BE> =
248 GLWESwitchingKeyPrepared::alloc_from_infos(module, &ksk);
249 ksk_prepared.prepare(module, &ksk, scratch.borrow());
250
251 glwe_out.keyswitch_inplace(module, &ksk_prepared, scratch.borrow());
252
253 let max_noise: f64 = var_noise_gglwe_product_v2(
254 module.n() as f64,
255 k_ksk,
256 dnum,
257 dsize,
258 base2k_key,
259 0.5,
260 0.5,
261 0f64,
262 SIGMA * SIGMA,
263 0f64,
264 rank as f64,
265 )
266 .sqrt()
267 .log2();
268
269 assert!(
270 glwe_out
271 .noise(module, &pt_want, &sk_out_prepared, scratch.borrow())
272 .std()
273 .log2()
274 <= max_noise + 1.0
275 )
276 }
277 }
278}