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 ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApply, SvpApplyInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare,
4 VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxAutomorphismInplace,
5 VecZnxBigAddInplace, VecZnxBigAddSmallInplace, VecZnxBigAlloc, VecZnxBigAllocBytes, VecZnxBigNormalize,
6 VecZnxBigNormalizeTmpBytes, VecZnxCopy, VecZnxDftAddInplace, VecZnxDftAlloc, VecZnxDftAllocBytes, VecZnxDftCopy,
7 VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume, VecZnxDftToVecZnxBigTmpA, VecZnxFillUniform, VecZnxNormalize,
8 VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace, VecZnxSwithcDegree, VmpApply,
9 VmpApplyAdd, VmpApplyTmpBytes, VmpPMatAlloc, 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 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 sigma: f64,
39) where
40 Module<B>: VecZnxDftAllocBytes
41 + VecZnxBigAllocBytes
42 + VecZnxDftFromVecZnx<B>
43 + SvpApplyInplace<B>
44 + VecZnxDftToVecZnxBigConsume<B>
45 + VecZnxBigAddInplace<B>
46 + VecZnxBigAddSmallInplace<B>
47 + VecZnxBigNormalize<B>
48 + VecZnxNormalizeTmpBytes
49 + VecZnxBigAlloc<B>
50 + VecZnxDftAlloc<B>
51 + VecZnxBigNormalizeTmpBytes
52 + VecZnxDftToVecZnxBigTmpA<B>
53 + SvpPrepare<B>
54 + SvpPPolAllocBytes
55 + SvpPPolAlloc<B>
56 + VecZnxAddScalarInplace
57 + VecZnxCopy
58 + VecZnxSubABInplace
59 + VmpPMatAlloc<B>
60 + VmpPrepare<B>
61 + VmpApplyTmpBytes
62 + VmpApply<B>
63 + VmpApplyAdd<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, n, basek, k_in, rank)
106 | GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, n, basek, k_ksk, rank)
107 | GGLWETensorKey::encrypt_sk_scratch_space(module, n, basek, k_tsk, rank)
108 | GGSWCiphertext::automorphism_scratch_space(
109 module, n, 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 sigma,
126 scratch.borrow(),
127 );
128 tensor_key.encrypt_sk(
129 module,
130 &sk,
131 &mut source_xa,
132 &mut source_xe,
133 sigma,
134 scratch.borrow(),
135 );
136
137 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
138
139 ct_in.encrypt_sk(
140 module,
141 &pt_scalar,
142 &sk_prepared,
143 &mut source_xa,
144 &mut source_xe,
145 sigma,
146 scratch.borrow(),
147 );
148
149 let mut auto_key_prepared: GGLWEAutomorphismKeyPrepared<Vec<u8>, B> =
150 GGLWEAutomorphismKeyPrepared::alloc(module, n, basek, k_ksk, rows, digits, rank);
151 auto_key_prepared.prepare(module, &auto_key, scratch.borrow());
152
153 let mut tsk_prepared: GGLWETensorKeyPrepared<Vec<u8>, B> =
154 GGLWETensorKeyPrepared::alloc(module, n, basek, k_tsk, rows, digits, rank);
155 tsk_prepared.prepare(module, &tensor_key, scratch.borrow());
156
157 ct_out.automorphism(
158 module,
159 &ct_in,
160 &auto_key_prepared,
161 &tsk_prepared,
162 scratch.borrow(),
163 );
164
165 module.vec_znx_automorphism_inplace(p, &mut pt_scalar.as_vec_znx_mut(), 0);
166
167 let max_noise = |col_j: usize| -> f64 {
168 noise_ggsw_keyswitch(
169 n as f64,
170 basek * digits,
171 col_j,
172 var_xs,
173 0f64,
174 sigma * sigma,
175 0f64,
176 rank as f64,
177 k_in,
178 k_ksk,
179 k_tsk,
180 ) + 0.5
181 };
182
183 ct_out.assert_noise(module, &sk_prepared, &pt_scalar, max_noise);
184}
185
186#[allow(clippy::too_many_arguments)]
187pub fn test_ggsw_automorphism_inplace<B>(
188 p: i64,
189 module: &Module<B>,
190 basek: usize,
191 k_ct: usize,
192 k_ksk: usize,
193 k_tsk: usize,
194 digits: usize,
195 rank: usize,
196 sigma: f64,
197) where
198 Module<B>: VecZnxDftAllocBytes
199 + VecZnxBigAllocBytes
200 + VecZnxDftFromVecZnx<B>
201 + SvpApplyInplace<B>
202 + VecZnxDftToVecZnxBigConsume<B>
203 + VecZnxBigAddInplace<B>
204 + VecZnxBigAddSmallInplace<B>
205 + VecZnxBigNormalize<B>
206 + VecZnxNormalizeTmpBytes
207 + VecZnxBigAlloc<B>
208 + VecZnxDftAlloc<B>
209 + VecZnxBigNormalizeTmpBytes
210 + VecZnxDftToVecZnxBigTmpA<B>
211 + SvpPrepare<B>
212 + SvpPPolAllocBytes
213 + SvpPPolAlloc<B>
214 + VecZnxAddScalarInplace
215 + VecZnxCopy
216 + VecZnxSubABInplace
217 + VmpPMatAlloc<B>
218 + VmpPrepare<B>
219 + VmpApplyTmpBytes
220 + VmpApply<B>
221 + VmpApplyAdd<B>
222 + VecZnxBigAddSmallInplace<B>
223 + VecZnxDftCopy<B>
224 + VecZnxDftAddInplace<B>
225 + VecZnxAddInplace
226 + VecZnxNormalizeInplace<B>
227 + VecZnxAddNormal
228 + VecZnxNormalize<B>
229 + VecZnxSub
230 + VecZnxFillUniform
231 + SvpApply<B>
232 + VecZnxSwithcDegree
233 + VecZnxAutomorphismInplace
234 + VecZnxAutomorphism,
235 B: Backend
236 + TakeVecZnxDftImpl<B>
237 + TakeVecZnxBigImpl<B>
238 + ScratchOwnedAllocImpl<B>
239 + ScratchOwnedBorrowImpl<B>
240 + ScratchAvailableImpl<B>
241 + TakeScalarZnxImpl<B>
242 + TakeVecZnxImpl<B>
243 + VecZnxDftAllocBytesImpl<B>
244 + VecZnxBigAllocBytesImpl<B>
245 + TakeSvpPPolImpl<B>,
246{
247 let n: usize = module.n();
248 let rows: usize = k_ct.div_ceil(digits * basek);
249 let rows_in: usize = k_ct.div_euclid(basek * digits);
250 let digits_in: usize = 1;
251
252 let mut ct: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_ct, rows_in, digits_in, rank);
253 let mut tensor_key: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(n, basek, k_tsk, rows, digits, rank);
254 let mut auto_key: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_ksk, rows, digits, rank);
255 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
256
257 let mut source_xs: Source = Source::new([0u8; 32]);
258 let mut source_xe: Source = Source::new([0u8; 32]);
259 let mut source_xa: Source = Source::new([0u8; 32]);
260
261 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
262 GGSWCiphertext::encrypt_sk_scratch_space(module, n, basek, k_ct, rank)
263 | GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, n, basek, k_ksk, rank)
264 | GGLWETensorKey::encrypt_sk_scratch_space(module, n, basek, k_tsk, rank)
265 | GGSWCiphertext::automorphism_inplace_scratch_space(module, n, basek, k_ct, k_ksk, digits, k_tsk, digits, rank),
266 );
267
268 let var_xs: f64 = 0.5;
269
270 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
271 sk.fill_ternary_prob(var_xs, &mut source_xs);
272 let sk_prepared: GLWESecretPrepared<Vec<u8>, B> = sk.prepare_alloc(module, scratch.borrow());
273
274 auto_key.encrypt_sk(
275 module,
276 p,
277 &sk,
278 &mut source_xa,
279 &mut source_xe,
280 sigma,
281 scratch.borrow(),
282 );
283 tensor_key.encrypt_sk(
284 module,
285 &sk,
286 &mut source_xa,
287 &mut source_xe,
288 sigma,
289 scratch.borrow(),
290 );
291
292 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
293
294 ct.encrypt_sk(
295 module,
296 &pt_scalar,
297 &sk_prepared,
298 &mut source_xa,
299 &mut source_xe,
300 sigma,
301 scratch.borrow(),
302 );
303
304 let mut auto_key_prepared: GGLWEAutomorphismKeyPrepared<Vec<u8>, B> =
305 GGLWEAutomorphismKeyPrepared::alloc(module, n, basek, k_ksk, rows, digits, rank);
306 auto_key_prepared.prepare(module, &auto_key, scratch.borrow());
307
308 let mut tsk_prepared: GGLWETensorKeyPrepared<Vec<u8>, B> =
309 GGLWETensorKeyPrepared::alloc(module, n, basek, k_tsk, rows, digits, rank);
310 tsk_prepared.prepare(module, &tensor_key, scratch.borrow());
311
312 ct.automorphism_inplace(module, &auto_key_prepared, &tsk_prepared, scratch.borrow());
313
314 module.vec_znx_automorphism_inplace(p, &mut pt_scalar.as_vec_znx_mut(), 0);
315
316 let max_noise = |col_j: usize| -> f64 {
317 noise_ggsw_keyswitch(
318 n as f64,
319 basek * digits,
320 col_j,
321 var_xs,
322 0f64,
323 sigma * sigma,
324 0f64,
325 rank as f64,
326 k_ct,
327 k_ksk,
328 k_tsk,
329 ) + 0.5
330 };
331
332 ct.assert_noise(module, &sk_prepared, &pt_scalar, max_noise);
333}