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 in_base2k: usize = 17;
33 let key_base2k: usize = 13;
34 let out_base2k: usize = 15;
35 let k_in: usize = 102;
36 let max_dsize: usize = k_in.div_ceil(key_base2k);
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 + key_base2k * 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(key_base2k * dsize);
46
47 let glwe_in_infos: GLWELayout = GLWELayout {
48 n: n.into(),
49 base2k: in_base2k.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: out_base2k.into(),
57 k: k_out.into(),
58 rank: rank_out.into(),
59 };
60
61 let ksk: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
62 n: n.into(),
63 base2k: key_base2k.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(module, &sk_in, &sk_out, &mut source_xa, &mut source_xe, scratch.borrow());
102
103 glwe_in.encrypt_sk(
104 module,
105 &pt_in,
106 &sk_in_prepared,
107 &mut source_xa,
108 &mut source_xe,
109 scratch.borrow(),
110 );
111
112 let mut ksk_prepared: GLWESwitchingKeyPrepared<Vec<u8>, BE> =
113 GLWESwitchingKeyPrepared::alloc_from_infos(module, &ksk);
114 ksk_prepared.prepare(module, &ksk, scratch.borrow());
115
116 glwe_out.keyswitch(module, &glwe_in, &ksk_prepared, scratch.borrow());
117
118 let max_noise: f64 = var_noise_gglwe_product_v2(
119 module.n() as f64,
120 k_ksk,
121 dnum,
122 dsize,
123 key_base2k,
124 0.5,
125 0.5,
126 0f64,
127 SIGMA * SIGMA,
128 0f64,
129 rank_in as f64,
130 )
131 .sqrt()
132 .log2();
133
134 module.glwe_normalize(&mut pt_out, &pt_in, scratch.borrow());
135
136 assert!(
137 glwe_out
138 .noise(module, &pt_out, &sk_out_prepared, scratch.borrow())
139 .std()
140 .log2()
141 <= max_noise + 1.0
142 )
143 }
144 }
145 }
146}
147
148pub fn test_glwe_keyswitch_inplace<BE: Backend>(module: &Module<BE>)
149where
150 Module<BE>: VecZnxFillUniform
151 + GLWESwitchingKeyEncryptSk<BE>
152 + GLWEEncryptSk<BE>
153 + GLWEKeyswitch<BE>
154 + GLWESecretPreparedFactory<BE>
155 + GLWESwitchingKeyPreparedFactory<BE>
156 + GLWENoise<BE>,
157 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
158 Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
159{
160 let out_base2k: usize = 17;
161 let key_base2k: usize = 13;
162
163 let k_out: usize = 102;
164 let max_dsize: usize = k_out.div_ceil(key_base2k);
165
166 for rank in 1_usize..3 {
167 for dsize in 1..max_dsize + 1 {
168 let k_ksk: usize = k_out + key_base2k * dsize;
169
170 let n: usize = module.n();
171 let dnum: usize = k_out.div_ceil(key_base2k * dsize);
172 let glwe_out_infos: GLWELayout = GLWELayout {
173 n: n.into(),
174 base2k: out_base2k.into(),
175 k: k_out.into(),
176 rank: rank.into(),
177 };
178
179 let ksk_infos: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
180 n: n.into(),
181 base2k: key_base2k.into(),
182 k: k_ksk.into(),
183 dnum: dnum.into(),
184 dsize: dsize.into(),
185 rank_in: rank.into(),
186 rank_out: rank.into(),
187 };
188
189 let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&ksk_infos);
190 let mut glwe_out: GLWE<Vec<u8>> = GLWE::alloc_from_infos(&glwe_out_infos);
191 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc_from_infos(&glwe_out_infos);
192
193 let mut source_xs: Source = Source::new([0u8; 32]);
194 let mut source_xe: Source = Source::new([0u8; 32]);
195 let mut source_xa: Source = Source::new([0u8; 32]);
196
197 module.vec_znx_fill_uniform(pt_want.base2k().into(), &mut pt_want.data, 0, &mut source_xa);
198
199 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
200 GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &ksk_infos)
201 | GLWE::encrypt_sk_tmp_bytes(module, &glwe_out_infos)
202 | GLWE::keyswitch_tmp_bytes(module, &glwe_out_infos, &glwe_out_infos, &ksk_infos),
203 );
204
205 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
206 sk_in.fill_ternary_prob(0.5, &mut source_xs);
207
208 let mut sk_in_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
209 sk_in_prepared.prepare(module, &sk_in);
210
211 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
212 sk_out.fill_ternary_prob(0.5, &mut source_xs);
213
214 let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
215 sk_out_prepared.prepare(module, &sk_out);
216
217 ksk.encrypt_sk(module, &sk_in, &sk_out, &mut source_xa, &mut source_xe, scratch.borrow());
218
219 glwe_out.encrypt_sk(
220 module,
221 &pt_want,
222 &sk_in_prepared,
223 &mut source_xa,
224 &mut source_xe,
225 scratch.borrow(),
226 );
227
228 let mut ksk_prepared: GLWESwitchingKeyPrepared<Vec<u8>, BE> =
229 GLWESwitchingKeyPrepared::alloc_from_infos(module, &ksk);
230 ksk_prepared.prepare(module, &ksk, scratch.borrow());
231
232 glwe_out.keyswitch_inplace(module, &ksk_prepared, scratch.borrow());
233
234 let max_noise: f64 = var_noise_gglwe_product_v2(
235 module.n() as f64,
236 k_ksk,
237 dnum,
238 dsize,
239 key_base2k,
240 0.5,
241 0.5,
242 0f64,
243 SIGMA * SIGMA,
244 0f64,
245 rank as f64,
246 )
247 .sqrt()
248 .log2();
249
250 assert!(
251 glwe_out
252 .noise(module, &pt_want, &sk_out_prepared, scratch.borrow())
253 .std()
254 .log2()
255 <= max_noise + 1.0
256 )
257 }
258 }
259}