poulpy/poulpy-core/src/automorphism/glwe_ct.rs
poulpy-core/src/automorphism/glwe_ct.rsRaw
1use poulpy_hal::{
2 api::{
3 ScratchTakeBasic, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace, VecZnxBigAutomorphismInplace, VecZnxBigNormalize,
4 VecZnxBigSubSmallInplace, VecZnxBigSubSmallNegateInplace, VecZnxNormalize,
5 },
6 layouts::{Backend, DataMut, Module, Scratch, VecZnxBig},
7};
8
9use crate::{
10 GLWEKeySwitchInternal, GLWEKeyswitch, GLWENormalize, ScratchTakeCore,
11 layouts::{GGLWEInfos, GGLWEPreparedToRef, GLWE, GLWEInfos, GLWELayout, GLWEToMut, GLWEToRef, GetGaloisElement, LWEInfos},
12};
13
14impl GLWE<Vec<u8>> {
15 pub fn automorphism_tmp_bytes<M, R, A, K, BE: Backend>(module: &M, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
16 where
17 R: GLWEInfos,
18 A: GLWEInfos,
19 K: GGLWEInfos,
20 M: GLWEAutomorphism<BE>,
21 {
22 module.glwe_automorphism_tmp_bytes(res_infos, a_infos, key_infos)
23 }
24}
25
26impl<DataSelf: DataMut> GLWE<DataSelf> {
27 pub fn automorphism<M, A, K, BE: Backend>(&mut self, module: &M, a: &A, key: &K, scratch: &mut Scratch<BE>)
28 where
29 M: GLWEAutomorphism<BE>,
30 A: GLWEToRef + GLWEInfos,
31 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
32 Scratch<BE>: ScratchTakeCore<BE>,
33 {
34 module.glwe_automorphism(self, a, key, scratch);
35 }
36
37 pub fn automorphism_add<M, A, K, BE: Backend>(&mut self, module: &M, a: &A, key: &K, scratch: &mut Scratch<BE>)
38 where
39 M: GLWEAutomorphism<BE>,
40 A: GLWEToRef + GLWEInfos,
41 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
42 Scratch<BE>: ScratchTakeCore<BE>,
43 {
44 module.glwe_automorphism_add(self, a, key, scratch);
45 }
46
47 pub fn automorphism_sub<M, A, K, BE: Backend>(&mut self, module: &M, a: &A, key: &K, scratch: &mut Scratch<BE>)
48 where
49 M: GLWEAutomorphism<BE>,
50 A: GLWEToRef + GLWEInfos,
51 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
52 Scratch<BE>: ScratchTakeCore<BE>,
53 {
54 module.glwe_automorphism_sub(self, a, key, scratch);
55 }
56
57 pub fn automorphism_sub_negate<M, A, K, BE: Backend>(&mut self, module: &M, a: &A, key: &K, scratch: &mut Scratch<BE>)
58 where
59 M: GLWEAutomorphism<BE>,
60 A: GLWEToRef + GLWEInfos,
61 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
62 Scratch<BE>: ScratchTakeCore<BE>,
63 {
64 module.glwe_automorphism_sub_negate(self, a, key, scratch);
65 }
66
67 pub fn automorphism_inplace<M, K, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
68 where
69 M: GLWEAutomorphism<BE>,
70 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
71 Scratch<BE>: ScratchTakeCore<BE>,
72 {
73 module.glwe_automorphism_inplace(self, key, scratch);
74 }
75
76 pub fn automorphism_add_inplace<M, K, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
77 where
78 M: GLWEAutomorphism<BE>,
79 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
80 Scratch<BE>: ScratchTakeCore<BE>,
81 {
82 module.glwe_automorphism_add_inplace(self, key, scratch);
83 }
84
85 pub fn automorphism_sub_inplace<M, K, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
86 where
87 M: GLWEAutomorphism<BE>,
88 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
89 Scratch<BE>: ScratchTakeCore<BE>,
90 {
91 module.glwe_automorphism_sub_inplace(self, key, scratch);
92 }
93
94 pub fn automorphism_sub_negate_inplace<M, K, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
95 where
96 M: GLWEAutomorphism<BE>,
97 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
98 Scratch<BE>: ScratchTakeCore<BE>,
99 {
100 module.glwe_automorphism_sub_negate_inplace(self, key, scratch);
101 }
102}
103
104pub trait GLWEAutomorphism<BE: Backend> {
105 fn glwe_automorphism_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
106 where
107 R: GLWEInfos,
108 A: GLWEInfos,
109 K: GGLWEInfos;
110
111 fn glwe_automorphism<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
112 where
113 R: GLWEToMut + GLWEInfos,
114 A: GLWEToRef + GLWEInfos,
115 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
116
117 fn glwe_automorphism_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
118 where
119 R: GLWEToMut + GLWEInfos,
120 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
121
122 fn glwe_automorphism_add<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
123 where
124 R: GLWEToMut + GLWEInfos,
125 A: GLWEToRef + GLWEInfos,
126 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
127
128 fn glwe_automorphism_add_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
129 where
130 R: GLWEToMut + GLWEInfos,
131 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
132
133 fn glwe_automorphism_sub<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
134 where
135 R: GLWEToMut + GLWEInfos,
136 A: GLWEToRef + GLWEInfos,
137 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
138
139 fn glwe_automorphism_sub_negate<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
140 where
141 R: GLWEToMut + GLWEInfos,
142 A: GLWEToRef + GLWEInfos,
143 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
144
145 fn glwe_automorphism_sub_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
146 where
147 R: GLWEToMut + GLWEInfos,
148 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
149
150 fn glwe_automorphism_sub_negate_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
151 where
152 R: GLWEToMut + GLWEInfos,
153 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
154}
155
156impl<BE: Backend> GLWEAutomorphism<BE> for Module<BE>
157where
158 Self: Sized
159 + GLWEKeyswitch<BE>
160 + GLWEKeySwitchInternal<BE>
161 + VecZnxNormalize<BE>
162 + VecZnxAutomorphismInplace<BE>
163 + VecZnxBigAutomorphismInplace<BE>
164 + VecZnxBigSubSmallInplace<BE>
165 + VecZnxBigSubSmallNegateInplace<BE>
166 + VecZnxBigAddSmallInplace<BE>
167 + VecZnxBigNormalize<BE>
168 + GLWENormalize<BE>,
169 Scratch<BE>: ScratchTakeCore<BE>,
170{
171 fn glwe_automorphism_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
172 where
173 R: GLWEInfos,
174 A: GLWEInfos,
175 K: GGLWEInfos,
176 {
177 self.glwe_keyswitch_tmp_bytes(res_infos, a_infos, key_infos)
178 }
179
180 fn glwe_automorphism<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
181 where
182 R: GLWEToMut + GLWEInfos,
183 A: GLWEToRef + GLWEInfos,
184 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
185 Scratch<BE>: ScratchTakeCore<BE>,
186 {
187 self.glwe_keyswitch(res, a, key, scratch);
188
189 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
190
191 for i in 0..res.rank().as_usize() + 1 {
192 self.vec_znx_automorphism_inplace(key.p(), res.data_mut(), i, scratch);
193 }
194 }
195
196 fn glwe_automorphism_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
197 where
198 R: GLWEToMut + GLWEInfos,
199 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
200 Scratch<BE>: ScratchTakeCore<BE>,
201 {
202 self.glwe_keyswitch_inplace(res, key, scratch);
203
204 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
205
206 for i in 0..res.rank().as_usize() + 1 {
207 self.vec_znx_automorphism_inplace(key.p(), res.data_mut(), i, scratch);
208 }
209 }
210
211 fn glwe_automorphism_add<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
212 where
213 R: GLWEToMut + GLWEInfos,
214 A: GLWEToRef + GLWEInfos,
215 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
216 Scratch<BE>: ScratchTakeCore<BE>,
217 {
218 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
219 let a: &GLWE<&[u8]> = &a.to_ref();
220
221 let a_base2k: usize = a.base2k().into();
222 let key_base2k: usize = key.base2k().into();
223 let res_base2k: usize = res.base2k().into();
224
225 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
226
227 if a_base2k != key_base2k {
228 let (mut a_conv, scratch_2) = scratch_1.take_glwe(&GLWELayout {
229 n: a.n(),
230 base2k: key.base2k(),
231 k: a.k(),
232 rank: a.rank(),
233 });
234 self.glwe_normalize(&mut a_conv, a, scratch_2);
235 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, &a_conv, key, scratch_2);
236 for i in 0..res.rank().as_usize() + 1 {
237 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_2);
238 self.vec_znx_big_add_small_inplace(&mut res_big, i, a_conv.data(), i);
239 self.vec_znx_big_normalize(res.data_mut(), res_base2k, 0, i, &res_big, key_base2k, i, scratch_2);
240 }
241 } else {
242 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, a, key, scratch_1);
243 for i in 0..res.rank().as_usize() + 1 {
244 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
245 self.vec_znx_big_add_small_inplace(&mut res_big, i, a.data(), i);
246 self.vec_znx_big_normalize(res.data_mut(), res_base2k, 0, i, &res_big, key_base2k, i, scratch_1);
247 }
248 };
249 }
250
251 fn glwe_automorphism_add_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
252 where
253 R: GLWEToMut,
254 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
255 Scratch<BE>: ScratchTakeCore<BE>,
256 {
257 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
258
259 let key_base2k: usize = key.base2k().into();
260 let res_base2k: usize = res.base2k().into();
261
262 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
263
264 if res_base2k != key_base2k {
265 let (mut res_conv, scratch_2) = scratch_1.take_glwe(&GLWELayout {
266 n: res.n(),
267 base2k: key.base2k(),
268 k: res.k(),
269 rank: res.rank(),
270 });
271 self.glwe_normalize(&mut res_conv, res, scratch_2);
272 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, &res_conv, key, scratch_2);
273 for i in 0..res.rank().as_usize() + 1 {
274 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_2);
275 self.vec_znx_big_add_small_inplace(&mut res_big, i, res_conv.data(), i);
276 self.vec_znx_big_normalize(res.data_mut(), res_base2k, 0, i, &res_big, key_base2k, i, scratch_2);
277 }
278 } else {
279 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, res, key, scratch_1);
280 for i in 0..res.rank().as_usize() + 1 {
281 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
282 self.vec_znx_big_add_small_inplace(&mut res_big, i, res.data(), i);
283 self.vec_znx_big_normalize(res.data_mut(), res_base2k, 0, i, &res_big, key_base2k, i, scratch_1);
284 }
285 };
286 }
287
288 fn glwe_automorphism_sub<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
289 where
290 R: GLWEToMut,
291 A: GLWEToRef,
292 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
293 Scratch<BE>: ScratchTakeCore<BE>,
294 {
295 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
296 let a: &GLWE<&[u8]> = &a.to_ref();
297
298 let a_base2k: usize = a.base2k().into();
299 let key_base2k: usize = key.base2k().into();
300 let res_base2k: usize = res.base2k().into();
301
302 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
303
304 if a_base2k != key_base2k {
305 let (mut a_conv, scratch_2) = scratch_1.take_glwe(&GLWELayout {
306 n: a.n(),
307 base2k: key.base2k(),
308 k: a.k(),
309 rank: a.rank(),
310 });
311 self.glwe_normalize(&mut a_conv, a, scratch_2);
312 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, &a_conv, key, scratch_2);
313 for i in 0..res.rank().as_usize() + 1 {
314 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_2);
315 self.vec_znx_big_sub_small_inplace(&mut res_big, i, a_conv.data(), i);
316 self.vec_znx_big_normalize(res.data_mut(), res_base2k, 0, i, &res_big, key_base2k, i, scratch_2);
317 }
318 } else {
319 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, a, key, scratch_1);
320 for i in 0..res.rank().as_usize() + 1 {
321 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
322 self.vec_znx_big_sub_small_inplace(&mut res_big, i, a.data(), i);
323 self.vec_znx_big_normalize(res.data_mut(), res_base2k, 0, i, &res_big, key_base2k, i, scratch_1);
324 }
325 };
326 }
327
328 fn glwe_automorphism_sub_negate<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
329 where
330 R: GLWEToMut,
331 A: GLWEToRef,
332 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
333 Scratch<BE>: ScratchTakeCore<BE>,
334 {
335 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
336 let a: &GLWE<&[u8]> = &a.to_ref();
337
338 let a_base2k: usize = a.base2k().into();
339 let key_base2k: usize = key.base2k().into();
340 let res_base2k: usize = res.base2k().into();
341
342 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
343
344 if a_base2k != key_base2k {
345 let (mut a_conv, scratch_2) = scratch_1.take_glwe(&GLWELayout {
346 n: a.n(),
347 base2k: key.base2k(),
348 k: a.k(),
349 rank: a.rank(),
350 });
351 self.glwe_normalize(&mut a_conv, a, scratch_2);
352 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, &a_conv, key, scratch_2);
353 for i in 0..res.rank().as_usize() + 1 {
354 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_2);
355 self.vec_znx_big_sub_small_negate_inplace(&mut res_big, i, a_conv.data(), i);
356 self.vec_znx_big_normalize(res.data_mut(), res_base2k, 0, i, &res_big, key_base2k, i, scratch_2);
357 }
358 } else {
359 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, a, key, scratch_1);
360 for i in 0..res.rank().as_usize() + 1 {
361 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
362 self.vec_znx_big_sub_small_negate_inplace(&mut res_big, i, a.data(), i);
363 self.vec_znx_big_normalize(res.data_mut(), res_base2k, 0, i, &res_big, key_base2k, i, scratch_1);
364 }
365 };
366 }
367
368 fn glwe_automorphism_sub_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
369 where
370 R: GLWEToMut,
371 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
372 Scratch<BE>: ScratchTakeCore<BE>,
373 {
374 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
375
376 let key_base2k: usize = key.base2k().into();
377 let res_base2k: usize = res.base2k().into();
378
379 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
380
381 if res_base2k != key_base2k {
382 let (mut res_conv, scratch_2) = scratch_1.take_glwe(&GLWELayout {
383 n: res.n(),
384 base2k: key.base2k(),
385 k: res.k(),
386 rank: res.rank(),
387 });
388 self.glwe_normalize(&mut res_conv, res, scratch_2);
389 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, &res_conv, key, scratch_2);
390 for i in 0..res.rank().as_usize() + 1 {
391 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_2);
392 self.vec_znx_big_sub_small_inplace(&mut res_big, i, res_conv.data(), i);
393 self.vec_znx_big_normalize(res.data_mut(), res_base2k, 0, i, &res_big, key_base2k, i, scratch_2);
394 }
395 } else {
396 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, res, key, scratch_1);
397 for i in 0..res.rank().as_usize() + 1 {
398 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
399 self.vec_znx_big_sub_small_inplace(&mut res_big, i, res.data(), i);
400 self.vec_znx_big_normalize(res.data_mut(), res_base2k, 0, i, &res_big, key_base2k, i, scratch_1);
401 }
402 };
403 }
404
405 fn glwe_automorphism_sub_negate_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
406 where
407 R: GLWEToMut,
408 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
409 Scratch<BE>: ScratchTakeCore<BE>,
410 {
411 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
412
413 let key_base2k: usize = key.base2k().into();
414 let res_base2k: usize = res.base2k().into();
415
416 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
417
418 if res_base2k != key_base2k {
419 let (mut res_conv, scratch_2) = scratch_1.take_glwe(&GLWELayout {
420 n: res.n(),
421 base2k: key.base2k(),
422 k: res.k(),
423 rank: res.rank(),
424 });
425 self.glwe_normalize(&mut res_conv, res, scratch_2);
426 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, &res_conv, key, scratch_2);
427 for i in 0..res.rank().as_usize() + 1 {
428 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_2);
429 self.vec_znx_big_sub_small_negate_inplace(&mut res_big, i, res_conv.data(), i);
430 self.vec_znx_big_normalize(res.data_mut(), res_base2k, 0, i, &res_big, key_base2k, i, scratch_2);
431 }
432 } else {
433 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, res, key, scratch_1);
434 for i in 0..res.rank().as_usize() + 1 {
435 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
436 self.vec_znx_big_sub_small_negate_inplace(&mut res_big, i, res.data(), i);
437 self.vec_znx_big_normalize(res.data_mut(), res_base2k, 0, i, &res_big, key_base2k, i, scratch_1);
438 }
439 };
440 }
441}