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