poulpy/core/src/tests/generics/keyswitch/ggsw_ct.rs
core/src/tests/generics/keyswitch/ggsw_ct.rsRaw
1use backend::hal::{
2 api::{
3 ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApply, SvpApplyInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare,
4 VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace, VecZnxBigAlloc,
5 VecZnxBigAllocBytes, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftAddInplace, VecZnxDftAlloc,
6 VecZnxDftAllocBytes, VecZnxDftCopy, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume, VecZnxDftToVecZnxBigTmpA,
7 VecZnxFillUniform, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace,
8 VecZnxSwithcDegree, VmpApply, VmpApplyAdd, VmpApplyTmpBytes, VmpPMatAlloc, VmpPrepare,
9 },
10 layouts::{Backend, Module, ScalarZnx, ScratchOwned},
11 oep::{
12 ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
13 TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl,
14 },
15};
16use sampling::source::Source;
17
18use crate::{
19 layouts::{
20 GGLWESwitchingKey, GGLWETensorKey, GGSWCiphertext, GLWESecret,
21 prepared::{GGLWESwitchingKeyPrepared, GGLWETensorKeyPrepared, GLWESecretPrepared, PrepareAlloc},
22 },
23 noise::noise_ggsw_keyswitch,
24};
25
26pub fn test_ggsw_keyswitch<B: Backend>(
27 module: &Module<B>,
28 basek: usize,
29 k_out: usize,
30 k_in: usize,
31 k_ksk: usize,
32 k_tsk: usize,
33 digits: usize,
34 rank: 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 + SvpApply<B>
65 + VecZnxDftToVecZnxBigTmpA<B>
66 + VecZnxDftCopy<B>
67 + VecZnxDftAddInplace<B>
68 + VecZnxBigAlloc<B>
69 + VecZnxDftAlloc<B>,
70 B: TakeVecZnxDftImpl<B>
71 + TakeVecZnxBigImpl<B>
72 + TakeSvpPPolImpl<B>
73 + ScratchOwnedAllocImpl<B>
74 + ScratchOwnedBorrowImpl<B>
75 + ScratchAvailableImpl<B>
76 + TakeScalarZnxImpl<B>
77 + TakeVecZnxImpl<B>,
78{
79 let n: usize = module.n();
80 let rows: usize = k_in.div_ceil(digits * basek);
81
82 let digits_in: usize = 1;
83
84 let mut ct_in: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_in, rows, digits_in, rank);
85 let mut ct_out: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_out, rows, digits_in, rank);
86 let mut tsk: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(n, basek, k_ksk, rows, digits, rank);
87 let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank, rank);
88 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
89
90 let mut source_xs: Source = Source::new([0u8; 32]);
91 let mut source_xe: Source = Source::new([0u8; 32]);
92 let mut source_xa: Source = Source::new([0u8; 32]);
93
94 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
95 GGSWCiphertext::encrypt_sk_scratch_space(module, n, basek, k_in, rank)
96 | GGLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k_ksk, rank, rank)
97 | GGLWETensorKey::encrypt_sk_scratch_space(module, n, basek, k_tsk, rank)
98 | GGSWCiphertext::keyswitch_scratch_space(
99 module, n, basek, k_out, k_in, k_ksk, digits, k_tsk, digits, rank,
100 ),
101 );
102
103 let var_xs: f64 = 0.5;
104
105 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
106 sk_in.fill_ternary_prob(var_xs, &mut source_xs);
107 let sk_in_dft: GLWESecretPrepared<Vec<u8>, B> = sk_in.prepare_alloc(module, scratch.borrow());
108
109 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
110 sk_out.fill_ternary_prob(var_xs, &mut source_xs);
111 let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
112
113 ksk.encrypt_sk(
114 module,
115 &sk_in,
116 &sk_out,
117 &mut source_xa,
118 &mut source_xe,
119 sigma,
120 scratch.borrow(),
121 );
122 tsk.encrypt_sk(
123 module,
124 &sk_out,
125 &mut source_xa,
126 &mut source_xe,
127 sigma,
128 scratch.borrow(),
129 );
130
131 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
132
133 ct_in.encrypt_sk(
134 module,
135 &pt_scalar,
136 &sk_in_dft,
137 &mut source_xa,
138 &mut source_xe,
139 sigma,
140 scratch.borrow(),
141 );
142
143 let ksk_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
144 let tsk_prepared: GGLWETensorKeyPrepared<Vec<u8>, B> = tsk.prepare_alloc(module, scratch.borrow());
145
146 ct_out.keyswitch(
147 module,
148 &ct_in,
149 &ksk_prepared,
150 &tsk_prepared,
151 scratch.borrow(),
152 );
153
154 let max_noise = |col_j: usize| -> f64 {
155 noise_ggsw_keyswitch(
156 n as f64,
157 basek * digits,
158 col_j,
159 var_xs,
160 0f64,
161 sigma * sigma,
162 0f64,
163 rank as f64,
164 k_in,
165 k_ksk,
166 k_tsk,
167 ) + 0.5
168 };
169
170 ct_out.assert_noise(module, &sk_out_prepared, &pt_scalar, &max_noise);
171}
172
173pub fn test_ggsw_keyswitch_inplace<B: Backend>(
174 module: &Module<B>,
175 basek: usize,
176 k_ct: usize,
177 k_ksk: usize,
178 k_tsk: usize,
179 digits: usize,
180 rank: usize,
181 sigma: f64,
182) where
183 Module<B>: VecZnxDftAllocBytes
184 + VecZnxBigNormalize<B>
185 + VecZnxDftFromVecZnx<B>
186 + SvpApplyInplace<B>
187 + VecZnxDftToVecZnxBigConsume<B>
188 + VecZnxFillUniform
189 + VecZnxSubABInplace
190 + VecZnxAddInplace
191 + VecZnxNormalizeInplace<B>
192 + VecZnxAddNormal
193 + VecZnxNormalize<B>
194 + VecZnxSub
195 + SvpPrepare<B>
196 + SvpPPolAllocBytes
197 + SvpPPolAlloc<B>
198 + VecZnxBigAllocBytes
199 + VecZnxBigAddInplace<B>
200 + VecZnxBigAddSmallInplace<B>
201 + VecZnxNormalizeTmpBytes
202 + VecZnxAddScalarInplace
203 + VmpPMatAlloc<B>
204 + VmpPrepare<B>
205 + VmpApplyTmpBytes
206 + VmpApply<B>
207 + VmpApplyAdd<B>
208 + VecZnxBigNormalizeTmpBytes
209 + VecZnxSwithcDegree
210 + SvpApply<B>
211 + VecZnxDftToVecZnxBigTmpA<B>
212 + VecZnxDftCopy<B>
213 + VecZnxDftAddInplace<B>
214 + VecZnxBigAlloc<B>
215 + VecZnxDftAlloc<B>,
216 B: TakeVecZnxDftImpl<B>
217 + TakeVecZnxBigImpl<B>
218 + TakeSvpPPolImpl<B>
219 + ScratchOwnedAllocImpl<B>
220 + ScratchOwnedBorrowImpl<B>
221 + ScratchAvailableImpl<B>
222 + TakeScalarZnxImpl<B>
223 + TakeVecZnxImpl<B>,
224{
225 let n: usize = module.n();
226 let rows: usize = k_ct.div_ceil(digits * basek);
227
228 let digits_in: usize = 1;
229
230 let mut ct: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_ct, rows, digits_in, rank);
231 let mut tsk: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(n, basek, k_tsk, rows, digits, rank);
232 let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank, rank);
233 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
234
235 let mut source_xs: Source = Source::new([0u8; 32]);
236 let mut source_xe: Source = Source::new([0u8; 32]);
237 let mut source_xa: Source = Source::new([0u8; 32]);
238
239 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
240 GGSWCiphertext::encrypt_sk_scratch_space(module, n, basek, k_ct, rank)
241 | GGLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k_ksk, rank, rank)
242 | GGLWETensorKey::encrypt_sk_scratch_space(module, n, basek, k_tsk, rank)
243 | GGSWCiphertext::keyswitch_inplace_scratch_space(module, n, basek, k_ct, k_ksk, digits, k_tsk, digits, rank),
244 );
245
246 let var_xs: f64 = 0.5;
247
248 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
249 sk_in.fill_ternary_prob(var_xs, &mut source_xs);
250 let sk_in_dft: GLWESecretPrepared<Vec<u8>, B> = sk_in.prepare_alloc(module, scratch.borrow());
251
252 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
253 sk_out.fill_ternary_prob(var_xs, &mut source_xs);
254 let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
255
256 ksk.encrypt_sk(
257 module,
258 &sk_in,
259 &sk_out,
260 &mut source_xa,
261 &mut source_xe,
262 sigma,
263 scratch.borrow(),
264 );
265 tsk.encrypt_sk(
266 module,
267 &sk_out,
268 &mut source_xa,
269 &mut source_xe,
270 sigma,
271 scratch.borrow(),
272 );
273
274 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
275
276 ct.encrypt_sk(
277 module,
278 &pt_scalar,
279 &sk_in_dft,
280 &mut source_xa,
281 &mut source_xe,
282 sigma,
283 scratch.borrow(),
284 );
285
286 let ksk_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
287 let tsk_prepared: GGLWETensorKeyPrepared<Vec<u8>, B> = tsk.prepare_alloc(module, scratch.borrow());
288
289 ct.keyswitch_inplace(module, &ksk_prepared, &tsk_prepared, scratch.borrow());
290
291 let max_noise = |col_j: usize| -> f64 {
292 noise_ggsw_keyswitch(
293 n as f64,
294 basek * digits,
295 col_j,
296 var_xs,
297 0f64,
298 sigma * sigma,
299 0f64,
300 rank as f64,
301 k_ct,
302 k_ksk,
303 k_tsk,
304 ) + 0.5
305 };
306
307 ct.assert_noise(module, &sk_out_prepared, &pt_scalar, &max_noise);
308}