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