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