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