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