poulpy/core/src/tests/generics/keyswitch/glwe_ct.rs
core/src/tests/generics/keyswitch/glwe_ct.rsRaw
1use backend::hal::{
2 api::{
3 ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApplyInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare, VecZnxAddInplace,
4 VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace, VecZnxBigAllocBytes,
5 VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftAllocBytes, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume,
6 VecZnxFillUniform, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace,
7 VecZnxSwithcDegree, VmpApply, VmpApplyAdd, VmpApplyTmpBytes, VmpPMatAlloc, VmpPrepare,
8 },
9 layouts::{Backend, Module, ScratchOwned},
10 oep::{
11 ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
12 TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl,
13 },
14};
15use sampling::source::Source;
16
17use crate::{
18 layouts::{
19 GGLWESwitchingKey, GLWECiphertext, GLWEPlaintext, GLWESecret, Infos,
20 prepared::{GGLWESwitchingKeyPrepared, GLWESecretPrepared, PrepareAlloc},
21 },
22 noise::log2_std_noise_gglwe_product,
23};
24
25pub fn test_glwe_keyswitch<B: Backend>(
26 module: &Module<B>,
27 basek: usize,
28 k_out: usize,
29 k_in: usize,
30 k_ksk: usize,
31 digits: usize,
32 rank_in: usize,
33 rank_out: usize,
34 sigma: f64,
35) where
36 Module<B>: VecZnxDftAllocBytes
37 + VecZnxBigNormalize<B>
38 + VecZnxDftFromVecZnx<B>
39 + SvpApplyInplace<B>
40 + VecZnxDftToVecZnxBigConsume<B>
41 + VecZnxFillUniform
42 + VecZnxSubABInplace
43 + VecZnxAddInplace
44 + VecZnxNormalizeInplace<B>
45 + VecZnxAddNormal
46 + VecZnxNormalize<B>
47 + VecZnxSub
48 + SvpPrepare<B>
49 + SvpPPolAllocBytes
50 + SvpPPolAlloc<B>
51 + VecZnxBigAllocBytes
52 + VecZnxBigAddInplace<B>
53 + VecZnxBigAddSmallInplace<B>
54 + VecZnxNormalizeTmpBytes
55 + VecZnxAddScalarInplace
56 + VmpPMatAlloc<B>
57 + VmpPrepare<B>
58 + VmpApplyTmpBytes
59 + VmpApply<B>
60 + VmpApplyAdd<B>
61 + VecZnxBigNormalizeTmpBytes
62 + VecZnxSwithcDegree,
63 B: TakeVecZnxDftImpl<B>
64 + TakeVecZnxBigImpl<B>
65 + TakeSvpPPolImpl<B>
66 + ScratchOwnedAllocImpl<B>
67 + ScratchOwnedBorrowImpl<B>
68 + ScratchAvailableImpl<B>
69 + TakeScalarZnxImpl<B>
70 + TakeVecZnxImpl<B>,
71{
72 let n = module.n();
73 let rows: usize = k_in.div_ceil(basek * digits);
74
75 let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank_in, rank_out);
76 let mut ct_in: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_in, rank_in);
77 let mut ct_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_out, rank_out);
78 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_in);
79
80 let mut source_xs: Source = Source::new([0u8; 32]);
81 let mut source_xe: Source = Source::new([0u8; 32]);
82 let mut source_xa: Source = Source::new([0u8; 32]);
83
84 module.vec_znx_fill_uniform(basek, &mut pt_want.data, 0, k_in, &mut source_xa);
85
86 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
87 GGLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, ksk.k(), rank_in, rank_out)
88 | GLWECiphertext::encrypt_sk_scratch_space(module, n, basek, ct_in.k())
89 | GLWECiphertext::keyswitch_scratch_space(
90 module,
91 n,
92 basek,
93 ct_out.k(),
94 ct_in.k(),
95 ksk.k(),
96 digits,
97 rank_in,
98 rank_out,
99 ),
100 );
101
102 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_in);
103 sk_in.fill_ternary_prob(0.5, &mut source_xs);
104 let sk_in_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_in.prepare_alloc(module, scratch.borrow());
105
106 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out);
107 sk_out.fill_ternary_prob(0.5, &mut source_xs);
108 let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
109
110 ksk.encrypt_sk(
111 module,
112 &sk_in,
113 &sk_out,
114 &mut source_xa,
115 &mut source_xe,
116 sigma,
117 scratch.borrow(),
118 );
119
120 ct_in.encrypt_sk(
121 module,
122 &pt_want,
123 &sk_in_prepared,
124 &mut source_xa,
125 &mut source_xe,
126 sigma,
127 scratch.borrow(),
128 );
129
130 let ksk_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
131
132 ct_out.keyswitch(module, &ct_in, &ksk_prepared, scratch.borrow());
133
134 let max_noise: f64 = log2_std_noise_gglwe_product(
135 module.n() as f64,
136 basek * digits,
137 0.5,
138 0.5,
139 0f64,
140 sigma * sigma,
141 0f64,
142 rank_in as f64,
143 k_in,
144 k_ksk,
145 );
146
147 ct_out.assert_noise(module, &sk_out_prepared, &pt_want, max_noise + 0.5);
148}
149
150pub fn test_glwe_keyswitch_inplace<B: Backend>(
151 module: &Module<B>,
152 basek: usize,
153 k_ct: usize,
154 k_ksk: usize,
155 digits: usize,
156 rank: usize,
157 sigma: f64,
158) where
159 Module<B>: VecZnxDftAllocBytes
160 + VecZnxBigNormalize<B>
161 + VecZnxDftFromVecZnx<B>
162 + SvpApplyInplace<B>
163 + VecZnxDftToVecZnxBigConsume<B>
164 + VecZnxFillUniform
165 + VecZnxSubABInplace
166 + VecZnxAddInplace
167 + VecZnxNormalizeInplace<B>
168 + VecZnxAddNormal
169 + VecZnxNormalize<B>
170 + VecZnxSub
171 + SvpPrepare<B>
172 + SvpPPolAllocBytes
173 + SvpPPolAlloc<B>
174 + VecZnxBigAllocBytes
175 + VecZnxBigAddInplace<B>
176 + VecZnxBigAddSmallInplace<B>
177 + VecZnxNormalizeTmpBytes
178 + VecZnxAddScalarInplace
179 + VmpPMatAlloc<B>
180 + VmpPrepare<B>
181 + VmpApplyTmpBytes
182 + VmpApply<B>
183 + VmpApplyAdd<B>
184 + VecZnxBigNormalizeTmpBytes
185 + VecZnxSwithcDegree,
186 B: 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 n: usize = module.n();
196 let rows: usize = k_ct.div_ceil(basek * digits);
197
198 let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank, rank);
199 let mut ct_glwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_ct, rank);
200 let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_ct);
201
202 let mut source_xs: Source = Source::new([0u8; 32]);
203 let mut source_xe: Source = Source::new([0u8; 32]);
204 let mut source_xa: Source = Source::new([0u8; 32]);
205
206 module.vec_znx_fill_uniform(basek, &mut pt_want.data, 0, k_ct, &mut source_xa);
207
208 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
209 GGLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, ksk.k(), rank, rank)
210 | GLWECiphertext::encrypt_sk_scratch_space(module, n, basek, ct_glwe.k())
211 | GLWECiphertext::keyswitch_inplace_scratch_space(module, n, basek, ct_glwe.k(), ksk.k(), digits, rank),
212 );
213
214 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
215 sk_in.fill_ternary_prob(0.5, &mut source_xs);
216 let sk_in_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_in.prepare_alloc(module, scratch.borrow());
217
218 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
219 sk_out.fill_ternary_prob(0.5, &mut source_xs);
220 let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
221
222 ksk.encrypt_sk(
223 module,
224 &sk_in,
225 &sk_out,
226 &mut source_xa,
227 &mut source_xe,
228 sigma,
229 scratch.borrow(),
230 );
231
232 ct_glwe.encrypt_sk(
233 module,
234 &pt_want,
235 &sk_in_prepared,
236 &mut source_xa,
237 &mut source_xe,
238 sigma,
239 scratch.borrow(),
240 );
241
242 let ksk_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
243
244 ct_glwe.keyswitch_inplace(module, &ksk_prepared, scratch.borrow());
245
246 let max_noise: f64 = log2_std_noise_gglwe_product(
247 module.n() as f64,
248 basek * digits,
249 0.5,
250 0.5,
251 0f64,
252 sigma * sigma,
253 0f64,
254 rank as f64,
255 k_ct,
256 k_ksk,
257 );
258
259 ct_glwe.assert_noise(module, &sk_out_prepared, &pt_want, max_noise + 0.5);
260}