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