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 res_base2k: usize = res.base2k().as_usize();
114 let ggsw_base2k: 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 res_base2k != ggsw_base2k {
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(res.data_mut(), res_base2k, 0, j, &res_big, ggsw_base2k, j, scratch_1);
134 }
135 }
136
137 fn glwe_external_product<R, A, G>(&self, res: &mut R, a: &A, ggsw: &G, scratch: &mut Scratch<BE>)
138 where
139 R: GLWEToMut + GLWEInfos,
140 A: GLWEToRef + GLWEInfos,
141 G: GGSWPreparedToRef<BE> + GGSWInfos,
142 Scratch<BE>: ScratchTakeCore<BE>,
143 {
144 assert_eq!(ggsw.rank(), a.rank());
145 assert_eq!(ggsw.rank(), res.rank());
146 assert_eq!(ggsw.n(), res.n());
147 assert_eq!(a.n(), res.n());
148 assert!(scratch.available() >= self.glwe_external_product_tmp_bytes(res, a, ggsw));
149
150 let a_base2k: usize = a.base2k().into();
151 let ggsw_base2k: usize = ggsw.base2k().into();
152 let res_base2k: usize = res.base2k().into();
153
154 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), ggsw.size()); // Todo optimise
155
156 let res_big: VecZnxBig<&mut [u8], BE> = if a_base2k != ggsw_base2k {
157 let (mut a_conv, scratch_2) = scratch_1.take_glwe(&GLWELayout {
158 n: a.n(),
159 base2k: ggsw.base2k(),
160 k: a.k(),
161 rank: a.rank(),
162 });
163 self.glwe_normalize(&mut a_conv, a, scratch_2);
164 self.glwe_external_product_internal(res_dft, &a_conv, ggsw, scratch_2)
165 } else {
166 self.glwe_external_product_internal(res_dft, a, ggsw, scratch_1)
167 };
168
169 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
170 for j in 0..(res.rank() + 1).into() {
171 self.vec_znx_big_normalize(res.data_mut(), res_base2k, 0, j, &res_big, ggsw_base2k, j, scratch_1);
172 }
173 }
174}
175
176pub trait GLWEExternalProductInternal<BE: Backend> {
177 fn glwe_external_product_internal_tmp_bytes<R, A, B>(&self, res_infos: &R, a_infos: &A, b_infos: &B) -> usize
178 where
179 R: GLWEInfos,
180 A: GLWEInfos,
181 B: GGSWInfos;
182 fn glwe_external_product_internal<DR, A, G>(
183 &self,
184 res_dft: VecZnxDft<DR, BE>,
185 a: &A,
186 ggsw: &G,
187 scratch: &mut Scratch<BE>,
188 ) -> VecZnxBig<DR, BE>
189 where
190 DR: DataMut,
191 A: GLWEToRef,
192 G: GGSWPreparedToRef<BE>,
193 Scratch<BE>: ScratchTakeCore<BE>;
194}
195
196impl<BE: Backend> GLWEExternalProductInternal<BE> for Module<BE>
197where
198 Self: ModuleN
199 + VecZnxDftBytesOf
200 + VmpApplyDftToDftTmpBytes
201 + VecZnxNormalizeTmpBytes
202 + VecZnxDftApply<BE>
203 + VmpApplyDftToDft<BE>
204 + VmpApplyDftToDftAdd<BE>
205 + VecZnxIdftApplyConsume<BE>
206 + VecZnxBigNormalize<BE>
207 + VecZnxNormalize<BE>
208 + VecZnxDftBytesOf
209 + VmpApplyDftToDftTmpBytes
210 + VecZnxNormalizeTmpBytes,
211{
212 fn glwe_external_product_internal_tmp_bytes<R, A, B>(&self, res_infos: &R, a_infos: &A, b_infos: &B) -> usize
213 where
214 R: GLWEInfos,
215 A: GLWEInfos,
216 B: GGSWInfos,
217 {
218 let in_size: usize = a_infos.k().div_ceil(b_infos.base2k()).div_ceil(b_infos.dsize().into()) as usize;
219 let out_size: usize = res_infos.size();
220 let ggsw_size: usize = b_infos.size();
221 let a_dft: usize = self.bytes_of_vec_znx_dft((b_infos.rank() + 1).into(), in_size);
222 let vmp: usize = self.vmp_apply_dft_to_dft_tmp_bytes(
223 out_size,
224 in_size,
225 in_size, // rows
226 (b_infos.rank() + 1).into(), // cols in
227 (b_infos.rank() + 1).into(), // cols out
228 ggsw_size,
229 );
230 let normalize_big: usize = self.vec_znx_normalize_tmp_bytes();
231
232 a_dft + vmp.max(normalize_big)
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 assert_eq!(a.base2k(), ggsw.base2k());
252
253 let cols: usize = (ggsw.rank() + 1).into();
254 let dsize: usize = ggsw.dsize().into();
255 let a_size: usize = a.size();
256
257 let (mut a_dft, scratch_1) = scratch.take_vec_znx_dft(self, cols, a_size.div_ceil(dsize));
258 a_dft.data_mut().fill(0);
259
260 for di in 0..dsize {
261 // (lhs.size() + di) / dsize = (a - (digit - di - 1)).div_ceil(dsize)
262 a_dft.set_size((a.size() + di) / dsize);
263
264 // Small optimization for dsize > 2
265 // VMP produce some error e, and since we aggregate vmp * 2^{di * B}, then
266 // we also aggregate ei * 2^{di * B}, with the largest error being ei * 2^{(dsize-1) * B}.
267 // As such we can ignore the last dsize-2 limbs safely of the sum of vmp products.
268 // It is possible to further ignore the last dsize-1 limbs, but this introduce
269 // ~0.5 to 1 bit of additional noise, and thus not chosen here to ensure that the same
270 // noise is kept with respect to the ideal functionality.
271 res_dft.set_size(ggsw.size() - ((dsize - di) as isize - 2).max(0) as usize);
272
273 for j in 0..cols {
274 self.vec_znx_dft_apply(dsize, dsize - 1 - di, &mut a_dft, j, &a.data, j);
275 }
276
277 if di == 0 {
278 self.vmp_apply_dft_to_dft(&mut res_dft, &a_dft, &ggsw.data, scratch_1);
279 } else {
280 self.vmp_apply_dft_to_dft_add(&mut res_dft, &a_dft, &ggsw.data, di, scratch_1);
281 }
282 }
283
284 self.vec_znx_idft_apply_consume(res_dft)
285 }
286}