poulpy/poulpy-core/src/keyswitching/ggsw_ct.rs
poulpy-core/src/keyswitching/ggsw_ct.rsRaw
1use poulpy_hal::{
2 api::{
3 ScratchAvailable, TakeVecZnx, TakeVecZnxBig, TakeVecZnxDft, VecZnxBigAddSmallInplace, VecZnxBigBytesOf,
4 VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxCopy, VecZnxDftAddInplace, VecZnxDftApply, VecZnxDftBytesOf,
5 VecZnxDftCopy, VecZnxIdftApplyConsume, VecZnxIdftApplyTmpA, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft,
6 VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
7 },
8 layouts::{Backend, DataMut, DataRef, Module, Scratch, VecZnx, VmpPMat, ZnxInfos},
9};
10
11use crate::{
12 layouts::{
13 GGLWE, GGLWEInfos, GGSW, GGSWInfos, GLWE, GLWEInfos, LWEInfos,
14 prepared::{GLWESwitchingKeyPrepared, TensorKeyPrepared},
15 },
16 operations::GLWEOperations,
17};
18
19impl GGSW<Vec<u8>> {
20 pub(crate) fn expand_row_scratch_space<B: Backend, OUT, TSK>(module: &Module<B>, out_infos: &OUT, tsk_infos: &TSK) -> usize
21 where
22 OUT: GGSWInfos,
23 TSK: GGLWEInfos,
24 Module<B>: VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigBytesOf + VecZnxNormalizeTmpBytes,
25 {
26 let tsk_size: usize = tsk_infos.k().div_ceil(tsk_infos.base2k()) as usize;
27 let size_in: usize = out_infos
28 .k()
29 .div_ceil(tsk_infos.base2k())
30 .div_ceil(tsk_infos.dsize().into()) as usize;
31
32 let tmp_dft_i: usize = module.bytes_of_vec_znx_dft((tsk_infos.rank_out() + 1).into(), tsk_size);
33 let tmp_a: usize = module.bytes_of_vec_znx_dft(1, size_in);
34 let vmp: usize = module.vmp_apply_dft_to_dft_tmp_bytes(
35 tsk_size,
36 size_in,
37 size_in,
38 (tsk_infos.rank_in()).into(), // Verify if rank+1
39 (tsk_infos.rank_out()).into(), // Verify if rank+1
40 tsk_size,
41 );
42 let tmp_idft: usize = module.bytes_of_vec_znx_big(1, tsk_size);
43 let norm: usize = module.vec_znx_normalize_tmp_bytes();
44
45 tmp_dft_i + ((tmp_a + vmp) | (tmp_idft + norm))
46 }
47
48 #[allow(clippy::too_many_arguments)]
49 pub fn keyswitch_scratch_space<B: Backend, OUT, IN, KEY, TSK>(
50 module: &Module<B>,
51 out_infos: &OUT,
52 in_infos: &IN,
53 apply_infos: &KEY,
54 tsk_infos: &TSK,
55 ) -> usize
56 where
57 OUT: GGSWInfos,
58 IN: GGSWInfos,
59 KEY: GGLWEInfos,
60 TSK: GGLWEInfos,
61 Module<B>:
62 VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigBytesOf + VecZnxNormalizeTmpBytes + VecZnxBigNormalizeTmpBytes,
63 {
64 #[cfg(debug_assertions)]
65 {
66 assert_eq!(apply_infos.rank_in(), apply_infos.rank_out());
67 assert_eq!(tsk_infos.rank_in(), tsk_infos.rank_out());
68 assert_eq!(apply_infos.rank_in(), tsk_infos.rank_in());
69 }
70
71 let rank: usize = apply_infos.rank_out().into();
72
73 let size_out: usize = out_infos.k().div_ceil(out_infos.base2k()) as usize;
74 let res_znx: usize = VecZnx::bytes_of(module.n(), rank + 1, size_out);
75 let ci_dft: usize = module.bytes_of_vec_znx_dft(rank + 1, size_out);
76 let ks: usize = GLWE::keyswitch_scratch_space(module, out_infos, in_infos, apply_infos);
77 let expand_rows: usize = GGSW::expand_row_scratch_space(module, out_infos, tsk_infos);
78 let res_dft: usize = module.bytes_of_vec_znx_dft(rank + 1, size_out);
79
80 if in_infos.base2k() == tsk_infos.base2k() {
81 res_znx + ci_dft + (ks | expand_rows | res_dft)
82 } else {
83 let a_conv: usize = VecZnx::bytes_of(
84 module.n(),
85 1,
86 out_infos.k().div_ceil(tsk_infos.base2k()) as usize,
87 ) + module.vec_znx_normalize_tmp_bytes();
88 res_znx + ci_dft + (a_conv | ks | expand_rows | res_dft)
89 }
90 }
91
92 #[allow(clippy::too_many_arguments)]
93 pub fn keyswitch_inplace_scratch_space<B: Backend, OUT, KEY, TSK>(
94 module: &Module<B>,
95 out_infos: &OUT,
96 apply_infos: &KEY,
97 tsk_infos: &TSK,
98 ) -> usize
99 where
100 OUT: GGSWInfos,
101 KEY: GGLWEInfos,
102 TSK: GGLWEInfos,
103 Module<B>:
104 VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigBytesOf + VecZnxNormalizeTmpBytes + VecZnxBigNormalizeTmpBytes,
105 {
106 GGSW::keyswitch_scratch_space(module, out_infos, out_infos, apply_infos, tsk_infos)
107 }
108}
109
110impl<DataSelf: DataMut> GGSW<DataSelf> {
111 pub fn from_gglwe<DataA, DataTsk, B: Backend>(
112 &mut self,
113 module: &Module<B>,
114 a: &GGLWE<DataA>,
115 tsk: &TensorKeyPrepared<DataTsk, B>,
116 scratch: &mut Scratch<B>,
117 ) where
118 DataA: DataRef,
119 DataTsk: DataRef,
120 Module<B>: VecZnxCopy
121 + VecZnxDftBytesOf
122 + VmpApplyDftToDftTmpBytes
123 + VecZnxBigBytesOf
124 + VecZnxNormalizeTmpBytes
125 + VecZnxDftApply<B>
126 + VecZnxDftCopy<B>
127 + VmpApplyDftToDft<B>
128 + VmpApplyDftToDftAdd<B>
129 + VecZnxDftAddInplace<B>
130 + VecZnxBigNormalize<B>
131 + VecZnxIdftApplyTmpA<B>
132 + VecZnxNormalize<B>,
133 Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeVecZnx,
134 {
135 #[cfg(debug_assertions)]
136 {
137 use crate::layouts::{GLWEInfos, LWEInfos};
138
139 assert_eq!(self.rank(), a.rank_out());
140 assert_eq!(self.dnum(), a.dnum());
141 assert_eq!(self.n(), module.n() as u32);
142 assert_eq!(a.n(), module.n() as u32);
143 assert_eq!(tsk.n(), module.n() as u32);
144 }
145 (0..self.dnum().into()).for_each(|row_i| {
146 self.at_mut(row_i, 0).copy(module, &a.at(row_i, 0));
147 });
148 self.expand_row(module, tsk, scratch);
149 }
150
151 pub fn keyswitch<DataLhs: DataRef, DataKsk: DataRef, DataTsk: DataRef, B: Backend>(
152 &mut self,
153 module: &Module<B>,
154 lhs: &GGSW<DataLhs>,
155 ksk: &GLWESwitchingKeyPrepared<DataKsk, B>,
156 tsk: &TensorKeyPrepared<DataTsk, B>,
157 scratch: &mut Scratch<B>,
158 ) where
159 Module<B>: VecZnxDftBytesOf
160 + VmpApplyDftToDftTmpBytes
161 + VecZnxBigNormalizeTmpBytes
162 + VmpApplyDftToDft<B>
163 + VmpApplyDftToDftAdd<B>
164 + VecZnxDftApply<B>
165 + VecZnxIdftApplyConsume<B>
166 + VecZnxBigAddSmallInplace<B>
167 + VecZnxBigNormalize<B>
168 + VecZnxDftBytesOf
169 + VecZnxBigBytesOf
170 + VecZnxNormalizeTmpBytes
171 + VecZnxDftCopy<B>
172 + VecZnxDftAddInplace<B>
173 + VecZnxIdftApplyTmpA<B>
174 + VecZnxNormalize<B>,
175 Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeVecZnx,
176 {
177 (0..lhs.dnum().into()).for_each(|row_i| {
178 // Key-switch column 0, i.e.
179 // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0s0' + a1s1' + a2s2') + M[i], a0, a1, a2)
180 self.at_mut(row_i, 0)
181 .keyswitch(module, &lhs.at(row_i, 0), ksk, scratch);
182 });
183 self.expand_row(module, tsk, scratch);
184 }
185
186 pub fn keyswitch_inplace<DataKsk: DataRef, DataTsk: DataRef, B: Backend>(
187 &mut self,
188 module: &Module<B>,
189 ksk: &GLWESwitchingKeyPrepared<DataKsk, B>,
190 tsk: &TensorKeyPrepared<DataTsk, B>,
191 scratch: &mut Scratch<B>,
192 ) where
193 Module<B>: VecZnxDftBytesOf
194 + VmpApplyDftToDftTmpBytes
195 + VecZnxBigNormalizeTmpBytes
196 + VmpApplyDftToDft<B>
197 + VmpApplyDftToDftAdd<B>
198 + VecZnxDftApply<B>
199 + VecZnxIdftApplyConsume<B>
200 + VecZnxBigAddSmallInplace<B>
201 + VecZnxBigNormalize<B>
202 + VecZnxDftBytesOf
203 + VecZnxBigBytesOf
204 + VecZnxNormalizeTmpBytes
205 + VecZnxDftCopy<B>
206 + VecZnxDftAddInplace<B>
207 + VecZnxIdftApplyTmpA<B>
208 + VecZnxNormalize<B>,
209 Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeVecZnx,
210 {
211 (0..self.dnum().into()).for_each(|row_i| {
212 // Key-switch column 0, i.e.
213 // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0s0' + a1s1' + a2s2') + M[i], a0, a1, a2)
214 self.at_mut(row_i, 0)
215 .keyswitch_inplace(module, ksk, scratch);
216 });
217 self.expand_row(module, tsk, scratch);
218 }
219
220 pub fn expand_row<DataTsk: DataRef, B: Backend>(
221 &mut self,
222 module: &Module<B>,
223 tsk: &TensorKeyPrepared<DataTsk, B>,
224 scratch: &mut Scratch<B>,
225 ) where
226 Module<B>: VecZnxDftBytesOf
227 + VmpApplyDftToDftTmpBytes
228 + VecZnxBigBytesOf
229 + VecZnxNormalizeTmpBytes
230 + VecZnxDftApply<B>
231 + VecZnxDftCopy<B>
232 + VmpApplyDftToDft<B>
233 + VmpApplyDftToDftAdd<B>
234 + VecZnxDftAddInplace<B>
235 + VecZnxBigNormalize<B>
236 + VecZnxIdftApplyTmpA<B>
237 + VecZnxNormalize<B>,
238 Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeVecZnx,
239 {
240 let basek_in: usize = self.base2k().into();
241 let basek_tsk: usize = tsk.base2k().into();
242
243 assert!(scratch.available() >= GGSW::expand_row_scratch_space(module, self, tsk));
244
245 let n: usize = self.n().into();
246 let rank: usize = self.rank().into();
247 let cols: usize = rank + 1;
248
249 let a_size: usize = (self.size() * basek_in).div_ceil(basek_tsk);
250
251 // Keyswitch the j-th row of the col 0
252 for row_i in 0..self.dnum().into() {
253 let a = &self.at(row_i, 0).data;
254
255 // Pre-compute DFT of (a0, a1, a2)
256 let (mut ci_dft, scratch_1) = scratch.take_vec_znx_dft(n, cols, a_size);
257
258 if basek_in == basek_tsk {
259 for i in 0..cols {
260 module.vec_znx_dft_apply(1, 0, &mut ci_dft, i, a, i);
261 }
262 } else {
263 let (mut a_conv, scratch_2) = scratch_1.take_vec_znx(n, 1, a_size);
264 for i in 0..cols {
265 module.vec_znx_normalize(basek_tsk, &mut a_conv, 0, basek_in, a, i, scratch_2);
266 module.vec_znx_dft_apply(1, 0, &mut ci_dft, i, &a_conv, 0);
267 }
268 }
269
270 for col_j in 1..cols {
271 // Example for rank 3:
272 //
273 // Note: M is a vector (m, Bm, B^2m, B^3m, ...), so each column is
274 // actually composed of that many dnum and we focus on a specific row here
275 // implicitely given ci_dft.
276 //
277 // # Input
278 //
279 // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0 , a1 , a2 )
280 // col 1: (0, 0, 0, 0)
281 // col 2: (0, 0, 0, 0)
282 // col 3: (0, 0, 0, 0)
283 //
284 // # Output
285 //
286 // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0 , a1 , a2 )
287 // col 1: (-(b0s0 + b1s1 + b2s2) , b0 + M[i], b1 , b2 )
288 // col 2: (-(c0s0 + c1s1 + c2s2) , c0 , c1 + M[i], c2 )
289 // col 3: (-(d0s0 + d1s1 + d2s2) , d0 , d1 , d2 + M[i])
290
291 let dsize: usize = tsk.dsize().into();
292
293 let (mut tmp_dft_i, scratch_2) = scratch_1.take_vec_znx_dft(n, cols, tsk.size());
294 let (mut tmp_a, scratch_3) = scratch_2.take_vec_znx_dft(n, 1, ci_dft.size().div_ceil(dsize));
295
296 {
297 // Performs a key-switch for each combination of s[i]*s[j], i.e. for a0, a1, a2
298 //
299 // # Example for col=1
300 //
301 // a0 * (-(f0s0 + f1s1 + f1s2) + s0^2, f0, f1, f2) = (-(a0f0s0 + a0f1s1 + a0f1s2) + a0s0^2, a0f0, a0f1, a0f2)
302 // +
303 // a1 * (-(g0s0 + g1s1 + g1s2) + s0s1, g0, g1, g2) = (-(a1g0s0 + a1g1s1 + a1g1s2) + a1s0s1, a1g0, a1g1, a1g2)
304 // +
305 // a2 * (-(h0s0 + h1s1 + h1s2) + s0s2, h0, h1, h2) = (-(a2h0s0 + a2h1s1 + a2h1s2) + a2s0s2, a2h0, a2h1, a2h2)
306 // =
307 // (-(x0s0 + x1s1 + x2s2) + s0(a0s0 + a1s1 + a2s2), x0, x1, x2)
308 for col_i in 1..cols {
309 let pmat: &VmpPMat<DataTsk, B> = &tsk.at(col_i - 1, col_j - 1).key.data; // Selects Enc(s[i]s[j])
310
311 // Extracts a[i] and multipies with Enc(s[i]s[j])
312 for di in 0..dsize {
313 tmp_a.set_size((ci_dft.size() + di) / dsize);
314
315 // Small optimization for dsize > 2
316 // VMP produce some error e, and since we aggregate vmp * 2^{di * B}, then
317 // we also aggregate ei * 2^{di * B}, with the largest error being ei * 2^{(dsize-1) * B}.
318 // As such we can ignore the last dsize-2 limbs safely of the sum of vmp products.
319 // It is possible to further ignore the last dsize-1 limbs, but this introduce
320 // ~0.5 to 1 bit of additional noise, and thus not chosen here to ensure that the same
321 // noise is kept with respect to the ideal functionality.
322 tmp_dft_i.set_size(tsk.size() - ((dsize - di) as isize - 2).max(0) as usize);
323
324 module.vec_znx_dft_copy(dsize, dsize - 1 - di, &mut tmp_a, 0, &ci_dft, col_i);
325 if di == 0 && col_i == 1 {
326 module.vmp_apply_dft_to_dft(&mut tmp_dft_i, &tmp_a, pmat, scratch_3);
327 } else {
328 module.vmp_apply_dft_to_dft_add(&mut tmp_dft_i, &tmp_a, pmat, di, scratch_3);
329 }
330 }
331 }
332 }
333
334 // Adds -(sum a[i] * s[i]) + m) on the i-th column of tmp_idft_i
335 //
336 // (-(x0s0 + x1s1 + x2s2) + a0s0s0 + a1s0s1 + a2s0s2, x0, x1, x2)
337 // +
338 // (0, -(a0s0 + a1s1 + a2s2) + M[i], 0, 0)
339 // =
340 // (-(x0s0 + x1s1 + x2s2) + s0(a0s0 + a1s1 + a2s2), x0 -(a0s0 + a1s1 + a2s2) + M[i], x1, x2)
341 // =
342 // (-(x0s0 + x1s1 + x2s2), x0 + M[i], x1, x2)
343 module.vec_znx_dft_add_inplace(&mut tmp_dft_i, col_j, &ci_dft, 0);
344 let (mut tmp_idft, scratch_3) = scratch_2.take_vec_znx_big(n, 1, tsk.size());
345 for i in 0..cols {
346 module.vec_znx_idft_apply_tmpa(&mut tmp_idft, 0, &mut tmp_dft_i, i);
347 module.vec_znx_big_normalize(
348 basek_in,
349 &mut self.at_mut(row_i, col_j).data,
350 i,
351 basek_tsk,
352 &tmp_idft,
353 0,
354 scratch_3,
355 );
356 }
357 }
358 }
359 }
360}