poulpy/poulpy-core/src/external_product/glwe.rs
poulpy-core/src/external_product/glwe.rsRaw
1use poulpy_hal::{
2 api::{
3 ModuleN, ScratchTakeBasic, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftApply, VecZnxDftBytesOf,
4 VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft, VmpApplyDftToDftAdd,
5 VmpApplyDftToDftTmpBytes,
6 },
7 layouts::{Backend, DataMut, DataViewMut, Module, Scratch, VecZnx, VecZnxBig, VecZnxDft},
8};
9
10use crate::{
11 ScratchTakeCore,
12 layouts::{
13 GGSWInfos, GLWE, GLWEInfos, 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}
71
72impl<BE: Backend> GLWEExternalProduct<BE> for Module<BE>
73where
74 Self: GLWEExternalProductInternal<BE> + VecZnxDftBytesOf + VecZnxBigNormalize<BE> + VecZnxBigNormalizeTmpBytes,
75{
76 fn glwe_external_product_tmp_bytes<R, A, B>(&self, res_infos: &R, a_infos: &A, b_infos: &B) -> usize
77 where
78 R: GLWEInfos,
79 A: GLWEInfos,
80 B: GGSWInfos,
81 {
82 let res_dft: usize = self.bytes_of_vec_znx_dft((b_infos.rank() + 1).into(), b_infos.size());
83 res_dft
84 + self
85 .glwe_external_product_internal_tmp_bytes(res_infos, a_infos, b_infos)
86 .max(self.vec_znx_big_normalize_tmp_bytes())
87 }
88
89 fn glwe_external_product_inplace<R, D>(&self, res: &mut R, a: &D, scratch: &mut Scratch<BE>)
90 where
91 R: GLWEToMut,
92 D: GGSWPreparedToRef<BE> + GGSWInfos,
93 Scratch<BE>: ScratchTakeCore<BE>,
94 {
95 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
96 let rhs: &GGSWPrepared<&[u8], BE> = &a.to_ref();
97
98 let basek_in: usize = res.base2k().into();
99 let basek_ggsw: usize = rhs.base2k().into();
100
101 #[cfg(debug_assertions)]
102 {
103 use poulpy_hal::api::ScratchAvailable;
104
105 assert_eq!(rhs.rank(), res.rank());
106 assert_eq!(rhs.n(), res.n());
107 assert!(scratch.available() >= self.glwe_external_product_tmp_bytes(res, res, rhs));
108 }
109
110 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), a.size()); // Todo optimise
111 let res_big = self.glwe_external_product_internal(res_dft, res, a, scratch_1);
112 for j in 0..(res.rank() + 1).into() {
113 self.vec_znx_big_normalize(
114 basek_in,
115 &mut res.data,
116 j,
117 basek_ggsw,
118 &res_big,
119 j,
120 scratch_1,
121 );
122 }
123 }
124
125 fn glwe_external_product<R, A, D>(&self, res: &mut R, lhs: &A, rhs: &D, scratch: &mut Scratch<BE>)
126 where
127 R: GLWEToMut,
128 A: GLWEToRef,
129 D: GGSWPreparedToRef<BE>,
130 Scratch<BE>: ScratchTakeCore<BE>,
131 {
132 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
133 let lhs: &GLWE<&[u8]> = &lhs.to_ref();
134
135 let rhs: &GGSWPrepared<&[u8], BE> = &rhs.to_ref();
136
137 let basek_ggsw: usize = rhs.base2k().into();
138 let basek_out: usize = res.base2k().into();
139
140 #[cfg(debug_assertions)]
141 {
142 use poulpy_hal::api::ScratchAvailable;
143
144 assert_eq!(rhs.rank(), lhs.rank());
145 assert_eq!(rhs.rank(), res.rank());
146 assert_eq!(rhs.n(), res.n());
147 assert_eq!(lhs.n(), res.n());
148 assert!(scratch.available() >= self.glwe_external_product_tmp_bytes(res, lhs, rhs));
149 }
150
151 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), rhs.size()); // Todo optimise
152 let res_big = self.glwe_external_product_internal(res_dft, lhs, rhs, scratch_1);
153
154 for j in 0..(res.rank() + 1).into() {
155 self.vec_znx_big_normalize(
156 basek_out,
157 &mut res.data,
158 j,
159 basek_ggsw,
160 &res_big,
161 j,
162 scratch_1,
163 );
164 }
165 }
166}
167
168pub trait GLWEExternalProductInternal<BE: Backend> {
169 fn glwe_external_product_internal_tmp_bytes<R, A, B>(&self, res_infos: &R, a_infos: &A, b_infos: &B) -> usize
170 where
171 R: GLWEInfos,
172 A: GLWEInfos,
173 B: GGSWInfos;
174 fn glwe_external_product_internal<DR, A, G>(
175 &self,
176 res_dft: VecZnxDft<DR, BE>,
177 a: &A,
178 ggsw: &G,
179 scratch: &mut Scratch<BE>,
180 ) -> VecZnxBig<DR, BE>
181 where
182 DR: DataMut,
183 A: GLWEToRef,
184 G: GGSWPreparedToRef<BE>,
185 Scratch<BE>: ScratchTakeCore<BE>;
186}
187
188impl<BE: Backend> GLWEExternalProductInternal<BE> for Module<BE>
189where
190 Self: ModuleN
191 + VecZnxDftBytesOf
192 + VmpApplyDftToDftTmpBytes
193 + VecZnxNormalizeTmpBytes
194 + VecZnxDftApply<BE>
195 + VmpApplyDftToDft<BE>
196 + VmpApplyDftToDftAdd<BE>
197 + VecZnxIdftApplyConsume<BE>
198 + VecZnxBigNormalize<BE>
199 + VecZnxNormalize<BE>
200 + VecZnxDftBytesOf
201 + VmpApplyDftToDftTmpBytes
202 + VecZnxNormalizeTmpBytes,
203{
204 fn glwe_external_product_internal_tmp_bytes<R, A, B>(&self, res_infos: &R, a_infos: &A, b_infos: &B) -> usize
205 where
206 R: GLWEInfos,
207 A: GLWEInfos,
208 B: GGSWInfos,
209 {
210 let in_size: usize = a_infos
211 .k()
212 .div_ceil(b_infos.base2k())
213 .div_ceil(b_infos.dsize().into()) as usize;
214 let out_size: usize = res_infos.size();
215 let ggsw_size: usize = b_infos.size();
216 let a_dft: usize = self.bytes_of_vec_znx_dft((b_infos.rank() + 1).into(), in_size);
217 let vmp: usize = self.vmp_apply_dft_to_dft_tmp_bytes(
218 out_size,
219 in_size,
220 in_size, // rows
221 (b_infos.rank() + 1).into(), // cols in
222 (b_infos.rank() + 1).into(), // cols out
223 ggsw_size,
224 );
225 let normalize_big: usize = self.vec_znx_normalize_tmp_bytes();
226
227 if a_infos.base2k() == b_infos.base2k() {
228 a_dft + (vmp | normalize_big)
229 } else {
230 let normalize_conv: usize = VecZnx::bytes_of(self.n(), (b_infos.rank() + 1).into(), in_size);
231 (a_dft + normalize_conv + (self.vec_znx_normalize_tmp_bytes() | vmp)) | normalize_big
232 }
233 }
234
235 fn glwe_external_product_internal<DR, A, G>(
236 &self,
237 mut res_dft: VecZnxDft<DR, BE>,
238 a: &A,
239 ggsw: &G,
240 scratch: &mut Scratch<BE>,
241 ) -> VecZnxBig<DR, BE>
242 where
243 DR: DataMut,
244 A: GLWEToRef,
245 G: GGSWPreparedToRef<BE>,
246 Scratch<BE>: ScratchTakeCore<BE>,
247 {
248 let a: &GLWE<&[u8]> = &a.to_ref();
249 let ggsw: &GGSWPrepared<&[u8], BE> = &ggsw.to_ref();
250
251 let basek_in: usize = a.base2k().into();
252 let basek_ggsw: usize = ggsw.base2k().into();
253
254 let cols: usize = (ggsw.rank() + 1).into();
255 let dsize: usize = ggsw.dsize().into();
256 let a_size: usize = (a.size() * basek_in).div_ceil(basek_ggsw);
257
258 let (mut a_dft, scratch_1) = scratch.take_vec_znx_dft(self, cols, a_size.div_ceil(dsize));
259 a_dft.data_mut().fill(0);
260
261 if basek_in == basek_ggsw {
262 for di in 0..dsize {
263 // (lhs.size() + di) / dsize = (a - (digit - di - 1)).div_ceil(dsize)
264 a_dft.set_size((a.size() + di) / dsize);
265
266 // Small optimization for dsize > 2
267 // VMP produce some error e, and since we aggregate vmp * 2^{di * B}, then
268 // we also aggregate ei * 2^{di * B}, with the largest error being ei * 2^{(dsize-1) * B}.
269 // As such we can ignore the last dsize-2 limbs safely of the sum of vmp products.
270 // It is possible to further ignore the last dsize-1 limbs, but this introduce
271 // ~0.5 to 1 bit of additional noise, and thus not chosen here to ensure that the same
272 // noise is kept with respect to the ideal functionality.
273 res_dft.set_size(ggsw.size() - ((dsize - di) as isize - 2).max(0) as usize);
274
275 for j in 0..cols {
276 self.vec_znx_dft_apply(dsize, dsize - 1 - di, &mut a_dft, j, &a.data, j);
277 }
278
279 if di == 0 {
280 self.vmp_apply_dft_to_dft(&mut res_dft, &a_dft, &ggsw.data, scratch_1);
281 } else {
282 self.vmp_apply_dft_to_dft_add(&mut res_dft, &a_dft, &ggsw.data, di, scratch_1);
283 }
284 }
285 } else {
286 let (mut a_conv, scratch_3) = scratch_1.take_vec_znx(self.n(), cols, a_size);
287
288 for j in 0..cols {
289 self.vec_znx_normalize(basek_ggsw, &mut a_conv, j, basek_in, &a.data, j, scratch_3);
290 }
291
292 for di in 0..dsize {
293 // (lhs.size() + di) / dsize = (a - (digit - di - 1)).div_ceil(dsize)
294 a_dft.set_size((a.size() + di) / dsize);
295
296 // Small optimization for dsize > 2
297 // VMP produce some error e, and since we aggregate vmp * 2^{di * B}, then
298 // we also aggregate ei * 2^{di * B}, with the largest error being ei * 2^{(dsize-1) * B}.
299 // As such we can ignore the last dsize-2 limbs safely of the sum of vmp products.
300 // It is possible to further ignore the last dsize-1 limbs, but this introduce
301 // ~0.5 to 1 bit of additional noise, and thus not chosen here to ensure that the same
302 // noise is kept with respect to the ideal functionality.
303 res_dft.set_size(ggsw.size() - ((dsize - di) as isize - 2).max(0) as usize);
304
305 for j in 0..cols {
306 self.vec_znx_dft_apply(dsize, dsize - 1 - di, &mut a_dft, j, &a.data, j);
307 }
308
309 if di == 0 {
310 self.vmp_apply_dft_to_dft(&mut res_dft, &a_dft, &ggsw.data, scratch_1);
311 } else {
312 self.vmp_apply_dft_to_dft_add(&mut res_dft, &a_dft, &ggsw.data, di, scratch_1);
313 }
314 }
315 }
316
317 self.vec_znx_idft_apply_consume(res_dft)
318 }
319}