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