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