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