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 VecZnxSubABInplace, 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 GGLWESwitchingKey, GLWECiphertext, GLWEPlaintext, GLWESecret, Infos,
22 prepared::{GGLWESwitchingKeyPrepared, GLWESecretPrepared, 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 + VecZnxSubABInplace
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 basek: usize = 12;
68 let k_in: usize = 45;
69 let digits: usize = k_in.div_ceil(basek);
70
71 (1..3).for_each(|rank_in| {
72 (1..3).for_each(|rank_out| {
73 (1..digits + 1).for_each(|di| {
74 let k_ksk: usize = k_in + basek * di;
75 let k_out: usize = k_ksk; // better capture noise
76
77 let n: usize = module.n();
78 let rows: usize = k_in.div_ceil(basek * digits);
79
80 let mut ksk: GGLWESwitchingKey<Vec<u8>> =
81 GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank_in, rank_out);
82 let mut ct_in: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_in, rank_in);
83 let mut ct_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_out, rank_out);
84 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_in);
85
86 let mut source_xs: Source = Source::new([0u8; 32]);
87 let mut source_xe: Source = Source::new([0u8; 32]);
88 let mut source_xa: Source = Source::new([0u8; 32]);
89
90 module.vec_znx_fill_uniform(basek, &mut pt_want.data, 0, &mut source_xa);
91
92 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
93 GGLWESwitchingKey::encrypt_sk_scratch_space(module, basek, ksk.k(), rank_in, rank_out)
94 | GLWECiphertext::encrypt_sk_scratch_space(module, basek, ct_in.k())
95 | GLWECiphertext::keyswitch_scratch_space(
96 module,
97 basek,
98 ct_out.k(),
99 ct_in.k(),
100 ksk.k(),
101 digits,
102 rank_in,
103 rank_out,
104 ),
105 );
106
107 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_in);
108 sk_in.fill_ternary_prob(0.5, &mut source_xs);
109 let sk_in_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_in.prepare_alloc(module, scratch.borrow());
110
111 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out);
112 sk_out.fill_ternary_prob(0.5, &mut source_xs);
113 let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
114
115 ksk.encrypt_sk(
116 module,
117 &sk_in,
118 &sk_out,
119 &mut source_xa,
120 &mut source_xe,
121 scratch.borrow(),
122 );
123
124 ct_in.encrypt_sk(
125 module,
126 &pt_want,
127 &sk_in_prepared,
128 &mut source_xa,
129 &mut source_xe,
130 scratch.borrow(),
131 );
132
133 let ksk_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
134
135 ct_out.keyswitch(module, &ct_in, &ksk_prepared, scratch.borrow());
136
137 let max_noise: f64 = log2_std_noise_gglwe_product(
138 module.n() as f64,
139 basek * digits,
140 0.5,
141 0.5,
142 0f64,
143 SIGMA * SIGMA,
144 0f64,
145 rank_in as f64,
146 k_in,
147 k_ksk,
148 );
149
150 ct_out.assert_noise(module, &sk_out_prepared, &pt_want, max_noise + 0.5);
151 })
152 });
153 });
154}
155
156pub fn test_glwe_keyswitch_inplace<B>(module: &Module<B>)
157where
158 Module<B>: VecZnxDftAllocBytes
159 + VecZnxBigNormalize<B>
160 + VecZnxDftApply<B>
161 + SvpApplyDftToDftInplace<B>
162 + VecZnxIdftApplyConsume<B>
163 + VecZnxFillUniform
164 + VecZnxSubABInplace
165 + VecZnxAddInplace
166 + VecZnxNormalizeInplace<B>
167 + VecZnxAddNormal
168 + VecZnxNormalize<B>
169 + VecZnxSub
170 + SvpPrepare<B>
171 + SvpPPolAllocBytes
172 + SvpPPolAlloc<B>
173 + VecZnxBigAllocBytes
174 + VecZnxBigAddInplace<B>
175 + VecZnxBigAddSmallInplace<B>
176 + VecZnxNormalizeTmpBytes
177 + VecZnxAddScalarInplace
178 + VmpPMatAlloc<B>
179 + VmpPrepare<B>
180 + VmpApplyDftToDftTmpBytes
181 + VmpApplyDftToDft<B>
182 + VmpApplyDftToDftAdd<B>
183 + VecZnxBigNormalizeTmpBytes
184 + VecZnxSwitchRing,
185 B: Backend
186 + TakeVecZnxDftImpl<B>
187 + TakeVecZnxBigImpl<B>
188 + TakeSvpPPolImpl<B>
189 + ScratchOwnedAllocImpl<B>
190 + ScratchOwnedBorrowImpl<B>
191 + ScratchAvailableImpl<B>
192 + TakeScalarZnxImpl<B>
193 + TakeVecZnxImpl<B>,
194{
195 let basek: usize = 12;
196 let k_ct: usize = 45;
197 let digits: usize = k_ct.div_ceil(basek);
198
199 (1..3).for_each(|rank| {
200 (1..digits + 1).for_each(|di| {
201 let k_ksk: usize = k_ct + basek * di;
202
203 let n: usize = module.n();
204 let rows: usize = k_ct.div_ceil(basek * digits);
205
206 let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank, rank);
207 let mut ct_glwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_ct, rank);
208 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_ct);
209
210 let mut source_xs: Source = Source::new([0u8; 32]);
211 let mut source_xe: Source = Source::new([0u8; 32]);
212 let mut source_xa: Source = Source::new([0u8; 32]);
213
214 module.vec_znx_fill_uniform(basek, &mut pt_want.data, 0, &mut source_xa);
215
216 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
217 GGLWESwitchingKey::encrypt_sk_scratch_space(module, basek, ksk.k(), rank, rank)
218 | GLWECiphertext::encrypt_sk_scratch_space(module, basek, ct_glwe.k())
219 | GLWECiphertext::keyswitch_inplace_scratch_space(module, basek, ct_glwe.k(), ksk.k(), digits, rank),
220 );
221
222 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
223 sk_in.fill_ternary_prob(0.5, &mut source_xs);
224 let sk_in_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_in.prepare_alloc(module, scratch.borrow());
225
226 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
227 sk_out.fill_ternary_prob(0.5, &mut source_xs);
228 let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
229
230 ksk.encrypt_sk(
231 module,
232 &sk_in,
233 &sk_out,
234 &mut source_xa,
235 &mut source_xe,
236 scratch.borrow(),
237 );
238
239 ct_glwe.encrypt_sk(
240 module,
241 &pt_want,
242 &sk_in_prepared,
243 &mut source_xa,
244 &mut source_xe,
245 scratch.borrow(),
246 );
247
248 let ksk_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
249
250 ct_glwe.keyswitch_inplace(module, &ksk_prepared, scratch.borrow());
251
252 let max_noise: f64 = log2_std_noise_gglwe_product(
253 module.n() as f64,
254 basek * digits,
255 0.5,
256 0.5,
257 0f64,
258 SIGMA * SIGMA,
259 0f64,
260 rank as f64,
261 k_ct,
262 k_ksk,
263 );
264
265 ct_glwe.assert_noise(module, &sk_out_prepared, &pt_want, max_noise + 0.5);
266 });
267 });
268}