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