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