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