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