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