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 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>, rank: usize, size: usize) -> usize {
83 GGLWECiphertext::generate_from_sk_scratch_space(module, rank, size)
84 }
85
86 pub fn encrypt_pk_scratch_space(module: &Module<FFT64>, rank: usize, pk_size: usize) -> usize {
87 GGLWECiphertext::generate_from_pk_scratch_space(module, rank, pk_size)
88 }
89
90 pub fn keyswitch_scratch_space(
91 module: &Module<FFT64>,
92 out_size: usize,
93 out_rank: usize,
94 in_size: usize,
95 in_rank: usize,
96 ksk_size: usize,
97 ) -> usize {
98 let tmp_in: usize = module.bytes_of_vec_znx_dft(in_rank + 1, in_size);
99 let tmp_out: usize = module.bytes_of_vec_znx_dft(out_rank + 1, out_size);
100 let ksk: usize = GLWECiphertextFourier::keyswitch_scratch_space(module, out_size, out_rank, in_size, in_rank, ksk_size);
101 tmp_in + tmp_out + ksk
102 }
103
104 pub fn keyswitch_inplace_scratch_space(module: &Module<FFT64>, out_size: usize, out_rank: usize, ksk_size: usize) -> usize {
105 let tmp: usize = module.bytes_of_vec_znx_dft(out_rank + 1, out_size);
106 let ksk: usize = GLWECiphertextFourier::keyswitch_inplace_scratch_space(module, out_size, out_rank, ksk_size);
107 tmp + ksk
108 }
109
110 pub fn external_product_scratch_space(
111 module: &Module<FFT64>,
112 out_size: usize,
113 in_size: usize,
114 ggsw_size: usize,
115 rank: usize,
116 ) -> usize {
117 let tmp_in: usize = module.bytes_of_vec_znx_dft(rank + 1, in_size);
118 let tmp_out: usize = module.bytes_of_vec_znx_dft(rank + 1, out_size);
119 let ggsw: usize = GLWECiphertextFourier::external_product_scratch_space(module, out_size, in_size, ggsw_size, rank);
120 tmp_in + tmp_out + ggsw
121 }
122
123 pub fn external_product_inplace_scratch_space(
124 module: &Module<FFT64>,
125 out_size: usize,
126 ggsw_size: usize,
127 rank: usize,
128 ) -> usize {
129 let tmp: usize = module.bytes_of_vec_znx_dft(rank + 1, out_size);
130 let ggsw: usize = GLWECiphertextFourier::external_product_inplace_scratch_space(module, out_size, ggsw_size, rank);
131 tmp + ggsw
132 }
133}
134impl<DataSelf: AsMut<[u8]> + AsRef<[u8]>> GLWESwitchingKey<DataSelf, FFT64> {
135 pub fn generate_from_sk<DataSkIn: AsRef<[u8]>, DataSkOut: AsRef<[u8]>>(
136 &mut self,
137 module: &Module<FFT64>,
138 sk_in: &SecretKey<DataSkIn>,
139 sk_out_dft: &SecretKeyFourier<DataSkOut, FFT64>,
140 source_xa: &mut Source,
141 source_xe: &mut Source,
142 sigma: f64,
143 scratch: &mut Scratch,
144 ) {
145 self.0.encrypt_sk(
146 module,
147 &sk_in.data,
148 sk_out_dft,
149 source_xa,
150 source_xe,
151 sigma,
152 scratch,
153 );
154 }
155
156 pub fn keyswitch<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
157 &mut self,
158 module: &Module<FFT64>,
159 lhs: &GLWESwitchingKey<DataLhs, FFT64>,
160 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
161 scratch: &mut Scratch,
162 ) {
163 #[cfg(debug_assertions)]
164 {
165 assert_eq!(
166 self.rank_in(),
167 lhs.rank_in(),
168 "ksk_out input rank: {} != ksk_in input rank: {}",
169 self.rank_in(),
170 lhs.rank_in()
171 );
172 assert_eq!(
173 lhs.rank_out(),
174 rhs.rank_in(),
175 "ksk_in output rank: {} != ksk_apply input rank: {}",
176 self.rank_out(),
177 rhs.rank_in()
178 );
179 assert_eq!(
180 self.rank_out(),
181 rhs.rank_out(),
182 "ksk_out output rank: {} != ksk_apply output rank: {}",
183 self.rank_out(),
184 rhs.rank_out()
185 );
186 }
187
188 let (mut tmp_in, scratch1) = scratch.tmp_glwe_fourier(module, lhs.basek(), lhs.k(), lhs.rank());
189 let (mut tmp_out, scratch2) = scratch1.tmp_glwe_fourier(module, self.basek(), self.k(), self.rank());
190
191 (0..self.rank_in()).for_each(|col_i| {
192 (0..self.rows()).for_each(|row_j| {
193 lhs.get_row(module, row_j, col_i, &mut tmp_in);
194 tmp_out.keyswitch(module, &tmp_in, rhs, scratch2);
195 self.set_row(module, row_j, col_i, &tmp_out);
196 });
197 });
198
199 tmp_out.data.zero();
200
201 (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
202 (0..self.rank_in()).for_each(|col_j| {
203 self.set_row(module, row_i, col_j, &tmp_out);
204 });
205 });
206 }
207
208 pub fn keyswitch_inplace<DataRhs: AsRef<[u8]>>(
209 &mut self,
210 module: &Module<FFT64>,
211 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
212 scratch: &mut Scratch,
213 ) {
214 #[cfg(debug_assertions)]
215 {
216 assert_eq!(
217 self.rank_out(),
218 rhs.rank_out(),
219 "ksk_out output rank: {} != ksk_apply output rank: {}",
220 self.rank_out(),
221 rhs.rank_out()
222 );
223 }
224
225 let (mut tmp, scratch1) = scratch.tmp_glwe_fourier(module, self.basek(), self.k(), self.rank());
226
227 (0..self.rank_in()).for_each(|col_i| {
228 (0..self.rows()).for_each(|row_j| {
229 self.get_row(module, row_j, col_i, &mut tmp);
230 tmp.keyswitch_inplace(module, rhs, scratch1);
231 self.set_row(module, row_j, col_i, &tmp);
232 });
233 });
234 }
235
236 pub fn external_product<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
237 &mut self,
238 module: &Module<FFT64>,
239 lhs: &GLWESwitchingKey<DataLhs, FFT64>,
240 rhs: &GGSWCiphertext<DataRhs, FFT64>,
241 scratch: &mut Scratch,
242 ) {
243 #[cfg(debug_assertions)]
244 {
245 assert_eq!(
246 self.rank_in(),
247 lhs.rank_in(),
248 "ksk_out input rank: {} != ksk_in input rank: {}",
249 self.rank_in(),
250 lhs.rank_in()
251 );
252 assert_eq!(
253 lhs.rank_out(),
254 rhs.rank(),
255 "ksk_in output rank: {} != ggsw rank: {}",
256 self.rank_out(),
257 rhs.rank()
258 );
259 assert_eq!(
260 self.rank_out(),
261 rhs.rank(),
262 "ksk_out output rank: {} != ggsw rank: {}",
263 self.rank_out(),
264 rhs.rank()
265 );
266 }
267
268 let (mut tmp_in, scratch1) = scratch.tmp_glwe_fourier(module, lhs.basek(), lhs.k(), lhs.rank());
269 let (mut tmp_out, scratch2) = scratch1.tmp_glwe_fourier(module, self.basek(), self.k(), self.rank());
270
271 (0..self.rank_in()).for_each(|col_i| {
272 (0..self.rows()).for_each(|row_j| {
273 lhs.get_row(module, row_j, col_i, &mut tmp_in);
274 tmp_out.external_product(module, &tmp_in, rhs, scratch2);
275 self.set_row(module, row_j, col_i, &tmp_out);
276 });
277 });
278
279 tmp_out.data.zero();
280
281 (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
282 (0..self.rank_in()).for_each(|col_j| {
283 self.set_row(module, row_i, col_j, &tmp_out);
284 });
285 });
286 }
287
288 pub fn external_product_inplace<DataRhs: AsRef<[u8]>>(
289 &mut self,
290 module: &Module<FFT64>,
291 rhs: &GGSWCiphertext<DataRhs, FFT64>,
292 scratch: &mut Scratch,
293 ) {
294 #[cfg(debug_assertions)]
295 {
296 assert_eq!(
297 self.rank_out(),
298 rhs.rank(),
299 "ksk_out output rank: {} != ggsw rank: {}",
300 self.rank_out(),
301 rhs.rank()
302 );
303 }
304
305 let (mut tmp, scratch1) = scratch.tmp_glwe_fourier(module, self.basek(), self.k(), self.rank());
306
307 (0..self.rank_in()).for_each(|col_i| {
308 (0..self.rows()).for_each(|row_j| {
309 self.get_row(module, row_j, col_i, &mut tmp);
310 tmp.external_product_inplace(module, rhs, scratch1);
311 self.set_row(module, row_j, col_i, &tmp);
312 });
313 });
314 }
315}