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