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