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