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