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