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