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