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