poulpy/core/src/gglwe/encryption.rs
core/src/gglwe/encryption.rsRaw
1use backend::hal::{
2 api::{
3 ScalarZnxAllocBytes, ScalarZnxAutomorphism, ScratchAvailable, SvpApply, TakeScalarZnx, TakeVecZnx, TakeVecZnxBig,
4 TakeVecZnxDft, VecZnxAddScalarInplace, VecZnxAllocBytes, VecZnxBigAllocBytes, VecZnxDftToVecZnxBigTmpA,
5 VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSwithcDegree, ZnxZero,
6 },
7 layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, Scratch},
8};
9use sampling::source::Source;
10
11use crate::{
12 AutomorphismKey, GGLWECiphertext, GLWECiphertext, GLWEDecryptFamily, GLWEEncryptSkFamily, GLWEPlaintext, GLWESecret,
13 GLWESecretExec, GLWESecretFamily, GLWESwitchingKey, GLWETensorKey, Infos, TakeGLWEPt, TakeGLWESecret, TakeGLWESecretExec,
14};
15
16pub trait GGLWEEncryptSkFamily<B: Backend> = GLWEEncryptSkFamily<B> + GLWESecretFamily<B>;
17
18impl GGLWECiphertext<Vec<u8>> {
19 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, basek: usize, k: usize) -> usize
20 where
21 Module<B>: GGLWEEncryptSkFamily<B> + VecZnxAllocBytes,
22 {
23 GLWECiphertext::encrypt_sk_scratch_space(module, basek, k)
24 + (GLWEPlaintext::byte_of(module, basek, k) | module.vec_znx_normalize_tmp_bytes(module.n()))
25 }
26
27 pub fn encrypt_pk_scratch_space<B: Backend>(_module: &Module<B>, _basek: usize, _k: usize, _rank: usize) -> usize {
28 unimplemented!()
29 }
30}
31
32impl<DataSelf: DataMut> GGLWECiphertext<DataSelf> {
33 pub fn encrypt_sk<DataPt: DataRef, DataSk: DataRef, B: Backend>(
34 &mut self,
35 module: &Module<B>,
36 pt: &ScalarZnx<DataPt>,
37 sk: &GLWESecretExec<DataSk, B>,
38 source_xa: &mut Source,
39 source_xe: &mut Source,
40 sigma: f64,
41 scratch: &mut Scratch<B>,
42 ) where
43 Module<B>: GGLWEEncryptSkFamily<B> + VecZnxAllocBytes + VecZnxAddScalarInplace,
44 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx<B>,
45 {
46 #[cfg(debug_assertions)]
47 {
48 use backend::hal::api::ZnxInfos;
49
50 assert_eq!(
51 self.rank_in(),
52 pt.cols(),
53 "self.rank_in(): {} != pt.cols(): {}",
54 self.rank_in(),
55 pt.cols()
56 );
57 assert_eq!(
58 self.rank_out(),
59 sk.rank(),
60 "self.rank_out(): {} != sk.rank(): {}",
61 self.rank_out(),
62 sk.rank()
63 );
64 assert_eq!(self.n(), module.n());
65 assert_eq!(sk.n(), module.n());
66 assert_eq!(pt.n(), module.n());
67 assert!(
68 scratch.available() >= GGLWECiphertext::encrypt_sk_scratch_space(module, self.basek(), self.k()),
69 "scratch.available: {} < GGLWECiphertext::encrypt_sk_scratch_space(module, self.rank()={}, self.size()={}): {}",
70 scratch.available(),
71 self.rank(),
72 self.size(),
73 GGLWECiphertext::encrypt_sk_scratch_space(module, self.basek(), self.k())
74 );
75 assert!(
76 self.rows() * self.digits() * self.basek() <= self.k(),
77 "self.rows() : {} * self.digits() : {} * self.basek() : {} = {} >= self.k() = {}",
78 self.rows(),
79 self.digits(),
80 self.basek(),
81 self.rows() * self.digits() * self.basek(),
82 self.k()
83 );
84 }
85
86 let rows: usize = self.rows();
87 let digits: usize = self.digits();
88 let basek: usize = self.basek();
89 let k: usize = self.k();
90 let rank_in: usize = self.rank_in();
91
92 let (mut tmp_pt, scrach_1) = scratch.take_glwe_pt(module, basek, k);
93 // For each input column (i.e. rank) produces a GGLWE ciphertext of rank_out+1 columns
94 //
95 // Example for ksk rank 2 to rank 3:
96 //
97 // (-(a0*s0 + a1*s1 + a2*s2) + s0', a0, a1, a2)
98 // (-(b0*s0 + b1*s1 + b2*s2) + s0', b0, b1, b2)
99 //
100 // Example ksk rank 2 to rank 1
101 //
102 // (-(a*s) + s0, a)
103 // (-(b*s) + s1, b)
104 (0..rank_in).for_each(|col_i| {
105 (0..rows).for_each(|row_i| {
106 // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
107 tmp_pt.data.zero(); // zeroes for next iteration
108 module.vec_znx_add_scalar_inplace(
109 &mut tmp_pt.data,
110 0,
111 (digits - 1) + row_i * digits,
112 pt,
113 col_i,
114 );
115 module.vec_znx_normalize_inplace(basek, &mut tmp_pt.data, 0, scrach_1);
116
117 // rlwe encrypt of vec_znx_pt into vec_znx_ct
118 self.at_mut(row_i, col_i)
119 .encrypt_sk(module, &tmp_pt, sk, source_xa, source_xe, sigma, scrach_1);
120 });
121 });
122 }
123}
124
125pub trait GLWESwitchingKeyEncryptSkFamily<B: Backend> = GGLWEEncryptSkFamily<B>;
126
127impl GLWESwitchingKey<Vec<u8>> {
128 pub fn encrypt_sk_scratch_space<B: Backend>(
129 module: &Module<B>,
130 basek: usize,
131 k: usize,
132 rank_in: usize,
133 rank_out: usize,
134 ) -> usize
135 where
136 Module<B>: GLWESwitchingKeyEncryptSkFamily<B> + ScalarZnxAllocBytes + VecZnxAllocBytes,
137 {
138 (GGLWECiphertext::encrypt_sk_scratch_space(module, basek, k) | module.scalar_znx_alloc_bytes(1))
139 + module.scalar_znx_alloc_bytes(rank_in)
140 + GLWESecretExec::bytes_of(module, rank_out)
141 }
142
143 pub fn encrypt_pk_scratch_space<B: Backend>(
144 module: &Module<B>,
145 _basek: usize,
146 _k: usize,
147 _rank_in: usize,
148 _rank_out: usize,
149 ) -> usize {
150 GGLWECiphertext::encrypt_pk_scratch_space(module, _basek, _k, _rank_out)
151 }
152}
153
154impl<DataSelf: DataMut> GLWESwitchingKey<DataSelf> {
155 pub fn encrypt_sk<DataSkIn: DataRef, DataSkOut: DataRef, B: Backend>(
156 &mut self,
157 module: &Module<B>,
158 sk_in: &GLWESecret<DataSkIn>,
159 sk_out: &GLWESecret<DataSkOut>,
160 source_xa: &mut Source,
161 source_xe: &mut Source,
162 sigma: f64,
163 scratch: &mut Scratch<B>,
164 ) where
165 Module<B>: GLWESwitchingKeyEncryptSkFamily<B>
166 + ScalarZnxAllocBytes
167 + VecZnxSwithcDegree
168 + VecZnxAllocBytes
169 + VecZnxAddScalarInplace,
170 Scratch<B>:
171 ScratchAvailable + TakeScalarZnx<B> + TakeVecZnxDft<B> + TakeGLWESecretExec<B> + ScratchAvailable + TakeVecZnx<B>,
172 {
173 #[cfg(debug_assertions)]
174 {
175 assert!(sk_in.n() <= module.n());
176 assert!(sk_out.n() <= module.n());
177 assert!(
178 scratch.available()
179 >= GLWESwitchingKey::encrypt_sk_scratch_space(
180 module,
181 self.basek(),
182 self.k(),
183 self.rank_in(),
184 self.rank_out()
185 ),
186 "scratch.available()={} < GLWESwitchingKey::encrypt_sk_scratch_space={}",
187 scratch.available(),
188 GLWESwitchingKey::encrypt_sk_scratch_space(
189 module,
190 self.basek(),
191 self.k(),
192 self.rank_in(),
193 self.rank_out()
194 )
195 )
196 }
197
198 let (mut sk_in_tmp, scratch1) = scratch.take_scalar_znx(module, sk_in.rank());
199 (0..sk_in.rank()).for_each(|i| {
200 module.vec_znx_switch_degree(&mut sk_in_tmp, i, &sk_in.data, i);
201 });
202
203 let (mut sk_out_tmp, scratch2) = scratch1.take_glwe_secret_exec(module, sk_out.rank());
204 {
205 let (mut tmp, _) = scratch2.take_scalar_znx(module, 1);
206 (0..sk_out.rank()).for_each(|i| {
207 module.vec_znx_switch_degree(&mut tmp, 0, &sk_out.data, i);
208 module.svp_prepare(&mut sk_out_tmp.data, i, &tmp, 0);
209 });
210 }
211
212 self.key.encrypt_sk(
213 module,
214 &sk_in_tmp,
215 &sk_out_tmp,
216 source_xa,
217 source_xe,
218 sigma,
219 scratch2,
220 );
221 self.sk_in_n = sk_in.n();
222 self.sk_out_n = sk_out.n();
223 }
224}
225
226pub trait AutomorphismKeyEncryptSkFamily<B: Backend> = GGLWEEncryptSkFamily<B>;
227
228impl AutomorphismKey<Vec<u8>> {
229 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, basek: usize, k: usize, rank: usize) -> usize
230 where
231 Module<B>: AutomorphismKeyEncryptSkFamily<B> + ScalarZnxAllocBytes + VecZnxAllocBytes,
232 {
233 GLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k, rank, rank) + GLWESecret::bytes_of(module, rank)
234 }
235
236 pub fn encrypt_pk_scratch_space<B: Backend>(module: &Module<B>, _basek: usize, _k: usize, _rank: usize) -> usize {
237 GLWESwitchingKey::encrypt_pk_scratch_space(module, _basek, _k, _rank, _rank)
238 }
239}
240
241impl<DataSelf: DataMut> AutomorphismKey<DataSelf> {
242 pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
243 &mut self,
244 module: &Module<B>,
245 p: i64,
246 sk: &GLWESecret<DataSk>,
247 source_xa: &mut Source,
248 source_xe: &mut Source,
249 sigma: f64,
250 scratch: &mut Scratch<B>,
251 ) where
252 Module<B>: AutomorphismKeyEncryptSkFamily<B>
253 + ScalarZnxAutomorphism
254 + ScalarZnxAllocBytes
255 + VecZnxAllocBytes
256 + VecZnxSwithcDegree
257 + VecZnxAddScalarInplace,
258 Scratch<B>: ScratchAvailable + TakeScalarZnx<B> + TakeVecZnxDft<B> + TakeGLWESecretExec<B> + TakeVecZnx<B>,
259 {
260 #[cfg(debug_assertions)]
261 {
262 assert_eq!(self.n(), module.n());
263 assert_eq!(sk.n(), module.n());
264 assert_eq!(self.rank_out(), self.rank_in());
265 assert_eq!(sk.rank(), self.rank());
266 assert!(
267 scratch.available() >= AutomorphismKey::encrypt_sk_scratch_space(module, self.basek(), self.k(), self.rank()),
268 "scratch.available(): {} < AutomorphismKey::encrypt_sk_scratch_space(module, self.rank()={}, self.size()={}): {}",
269 scratch.available(),
270 self.rank(),
271 self.size(),
272 AutomorphismKey::encrypt_sk_scratch_space(module, self.basek(), self.k(), self.rank())
273 )
274 }
275
276 let (mut sk_out, scratch_1) = scratch.take_glwe_secret(module, sk.rank());
277
278 {
279 (0..self.rank()).for_each(|i| {
280 module.scalar_znx_automorphism(
281 module.galois_element_inv(p),
282 &mut sk_out.data,
283 i,
284 &sk.data,
285 i,
286 );
287 });
288 }
289
290 self.key
291 .encrypt_sk(module, &sk, &sk_out, source_xa, source_xe, sigma, scratch_1);
292
293 self.p = p;
294 }
295}
296
297pub trait GLWETensorKeyEncryptSkFamily<B: Backend> =
298 GGLWEEncryptSkFamily<B> + VecZnxBigAllocBytes + VecZnxDftToVecZnxBigTmpA<B> + SvpApply<B>;
299
300impl GLWETensorKey<Vec<u8>> {
301 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, basek: usize, k: usize, rank: usize) -> usize
302 where
303 Module<B>: GLWETensorKeyEncryptSkFamily<B> + ScalarZnxAllocBytes + VecZnxAllocBytes,
304 {
305 GLWESecretExec::bytes_of(module, rank)
306 + module.vec_znx_dft_alloc_bytes(rank, 1)
307 + module.vec_znx_big_alloc_bytes(1, 1)
308 + module.vec_znx_dft_alloc_bytes(1, 1)
309 + GLWESecret::bytes_of(module, 1)
310 + GLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k, rank, rank)
311 }
312}
313
314impl<DataSelf: DataMut> GLWETensorKey<DataSelf> {
315 pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
316 &mut self,
317 module: &Module<B>,
318 sk: &GLWESecret<DataSk>,
319 source_xa: &mut Source,
320 source_xe: &mut Source,
321 sigma: f64,
322 scratch: &mut Scratch<B>,
323 ) where
324 Module<B>: GLWETensorKeyEncryptSkFamily<B>
325 + ScalarZnxAllocBytes
326 + VecZnxSwithcDegree
327 + VecZnxAllocBytes
328 + VecZnxAddScalarInplace,
329 Scratch<B>:
330 ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeGLWESecretExec<B> + TakeScalarZnx<B> + TakeVecZnx<B>,
331 {
332 #[cfg(debug_assertions)]
333 {
334 assert_eq!(self.rank(), sk.rank());
335 assert_eq!(self.n(), module.n());
336 assert_eq!(sk.n(), module.n());
337 }
338
339 let rank: usize = self.rank();
340
341 let (mut sk_dft_prep, scratch1) = scratch.take_glwe_secret_exec(module, rank);
342 sk_dft_prep.prepare(module, &sk);
343
344 let (mut sk_dft, scratch2) = scratch1.take_vec_znx_dft(module, rank, 1);
345
346 (0..rank).for_each(|i| {
347 module.vec_znx_dft_from_vec_znx(1, 0, &mut sk_dft, i, &sk.data, i);
348 });
349
350 let (mut sk_ij_big, scratch3) = scratch2.take_vec_znx_big(module, 1, 1);
351 let (mut sk_ij, scratch4) = scratch3.take_glwe_secret(module, 1);
352 let (mut sk_ij_dft, scratch5) = scratch4.take_vec_znx_dft(module, 1, 1);
353
354 (0..rank).for_each(|i| {
355 (i..rank).for_each(|j| {
356 module.svp_apply(&mut sk_ij_dft, 0, &sk_dft_prep.data, j, &sk_dft, i);
357
358 module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
359 module.vec_znx_big_normalize(self.basek(), &mut sk_ij.data, 0, &sk_ij_big, 0, scratch5);
360
361 self.at_mut(i, j)
362 .encrypt_sk(module, &sk_ij, sk, source_xa, source_xe, sigma, scratch5);
363 });
364 })
365 }
366}