poulpy/core/src/automorphism.rs
core/src/automorphism.rsRaw
1use base2k::{
2 Backend, FFT64, MatZnxDft, MatZnxDftOps, MatZnxDftToMut, MatZnxDftToRef, Module, ScalarZnx, ScalarZnxDftOps, ScalarZnxOps,
3 ScalarZnxToRef, Scratch, VecZnx, VecZnxBigAlloc, VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, VecZnxDftToMut, VecZnxDftToRef,
4 VecZnxOps, ZnxZero,
5};
6use sampling::source::Source;
7
8use crate::{
9 elem::{GetRow, Infos, SetRow},
10 gglwe_ciphertext::GGLWECiphertext,
11 ggsw_ciphertext::GGSWCiphertext,
12 glwe_ciphertext::GLWECiphertext,
13 glwe_ciphertext_fourier::GLWECiphertextFourier,
14 keys::{SecretKey, SecretKeyFourier},
15 keyswitch_key::GLWESwitchingKey,
16};
17
18pub struct AutomorphismKey<Data, B: Backend> {
19 pub(crate) key: GLWESwitchingKey<Data, B>,
20 pub(crate) p: i64,
21}
22
23impl AutomorphismKey<Vec<u8>, FFT64> {
24 pub fn new(module: &Module<FFT64>, basek: usize, k: usize, rows: usize, rank: usize) -> Self {
25 AutomorphismKey {
26 key: GLWESwitchingKey::new(module, basek, k, rows, rank, rank),
27 p: 0,
28 }
29 }
30}
31
32impl<T, B: Backend> Infos for AutomorphismKey<T, B> {
33 type Inner = MatZnxDft<T, B>;
34
35 fn inner(&self) -> &Self::Inner {
36 &self.key.inner()
37 }
38
39 fn basek(&self) -> usize {
40 self.key.basek()
41 }
42
43 fn k(&self) -> usize {
44 self.key.k()
45 }
46}
47
48impl<T, B: Backend> AutomorphismKey<T, B> {
49 pub fn p(&self) -> i64 {
50 self.p
51 }
52
53 pub fn rank(&self) -> usize {
54 self.key.rank()
55 }
56
57 pub fn rank_in(&self) -> usize {
58 self.key.rank_in()
59 }
60
61 pub fn rank_out(&self) -> usize {
62 self.key.rank_out()
63 }
64}
65
66impl<DataSelf, B: Backend> MatZnxDftToMut<B> for AutomorphismKey<DataSelf, B>
67where
68 MatZnxDft<DataSelf, B>: MatZnxDftToMut<B>,
69{
70 fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
71 self.key.to_mut()
72 }
73}
74
75impl<DataSelf, B: Backend> MatZnxDftToRef<B> for AutomorphismKey<DataSelf, B>
76where
77 MatZnxDft<DataSelf, B>: MatZnxDftToRef<B>,
78{
79 fn to_ref(&self) -> MatZnxDft<&[u8], B> {
80 self.key.to_ref()
81 }
82}
83
84impl<C> GetRow<FFT64> for AutomorphismKey<C, FFT64>
85where
86 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
87{
88 fn get_row<R>(&self, module: &Module<FFT64>, row_i: usize, col_j: usize, res: &mut GLWECiphertextFourier<R, FFT64>)
89 where
90 VecZnxDft<R, FFT64>: VecZnxDftToMut<FFT64>,
91 {
92 module.vmp_extract_row(res, self, row_i, col_j);
93 }
94}
95
96impl<C> SetRow<FFT64> for AutomorphismKey<C, FFT64>
97where
98 MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
99{
100 fn set_row<R>(&mut self, module: &Module<FFT64>, row_i: usize, col_j: usize, a: &GLWECiphertextFourier<R, FFT64>)
101 where
102 VecZnxDft<R, FFT64>: VecZnxDftToRef<FFT64>,
103 {
104 module.vmp_prepare_row(self, row_i, col_j, a);
105 }
106}
107
108impl AutomorphismKey<Vec<u8>, FFT64> {
109 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, rank: usize, size: usize) -> usize {
110 GGLWECiphertext::encrypt_sk_scratch_space(module, rank, size)
111 }
112
113 pub fn encrypt_pk_scratch_space(module: &Module<FFT64>, rank: usize, pk_size: usize) -> usize {
114 GGLWECiphertext::encrypt_pk_scratch_space(module, rank, pk_size)
115 }
116
117 pub fn keyswitch_scratch_space(
118 module: &Module<FFT64>,
119 out_size: usize,
120 in_size: usize,
121 ksk_size: usize,
122 rank: usize,
123 ) -> usize {
124 GLWESwitchingKey::keyswitch_scratch_space(module, out_size, rank, in_size, rank, ksk_size)
125 }
126
127 pub fn keyswitch_inplace_scratch_space(module: &Module<FFT64>, out_size: usize, out_rank: usize, ksk_size: usize) -> usize {
128 GLWESwitchingKey::keyswitch_inplace_scratch_space(module, out_size, out_rank, ksk_size)
129 }
130
131 pub fn automorphism_scratch_space(
132 module: &Module<FFT64>,
133 out_size: usize,
134 in_size: usize,
135 ksk_size: usize,
136 rank: usize,
137 ) -> usize {
138 let tmp_dft: usize = module.bytes_of_vec_znx_dft(rank + 1, in_size);
139 let tmp_idft: usize = module.bytes_of_vec_znx_big(rank + 1, out_size);
140 let idft: usize = module.vec_znx_idft_tmp_bytes();
141 let keyswitch: usize = GLWECiphertext::keyswitch_inplace_scratch_space(module, out_size, rank, ksk_size);
142 tmp_dft + tmp_idft + idft + keyswitch
143 }
144
145 pub fn automorphism_inplace_scratch_space(module: &Module<FFT64>, out_size: usize, ksk_size: usize, rank: usize) -> usize {
146 AutomorphismKey::automorphism_scratch_space(module, out_size, out_size, ksk_size, rank)
147 }
148
149 pub fn external_product_scratch_space(
150 module: &Module<FFT64>,
151 out_size: usize,
152 in_size: usize,
153 ggsw_size: usize,
154 rank: usize,
155 ) -> usize {
156 GLWESwitchingKey::external_product_scratch_space(module, out_size, in_size, ggsw_size, rank)
157 }
158
159 pub fn external_product_inplace_scratch_space(
160 module: &Module<FFT64>,
161 out_size: usize,
162 ggsw_size: usize,
163 rank: usize,
164 ) -> usize {
165 GLWESwitchingKey::external_product_inplace_scratch_space(module, out_size, ggsw_size, rank)
166 }
167}
168
169impl<DataSelf> AutomorphismKey<DataSelf, FFT64>
170where
171 MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64>,
172{
173 pub fn encrypt_sk<DataSk>(
174 &mut self,
175 module: &Module<FFT64>,
176 p: i64,
177 sk: &SecretKey<DataSk>,
178 source_xa: &mut Source,
179 source_xe: &mut Source,
180 sigma: f64,
181 scratch: &mut Scratch,
182 ) where
183 ScalarZnx<DataSk>: ScalarZnxToRef,
184 {
185 #[cfg(debug_assertions)]
186 {
187 assert_eq!(self.n(), module.n());
188 assert_eq!(sk.n(), module.n());
189 assert_eq!(self.rank_out(), self.rank_in());
190 assert_eq!(sk.rank(), self.rank());
191 }
192
193 let (sk_out_dft_data, scratch_1) = scratch.tmp_scalar_znx_dft(module, sk.rank());
194
195 let mut sk_out_dft: SecretKeyFourier<&mut [u8], FFT64> = SecretKeyFourier {
196 data: sk_out_dft_data,
197 dist: sk.dist,
198 };
199
200 {
201 (0..self.rank()).for_each(|i| {
202 let (mut sk_inv_auto, _) = scratch_1.tmp_scalar_znx(module, 1);
203 module.scalar_znx_automorphism(module.galois_element_inv(p), &mut sk_inv_auto, 0, sk, i);
204 module.svp_prepare(&mut sk_out_dft, i, &sk_inv_auto, 0);
205 });
206 }
207
208 self.key.encrypt_sk(
209 module,
210 &sk,
211 &sk_out_dft,
212 source_xa,
213 source_xe,
214 sigma,
215 scratch_1,
216 );
217
218 self.p = p;
219 }
220}
221
222impl<DataSelf> AutomorphismKey<DataSelf, FFT64>
223where
224 MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64>,
225{
226 pub fn automorphism<DataLhs, DataRhs>(
227 &mut self,
228 module: &Module<FFT64>,
229 lhs: &AutomorphismKey<DataLhs, FFT64>,
230 rhs: &AutomorphismKey<DataRhs, FFT64>,
231 scratch: &mut base2k::Scratch,
232 ) where
233 MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
234 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
235 {
236 #[cfg(debug_assertions)]
237 {
238 assert_eq!(
239 self.rank_in(),
240 lhs.rank_in(),
241 "ksk_out input rank: {} != ksk_in input rank: {}",
242 self.rank_in(),
243 lhs.rank_in()
244 );
245 assert_eq!(
246 lhs.rank_out(),
247 rhs.rank_in(),
248 "ksk_in output rank: {} != ksk_apply input rank: {}",
249 self.rank_out(),
250 rhs.rank_in()
251 );
252 assert_eq!(
253 self.rank_out(),
254 rhs.rank_out(),
255 "ksk_out output rank: {} != ksk_apply output rank: {}",
256 self.rank_out(),
257 rhs.rank_out()
258 );
259 }
260
261 let cols_out: usize = rhs.rank_out() + 1;
262
263 let (tmp_dft_data, scratch1) = scratch.tmp_vec_znx_dft(module, cols_out, lhs.size());
264
265 let mut tmp_dft: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
266 data: tmp_dft_data,
267 basek: lhs.basek(),
268 k: lhs.k(),
269 };
270
271 (0..self.rank_in()).for_each(|col_i| {
272 (0..self.rows()).for_each(|row_j| {
273 // Extracts relevant row
274 lhs.get_row(module, row_j, col_i, &mut tmp_dft);
275
276 // Get a VecZnxBig from scratch space
277 let (mut tmp_idft_data, scratch2) = scratch1.tmp_vec_znx_big(module, cols_out, self.size());
278
279 // Switches input outside of DFT
280 (0..cols_out).for_each(|i| {
281 module.vec_znx_idft(&mut tmp_idft_data, i, &tmp_dft.data, i, scratch2);
282 });
283
284 // Consumes to small vec znx
285 let mut tmp_idft_small_data: VecZnx<&mut [u8]> = tmp_idft_data.to_vec_znx_small();
286
287 // Reverts the automorphis key from (-pi^{-1}_{k}(s)a + s, a) to (-sa + pi_{k}(s), a)
288 (0..cols_out).for_each(|i| {
289 module.vec_znx_automorphism_inplace(lhs.p(), &mut tmp_idft_small_data, i);
290 });
291
292 // Wraps into ciphertext
293 let mut tmp_idft: GLWECiphertext<&mut [u8]> = GLWECiphertext::<&mut [u8]> {
294 data: tmp_idft_small_data,
295 basek: self.basek(),
296 k: self.k(),
297 };
298
299 // Key-switch (-sa + pi_{k}(s), a) to (-pi^{-1}_{k'}(s)a + pi_{k}(s), a)
300 tmp_idft.keyswitch_inplace(module, &rhs.key, scratch2);
301
302 // Applies back the automorphism X^{k}: (-pi^{-1}_{k'}(s)a + pi_{k}(s), a) -> (-pi^{-1}_{k'+k}(s)a + s, a)
303 // and switches back to DFT domain
304 (0..self.rank_out() + 1).for_each(|i| {
305 module.vec_znx_automorphism_inplace(lhs.p(), &mut tmp_idft, i);
306 module.vec_znx_dft(&mut tmp_dft, i, &tmp_idft, i);
307 });
308
309 // Sets back the relevant row
310 self.set_row(module, row_j, col_i, &tmp_dft);
311 });
312 });
313
314 tmp_dft.data.zero();
315
316 (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
317 (0..self.rank_in()).for_each(|col_j| {
318 self.set_row(module, row_i, col_j, &tmp_dft);
319 });
320 });
321
322 self.p = (lhs.p * rhs.p) % (module.cyclotomic_order() as i64);
323 }
324
325 pub fn automorphism_inplace<DataRhs>(
326 &mut self,
327 module: &Module<FFT64>,
328 rhs: &AutomorphismKey<DataRhs, FFT64>,
329 scratch: &mut Scratch,
330 ) where
331 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
332 {
333 unsafe {
334 let self_ptr: *mut AutomorphismKey<DataSelf, FFT64> = self as *mut AutomorphismKey<DataSelf, FFT64>;
335 self.automorphism(&module, &*self_ptr, rhs, scratch);
336 }
337 }
338
339 pub fn keyswitch<DataLhs, DataRhs>(
340 &mut self,
341 module: &Module<FFT64>,
342 lhs: &AutomorphismKey<DataLhs, FFT64>,
343 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
344 scratch: &mut base2k::Scratch,
345 ) where
346 MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
347 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
348 {
349 self.key.keyswitch(module, &lhs.key, rhs, scratch);
350 }
351
352 pub fn keyswitch_inplace<DataRhs>(
353 &mut self,
354 module: &Module<FFT64>,
355 rhs: &AutomorphismKey<DataRhs, FFT64>,
356 scratch: &mut base2k::Scratch,
357 ) where
358 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
359 {
360 self.key.keyswitch_inplace(module, &rhs.key, scratch);
361 }
362
363 pub fn external_product<DataLhs, DataRhs>(
364 &mut self,
365 module: &Module<FFT64>,
366 lhs: &AutomorphismKey<DataLhs, FFT64>,
367 rhs: &GGSWCiphertext<DataRhs, FFT64>,
368 scratch: &mut Scratch,
369 ) where
370 MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
371 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
372 {
373 self.key.external_product(module, &lhs.key, rhs, scratch);
374 }
375
376 pub fn external_product_inplace<DataRhs>(
377 &mut self,
378 module: &Module<FFT64>,
379 rhs: &GGSWCiphertext<DataRhs, FFT64>,
380 scratch: &mut Scratch,
381 ) where
382 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
383 {
384 self.key.external_product_inplace(module, rhs, scratch);
385 }
386}