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