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