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