poulpy/poulpy-core/src/tests/generics/automorphism/ggsw_ct.rs
poulpy-core/src/tests/generics/automorphism/ggsw_ct.rsRaw
1use poulpy_hal::{
2 api::{
3 DFT, IDFTConsume, IDFTTmpA, ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApply, SvpApplyInplace, SvpPPolAlloc,
4 SvpPPolAllocBytes, SvpPrepare, VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxAutomorphism,
5 VecZnxAutomorphismInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace, VecZnxBigAlloc, VecZnxBigAllocBytes,
6 VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxCopy, VecZnxDftAddInplace, VecZnxDftAlloc, VecZnxDftAllocBytes,
7 VecZnxDftCopy, VecZnxFillUniform, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub,
8 VecZnxSubABInplace, VecZnxSwithcDegree, VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes, VmpPMatAlloc,
9 VmpPrepare,
10 },
11 layouts::{Backend, Module, ScalarZnx, ScratchOwned},
12 oep::{
13 ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
14 TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl, VecZnxBigAllocBytesImpl, VecZnxDftAllocBytesImpl,
15 },
16 source::Source,
17};
18
19use crate::{
20 encryption::SIGMA,
21 layouts::{
22 GGLWEAutomorphismKey, GGLWETensorKey, GGSWCiphertext, GLWESecret,
23 prepared::{GGLWEAutomorphismKeyPrepared, GGLWETensorKeyPrepared, GLWESecretPrepared, Prepare, PrepareAlloc},
24 },
25 noise::noise_ggsw_keyswitch,
26};
27
28#[allow(clippy::too_many_arguments)]
29pub fn test_ggsw_automorphism<B>(
30 p: i64,
31 module: &Module<B>,
32 basek: usize,
33 k_out: usize,
34 k_in: usize,
35 k_ksk: usize,
36 k_tsk: usize,
37 digits: usize,
38 rank: usize,
39) where
40 Module<B>: VecZnxDftAllocBytes
41 + VecZnxBigAllocBytes
42 + DFT<B>
43 + SvpApplyInplace<B>
44 + IDFTConsume<B>
45 + VecZnxBigAddInplace<B>
46 + VecZnxBigAddSmallInplace<B>
47 + VecZnxBigNormalize<B>
48 + VecZnxNormalizeTmpBytes
49 + VecZnxBigAlloc<B>
50 + VecZnxDftAlloc<B>
51 + VecZnxBigNormalizeTmpBytes
52 + IDFTTmpA<B>
53 + SvpPrepare<B>
54 + SvpPPolAllocBytes
55 + SvpPPolAlloc<B>
56 + VecZnxAddScalarInplace
57 + VecZnxCopy
58 + VecZnxSubABInplace
59 + VmpPMatAlloc<B>
60 + VmpPrepare<B>
61 + VmpApplyDftToDftTmpBytes
62 + VmpApplyDftToDft<B>
63 + VmpApplyDftToDftAdd<B>
64 + VecZnxDftCopy<B>
65 + VecZnxDftAddInplace<B>
66 + VecZnxFillUniform
67 + VecZnxAddInplace
68 + VecZnxNormalizeInplace<B>
69 + VecZnxAddNormal
70 + VecZnxNormalize<B>
71 + VecZnxSub
72 + SvpApply<B>
73 + VecZnxSwithcDegree
74 + VecZnxAutomorphismInplace
75 + VecZnxAutomorphism,
76 B: Backend
77 + TakeVecZnxDftImpl<B>
78 + TakeVecZnxBigImpl<B>
79 + ScratchOwnedAllocImpl<B>
80 + ScratchOwnedBorrowImpl<B>
81 + ScratchAvailableImpl<B>
82 + TakeScalarZnxImpl<B>
83 + TakeVecZnxImpl<B>
84 + VecZnxDftAllocBytesImpl<B>
85 + VecZnxBigAllocBytesImpl<B>
86 + TakeSvpPPolImpl<B>,
87{
88 let n: usize = module.n();
89 let rows: usize = k_in.div_ceil(basek * digits);
90 let rows_in: usize = k_in.div_euclid(basek * digits);
91
92 let digits_in: usize = 1;
93
94 let mut ct_in: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_in, rows_in, digits_in, rank);
95 let mut ct_out: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_out, rows_in, digits_in, rank);
96 let mut tensor_key: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(n, basek, k_tsk, rows, digits, rank);
97 let mut auto_key: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_ksk, rows, digits, rank);
98 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
99
100 let mut source_xs: Source = Source::new([0u8; 32]);
101 let mut source_xe: Source = Source::new([0u8; 32]);
102 let mut source_xa: Source = Source::new([0u8; 32]);
103
104 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
105 GGSWCiphertext::encrypt_sk_scratch_space(module, basek, k_in, rank)
106 | GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, basek, k_ksk, rank)
107 | GGLWETensorKey::encrypt_sk_scratch_space(module, basek, k_tsk, rank)
108 | GGSWCiphertext::automorphism_scratch_space(
109 module, basek, k_out, k_in, k_ksk, digits, k_tsk, digits, rank,
110 ),
111 );
112
113 let var_xs: f64 = 0.5;
114
115 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
116 sk.fill_ternary_prob(var_xs, &mut source_xs);
117 let sk_prepared: GLWESecretPrepared<Vec<u8>, B> = sk.prepare_alloc(module, scratch.borrow());
118
119 auto_key.encrypt_sk(
120 module,
121 p,
122 &sk,
123 &mut source_xa,
124 &mut source_xe,
125 scratch.borrow(),
126 );
127 tensor_key.encrypt_sk(
128 module,
129 &sk,
130 &mut source_xa,
131 &mut source_xe,
132 scratch.borrow(),
133 );
134
135 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
136
137 ct_in.encrypt_sk(
138 module,
139 &pt_scalar,
140 &sk_prepared,
141 &mut source_xa,
142 &mut source_xe,
143 scratch.borrow(),
144 );
145
146 let mut auto_key_prepared: GGLWEAutomorphismKeyPrepared<Vec<u8>, B> =
147 GGLWEAutomorphismKeyPrepared::alloc(module, basek, k_ksk, rows, digits, rank);
148 auto_key_prepared.prepare(module, &auto_key, scratch.borrow());
149
150 let mut tsk_prepared: GGLWETensorKeyPrepared<Vec<u8>, B> =
151 GGLWETensorKeyPrepared::alloc(module, basek, k_tsk, rows, digits, rank);
152 tsk_prepared.prepare(module, &tensor_key, scratch.borrow());
153
154 ct_out.automorphism(
155 module,
156 &ct_in,
157 &auto_key_prepared,
158 &tsk_prepared,
159 scratch.borrow(),
160 );
161
162 module.vec_znx_automorphism_inplace(p, &mut pt_scalar.as_vec_znx_mut(), 0);
163
164 let max_noise = |col_j: usize| -> f64 {
165 noise_ggsw_keyswitch(
166 n as f64,
167 basek * digits,
168 col_j,
169 var_xs,
170 0f64,
171 SIGMA * SIGMA,
172 0f64,
173 rank as f64,
174 k_in,
175 k_ksk,
176 k_tsk,
177 ) + 0.5
178 };
179
180 ct_out.assert_noise(module, &sk_prepared, &pt_scalar, max_noise);
181}
182
183#[allow(clippy::too_many_arguments)]
184pub fn test_ggsw_automorphism_inplace<B>(
185 p: i64,
186 module: &Module<B>,
187 basek: usize,
188 k_ct: usize,
189 k_ksk: usize,
190 k_tsk: usize,
191 digits: usize,
192 rank: usize,
193) where
194 Module<B>: VecZnxDftAllocBytes
195 + VecZnxBigAllocBytes
196 + DFT<B>
197 + SvpApplyInplace<B>
198 + IDFTConsume<B>
199 + VecZnxBigAddInplace<B>
200 + VecZnxBigAddSmallInplace<B>
201 + VecZnxBigNormalize<B>
202 + VecZnxNormalizeTmpBytes
203 + VecZnxBigAlloc<B>
204 + VecZnxDftAlloc<B>
205 + VecZnxBigNormalizeTmpBytes
206 + IDFTTmpA<B>
207 + SvpPrepare<B>
208 + SvpPPolAllocBytes
209 + SvpPPolAlloc<B>
210 + VecZnxAddScalarInplace
211 + VecZnxCopy
212 + VecZnxSubABInplace
213 + VmpPMatAlloc<B>
214 + VmpPrepare<B>
215 + VmpApplyDftToDftTmpBytes
216 + VmpApplyDftToDft<B>
217 + VmpApplyDftToDftAdd<B>
218 + VecZnxBigAddSmallInplace<B>
219 + VecZnxDftCopy<B>
220 + VecZnxDftAddInplace<B>
221 + VecZnxAddInplace
222 + VecZnxNormalizeInplace<B>
223 + VecZnxAddNormal
224 + VecZnxNormalize<B>
225 + VecZnxSub
226 + VecZnxFillUniform
227 + SvpApply<B>
228 + VecZnxSwithcDegree
229 + VecZnxAutomorphismInplace
230 + VecZnxAutomorphism,
231 B: Backend
232 + TakeVecZnxDftImpl<B>
233 + TakeVecZnxBigImpl<B>
234 + ScratchOwnedAllocImpl<B>
235 + ScratchOwnedBorrowImpl<B>
236 + ScratchAvailableImpl<B>
237 + TakeScalarZnxImpl<B>
238 + TakeVecZnxImpl<B>
239 + VecZnxDftAllocBytesImpl<B>
240 + VecZnxBigAllocBytesImpl<B>
241 + TakeSvpPPolImpl<B>,
242{
243 let n: usize = module.n();
244 let rows: usize = k_ct.div_ceil(digits * basek);
245 let rows_in: usize = k_ct.div_euclid(basek * digits);
246 let digits_in: usize = 1;
247
248 let mut ct: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_ct, rows_in, digits_in, rank);
249 let mut tensor_key: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(n, basek, k_tsk, rows, digits, rank);
250 let mut auto_key: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_ksk, rows, digits, rank);
251 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
252
253 let mut source_xs: Source = Source::new([0u8; 32]);
254 let mut source_xe: Source = Source::new([0u8; 32]);
255 let mut source_xa: Source = Source::new([0u8; 32]);
256
257 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
258 GGSWCiphertext::encrypt_sk_scratch_space(module, basek, k_ct, rank)
259 | GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, basek, k_ksk, rank)
260 | GGLWETensorKey::encrypt_sk_scratch_space(module, basek, k_tsk, rank)
261 | GGSWCiphertext::automorphism_inplace_scratch_space(module, basek, k_ct, k_ksk, digits, k_tsk, digits, rank),
262 );
263
264 let var_xs: f64 = 0.5;
265
266 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
267 sk.fill_ternary_prob(var_xs, &mut source_xs);
268 let sk_prepared: GLWESecretPrepared<Vec<u8>, B> = sk.prepare_alloc(module, scratch.borrow());
269
270 auto_key.encrypt_sk(
271 module,
272 p,
273 &sk,
274 &mut source_xa,
275 &mut source_xe,
276 scratch.borrow(),
277 );
278 tensor_key.encrypt_sk(
279 module,
280 &sk,
281 &mut source_xa,
282 &mut source_xe,
283 scratch.borrow(),
284 );
285
286 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
287
288 ct.encrypt_sk(
289 module,
290 &pt_scalar,
291 &sk_prepared,
292 &mut source_xa,
293 &mut source_xe,
294 scratch.borrow(),
295 );
296
297 let mut auto_key_prepared: GGLWEAutomorphismKeyPrepared<Vec<u8>, B> =
298 GGLWEAutomorphismKeyPrepared::alloc(module, basek, k_ksk, rows, digits, rank);
299 auto_key_prepared.prepare(module, &auto_key, scratch.borrow());
300
301 let mut tsk_prepared: GGLWETensorKeyPrepared<Vec<u8>, B> =
302 GGLWETensorKeyPrepared::alloc(module, basek, k_tsk, rows, digits, rank);
303 tsk_prepared.prepare(module, &tensor_key, scratch.borrow());
304
305 ct.automorphism_inplace(module, &auto_key_prepared, &tsk_prepared, scratch.borrow());
306
307 module.vec_znx_automorphism_inplace(p, &mut pt_scalar.as_vec_znx_mut(), 0);
308
309 let max_noise = |col_j: usize| -> f64 {
310 noise_ggsw_keyswitch(
311 n as f64,
312 basek * digits,
313 col_j,
314 var_xs,
315 0f64,
316 SIGMA * SIGMA,
317 0f64,
318 rank as f64,
319 k_ct,
320 k_ksk,
321 k_tsk,
322 ) + 0.5
323 };
324
325 ct.assert_noise(module, &sk_prepared, &pt_scalar, max_noise);
326}