poulpy/poulpy-core/src/keyswitching/glwe.rs
poulpy-core/src/keyswitching/glwe.rsRaw
1use poulpy_hal::{
2 api::{
3 ModuleN, ScratchAvailable, ScratchTakeBasic, VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes,
4 VecZnxDftApply, VecZnxDftBytesOf, VecZnxDftCopy, VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeTmpBytes,
5 VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
6 },
7 layouts::{Backend, DataMut, DataViewMut, Module, Scratch, VecZnxBig, VecZnxDft, VecZnxDftToRef, VmpPMat, ZnxInfos},
8};
9
10use crate::{
11 GLWENormalize, ScratchTakeCore,
12 layouts::{GGLWEInfos, GGLWEPrepared, GGLWEPreparedToRef, GLWE, GLWEInfos, GLWELayout, GLWEToMut, GLWEToRef, LWEInfos},
13};
14
15impl GLWE<Vec<u8>> {
16 pub fn keyswitch_tmp_bytes<M, R, A, B, BE: Backend>(module: &M, res_infos: &R, a_infos: &A, key_infos: &B) -> usize
17 where
18 R: GLWEInfos,
19 A: GLWEInfos,
20 B: GGLWEInfos,
21 M: GLWEKeyswitch<BE>,
22 {
23 module.glwe_keyswitch_tmp_bytes(res_infos, a_infos, key_infos)
24 }
25}
26
27impl<D: DataMut> GLWE<D> {
28 pub fn keyswitch<A, B, M, BE: Backend>(&mut self, module: &M, a: &A, b: &B, scratch: &mut Scratch<BE>)
29 where
30 A: GLWEToRef + GLWEInfos,
31 B: GGLWEPreparedToRef<BE> + GGLWEInfos,
32 M: GLWEKeyswitch<BE>,
33 Scratch<BE>: ScratchTakeCore<BE>,
34 {
35 module.glwe_keyswitch(self, a, b, scratch);
36 }
37
38 pub fn keyswitch_inplace<A, M, BE: Backend>(&mut self, module: &M, a: &A, scratch: &mut Scratch<BE>)
39 where
40 A: GGLWEPreparedToRef<BE> + GGLWEInfos,
41 M: GLWEKeyswitch<BE>,
42 Scratch<BE>: ScratchTakeCore<BE>,
43 {
44 module.glwe_keyswitch_inplace(self, a, scratch);
45 }
46}
47
48impl<BE: Backend> GLWEKeyswitch<BE> for Module<BE>
49where
50 Self: Sized + GLWEKeySwitchInternal<BE> + VecZnxBigNormalizeTmpBytes + VecZnxBigNormalize<BE> + GLWENormalize<BE>,
51 Scratch<BE>: ScratchTakeCore<BE>,
52{
53 fn glwe_keyswitch_tmp_bytes<R, A, B>(&self, res_infos: &R, a_infos: &A, key_infos: &B) -> usize
54 where
55 R: GLWEInfos,
56 A: GLWEInfos,
57 B: GGLWEInfos,
58 {
59 let cols: usize = res_infos.rank().as_usize() + 1;
60 let size: usize = if a_infos.base2k() != key_infos.base2k() {
61 let a_conv_infos = &GLWELayout {
62 n: a_infos.n(),
63 base2k: key_infos.base2k(),
64 k: a_infos.k(),
65 rank: a_infos.rank(),
66 };
67 self.glwe_keyswitch_internal_tmp_bytes(res_infos, a_conv_infos, key_infos) + GLWE::bytes_of_from_infos(a_conv_infos)
68 } else {
69 self.glwe_keyswitch_internal_tmp_bytes(res_infos, a_infos, key_infos)
70 };
71
72 size.max(self.vec_znx_big_normalize_tmp_bytes()) + self.bytes_of_vec_znx_dft(cols, key_infos.size())
73 }
74
75 fn glwe_keyswitch<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
76 where
77 R: GLWEToMut + GLWEInfos,
78 A: GLWEToRef + GLWEInfos,
79 K: GGLWEPreparedToRef<BE> + GGLWEInfos,
80 {
81 assert_eq!(
82 a.rank(),
83 key.rank_in(),
84 "a.rank(): {} != b.rank_in(): {}",
85 a.rank(),
86 key.rank_in()
87 );
88 assert_eq!(
89 res.rank(),
90 key.rank_out(),
91 "res.rank(): {} != b.rank_out(): {}",
92 res.rank(),
93 key.rank_out()
94 );
95
96 assert_eq!(res.n(), self.n() as u32);
97 assert_eq!(a.n(), self.n() as u32);
98 assert_eq!(key.n(), self.n() as u32);
99
100 let scrach_needed: usize = self.glwe_keyswitch_tmp_bytes(res, a, key);
101
102 assert!(
103 scratch.available() >= scrach_needed,
104 "scratch.available()={} < glwe_keyswitch_tmp_bytes={scrach_needed}",
105 scratch.available(),
106 );
107
108 let a_base2k: usize = a.base2k().into();
109 let key_base2k: usize = key.base2k().into();
110 let res_base2k: usize = res.base2k().into();
111
112 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // Todo optimise
113
114 let res_big: VecZnxBig<&mut [u8], BE> = if a_base2k != key_base2k {
115 let (mut a_conv, scratch_2) = scratch_1.take_glwe(&GLWELayout {
116 n: a.n(),
117 base2k: key.base2k(),
118 k: a.k(),
119 rank: a.rank(),
120 });
121 self.glwe_normalize(&mut a_conv, a, scratch_2);
122 self.glwe_keyswitch_internal(res_dft, &a_conv, key, scratch_2)
123 } else {
124 self.glwe_keyswitch_internal(res_dft, a, key, scratch_1)
125 };
126
127 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
128 for i in 0..(res.rank() + 1).into() {
129 self.vec_znx_big_normalize(res.data_mut(), res_base2k, 0, i, &res_big, key_base2k, i, scratch_1);
130 }
131 }
132
133 fn glwe_keyswitch_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
134 where
135 R: GLWEToMut + GLWEInfos,
136 K: GGLWEPreparedToRef<BE> + GGLWEInfos,
137 {
138 assert_eq!(
139 res.rank(),
140 key.rank_in(),
141 "res.rank(): {} != a.rank_in(): {}",
142 res.rank(),
143 key.rank_in()
144 );
145 assert_eq!(
146 res.rank(),
147 key.rank_out(),
148 "res.rank(): {} != b.rank_out(): {}",
149 res.rank(),
150 key.rank_out()
151 );
152
153 assert_eq!(res.n(), self.n() as u32);
154 assert_eq!(key.n(), self.n() as u32);
155
156 let scrach_needed: usize = self.glwe_keyswitch_tmp_bytes(res, res, key);
157
158 assert!(
159 scratch.available() >= scrach_needed,
160 "scratch.available()={} < glwe_keyswitch_tmp_bytes={scrach_needed}",
161 scratch.available(),
162 );
163
164 let res_base2k: usize = res.base2k().as_usize();
165 let key_base2k: usize = key.base2k().as_usize();
166
167 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // Todo optimise
168
169 let res_big: VecZnxBig<&mut [u8], BE> = if res_base2k != key_base2k {
170 let (mut res_conv, scratch_2) = scratch_1.take_glwe(&GLWELayout {
171 n: res.n(),
172 base2k: key.base2k(),
173 k: res.k(),
174 rank: res.rank(),
175 });
176 self.glwe_normalize(&mut res_conv, res, scratch_2);
177
178 self.glwe_keyswitch_internal(res_dft, &res_conv, key, scratch_2)
179 } else {
180 self.glwe_keyswitch_internal(res_dft, res, key, scratch_1)
181 };
182
183 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
184 for i in 0..(res.rank() + 1).into() {
185 self.vec_znx_big_normalize(res.data_mut(), res_base2k, 0, i, &res_big, key_base2k, i, scratch_1);
186 }
187 }
188}
189
190pub trait GLWEKeyswitch<BE: Backend> {
191 fn glwe_keyswitch_tmp_bytes<R, A, B>(&self, res_infos: &R, a_infos: &A, key_infos: &B) -> usize
192 where
193 R: GLWEInfos,
194 A: GLWEInfos,
195 B: GGLWEInfos;
196
197 fn glwe_keyswitch<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
198 where
199 R: GLWEToMut + GLWEInfos,
200 A: GLWEToRef + GLWEInfos,
201 K: GGLWEPreparedToRef<BE> + GGLWEInfos;
202
203 fn glwe_keyswitch_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
204 where
205 R: GLWEToMut + GLWEInfos,
206 K: GGLWEPreparedToRef<BE> + GGLWEInfos;
207}
208
209impl<BE: Backend> GLWEKeySwitchInternal<BE> for Module<BE> where
210 Self: GGLWEProduct<BE>
211 + VecZnxDftApply<BE>
212 + VecZnxNormalize<BE>
213 + VecZnxIdftApplyConsume<BE>
214 + VecZnxBigAddSmallInplace<BE>
215 + VecZnxNormalizeTmpBytes
216{
217}
218
219pub(crate) trait GLWEKeySwitchInternal<BE: Backend>
220where
221 Self: GGLWEProduct<BE>
222 + VecZnxDftApply<BE>
223 + VecZnxNormalize<BE>
224 + VecZnxIdftApplyConsume<BE>
225 + VecZnxBigAddSmallInplace<BE>
226 + VecZnxNormalizeTmpBytes,
227{
228 fn glwe_keyswitch_internal_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
229 where
230 R: GLWEInfos,
231 A: GLWEInfos,
232 K: GGLWEInfos,
233 {
234 let cols: usize = (a_infos.rank() + 1).into();
235 let a_size: usize = a_infos.size();
236 self.gglwe_product_dft_tmp_bytes(res_infos.size(), a_size, key_infos) + self.bytes_of_vec_znx_dft(cols - 1, a_size)
237 }
238
239 fn glwe_keyswitch_internal<DR, A, K>(
240 &self,
241 mut res: VecZnxDft<DR, BE>,
242 a: &A,
243 key: &K,
244 scratch: &mut Scratch<BE>,
245 ) -> VecZnxBig<DR, BE>
246 where
247 DR: DataMut,
248 A: GLWEToRef,
249 K: GGLWEPreparedToRef<BE>,
250 Scratch<BE>: ScratchTakeCore<BE>,
251 {
252 let a: &GLWE<&[u8]> = &a.to_ref();
253 let key: &GGLWEPrepared<&[u8], BE> = &key.to_ref();
254 assert_eq!(a.base2k(), key.base2k());
255 let cols: usize = (a.rank() + 1).into();
256 let a_size: usize = a.size();
257 let (mut a_dft, scratch_1) = scratch.take_vec_znx_dft(self, cols - 1, a_size);
258 for col_i in 0..cols - 1 {
259 self.vec_znx_dft_apply(1, 0, &mut a_dft, col_i, a.data(), col_i + 1);
260 }
261 self.gglwe_product_dft(&mut res, &a_dft, key, scratch_1);
262 let mut res_big: VecZnxBig<DR, BE> = self.vec_znx_idft_apply_consume(res);
263 self.vec_znx_big_add_small_inplace(&mut res_big, 0, a.data(), 0);
264 res_big
265 }
266}
267
268impl<BE: Backend> GGLWEProduct<BE> for Module<BE> where
269 Self: Sized
270 + ModuleN
271 + VecZnxDftBytesOf
272 + VmpApplyDftToDftTmpBytes
273 + VmpApplyDftToDft<BE>
274 + VmpApplyDftToDftAdd<BE>
275 + VecZnxDftCopy<BE>
276{
277}
278
279pub(crate) trait GGLWEProduct<BE: Backend>
280where
281 Self: Sized
282 + ModuleN
283 + VecZnxDftBytesOf
284 + VmpApplyDftToDftTmpBytes
285 + VmpApplyDftToDft<BE>
286 + VmpApplyDftToDftAdd<BE>
287 + VecZnxDftCopy<BE>,
288{
289 fn gglwe_product_dft_tmp_bytes<K>(&self, res_size: usize, a_size: usize, key_infos: &K) -> usize
290 where
291 K: GGLWEInfos,
292 {
293 let dsize: usize = key_infos.dsize().as_usize();
294
295 if dsize == 1 {
296 self.vmp_apply_dft_to_dft_tmp_bytes(
297 res_size,
298 a_size,
299 key_infos.dnum().into(),
300 (key_infos.rank_in()).into(),
301 (key_infos.rank_out() + 1).into(),
302 key_infos.size(),
303 )
304 } else {
305 let dnum: usize = key_infos.dnum().into();
306 let a_size: usize = a_size.div_ceil(dsize).min(dnum);
307 let ai_dft: usize = self.bytes_of_vec_znx_dft(key_infos.rank_in().into(), a_size);
308
309 let vmp: usize = self.vmp_apply_dft_to_dft_tmp_bytes(
310 res_size,
311 a_size,
312 dnum,
313 (key_infos.rank_in()).into(),
314 (key_infos.rank_out() + 1).into(),
315 key_infos.size(),
316 );
317
318 ai_dft + vmp
319 }
320 }
321
322 fn gglwe_product_dft<DR, A, K>(&self, res: &mut VecZnxDft<DR, BE>, a: &A, key: &K, scratch: &mut Scratch<BE>)
323 where
324 DR: DataMut,
325 A: VecZnxDftToRef<BE>,
326 K: GGLWEPreparedToRef<BE>,
327 Scratch<BE>: ScratchTakeCore<BE>,
328 {
329 let a: &VecZnxDft<&[u8], BE> = &a.to_ref();
330 let key: &GGLWEPrepared<&[u8], BE> = &key.to_ref();
331
332 let cols: usize = a.cols();
333 let a_size: usize = a.size();
334 let pmat: &VmpPMat<&[u8], BE> = &key.data;
335
336 // If dsize == 1, then the digit decomposition is equal to Base2K and we can simply
337 // can the vmp API.
338 if key.dsize() == 1 {
339 self.vmp_apply_dft_to_dft(res, a, pmat, scratch);
340 // If dsize != 1, then the digit decomposition is k * Base2K with k > 1.
341 // As such we need to perform a bivariate polynomial convolution in (X, Y) / (X^{N}+1) with Y = 2^-K
342 // (instead of yn univariate one in X).
343 //
344 // Since the basis in Y is small (in practice degree 6-7 max), we perform it naiveley.
345 // To do so, we group the different limbs of ai_dft by their respective degree in Y
346 // which are multiples of the current digit.
347 // For example if dsize = 3, with ai_dft = [a0, a1, a2, a3, a4, a5, a6],
348 // we group them as [[a0, a3, a5], [a1, a4, a6], [a2, a5, 0]]
349 // and evaluate sum(a_di * pmat * 2^{di*Base2k})
350 } else {
351 let dsize: usize = key.dsize().into();
352 let dnum: usize = key.dnum().into();
353
354 // We bound ai_dft size by the number of rows of the matrix
355 let (mut ai_dft, scratch_1) = scratch.take_vec_znx_dft(self, cols, a_size.div_ceil(dsize).min(dnum));
356 ai_dft.data_mut().fill(0);
357
358 for di in 0..dsize {
359 // Sets ai_dft size according to the current digit (if dsize does not divides a_size),
360 // bounded by the number of rows (digits) in the prepared matrix.
361 ai_dft.set_size(((a_size + di) / dsize).min(dnum));
362
363 // Small optimization for dsize > 2
364 // VMP produce some error e, and since we aggregate vmp * 2^{di * Base2k}, then
365 // we also aggregate ei * 2^{di * Base2k}, with the largest error being ei * 2^{(dsize-1) * Base2k}.
366 // As such we can ignore the last dsize-2 limbs safely of the sum of vmp products.
367 // It is possible to further ignore the last dsize-1 limbs, but this introduce
368 // ~0.5 to 1 bit of additional noise, and thus not chosen here to ensure that the same
369 // noise is kept with respect to the ideal functionality.
370 res.set_size(pmat.size() - ((dsize - di) as isize - 2).max(0) as usize);
371
372 for j in 0..cols {
373 self.vec_znx_dft_copy(dsize, dsize - di - 1, &mut ai_dft, j, a, j);
374 }
375
376 if di == 0 {
377 // res = pmat * ai_dft
378 self.vmp_apply_dft_to_dft(res, &ai_dft, pmat, scratch_1);
379 } else {
380 // res = (pmat * ai_dft) * 2^{di * Base2k}
381 self.vmp_apply_dft_to_dft_add(res, &ai_dft, pmat, di, scratch_1);
382 }
383 }
384
385 res.set_size(res.max_size());
386 }
387 }
388}