poulpy/poulpy-core/src/automorphism/glwe_ct.rs
poulpy-core/src/automorphism/glwe_ct.rsRaw
1use poulpy_hal::{
2 api::{
3 ScratchTakeBasic, VecZnxAutomorphismInplace, VecZnxBigAutomorphismInplace, VecZnxBigSubSmallInplace,
4 VecZnxBigSubSmallNegateInplace,
5 },
6 layouts::{Backend, DataMut, Module, Scratch, VecZnxBig},
7};
8
9use crate::{
10 GLWEKeyswitch, ScratchTakeCore, keyswitch_internal,
11 layouts::{
12 GGLWEInfos, GLWE, GLWEInfos, GLWEToMut, GLWEToRef, LWEInfos,
13 prepared::{AutomorphismKeyPrepared, AutomorphismKeyPreparedToRef, GetAutomorphismGaloisElement},
14 },
15};
16
17impl GLWE<Vec<u8>> {
18 pub fn automorphism_tmp_bytes<M, R, A, K, BE: Backend>(module: &M, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
19 where
20 R: GLWEInfos,
21 A: GLWEInfos,
22 K: GGLWEInfos,
23 M: GLWEAutomorphism<BE>,
24 {
25 module.glwe_automorphism_tmp_bytes(res_infos, a_infos, key_infos)
26 }
27}
28
29impl<DataSelf: DataMut> GLWE<DataSelf> {
30 pub fn automorphism<M, A, K, BE: Backend>(&mut self, module: &M, a: &A, key: &K, scratch: &mut Scratch<BE>)
31 where
32 M: GLWEAutomorphism<BE>,
33 A: GLWEToRef,
34 K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
35 Scratch<BE>: ScratchTakeCore<BE>,
36 {
37 module.glwe_automorphism(self, a, key, scratch);
38 }
39
40 pub fn automorphism_add<M, A, K, BE: Backend>(&mut self, module: &M, a: &A, key: &K, scratch: &mut Scratch<BE>)
41 where
42 M: GLWEAutomorphism<BE>,
43 A: GLWEToRef,
44 K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
45 Scratch<BE>: ScratchTakeCore<BE>,
46 {
47 module.glwe_automorphism_add(self, a, key, scratch);
48 }
49
50 pub fn automorphism_sub<M, A, K, BE: Backend>(&mut self, module: &M, a: &A, key: &K, scratch: &mut Scratch<BE>)
51 where
52 M: GLWEAutomorphism<BE>,
53 A: GLWEToRef,
54 K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
55 Scratch<BE>: ScratchTakeCore<BE>,
56 {
57 module.glwe_automorphism_sub(self, a, key, scratch);
58 }
59
60 pub fn glwe_automorphism_sub_negate<M, A, K, BE: Backend>(&mut self, module: &M, a: &A, key: &K, scratch: &mut Scratch<BE>)
61 where
62 M: GLWEAutomorphism<BE>,
63 A: GLWEToRef,
64 K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
65 Scratch<BE>: ScratchTakeCore<BE>,
66 {
67 module.glwe_automorphism_sub_negate(self, a, key, scratch);
68 }
69
70 pub fn automorphism_inplace<M, K, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
71 where
72 M: GLWEAutomorphism<BE>,
73 K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
74 Scratch<BE>: ScratchTakeCore<BE>,
75 {
76 module.glwe_automorphism_inplace(self, key, scratch);
77 }
78
79 pub fn automorphism_add_inplace<M, K, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
80 where
81 M: GLWEAutomorphism<BE>,
82 K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
83 Scratch<BE>: ScratchTakeCore<BE>,
84 {
85 module.glwe_automorphism_add_inplace(self, key, scratch);
86 }
87
88 pub fn automorphism_sub_inplace<M, K, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
89 where
90 M: GLWEAutomorphism<BE>,
91 K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
92 Scratch<BE>: ScratchTakeCore<BE>,
93 {
94 module.glwe_automorphism_sub_inplace(self, key, scratch);
95 }
96
97 pub fn automorphism_sub_negate_inplace<M, K, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
98 where
99 M: GLWEAutomorphism<BE>,
100 K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
101 Scratch<BE>: ScratchTakeCore<BE>,
102 {
103 module.glwe_automorphism_sub_negate_inplace(self, key, scratch);
104 }
105}
106
107pub trait GLWEAutomorphism<BE: Backend>
108where
109 Self: GLWEKeyswitch<BE>
110 + VecZnxAutomorphismInplace<BE>
111 + VecZnxBigAutomorphismInplace<BE>
112 + VecZnxBigSubSmallInplace<BE>
113 + VecZnxBigSubSmallNegateInplace<BE>,
114{
115 fn glwe_automorphism_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
116 where
117 R: GLWEInfos,
118 A: GLWEInfos,
119 K: GGLWEInfos,
120 {
121 self.glwe_keyswitch_tmp_bytes(res_infos, a_infos, key_infos)
122 }
123
124 fn glwe_automorphism<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
125 where
126 R: GLWEToMut,
127 A: GLWEToRef,
128 K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
129 Scratch<BE>: ScratchTakeCore<BE>,
130 {
131 self.glwe_keyswitch(res, a, &key.to_ref().key, scratch);
132
133 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
134
135 for i in 0..res.rank().as_usize() + 1 {
136 self.vec_znx_automorphism_inplace(key.p(), res.data_mut(), i, scratch);
137 }
138 }
139
140 fn glwe_automorphism_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
141 where
142 R: GLWEToMut,
143 K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
144 Scratch<BE>: ScratchTakeCore<BE>,
145 {
146 self.glwe_keyswitch_inplace(res, &key.to_ref().key, scratch);
147
148 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
149
150 for i in 0..res.rank().as_usize() + 1 {
151 self.vec_znx_automorphism_inplace(key.p(), res.data_mut(), i, scratch);
152 }
153 }
154
155 fn glwe_automorphism_add<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
156 where
157 R: GLWEToMut,
158 A: GLWEToRef,
159 K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
160 Scratch<BE>: ScratchTakeCore<BE>,
161 {
162 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
163 let a: &GLWE<&[u8]> = &a.to_ref();
164 let key: &AutomorphismKeyPrepared<&[u8], BE> = &key.to_ref();
165
166 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
167 let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, a, &key.key, scratch_1);
168
169 for i in 0..res.rank().as_usize() + 1 {
170 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
171 self.vec_znx_big_add_small_inplace(&mut res_big, i, a.data(), i);
172 self.vec_znx_big_normalize(
173 res.base2k().into(),
174 res.data_mut(),
175 i,
176 key.base2k().into(),
177 &res_big,
178 i,
179 scratch_1,
180 );
181 }
182 }
183
184 fn glwe_automorphism_add_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
185 where
186 R: GLWEToMut,
187 K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
188 Scratch<BE>: ScratchTakeCore<BE>,
189 {
190 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
191 let key: &AutomorphismKeyPrepared<&[u8], BE> = &key.to_ref();
192
193 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
194 let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, res, &key.key, scratch_1);
195
196 for i in 0..res.rank().as_usize() + 1 {
197 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
198 self.vec_znx_big_add_small_inplace(&mut res_big, i, res.data(), i);
199 self.vec_znx_big_normalize(
200 res.base2k().into(),
201 res.data_mut(),
202 i,
203 key.base2k().into(),
204 &res_big,
205 i,
206 scratch_1,
207 );
208 }
209 }
210
211 fn glwe_automorphism_sub<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
212 where
213 R: GLWEToMut,
214 A: GLWEToRef,
215 K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
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 let key: &AutomorphismKeyPrepared<&[u8], BE> = &key.to_ref();
221
222 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
223 let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, a, &key.key, scratch_1);
224
225 for i in 0..res.rank().as_usize() + 1 {
226 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
227 self.vec_znx_big_sub_small_inplace(&mut res_big, i, a.data(), i);
228 self.vec_znx_big_normalize(
229 res.base2k().into(),
230 res.data_mut(),
231 i,
232 key.base2k().into(),
233 &res_big,
234 i,
235 scratch_1,
236 );
237 }
238 }
239
240 fn glwe_automorphism_sub_negate<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
241 where
242 R: GLWEToMut,
243 A: GLWEToRef,
244 K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
245 Scratch<BE>: ScratchTakeCore<BE>,
246 {
247 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
248 let a: &GLWE<&[u8]> = &a.to_ref();
249 let key: &AutomorphismKeyPrepared<&[u8], BE> = &key.to_ref();
250
251 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
252 let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, a, &key.key, scratch_1);
253
254 for i in 0..res.rank().as_usize() + 1 {
255 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
256 self.vec_znx_big_sub_small_negate_inplace(&mut res_big, i, a.data(), i);
257 self.vec_znx_big_normalize(
258 res.base2k().into(),
259 res.data_mut(),
260 i,
261 key.base2k().into(),
262 &res_big,
263 i,
264 scratch_1,
265 );
266 }
267 }
268
269 fn glwe_automorphism_sub_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
270 where
271 R: GLWEToMut,
272 K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
273 Scratch<BE>: ScratchTakeCore<BE>,
274 {
275 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
276 let key: &AutomorphismKeyPrepared<&[u8], BE> = &key.to_ref();
277
278 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
279 let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, res, &key.key, scratch_1);
280
281 for i in 0..res.rank().as_usize() + 1 {
282 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
283 self.vec_znx_big_sub_small_inplace(&mut res_big, i, res.data(), i);
284 self.vec_znx_big_normalize(
285 res.base2k().into(),
286 res.data_mut(),
287 i,
288 key.base2k().into(),
289 &res_big,
290 i,
291 scratch_1,
292 );
293 }
294 }
295
296 fn glwe_automorphism_sub_negate_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
297 where
298 R: GLWEToMut,
299 K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
300 Scratch<BE>: ScratchTakeCore<BE>,
301 {
302 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
303 let key: &AutomorphismKeyPrepared<&[u8], BE> = &key.to_ref();
304
305 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
306 let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, res, &key.key, scratch_1);
307
308 for i in 0..res.rank().as_usize() + 1 {
309 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
310 self.vec_znx_big_sub_small_negate_inplace(&mut res_big, i, res.data(), i);
311 self.vec_znx_big_normalize(
312 res.base2k().into(),
313 res.data_mut(),
314 i,
315 key.base2k().into(),
316 &res_big,
317 i,
318 scratch_1,
319 );
320 }
321 }
322}
323
324impl<BE: Backend> GLWEAutomorphism<BE> for Module<BE> where
325 Self: GLWEKeyswitch<BE>
326 + VecZnxAutomorphismInplace<BE>
327 + VecZnxBigAutomorphismInplace<BE>
328 + VecZnxBigSubSmallInplace<BE>
329 + VecZnxBigSubSmallNegateInplace<BE>
330{
331}