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