poulpy/core/src/lwe/keyswtich.rs
core/src/lwe/keyswtich.rsRaw
1use backend::hal::{
2 api::{
3 MatZnxAlloc, ScalarZnxAllocBytes, ScratchAvailable, TakeScalarZnx, TakeVecZnx, TakeVecZnxDft, VecZnxAddScalarInplace,
4 VecZnxAllocBytes, VecZnxAutomorphismInplace, VecZnxSwithcDegree, ZnxView, ZnxViewMut, ZnxZero,
5 },
6 layouts::{Backend, Data, DataMut, DataRef, Module, ReaderFrom, Scratch, WriterTo},
7};
8use sampling::source::Source;
9
10use crate::{
11 GGLWEEncryptSkFamily, GGLWEExecLayoutFamily, GLWECiphertext, GLWEKeyswitchFamily, GLWESecret, GLWESecretExec,
12 GLWESwitchingKey, GLWESwitchingKeyExec, Infos, LWECiphertext, LWESecret, TakeGLWECt, TakeGLWESecret, TakeGLWESecretExec,
13};
14
15/// A special [GLWESwitchingKey] required to for the conversion from [GLWECiphertext] to [LWECiphertext].
16#[derive(PartialEq, Eq)]
17pub struct GLWEToLWESwitchingKey<D: Data>(GLWESwitchingKey<D>);
18
19impl<D: DataMut> ReaderFrom for GLWEToLWESwitchingKey<D> {
20 fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
21 self.0.read_from(reader)
22 }
23}
24
25impl<D: DataRef> WriterTo for GLWEToLWESwitchingKey<D> {
26 fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
27 self.0.write_to(writer)
28 }
29}
30
31#[derive(PartialEq, Eq)]
32pub struct GLWEToLWESwitchingKeyExec<D: Data, B: Backend>(GLWESwitchingKeyExec<D, B>);
33
34impl<B: Backend> GLWEToLWESwitchingKeyExec<Vec<u8>, B> {
35 pub fn alloc(module: &Module<B>, basek: usize, k: usize, rows: usize, rank_in: usize) -> Self
36 where
37 Module<B>: GGLWEExecLayoutFamily<B>,
38 {
39 Self(GLWESwitchingKeyExec::alloc(
40 module, basek, k, rows, 1, rank_in, 1,
41 ))
42 }
43
44 pub fn bytes_of(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank_in: usize) -> usize
45 where
46 Module<B>: GGLWEExecLayoutFamily<B>,
47 {
48 GLWESwitchingKeyExec::<Vec<u8>, B>::bytes_of(module, basek, k, rows, digits, rank_in, 1)
49 }
50
51 pub fn from<DataOther: DataRef>(
52 module: &Module<B>,
53 other: &GLWEToLWESwitchingKey<DataOther>,
54 scratch: &mut Scratch<B>,
55 ) -> Self
56 where
57 Module<B>: GGLWEExecLayoutFamily<B>,
58 {
59 let mut ksk_exec: GLWEToLWESwitchingKeyExec<Vec<u8>, B> = Self::alloc(
60 module,
61 other.0.basek(),
62 other.0.k(),
63 other.0.rows(),
64 other.0.rank_in(),
65 );
66 ksk_exec.prepare(module, other, scratch);
67 ksk_exec
68 }
69}
70
71impl<D: DataMut, B: Backend> GLWEToLWESwitchingKeyExec<D, B> {
72 pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &GLWEToLWESwitchingKey<DataOther>, scratch: &mut Scratch<B>)
73 where
74 DataOther: DataRef,
75 Module<B>: GGLWEExecLayoutFamily<B>,
76 {
77 self.0.prepare(module, &other.0, scratch);
78 }
79}
80
81impl GLWEToLWESwitchingKey<Vec<u8>> {
82 pub fn alloc<B: Backend>(module: &Module<B>, basek: usize, k: usize, rows: usize, rank_in: usize) -> Self
83 where
84 Module<B>: MatZnxAlloc,
85 {
86 Self(GLWESwitchingKey::alloc(
87 module, basek, k, rows, 1, rank_in, 1,
88 ))
89 }
90
91 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, basek: usize, k: usize, rank_in: usize) -> usize
92 where
93 Module<B>: GGLWEEncryptSkFamily<B> + ScalarZnxAllocBytes + VecZnxAllocBytes,
94 {
95 GLWESecretExec::bytes_of(module, rank_in)
96 + (GLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k, rank_in, 1) | GLWESecret::bytes_of(module, rank_in))
97 }
98}
99
100impl<D: DataMut> GLWEToLWESwitchingKey<D> {
101 pub fn encrypt_sk<DLwe, DGlwe, B: Backend>(
102 &mut self,
103 module: &Module<B>,
104 sk_lwe: &LWESecret<DLwe>,
105 sk_glwe: &GLWESecret<DGlwe>,
106 source_xa: &mut Source,
107 source_xe: &mut Source,
108 sigma: f64,
109 scratch: &mut Scratch<B>,
110 ) where
111 DLwe: DataRef,
112 DGlwe: DataRef,
113 Module<B>: GGLWEEncryptSkFamily<B>
114 + VecZnxAutomorphismInplace
115 + ScalarZnxAllocBytes
116 + VecZnxSwithcDegree
117 + VecZnxAllocBytes
118 + VecZnxAddScalarInplace,
119 Scratch<B>: ScratchAvailable + TakeScalarZnx<B> + TakeVecZnxDft<B> + TakeGLWESecretExec<B> + TakeVecZnx<B>,
120 {
121 #[cfg(debug_assertions)]
122 {
123 assert!(sk_lwe.n() <= module.n());
124 }
125
126 let (mut sk_lwe_as_glwe, scratch1) = scratch.take_glwe_secret(module, 1);
127 sk_lwe_as_glwe.data.zero();
128 sk_lwe_as_glwe.data.at_mut(0, 0)[..sk_lwe.n()].copy_from_slice(sk_lwe.data.at(0, 0));
129 module.vec_znx_automorphism_inplace(-1, &mut sk_lwe_as_glwe.data, 0);
130
131 self.0.encrypt_sk(
132 module,
133 sk_glwe,
134 &sk_lwe_as_glwe,
135 source_xa,
136 source_xe,
137 sigma,
138 scratch1,
139 );
140 }
141}
142
143/// A special [GLWESwitchingKey] required to for the conversion from [LWECiphertext] to [GLWECiphertext].
144#[derive(PartialEq, Eq)]
145pub struct LWEToGLWESwitchingKeyExec<D: Data, B: Backend>(GLWESwitchingKeyExec<D, B>);
146
147impl<B: Backend> LWEToGLWESwitchingKeyExec<Vec<u8>, B> {
148 pub fn alloc(module: &Module<B>, basek: usize, k: usize, rows: usize, rank_out: usize) -> Self
149 where
150 Module<B>: GGLWEExecLayoutFamily<B>,
151 {
152 Self(GLWESwitchingKeyExec::alloc(
153 module, basek, k, rows, 1, 1, rank_out,
154 ))
155 }
156
157 pub fn bytes_of(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank_out: usize) -> usize
158 where
159 Module<B>: GGLWEExecLayoutFamily<B>,
160 {
161 GLWESwitchingKeyExec::<Vec<u8>, B>::bytes_of(module, basek, k, rows, digits, 1, rank_out)
162 }
163
164 pub fn from<DataOther: DataRef>(
165 module: &Module<B>,
166 other: &LWEToGLWESwitchingKey<DataOther>,
167 scratch: &mut Scratch<B>,
168 ) -> Self
169 where
170 Module<B>: GGLWEExecLayoutFamily<B>,
171 {
172 let mut ksk_exec: LWEToGLWESwitchingKeyExec<Vec<u8>, B> = Self::alloc(
173 module,
174 other.0.basek(),
175 other.0.k(),
176 other.0.rows(),
177 other.0.rank(),
178 );
179 ksk_exec.prepare(module, other, scratch);
180 ksk_exec
181 }
182}
183
184impl<D: DataMut, B: Backend> LWEToGLWESwitchingKeyExec<D, B> {
185 pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &LWEToGLWESwitchingKey<DataOther>, scratch: &mut Scratch<B>)
186 where
187 DataOther: DataRef,
188 Module<B>: GGLWEExecLayoutFamily<B>,
189 {
190 self.0.prepare(module, &other.0, scratch);
191 }
192}
193#[derive(PartialEq, Eq)]
194pub struct LWEToGLWESwitchingKey<D: Data>(GLWESwitchingKey<D>);
195
196impl<D: DataMut> ReaderFrom for LWEToGLWESwitchingKey<D> {
197 fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
198 self.0.read_from(reader)
199 }
200}
201
202impl<D: DataRef> WriterTo for LWEToGLWESwitchingKey<D> {
203 fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
204 self.0.write_to(writer)
205 }
206}
207
208impl LWEToGLWESwitchingKey<Vec<u8>> {
209 pub fn alloc<B: Backend>(module: &Module<B>, basek: usize, k: usize, rows: usize, rank_out: usize) -> Self
210 where
211 Module<B>: MatZnxAlloc,
212 {
213 Self(GLWESwitchingKey::alloc(
214 module, basek, k, rows, 1, 1, rank_out,
215 ))
216 }
217
218 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, basek: usize, k: usize, rank_out: usize) -> usize
219 where
220 Module<B>: GGLWEEncryptSkFamily<B> + ScalarZnxAllocBytes + VecZnxAllocBytes,
221 {
222 GLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k, 1, rank_out) + GLWESecret::bytes_of(module, 1)
223 }
224}
225
226impl<D: DataMut> LWEToGLWESwitchingKey<D> {
227 pub fn encrypt_sk<DLwe, DGlwe, B: Backend>(
228 &mut self,
229 module: &Module<B>,
230 sk_lwe: &LWESecret<DLwe>,
231 sk_glwe: &GLWESecret<DGlwe>,
232 source_xa: &mut Source,
233 source_xe: &mut Source,
234 sigma: f64,
235 scratch: &mut Scratch<B>,
236 ) where
237 DLwe: DataRef,
238 DGlwe: DataRef,
239 Module<B>: GGLWEEncryptSkFamily<B>
240 + VecZnxAutomorphismInplace
241 + ScalarZnxAllocBytes
242 + VecZnxSwithcDegree
243 + VecZnxAllocBytes
244 + VecZnxAddScalarInplace,
245 Scratch<B>: ScratchAvailable + TakeScalarZnx<B> + TakeVecZnxDft<B> + TakeGLWESecretExec<B> + TakeVecZnx<B>,
246 {
247 #[cfg(debug_assertions)]
248 {
249 assert!(sk_lwe.n() <= module.n());
250 }
251
252 let (mut sk_lwe_as_glwe, scratch1) = scratch.take_glwe_secret(module, 1);
253 sk_lwe_as_glwe.data.at_mut(0, 0)[..sk_lwe.n()].copy_from_slice(sk_lwe.data.at(0, 0));
254 sk_lwe_as_glwe.data.at_mut(0, 0)[sk_lwe.n()..].fill(0);
255 module.vec_znx_automorphism_inplace(-1, &mut sk_lwe_as_glwe.data, 0);
256
257 self.0.encrypt_sk(
258 module,
259 &sk_lwe_as_glwe,
260 &sk_glwe,
261 source_xa,
262 source_xe,
263 sigma,
264 scratch1,
265 );
266 }
267}
268
269#[derive(PartialEq, Eq)]
270pub struct LWESwitchingKeyExec<D: Data, B: Backend>(GLWESwitchingKeyExec<D, B>);
271
272impl<B: Backend> LWESwitchingKeyExec<Vec<u8>, B> {
273 pub fn alloc(module: &Module<B>, basek: usize, k: usize, rows: usize) -> Self
274 where
275 Module<B>: GGLWEExecLayoutFamily<B>,
276 {
277 Self(GLWESwitchingKeyExec::alloc(module, basek, k, rows, 1, 1, 1))
278 }
279
280 pub fn bytes_of(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize) -> usize
281 where
282 Module<B>: GGLWEExecLayoutFamily<B>,
283 {
284 GLWESwitchingKeyExec::<Vec<u8>, B>::bytes_of(module, basek, k, rows, digits, 1, 1)
285 }
286
287 pub fn from<DataOther: DataRef>(module: &Module<B>, other: &LWESwitchingKey<DataOther>, scratch: &mut Scratch<B>) -> Self
288 where
289 Module<B>: GGLWEExecLayoutFamily<B>,
290 {
291 let mut ksk_exec: LWESwitchingKeyExec<Vec<u8>, B> = Self::alloc(module, other.0.basek(), other.0.k(), other.0.rows());
292 ksk_exec.prepare(module, other, scratch);
293 ksk_exec
294 }
295}
296
297impl<D: DataMut, B: Backend> LWESwitchingKeyExec<D, B> {
298 pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &LWESwitchingKey<DataOther>, scratch: &mut Scratch<B>)
299 where
300 DataOther: DataRef,
301 Module<B>: GGLWEExecLayoutFamily<B>,
302 {
303 self.0.prepare(module, &other.0, scratch);
304 }
305}
306#[derive(PartialEq, Eq)]
307pub struct LWESwitchingKey<D: Data>(GLWESwitchingKey<D>);
308
309impl<D: DataMut> ReaderFrom for LWESwitchingKey<D> {
310 fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
311 self.0.read_from(reader)
312 }
313}
314
315impl<D: DataRef> WriterTo for LWESwitchingKey<D> {
316 fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
317 self.0.write_to(writer)
318 }
319}
320
321impl LWESwitchingKey<Vec<u8>> {
322 pub fn alloc<B: Backend>(module: &Module<B>, basek: usize, k: usize, rows: usize) -> Self
323 where
324 Module<B>: MatZnxAlloc,
325 {
326 Self(GLWESwitchingKey::alloc(module, basek, k, rows, 1, 1, 1))
327 }
328
329 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, basek: usize, k: usize) -> usize
330 where
331 Module<B>: GGLWEEncryptSkFamily<B> + ScalarZnxAllocBytes + VecZnxAllocBytes,
332 {
333 GLWESecret::bytes_of(module, 1)
334 + GLWESecretExec::bytes_of(module, 1)
335 + GLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k, 1, 1)
336 }
337}
338
339impl<D: DataMut> LWESwitchingKey<D> {
340 pub fn encrypt_sk<DIn, DOut, B: Backend>(
341 &mut self,
342 module: &Module<B>,
343 sk_lwe_in: &LWESecret<DIn>,
344 sk_lwe_out: &LWESecret<DOut>,
345 source_xa: &mut Source,
346 source_xe: &mut Source,
347 sigma: f64,
348 scratch: &mut Scratch<B>,
349 ) where
350 DIn: DataRef,
351 DOut: DataRef,
352 Module<B>: GGLWEEncryptSkFamily<B>
353 + VecZnxAutomorphismInplace
354 + ScalarZnxAllocBytes
355 + VecZnxSwithcDegree
356 + VecZnxAllocBytes
357 + VecZnxAddScalarInplace,
358 Scratch<B>: ScratchAvailable + TakeScalarZnx<B> + TakeVecZnxDft<B> + TakeGLWESecretExec<B> + TakeVecZnx<B>,
359 {
360 #[cfg(debug_assertions)]
361 {
362 assert!(sk_lwe_in.n() <= module.n());
363 assert!(sk_lwe_out.n() <= module.n());
364 }
365
366 let (mut sk_in_glwe, scratch1) = scratch.take_glwe_secret(module, 1);
367 let (mut sk_out_glwe, scratch2) = scratch1.take_glwe_secret(module, 1);
368
369 sk_out_glwe.data.at_mut(0, 0)[..sk_lwe_out.n()].copy_from_slice(sk_lwe_out.data.at(0, 0));
370 sk_out_glwe.data.at_mut(0, 0)[sk_lwe_out.n()..].fill(0);
371 module.vec_znx_automorphism_inplace(-1, &mut sk_out_glwe.data, 0);
372
373 sk_in_glwe.data.at_mut(0, 0)[..sk_lwe_in.n()].copy_from_slice(sk_lwe_in.data.at(0, 0));
374 sk_in_glwe.data.at_mut(0, 0)[sk_lwe_in.n()..].fill(0);
375 module.vec_znx_automorphism_inplace(-1, &mut sk_in_glwe.data, 0);
376
377 self.0.encrypt_sk(
378 module,
379 &sk_in_glwe,
380 &sk_out_glwe,
381 source_xa,
382 source_xe,
383 sigma,
384 scratch2,
385 );
386 }
387}
388
389impl LWECiphertext<Vec<u8>> {
390 pub fn from_glwe_scratch_space<B: Backend>(
391 module: &Module<B>,
392 basek: usize,
393 k_lwe: usize,
394 k_glwe: usize,
395 k_ksk: usize,
396 rank: usize,
397 ) -> usize
398 where
399 Module<B>: GLWEKeyswitchFamily<B> + VecZnxAllocBytes,
400 {
401 GLWECiphertext::bytes_of(module, basek, k_lwe, 1)
402 + GLWECiphertext::keyswitch_scratch_space(module, basek, k_lwe, k_glwe, k_ksk, 1, rank, 1)
403 }
404
405 pub fn keyswitch_scratch_space<B: Backend>(
406 module: &Module<B>,
407 basek: usize,
408 k_lwe_out: usize,
409 k_lwe_in: usize,
410 k_ksk: usize,
411 ) -> usize
412 where
413 Module<B>: GLWEKeyswitchFamily<B> + ScalarZnxAllocBytes + VecZnxAllocBytes,
414 {
415 GLWECiphertext::bytes_of(module, basek, k_lwe_out.max(k_lwe_in), 1)
416 + GLWECiphertext::keyswitch_inplace_scratch_space(module, basek, k_lwe_out, k_ksk, 1, 1)
417 }
418}
419
420impl<DLwe: DataMut> LWECiphertext<DLwe> {
421 pub fn sample_extract<DGlwe: DataRef>(&mut self, a: &GLWECiphertext<DGlwe>) {
422 #[cfg(debug_assertions)]
423 {
424 assert!(self.n() <= a.n());
425 }
426
427 let min_size: usize = self.size().min(a.size());
428 let n: usize = self.n();
429
430 self.data.zero();
431 (0..min_size).for_each(|i| {
432 let data_lwe: &mut [i64] = self.data.at_mut(0, i);
433 data_lwe[0] = a.data.at(0, i)[0];
434 data_lwe[1..].copy_from_slice(&a.data.at(1, i)[..n]);
435 });
436 }
437
438 pub fn from_glwe<DGlwe, DKs, B: Backend>(
439 &mut self,
440 module: &Module<B>,
441 a: &GLWECiphertext<DGlwe>,
442 ks: &GLWEToLWESwitchingKeyExec<DKs, B>,
443 scratch: &mut Scratch<B>,
444 ) where
445 DGlwe: DataRef,
446 DKs: DataRef,
447 Module<B>: GLWEKeyswitchFamily<B>,
448 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx<B>,
449 {
450 #[cfg(debug_assertions)]
451 {
452 assert_eq!(self.basek(), a.basek());
453 }
454 let (mut tmp_glwe, scratch1) = scratch.take_glwe_ct(module, a.basek(), self.k(), 1);
455 tmp_glwe.keyswitch(module, a, &ks.0, scratch1);
456 self.sample_extract(&tmp_glwe);
457 }
458
459 pub fn keyswitch<A, DKs, B: Backend>(
460 &mut self,
461 module: &Module<B>,
462 a: &LWECiphertext<A>,
463 ksk: &LWESwitchingKeyExec<DKs, B>,
464 scratch: &mut Scratch<B>,
465 ) where
466 A: DataRef,
467 DKs: DataRef,
468 Module<B>: GLWEKeyswitchFamily<B>,
469 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx<B>,
470 {
471 #[cfg(debug_assertions)]
472 {
473 assert!(self.n() <= module.n());
474 assert!(a.n() <= module.n());
475 assert_eq!(self.basek(), a.basek());
476 }
477
478 let max_k: usize = self.k().max(a.k());
479 let basek: usize = self.basek();
480
481 let (mut glwe, scratch1) = scratch.take_glwe_ct(&module, basek, max_k, 1);
482 glwe.data.zero();
483
484 let n_lwe: usize = a.n();
485
486 (0..a.size()).for_each(|i| {
487 let data_lwe: &[i64] = a.data.at(0, i);
488 glwe.data.at_mut(0, i)[0] = data_lwe[0];
489 glwe.data.at_mut(1, i)[..n_lwe].copy_from_slice(&data_lwe[1..]);
490 });
491
492 glwe.keyswitch_inplace(module, &ksk.0, scratch1);
493
494 self.sample_extract(&glwe);
495 }
496}
497
498impl GLWECiphertext<Vec<u8>> {
499 pub fn from_lwe_scratch_space<B: Backend>(
500 module: &Module<B>,
501 basek: usize,
502 k_lwe: usize,
503 k_glwe: usize,
504 k_ksk: usize,
505 rank: usize,
506 ) -> usize
507 where
508 Module<B>: GLWEKeyswitchFamily<B> + VecZnxAllocBytes,
509 {
510 GLWECiphertext::keyswitch_scratch_space(module, basek, k_glwe, k_lwe, k_ksk, 1, 1, rank)
511 + GLWECiphertext::bytes_of(module, basek, k_lwe, 1)
512 }
513}
514
515impl<D: DataMut> GLWECiphertext<D> {
516 pub fn from_lwe<DLwe, DKsk, B: Backend>(
517 &mut self,
518 module: &Module<B>,
519 lwe: &LWECiphertext<DLwe>,
520 ksk: &LWEToGLWESwitchingKeyExec<DKsk, B>,
521 scratch: &mut Scratch<B>,
522 ) where
523 DLwe: DataRef,
524 DKsk: DataRef,
525 Module<B>: GLWEKeyswitchFamily<B>,
526 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx<B>,
527 {
528 #[cfg(debug_assertions)]
529 {
530 assert!(lwe.n() <= self.n());
531 assert_eq!(self.basek(), self.basek());
532 }
533
534 let (mut glwe, scratch1) = scratch.take_glwe_ct(module, lwe.basek(), lwe.k(), 1);
535 glwe.data.zero();
536
537 let n_lwe: usize = lwe.n();
538
539 (0..lwe.size()).for_each(|i| {
540 let data_lwe: &[i64] = lwe.data.at(0, i);
541 glwe.data.at_mut(0, i)[0] = data_lwe[0];
542 glwe.data.at_mut(1, i)[..n_lwe].copy_from_slice(&data_lwe[1..]);
543 });
544
545 self.keyswitch(module, &glwe, &ksk.0, scratch1);
546 }
547}