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, VecZnxDft, 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>, base2k: usize, k: usize, rows: usize, rank_in: usize, rank_out: usize) -> Self {
19 GLWESwitchingKey(GGLWECiphertext::new(
20 module, base2k, 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 GLWECiphertextFourier<R, FFT64>)
78 where
79 VecZnxDft<R, FFT64>: VecZnxDftToMut<FFT64>,
80 {
81 #[cfg(debug_assertions)]
82 {
83 assert_eq!(col_j, 0);
84 }
85 module.vmp_extract_row(res, self, row_i, col_j);
86 }
87}
88
89impl<C> SetRow<FFT64> for GLWESwitchingKey<C, FFT64>
90where
91 MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
92{
93 fn set_row<R>(&mut self, module: &Module<FFT64>, row_i: usize, col_j: usize, a: &GLWECiphertextFourier<R, FFT64>)
94 where
95 VecZnxDft<R, FFT64>: VecZnxDftToRef<FFT64>,
96 {
97 #[cfg(debug_assertions)]
98 {
99 assert_eq!(col_j, 0);
100 }
101 module.vmp_prepare_row(self, row_i, col_j, a);
102 }
103}
104
105impl GLWESwitchingKey<Vec<u8>, FFT64> {
106 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, rank: usize, size: usize) -> usize {
107 GGLWECiphertext::encrypt_sk_scratch_space(module, rank, size)
108 }
109
110 pub fn encrypt_pk_scratch_space(module: &Module<FFT64>, rank: usize, pk_size: usize) -> usize {
111 GGLWECiphertext::encrypt_pk_scratch_space(module, rank, pk_size)
112 }
113}
114
115impl<DataSelf> GLWESwitchingKey<DataSelf, FFT64>
116where
117 MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64>,
118{
119 pub fn encrypt_sk<DataSkIn, DataSkOut>(
120 &mut self,
121 module: &Module<FFT64>,
122 sk_in: &SecretKey<DataSkIn>,
123 sk_out_dft: &SecretKeyFourier<DataSkOut, FFT64>,
124 source_xa: &mut Source,
125 source_xe: &mut Source,
126 sigma: f64,
127 scratch: &mut Scratch,
128 ) where
129 ScalarZnx<DataSkIn>: ScalarZnxToRef,
130 ScalarZnxDft<DataSkOut, FFT64>: ScalarZnxDftToRef<FFT64>,
131 {
132 self.0.encrypt_sk(
133 module,
134 &sk_in.data,
135 sk_out_dft,
136 source_xa,
137 source_xe,
138 sigma,
139 scratch,
140 );
141 }
142}
143
144impl GLWESwitchingKey<Vec<u8>, FFT64> {
145 pub fn keyswitch_scratch_space(
146 module: &Module<FFT64>,
147 out_size: usize,
148 out_rank: usize,
149 in_size: usize,
150 in_rank: usize,
151 ksk_size: usize,
152 ) -> usize {
153 let tmp_in: usize = module.bytes_of_vec_znx_dft(in_rank + 1, in_size);
154 let tmp_out: usize = module.bytes_of_vec_znx_dft(out_rank + 1, out_size);
155 let ksk: usize = GLWECiphertextFourier::keyswitch_scratch_space(module, out_size, out_rank, in_size, in_rank, ksk_size);
156 tmp_in + tmp_out + ksk
157 }
158
159 pub fn keyswitch_inplace_scratch_space(module: &Module<FFT64>, out_size: usize, out_rank: usize, ksk_size: usize) -> usize {
160 let tmp: usize = module.bytes_of_vec_znx_dft(out_rank + 1, out_size);
161 let ksk: usize = GLWECiphertextFourier::keyswitch_inplace_scratch_space(module, out_size, out_rank, ksk_size);
162 tmp + ksk
163 }
164
165 pub fn external_product_scratch_space(
166 module: &Module<FFT64>,
167 out_size: usize,
168 in_size: usize,
169 ggsw_size: usize,
170 rank: usize,
171 ) -> usize {
172 let tmp_in: usize = module.bytes_of_vec_znx_dft(rank + 1, in_size);
173 let tmp_out: usize = module.bytes_of_vec_znx_dft(rank + 1, out_size);
174 let ggsw: usize = GLWECiphertextFourier::external_product_scratch_space(module, out_size, in_size, ggsw_size, rank);
175 tmp_in + tmp_out + ggsw
176 }
177
178 pub fn external_product_inplace_scratch_space(
179 module: &Module<FFT64>,
180 out_size: usize,
181 ggsw_size: usize,
182 rank: usize,
183 ) -> usize {
184 let tmp: usize = module.bytes_of_vec_znx_dft(rank + 1, out_size);
185 let ggsw: usize = GLWECiphertextFourier::external_product_inplace_scratch_space(module, out_size, ggsw_size, rank);
186 tmp + ggsw
187 }
188}
189
190impl<DataSelf> GLWESwitchingKey<DataSelf, FFT64>
191where
192 MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64>,
193{
194 pub fn keyswitch<DataLhs, DataRhs>(
195 &mut self,
196 module: &Module<FFT64>,
197 lhs: &GLWESwitchingKey<DataLhs, FFT64>,
198 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
199 scratch: &mut base2k::Scratch,
200 ) where
201 MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
202 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
203 {
204 #[cfg(debug_assertions)]
205 {
206 assert_eq!(
207 self.rank_in(),
208 lhs.rank_in(),
209 "ksk_out input rank: {} != ksk_in input rank: {}",
210 self.rank_in(),
211 lhs.rank_in()
212 );
213 assert_eq!(
214 lhs.rank_out(),
215 rhs.rank_in(),
216 "ksk_in output rank: {} != ksk_apply input rank: {}",
217 self.rank_out(),
218 rhs.rank_in()
219 );
220 assert_eq!(
221 self.rank_out(),
222 rhs.rank_out(),
223 "ksk_out output rank: {} != ksk_apply output rank: {}",
224 self.rank_out(),
225 rhs.rank_out()
226 );
227 }
228
229 let (tmp_in_data, scratch1) = scratch.tmp_vec_znx_dft(module, lhs.rank_out() + 1, lhs.size());
230
231 let mut tmp_in: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
232 data: tmp_in_data,
233 basek: lhs.basek(),
234 k: lhs.k(),
235 };
236
237 let (tmp_out_data, scratch2) = scratch1.tmp_vec_znx_dft(module, self.rank_out() + 1, self.size());
238
239 let mut tmp_out: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
240 data: tmp_out_data,
241 basek: self.basek(),
242 k: self.k(),
243 };
244
245 (0..self.rank_in()).for_each(|col_i| {
246 (0..self.rows()).for_each(|row_j| {
247 lhs.get_row(module, row_j, col_i, &mut tmp_in);
248 tmp_out.keyswitch(module, &tmp_in, rhs, scratch2);
249 self.set_row(module, row_j, col_i, &tmp_out);
250 });
251 });
252
253 tmp_out.data.zero();
254
255 (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
256 (0..self.rank_in()).for_each(|col_j| {
257 self.set_row(module, row_i, col_j, &tmp_out);
258 });
259 });
260 }
261
262 pub fn keyswitch_inplace<DataRhs>(
263 &mut self,
264 module: &Module<FFT64>,
265 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
266 scratch: &mut base2k::Scratch,
267 ) where
268 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
269 {
270 #[cfg(debug_assertions)]
271 {
272 assert_eq!(
273 self.rank_out(),
274 rhs.rank_out(),
275 "ksk_out output rank: {} != ksk_apply output rank: {}",
276 self.rank_out(),
277 rhs.rank_out()
278 );
279 }
280
281 let (tmp_data, scratch1) = scratch.tmp_vec_znx_dft(module, self.rank_out() + 1, self.size());
282
283 let mut tmp: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
284 data: tmp_data,
285 basek: self.basek(),
286 k: self.k(),
287 };
288
289 (0..self.rank_in()).for_each(|col_i| {
290 (0..self.rows()).for_each(|row_j| {
291 self.get_row(module, row_j, col_i, &mut tmp);
292 tmp.keyswitch_inplace(module, rhs, scratch1);
293 self.set_row(module, row_j, col_i, &tmp);
294 });
295 });
296 }
297
298 pub fn external_product<DataLhs, DataRhs>(
299 &mut self,
300 module: &Module<FFT64>,
301 lhs: &GLWESwitchingKey<DataLhs, FFT64>,
302 rhs: &GGSWCiphertext<DataRhs, FFT64>,
303 scratch: &mut Scratch,
304 ) where
305 MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
306 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
307 {
308 #[cfg(debug_assertions)]
309 {
310 assert_eq!(
311 self.rank_in(),
312 lhs.rank_in(),
313 "ksk_out input rank: {} != ksk_in input rank: {}",
314 self.rank_in(),
315 lhs.rank_in()
316 );
317 assert_eq!(
318 lhs.rank_out(),
319 rhs.rank(),
320 "ksk_in output rank: {} != ggsw rank: {}",
321 self.rank_out(),
322 rhs.rank()
323 );
324 assert_eq!(
325 self.rank_out(),
326 rhs.rank(),
327 "ksk_out output rank: {} != ggsw rank: {}",
328 self.rank_out(),
329 rhs.rank()
330 );
331 }
332
333 let (tmp_in_data, scratch1) = scratch.tmp_vec_znx_dft(module, lhs.rank_out() + 1, lhs.size());
334
335 let mut tmp_in: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
336 data: tmp_in_data,
337 basek: lhs.basek(),
338 k: lhs.k(),
339 };
340
341 let (tmp_out_data, scratch2) = scratch1.tmp_vec_znx_dft(module, self.rank_out() + 1, self.size());
342
343 let mut tmp_out: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
344 data: tmp_out_data,
345 basek: self.basek(),
346 k: self.k(),
347 };
348
349 (0..self.rank_in()).for_each(|col_i| {
350 (0..self.rows()).for_each(|row_j| {
351 lhs.get_row(module, row_j, col_i, &mut tmp_in);
352 tmp_out.external_product(module, &tmp_in, rhs, scratch2);
353 self.set_row(module, row_j, col_i, &tmp_out);
354 });
355 });
356
357 tmp_out.data.zero();
358
359 (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
360 (0..self.rank_in()).for_each(|col_j| {
361 self.set_row(module, row_i, col_j, &tmp_out);
362 });
363 });
364 }
365
366 pub fn external_product_inplace<DataRhs>(
367 &mut self,
368 module: &Module<FFT64>,
369 rhs: &GGSWCiphertext<DataRhs, FFT64>,
370 scratch: &mut Scratch,
371 ) where
372 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
373 {
374 #[cfg(debug_assertions)]
375 {
376 assert_eq!(
377 self.rank_out(),
378 rhs.rank(),
379 "ksk_out output rank: {} != ggsw rank: {}",
380 self.rank_out(),
381 rhs.rank()
382 );
383 }
384
385 let (tmp_data, scratch1) = scratch.tmp_vec_znx_dft(module, self.rank_out() + 1, self.size());
386
387 let mut tmp: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
388 data: tmp_data,
389 basek: self.basek(),
390 k: self.k(),
391 };
392
393 (0..self.rank_in()).for_each(|col_i| {
394 (0..self.rows()).for_each(|row_j| {
395 self.get_row(module, row_j, col_i, &mut tmp);
396 tmp.external_product_inplace(module, rhs, scratch1);
397 self.set_row(module, row_j, col_i, &tmp);
398 });
399 });
400 }
401}