poulpy/poulpy-core/src/keyswitching/glwe_ct.rs
poulpy-core/src/keyswitching/glwe_ct.rsRaw
1use poulpy_hal::{
2 api::{
3 ModuleN, ScratchAvailable, ScratchTakeBasic, VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes,
4 VecZnxDftApply, VecZnxDftBytesOf, VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft,
5 VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
6 },
7 layouts::{Backend, DataMut, DataRef, DataViewMut, Module, Scratch, VecZnx, VecZnxBig, VecZnxDft, VmpPMat, ZnxInfos},
8};
9
10use crate::{
11 ScratchTakeCore,
12 layouts::{
13 GGLWEInfos, GLWE, GLWEInfos, GLWEToMut, GLWEToRef, LWEInfos,
14 prepared::{GLWESwitchingKeyPrepared, GLWESwitchingKeyPreparedToRef},
15 },
16};
17
18impl GLWE<Vec<u8>> {
19 pub fn keyswitch_tmp_bytes<M, R, A, B, BE: Backend>(module: &M, res_infos: &R, a_infos: &A, b_infos: &B) -> usize
20 where
21 R: GLWEInfos,
22 A: GLWEInfos,
23 B: GGLWEInfos,
24 M: GLWEKeySwitch<BE>,
25 {
26 module.glwe_keyswitch_tmp_bytes(res_infos, a_infos, b_infos)
27 }
28}
29
30impl<D: DataMut> GLWE<D> {
31 pub fn keyswitch<A, B, M, BE: Backend>(&mut self, module: &M, a: &A, b: &B, scratch: &mut Scratch<BE>)
32 where
33 A: GLWEToRef,
34 B: GLWESwitchingKeyPreparedToRef<BE>,
35 M: GLWEKeySwitch<BE>,
36 Scratch<BE>: ScratchTakeCore<BE>,
37 {
38 module.glwe_keyswitch(self, a, b, scratch);
39 }
40
41 pub fn keyswitch_inplace<A, M, BE: Backend>(&mut self, module: &M, a: &A, scratch: &mut Scratch<BE>)
42 where
43 A: GLWESwitchingKeyPreparedToRef<BE>,
44 M: GLWEKeySwitch<BE>,
45 Scratch<BE>: ScratchTakeCore<BE>,
46 {
47 module.glwe_keyswitch_inplace(self, a, scratch);
48 }
49}
50
51impl<BE: Backend> GLWEKeySwitch<BE> for Module<BE> where
52 Self: Sized
53 + ModuleN
54 + VecZnxDftBytesOf
55 + VmpApplyDftToDftTmpBytes
56 + VecZnxBigNormalizeTmpBytes
57 + VecZnxNormalizeTmpBytes
58 + VecZnxDftBytesOf
59 + VmpApplyDftToDftTmpBytes
60 + VecZnxBigNormalizeTmpBytes
61 + VmpApplyDftToDft<BE>
62 + VmpApplyDftToDftAdd<BE>
63 + VecZnxDftApply<BE>
64 + VecZnxIdftApplyConsume<BE>
65 + VecZnxBigAddSmallInplace<BE>
66 + VecZnxBigNormalize<BE>
67 + VecZnxNormalize<BE>
68 + VecZnxNormalizeTmpBytes
69{
70}
71
72pub trait GLWEKeySwitch<BE: Backend>
73where
74 Self: Sized
75 + ModuleN
76 + VecZnxDftBytesOf
77 + VmpApplyDftToDftTmpBytes
78 + VecZnxBigNormalizeTmpBytes
79 + VecZnxNormalizeTmpBytes
80 + VecZnxDftBytesOf
81 + VmpApplyDftToDftTmpBytes
82 + VecZnxBigNormalizeTmpBytes
83 + VmpApplyDftToDft<BE>
84 + VmpApplyDftToDftAdd<BE>
85 + VecZnxDftApply<BE>
86 + VecZnxIdftApplyConsume<BE>
87 + VecZnxBigAddSmallInplace<BE>
88 + VecZnxBigNormalize<BE>
89 + VecZnxNormalize<BE>
90 + VecZnxNormalizeTmpBytes,
91{
92 fn glwe_keyswitch_tmp_bytes<R, A, B>(&self, res_infos: &R, a_infos: &A, b_infos: &B) -> usize
93 where
94 R: GLWEInfos,
95 A: GLWEInfos,
96 B: GGLWEInfos,
97 {
98 let in_size: usize = a_infos
99 .k()
100 .div_ceil(b_infos.base2k())
101 .div_ceil(b_infos.dsize().into()) as usize;
102 let out_size: usize = res_infos.size();
103 let ksk_size: usize = b_infos.size();
104 let res_dft: usize = self.bytes_of_vec_znx_dft((b_infos.rank_out() + 1).into(), ksk_size); // TODO OPTIMIZE
105 let ai_dft: usize = self.bytes_of_vec_znx_dft((b_infos.rank_in()).into(), in_size);
106 let vmp: usize = self.vmp_apply_dft_to_dft_tmp_bytes(
107 out_size,
108 in_size,
109 in_size,
110 (b_infos.rank_in()).into(),
111 (b_infos.rank_out() + 1).into(),
112 ksk_size,
113 ) + self.bytes_of_vec_znx_dft((b_infos.rank_in()).into(), in_size);
114 let normalize_big: usize = self.vec_znx_big_normalize_tmp_bytes();
115 if a_infos.base2k() == b_infos.base2k() {
116 res_dft + ((ai_dft + vmp) | normalize_big)
117 } else if b_infos.dsize() == 1 {
118 // In this case, we only need one column, temporary, that we can drop once a_dft is computed.
119 let normalize_conv: usize = VecZnx::bytes_of(self.n(), 1, in_size) + self.vec_znx_normalize_tmp_bytes();
120 res_dft + (((ai_dft + normalize_conv) | vmp) | normalize_big)
121 } else {
122 // Since we stride over a to get a_dft when dsize > 1, we need to store the full columns of a with in the base conversion.
123 let normalize_conv: usize = VecZnx::bytes_of(self.n(), (b_infos.rank_in()).into(), in_size);
124 res_dft + ((ai_dft + normalize_conv + (self.vec_znx_normalize_tmp_bytes() | vmp)) | normalize_big)
125 }
126 }
127
128 fn glwe_keyswitch<R, A, B>(&self, res: &mut R, a: &A, b: &B, scratch: &mut Scratch<BE>)
129 where
130 R: GLWEToMut,
131 A: GLWEToRef,
132 B: GLWESwitchingKeyPreparedToRef<BE>,
133 Scratch<BE>: ScratchTakeCore<BE>,
134 {
135 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
136 let a: &GLWE<&[u8]> = &a.to_ref();
137 let b: &GLWESwitchingKeyPrepared<&[u8], BE> = &b.to_ref();
138
139 assert_eq!(
140 a.rank(),
141 b.rank_in(),
142 "a.rank(): {} != b.rank_in(): {}",
143 a.rank(),
144 b.rank_in()
145 );
146 assert_eq!(
147 res.rank(),
148 b.rank_out(),
149 "res.rank(): {} != b.rank_out(): {}",
150 res.rank(),
151 b.rank_out()
152 );
153
154 assert_eq!(res.n(), self.n() as u32);
155 assert_eq!(a.n(), self.n() as u32);
156 assert_eq!(b.n(), self.n() as u32);
157
158 let scrach_needed: usize = self.glwe_keyswitch_tmp_bytes(res, a, b);
159
160 assert!(
161 scratch.available() >= scrach_needed,
162 "scratch.available()={} < glwe_keyswitch_tmp_bytes={scrach_needed}",
163 scratch.available(),
164 );
165
166 let basek_out: usize = res.base2k().into();
167 let base2k_out: usize = b.base2k().into();
168
169 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), b.size()); // Todo optimise
170 let res_big: VecZnxBig<&mut [u8], BE> = keyswitch_internal(self, res_dft, a, b, scratch_1);
171 (0..(res.rank() + 1).into()).for_each(|i| {
172 self.vec_znx_big_normalize(
173 basek_out,
174 &mut res.data,
175 i,
176 base2k_out,
177 &res_big,
178 i,
179 scratch_1,
180 );
181 })
182 }
183
184 fn glwe_keyswitch_inplace<R, A>(&self, res: &mut R, a: &A, scratch: &mut Scratch<BE>)
185 where
186 R: GLWEToMut,
187 A: GLWESwitchingKeyPreparedToRef<BE>,
188 Scratch<BE>: ScratchTakeCore<BE>,
189 {
190 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
191 let a: &GLWESwitchingKeyPrepared<&[u8], BE> = &a.to_ref();
192
193 assert_eq!(
194 res.rank(),
195 a.rank_in(),
196 "res.rank(): {} != a.rank_in(): {}",
197 res.rank(),
198 a.rank_in()
199 );
200 assert_eq!(
201 res.rank(),
202 a.rank_out(),
203 "res.rank(): {} != b.rank_out(): {}",
204 res.rank(),
205 a.rank_out()
206 );
207
208 assert_eq!(res.n(), self.n() as u32);
209 assert_eq!(a.n(), self.n() as u32);
210
211 let scrach_needed: usize = self.glwe_keyswitch_tmp_bytes(res, res, a);
212
213 assert!(
214 scratch.available() >= scrach_needed,
215 "scratch.available()={} < glwe_keyswitch_tmp_bytes={scrach_needed}",
216 scratch.available(),
217 );
218
219 let base2k_in: usize = res.base2k().into();
220 let base2k_out: usize = a.base2k().into();
221
222 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), a.size()); // Todo optimise
223 let res_big: VecZnxBig<&mut [u8], BE> = keyswitch_internal(self, res_dft, res, a, scratch_1);
224 (0..(res.rank() + 1).into()).for_each(|i| {
225 self.vec_znx_big_normalize(
226 base2k_in,
227 &mut res.data,
228 i,
229 base2k_out,
230 &res_big,
231 i,
232 scratch_1,
233 );
234 })
235 }
236}
237
238impl GLWE<Vec<u8>> {}
239
240impl<DataSelf: DataMut> GLWE<DataSelf> {}
241
242fn keyswitch_internal<BE: Backend, M, DR, DA, DB>(
243 module: &M,
244 mut res: VecZnxDft<DR, BE>,
245 a: &GLWE<DA>,
246 b: &GLWESwitchingKeyPrepared<DB, BE>,
247 scratch: &mut Scratch<BE>,
248) -> VecZnxBig<DR, BE>
249where
250 DR: DataMut,
251 DA: DataRef,
252 DB: DataRef,
253 M: ModuleN
254 + VecZnxDftBytesOf
255 + VmpApplyDftToDftTmpBytes
256 + VecZnxBigNormalizeTmpBytes
257 + VmpApplyDftToDftTmpBytes
258 + VmpApplyDftToDft<BE>
259 + VmpApplyDftToDftAdd<BE>
260 + VecZnxDftApply<BE>
261 + VecZnxIdftApplyConsume<BE>
262 + VecZnxBigAddSmallInplace<BE>
263 + VecZnxBigNormalize<BE>
264 + VecZnxNormalize<BE>,
265 Scratch<BE>: ScratchTakeCore<BE>,
266{
267 let base2k_in: usize = a.base2k().into();
268 let base2k_out: usize = b.base2k().into();
269 let cols: usize = (a.rank() + 1).into();
270 let a_size: usize = (a.size() * base2k_in).div_ceil(base2k_out);
271 let pmat: &VmpPMat<DB, BE> = &b.key.data;
272
273 if b.dsize() == 1 {
274 let (mut ai_dft, scratch_1) = scratch.take_vec_znx_dft(module, cols - 1, a.size());
275
276 if base2k_in == base2k_out {
277 (0..cols - 1).for_each(|col_i| {
278 module.vec_znx_dft_apply(1, 0, &mut ai_dft, col_i, a.data(), col_i + 1);
279 });
280 } else {
281 let (mut a_conv, scratch_2) = scratch_1.take_vec_znx(module, 1, a_size);
282 (0..cols - 1).for_each(|col_i| {
283 module.vec_znx_normalize(
284 base2k_out,
285 &mut a_conv,
286 0,
287 base2k_in,
288 a.data(),
289 col_i + 1,
290 scratch_2,
291 );
292 module.vec_znx_dft_apply(1, 0, &mut ai_dft, col_i, &a_conv, 0);
293 });
294 }
295
296 module.vmp_apply_dft_to_dft(&mut res, &ai_dft, pmat, scratch_1);
297 } else {
298 let dsize: usize = b.dsize().into();
299
300 let (mut ai_dft, scratch_1) = scratch.take_vec_znx_dft(module, cols - 1, a_size.div_ceil(dsize));
301 ai_dft.data_mut().fill(0);
302
303 if base2k_in == base2k_out {
304 for di in 0..dsize {
305 ai_dft.set_size((a_size + di) / dsize);
306
307 // Small optimization for dsize > 2
308 // VMP produce some error e, and since we aggregate vmp * 2^{di * B}, then
309 // we also aggregate ei * 2^{di * B}, with the largest error being ei * 2^{(dsize-1) * B}.
310 // As such we can ignore the last dsize-2 limbs safely of the sum of vmp products.
311 // It is possible to further ignore the last dsize-1 limbs, but this introduce
312 // ~0.5 to 1 bit of additional noise, and thus not chosen here to ensure that the same
313 // noise is kept with respect to the ideal functionality.
314 res.set_size(pmat.size() - ((dsize - di) as isize - 2).max(0) as usize);
315
316 for j in 0..cols - 1 {
317 module.vec_znx_dft_apply(dsize, dsize - di - 1, &mut ai_dft, j, a.data(), j + 1);
318 }
319
320 if di == 0 {
321 module.vmp_apply_dft_to_dft(&mut res, &ai_dft, pmat, scratch_1);
322 } else {
323 module.vmp_apply_dft_to_dft_add(&mut res, &ai_dft, pmat, di, scratch_1);
324 }
325 }
326 } else {
327 let (mut a_conv, scratch_2) = scratch_1.take_vec_znx(module, cols - 1, a_size);
328 for j in 0..cols - 1 {
329 module.vec_znx_normalize(
330 base2k_out,
331 &mut a_conv,
332 j,
333 base2k_in,
334 a.data(),
335 j + 1,
336 scratch_2,
337 );
338 }
339
340 for di in 0..dsize {
341 ai_dft.set_size((a_size + di) / dsize);
342
343 // Small optimization for dsize > 2
344 // VMP produce some error e, and since we aggregate vmp * 2^{di * B}, then
345 // we also aggregate ei * 2^{di * B}, with the largest error being ei * 2^{(dsize-1) * B}.
346 // As such we can ignore the last dsize-2 limbs safely of the sum of vmp products.
347 // It is possible to further ignore the last dsize-1 limbs, but this introduce
348 // ~0.5 to 1 bit of additional noise, and thus not chosen here to ensure that the same
349 // noise is kept with respect to the ideal functionality.
350 res.set_size(pmat.size() - ((dsize - di) as isize - 2).max(0) as usize);
351
352 for j in 0..cols - 1 {
353 module.vec_znx_dft_apply(dsize, dsize - di - 1, &mut ai_dft, j, &a_conv, j);
354 }
355
356 if di == 0 {
357 module.vmp_apply_dft_to_dft(&mut res, &ai_dft, pmat, scratch_2);
358 } else {
359 module.vmp_apply_dft_to_dft_add(&mut res, &ai_dft, pmat, di, scratch_2);
360 }
361 }
362 }
363
364 res.set_size(res.max_size());
365 }
366
367 let mut res_big: VecZnxBig<DR, BE> = module.vec_znx_idft_apply_consume(res);
368 module.vec_znx_big_add_small_inplace(&mut res_big, 0, a.data(), 0);
369 res_big
370}