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