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