poulpy/core/src/keyswitch_key.rs
core/src/keyswitch_key.rsRaw
1use base2k::{
2 Backend, FFT64, MatZnxDft, MatZnxDftOps, MatZnxDftToMut, MatZnxDftToRef, Module, ScalarZnx, ScalarZnxDft, ScalarZnxDftToRef,
3 ScalarZnxToRef, Scratch, VecZnxDftAlloc, VecZnxDftToMut, VecZnxDftToRef, ZnxZero,
4};
5use sampling::source::Source;
6
7use crate::{
8 elem::{GetRow, Infos, SetRow},
9 gglwe_ciphertext::GGLWECiphertext,
10 ggsw_ciphertext::GGSWCiphertext,
11 glwe_ciphertext_fourier::GLWECiphertextFourier,
12 keys::{SecretKey, SecretKeyFourier},
13};
14
15pub struct GLWESwitchingKey<Data, B: Backend>(pub(crate) GGLWECiphertext<Data, B>);
16
17impl GLWESwitchingKey<Vec<u8>, FFT64> {
18 pub fn new(module: &Module<FFT64>, basek: usize, k: usize, rows: usize, rank_in: usize, rank_out: usize) -> Self {
19 GLWESwitchingKey(GGLWECiphertext::new(
20 module, basek, k, rows, rank_in, rank_out,
21 ))
22 }
23}
24
25impl<T, B: Backend> Infos for GLWESwitchingKey<T, B> {
26 type Inner = MatZnxDft<T, B>;
27
28 fn inner(&self) -> &Self::Inner {
29 self.0.inner()
30 }
31
32 fn basek(&self) -> usize {
33 self.0.basek()
34 }
35
36 fn k(&self) -> usize {
37 self.0.k()
38 }
39}
40
41impl<T, B: Backend> GLWESwitchingKey<T, B> {
42 pub fn rank(&self) -> usize {
43 self.0.data.cols_out() - 1
44 }
45
46 pub fn rank_in(&self) -> usize {
47 self.0.data.cols_in()
48 }
49
50 pub fn rank_out(&self) -> usize {
51 self.0.data.cols_out() - 1
52 }
53}
54
55impl<DataSelf, B: Backend> MatZnxDftToMut<B> for GLWESwitchingKey<DataSelf, B>
56where
57 MatZnxDft<DataSelf, B>: MatZnxDftToMut<B>,
58{
59 fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
60 self.0.data.to_mut()
61 }
62}
63
64impl<DataSelf, B: Backend> MatZnxDftToRef<B> for GLWESwitchingKey<DataSelf, B>
65where
66 MatZnxDft<DataSelf, B>: MatZnxDftToRef<B>,
67{
68 fn to_ref(&self) -> MatZnxDft<&[u8], B> {
69 self.0.data.to_ref()
70 }
71}
72
73impl<C> GetRow<FFT64> for GLWESwitchingKey<C, FFT64>
74where
75 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
76{
77 fn get_row<R>(&self, module: &Module<FFT64>, row_i: usize, col_j: usize, res: &mut R)
78 where
79 R: VecZnxDftToMut<FFT64>,
80 {
81 module.vmp_extract_row(res, self, row_i, col_j);
82 }
83}
84
85impl<C> SetRow<FFT64> for GLWESwitchingKey<C, FFT64>
86where
87 MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
88{
89 fn set_row<R>(&mut self, module: &Module<FFT64>, row_i: usize, col_j: usize, a: &R)
90 where
91 R: VecZnxDftToRef<FFT64>,
92 {
93 module.vmp_prepare_row(self, row_i, col_j, a);
94 }
95}
96
97impl GLWESwitchingKey<Vec<u8>, FFT64> {
98 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, rank: usize, size: usize) -> usize {
99 GGLWECiphertext::encrypt_sk_scratch_space(module, rank, size)
100 }
101
102 pub fn encrypt_pk_scratch_space(module: &Module<FFT64>, rank: usize, pk_size: usize) -> usize {
103 GGLWECiphertext::encrypt_pk_scratch_space(module, rank, pk_size)
104 }
105
106 pub fn keyswitch_scratch_space(
107 module: &Module<FFT64>,
108 out_size: usize,
109 out_rank: usize,
110 in_size: usize,
111 in_rank: usize,
112 ksk_size: usize,
113 ) -> usize {
114 let tmp_in: usize = module.bytes_of_vec_znx_dft(in_rank + 1, in_size);
115 let tmp_out: usize = module.bytes_of_vec_znx_dft(out_rank + 1, out_size);
116 let ksk: usize = GLWECiphertextFourier::keyswitch_scratch_space(module, out_size, out_rank, in_size, in_rank, ksk_size);
117 tmp_in + tmp_out + ksk
118 }
119
120 pub fn keyswitch_inplace_scratch_space(module: &Module<FFT64>, out_size: usize, out_rank: usize, ksk_size: usize) -> usize {
121 let tmp: usize = module.bytes_of_vec_znx_dft(out_rank + 1, out_size);
122 let ksk: usize = GLWECiphertextFourier::keyswitch_inplace_scratch_space(module, out_size, out_rank, ksk_size);
123 tmp + ksk
124 }
125
126 pub fn external_product_scratch_space(
127 module: &Module<FFT64>,
128 out_size: usize,
129 in_size: usize,
130 ggsw_size: usize,
131 rank: usize,
132 ) -> usize {
133 let tmp_in: usize = module.bytes_of_vec_znx_dft(rank + 1, in_size);
134 let tmp_out: usize = module.bytes_of_vec_znx_dft(rank + 1, out_size);
135 let ggsw: usize = GLWECiphertextFourier::external_product_scratch_space(module, out_size, in_size, ggsw_size, rank);
136 tmp_in + tmp_out + ggsw
137 }
138
139 pub fn external_product_inplace_scratch_space(
140 module: &Module<FFT64>,
141 out_size: usize,
142 ggsw_size: usize,
143 rank: usize,
144 ) -> usize {
145 let tmp: usize = module.bytes_of_vec_znx_dft(rank + 1, out_size);
146 let ggsw: usize = GLWECiphertextFourier::external_product_inplace_scratch_space(module, out_size, ggsw_size, rank);
147 tmp + ggsw
148 }
149}
150impl<DataSelf> GLWESwitchingKey<DataSelf, FFT64>
151where
152 MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64>,
153{
154 pub fn encrypt_sk<DataSkIn, DataSkOut>(
155 &mut self,
156 module: &Module<FFT64>,
157 sk_in: &SecretKey<DataSkIn>,
158 sk_out_dft: &SecretKeyFourier<DataSkOut, FFT64>,
159 source_xa: &mut Source,
160 source_xe: &mut Source,
161 sigma: f64,
162 scratch: &mut Scratch,
163 ) where
164 ScalarZnx<DataSkIn>: ScalarZnxToRef,
165 ScalarZnxDft<DataSkOut, FFT64>: ScalarZnxDftToRef<FFT64>,
166 {
167 self.0.encrypt_sk(
168 module,
169 &sk_in.data,
170 sk_out_dft,
171 source_xa,
172 source_xe,
173 sigma,
174 scratch,
175 );
176 }
177
178 pub fn keyswitch<DataLhs, DataRhs>(
179 &mut self,
180 module: &Module<FFT64>,
181 lhs: &GLWESwitchingKey<DataLhs, FFT64>,
182 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
183 scratch: &mut base2k::Scratch,
184 ) where
185 MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
186 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
187 {
188 #[cfg(debug_assertions)]
189 {
190 assert_eq!(
191 self.rank_in(),
192 lhs.rank_in(),
193 "ksk_out input rank: {} != ksk_in input rank: {}",
194 self.rank_in(),
195 lhs.rank_in()
196 );
197 assert_eq!(
198 lhs.rank_out(),
199 rhs.rank_in(),
200 "ksk_in output rank: {} != ksk_apply input rank: {}",
201 self.rank_out(),
202 rhs.rank_in()
203 );
204 assert_eq!(
205 self.rank_out(),
206 rhs.rank_out(),
207 "ksk_out output rank: {} != ksk_apply output rank: {}",
208 self.rank_out(),
209 rhs.rank_out()
210 );
211 }
212
213 let (tmp_in_data, scratch1) = scratch.tmp_vec_znx_dft(module, lhs.rank_out() + 1, lhs.size());
214
215 let mut tmp_in: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
216 data: tmp_in_data,
217 basek: lhs.basek(),
218 k: lhs.k(),
219 };
220
221 let (tmp_out_data, scratch2) = scratch1.tmp_vec_znx_dft(module, self.rank_out() + 1, self.size());
222
223 let mut tmp_out: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
224 data: tmp_out_data,
225 basek: self.basek(),
226 k: self.k(),
227 };
228
229 (0..self.rank_in()).for_each(|col_i| {
230 (0..self.rows()).for_each(|row_j| {
231 lhs.get_row(module, row_j, col_i, &mut tmp_in);
232 tmp_out.keyswitch(module, &tmp_in, rhs, scratch2);
233 self.set_row(module, row_j, col_i, &tmp_out);
234 });
235 });
236
237 tmp_out.data.zero();
238
239 (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
240 (0..self.rank_in()).for_each(|col_j| {
241 self.set_row(module, row_i, col_j, &tmp_out);
242 });
243 });
244 }
245
246 pub fn keyswitch_inplace<DataRhs>(
247 &mut self,
248 module: &Module<FFT64>,
249 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
250 scratch: &mut base2k::Scratch,
251 ) where
252 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
253 {
254 #[cfg(debug_assertions)]
255 {
256 assert_eq!(
257 self.rank_out(),
258 rhs.rank_out(),
259 "ksk_out output rank: {} != ksk_apply output rank: {}",
260 self.rank_out(),
261 rhs.rank_out()
262 );
263 }
264
265 let (tmp_data, scratch1) = scratch.tmp_vec_znx_dft(module, self.rank_out() + 1, self.size());
266
267 let mut tmp: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
268 data: tmp_data,
269 basek: self.basek(),
270 k: self.k(),
271 };
272
273 (0..self.rank_in()).for_each(|col_i| {
274 (0..self.rows()).for_each(|row_j| {
275 self.get_row(module, row_j, col_i, &mut tmp);
276 tmp.keyswitch_inplace(module, rhs, scratch1);
277 self.set_row(module, row_j, col_i, &tmp);
278 });
279 });
280 }
281
282 pub fn external_product<DataLhs, DataRhs>(
283 &mut self,
284 module: &Module<FFT64>,
285 lhs: &GLWESwitchingKey<DataLhs, FFT64>,
286 rhs: &GGSWCiphertext<DataRhs, FFT64>,
287 scratch: &mut Scratch,
288 ) where
289 MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
290 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
291 {
292 #[cfg(debug_assertions)]
293 {
294 assert_eq!(
295 self.rank_in(),
296 lhs.rank_in(),
297 "ksk_out input rank: {} != ksk_in input rank: {}",
298 self.rank_in(),
299 lhs.rank_in()
300 );
301 assert_eq!(
302 lhs.rank_out(),
303 rhs.rank(),
304 "ksk_in output rank: {} != ggsw rank: {}",
305 self.rank_out(),
306 rhs.rank()
307 );
308 assert_eq!(
309 self.rank_out(),
310 rhs.rank(),
311 "ksk_out output rank: {} != ggsw rank: {}",
312 self.rank_out(),
313 rhs.rank()
314 );
315 }
316
317 let (tmp_in_data, scratch1) = scratch.tmp_vec_znx_dft(module, lhs.rank_out() + 1, lhs.size());
318
319 let mut tmp_in: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
320 data: tmp_in_data,
321 basek: lhs.basek(),
322 k: lhs.k(),
323 };
324
325 let (tmp_out_data, scratch2) = scratch1.tmp_vec_znx_dft(module, self.rank_out() + 1, self.size());
326
327 let mut tmp_out: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
328 data: tmp_out_data,
329 basek: self.basek(),
330 k: self.k(),
331 };
332
333 (0..self.rank_in()).for_each(|col_i| {
334 (0..self.rows()).for_each(|row_j| {
335 lhs.get_row(module, row_j, col_i, &mut tmp_in);
336 tmp_out.external_product(module, &tmp_in, rhs, scratch2);
337 self.set_row(module, row_j, col_i, &tmp_out);
338 });
339 });
340
341 tmp_out.data.zero();
342
343 (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
344 (0..self.rank_in()).for_each(|col_j| {
345 self.set_row(module, row_i, col_j, &tmp_out);
346 });
347 });
348 }
349
350 pub fn external_product_inplace<DataRhs>(
351 &mut self,
352 module: &Module<FFT64>,
353 rhs: &GGSWCiphertext<DataRhs, FFT64>,
354 scratch: &mut Scratch,
355 ) where
356 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
357 {
358 #[cfg(debug_assertions)]
359 {
360 assert_eq!(
361 self.rank_out(),
362 rhs.rank(),
363 "ksk_out output rank: {} != ggsw rank: {}",
364 self.rank_out(),
365 rhs.rank()
366 );
367 }
368
369 let (tmp_data, scratch1) = scratch.tmp_vec_znx_dft(module, self.rank_out() + 1, self.size());
370
371 let mut tmp: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
372 data: tmp_data,
373 basek: self.basek(),
374 k: self.k(),
375 };
376
377 (0..self.rank_in()).for_each(|col_i| {
378 (0..self.rows()).for_each(|row_j| {
379 self.get_row(module, row_j, col_i, &mut tmp);
380 tmp.external_product_inplace(module, rhs, scratch1);
381 self.set_row(module, row_j, col_i, &tmp);
382 });
383 });
384 }
385}