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