poulpy/rlwe/src/automorphism.rs
rlwe/src/automorphism.rsRaw
1use crate::{
2 ciphertext::{Ciphertext, new_gadget_ciphertext},
3 elem::ElemCommon,
4 encryptor::{encrypt_grlwe_sk, encrypt_grlwe_sk_tmp_bytes},
5 keys::SecretKey,
6 parameters::Parameters,
7};
8use base2k::{
9 Module, Scalar, ScalarOps, SvpPPol, SvpPPolOps, VecZnx, VecZnxBig, VecZnxBigOps, VecZnxDft,
10 VecZnxDftOps, VecZnxOps, VmpPMat, VmpPMatOps, assert_alignement,
11};
12use sampling::source::Source;
13use std::{cmp::min, collections::HashMap};
14
15/// Stores DFT([-A*AUTO(s, -p) + 2^{-K*i}*s + E, A]) where AUTO(X, p): X^{i} -> X^{i*p}
16pub struct AutomorphismKey {
17 pub value: Ciphertext<VmpPMat>,
18 pub p: i64,
19}
20
21pub fn automorphis_key_new_tmp_bytes(
22 module: &Module,
23 log_base2k: usize,
24 rows: usize,
25 log_q: usize,
26) -> usize {
27 module.bytes_of_scalar()
28 + module.bytes_of_svp_ppol()
29 + encrypt_grlwe_sk_tmp_bytes(module, log_base2k, rows, log_q)
30}
31
32impl Parameters {
33 pub fn automorphism_key_new_tmp_bytes(&self, rows: usize, log_q: usize) -> usize {
34 automorphis_key_new_tmp_bytes(self.module(), self.log_base2k(), rows, log_q)
35 }
36
37 pub fn automorphism_tmp_bytes(
38 &self,
39 res_logq: usize,
40 in_logq: usize,
41 gct_logq: usize,
42 ) -> usize {
43 automorphism_tmp_bytes(
44 self.module(),
45 self.log_base2k(),
46 res_logq,
47 in_logq,
48 gct_logq,
49 )
50 }
51}
52
53impl AutomorphismKey {
54 pub fn new(
55 module: &Module,
56 p: i64,
57 sk: &SecretKey,
58 log_base2k: usize,
59 rows: usize,
60 log_q: usize,
61 source_xa: &mut Source,
62 source_xe: &mut Source,
63 sigma: f64,
64 tmp_bytes: &mut [u8],
65 ) -> Self {
66 Self::new_many_core(
67 module,
68 &vec![p],
69 sk,
70 log_base2k,
71 rows,
72 log_q,
73 source_xa,
74 source_xe,
75 sigma,
76 tmp_bytes,
77 )
78 .into_iter()
79 .next()
80 .unwrap()
81 }
82
83 pub fn new_many(
84 module: &Module,
85 p: &Vec<i64>,
86 sk: &SecretKey,
87 log_base2k: usize,
88 rows: usize,
89 log_q: usize,
90 source_xa: &mut Source,
91 source_xe: &mut Source,
92 sigma: f64,
93 tmp_bytes: &mut [u8],
94 ) -> HashMap<i64, AutomorphismKey> {
95 Self::new_many_core(
96 module, p, sk, log_base2k, rows, log_q, source_xa, source_xe, sigma, tmp_bytes,
97 )
98 .into_iter()
99 .zip(p.iter().cloned())
100 .map(|(key, pi)| (pi, key))
101 .collect()
102 }
103
104 fn new_many_core(
105 module: &Module,
106 p: &Vec<i64>,
107 sk: &SecretKey,
108 log_base2k: usize,
109 rows: usize,
110 log_q: usize,
111 source_xa: &mut Source,
112 source_xe: &mut Source,
113 sigma: f64,
114 tmp_bytes: &mut [u8],
115 ) -> Vec<Self> {
116 let (sk_auto_bytes, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_scalar());
117 let (sk_out_bytes, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_svp_ppol());
118
119 let sk_auto: Scalar = module.new_scalar_from_bytes_borrow(sk_auto_bytes);
120 let mut sk_out: SvpPPol = module.new_svp_ppol_from_bytes_borrow(sk_out_bytes);
121
122 let mut keys: Vec<AutomorphismKey> = Vec::new();
123
124 p.iter().for_each(|pi| {
125 let mut value: Ciphertext<VmpPMat> =
126 new_gadget_ciphertext(module, log_base2k, rows, log_q);
127
128 let p_inv: i64 = module.galois_element_inv(*pi);
129
130 module.vec_znx_automorphism(p_inv, &mut sk_auto.as_vec_znx(), &sk.0.as_vec_znx());
131 module.svp_prepare(&mut sk_out, &sk_auto);
132 encrypt_grlwe_sk(
133 module, &mut value, &sk.0, &sk_out, source_xa, source_xe, sigma, tmp_bytes,
134 );
135
136 keys.push(Self {
137 value: value,
138 p: *pi,
139 })
140 });
141
142 keys
143 }
144}
145
146pub fn automorphism_tmp_bytes(
147 module: &Module,
148 log_base2k: usize,
149 res_logq: usize,
150 in_logq: usize,
151 gct_logq: usize,
152) -> usize {
153 let gct_cols: usize = (gct_logq + log_base2k - 1) / log_base2k;
154 let in_cols: usize = (in_logq + log_base2k - 1) / log_base2k;
155 let res_cols: usize = (res_logq + log_base2k - 1) / log_base2k;
156 return module.vmp_apply_dft_to_dft_tmp_bytes(res_cols, in_cols, in_cols, gct_cols)
157 + module.bytes_of_vec_znx_dft(std::cmp::min(res_cols, in_cols))
158 + module.bytes_of_vec_znx_dft(gct_cols);
159}
160
161pub fn automorphism(
162 module: &Module,
163 c: &mut Ciphertext<VecZnx>,
164 a: &Ciphertext<VecZnx>,
165 b: &AutomorphismKey,
166 b_cols: usize,
167 tmp_bytes: &mut [u8],
168) {
169 let cols: usize = min(min(c.cols(), a.cols()), b.value.rows());
170
171 #[cfg(debug_assertions)]
172 {
173 assert!(b_cols <= b.value.cols());
174 assert!(
175 tmp_bytes.len()
176 >= automorphism_tmp_bytes(
177 module,
178 c.cols(),
179 a.cols(),
180 b.value.rows(),
181 b.value.cols()
182 )
183 );
184 assert_alignement(tmp_bytes.as_ptr());
185 }
186
187 let (tmp_bytes_a1_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
188 let (tmp_bytes_res_dft, tmp_bytes) =
189 tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(b_cols));
190
191 let mut a1_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_a1_dft);
192 let mut res_dft: VecZnxDft =
193 module.new_vec_znx_dft_from_bytes_borrow(b_cols, tmp_bytes_res_dft);
194 let mut res_big: VecZnxBig = res_dft.as_vec_znx_big();
195
196 // a1_dft = DFT(a[1])
197 module.vec_znx_dft(&mut a1_dft, a.at(1));
198
199 // res_dft = IDFT(<DFT(a), DFT([-A*AUTO(s, -p) + 2^{-K*i}*s + E])>) = [-b*AUTO(s, -p) + a * s + e]
200 module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, b.value.at(0), tmp_bytes);
201 module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft);
202
203 // res_dft = [-b*AUTO(s, -p) + a * s + e] + [-a * s + m + e] = [-b*AUTO(s, -p) + m + e]
204 module.vec_znx_big_add_small_inplace(&mut res_big, a.at(0));
205
206 // c[0] = NORMALIZE([-b*AUTO(s, -p) + m + e])
207 module.vec_znx_big_normalize(c.log_base2k(), c.at_mut(0), &mut res_big, tmp_bytes);
208
209 // c[0] = AUTO([-b*AUTO(s, -p) + m + e], p) = [-AUTO(b, p)*s + AUTO(m, p) + AUTO(b, e)]
210 module.vec_znx_automorphism_inplace(b.p, c.at_mut(0));
211
212 // res_dft = IDFT(<DFT(a), DFT([A])>) = [b]
213 module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, b.value.at(1), tmp_bytes);
214 module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft);
215
216 // c[1] = b
217 module.vec_znx_big_normalize(c.log_base2k(), c.at_mut(1), &mut res_big, tmp_bytes);
218
219 // c[1] = AUTO(b, p)
220 module.vec_znx_automorphism_inplace(b.p, c.at_mut(1));
221}
222
223pub fn automorphism_inplace_tmp_bytes(
224 module: &Module,
225 c_cols: usize,
226 a_cols: usize,
227 b_rows: usize,
228 b_cols: usize,
229) -> usize {
230 return module.vmp_apply_dft_to_dft_tmp_bytes(c_cols, a_cols, b_rows, b_cols)
231 + 2 * module.bytes_of_vec_znx_dft(std::cmp::min(c_cols, a_cols));
232}
233
234pub fn automorphism_inplace(
235 module: &Module,
236 a: &mut Ciphertext<VecZnx>,
237 b: &AutomorphismKey,
238 b_cols: usize,
239 tmp_bytes: &mut [u8],
240) {
241 let cols: usize = min(a.cols(), b.value.rows());
242
243 #[cfg(debug_assertions)]
244 {
245 assert!(b_cols <= b.value.cols());
246 assert!(
247 tmp_bytes.len()
248 >= automorphism_inplace_tmp_bytes(
249 module,
250 a.cols(),
251 a.cols(),
252 b.value.rows(),
253 b.value.cols()
254 )
255 );
256 assert_alignement(tmp_bytes.as_ptr());
257 }
258
259 let (tmp_bytes_b1_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
260 let (tmp_bytes_res_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
261
262 let mut a1_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_b1_dft);
263 let mut res_dft: VecZnxDft =
264 module.new_vec_znx_dft_from_bytes_borrow(b_cols, tmp_bytes_res_dft);
265 let mut res_big: VecZnxBig = res_dft.as_vec_znx_big();
266
267 // a1_dft = DFT(a[1])
268 module.vec_znx_dft(&mut a1_dft, a.at(1));
269
270 // res_dft = IDFT(<DFT(a), DFT([-A*AUTO(s, -p) + 2^{-K*i}*s + E])>) = [-b*AUTO(s, -p) + a * s + e]
271 module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, b.value.at(0), tmp_bytes);
272 module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft);
273
274 // res_dft = [-b*AUTO(s, -p) + a * s + e] + [-a * s + m + e] = [-b*AUTO(s, -p) + m + e]
275 module.vec_znx_big_add_small_inplace(&mut res_big, a.at(0));
276
277 // a[0] = NORMALIZE([-b*AUTO(s, -p) + m + e])
278 module.vec_znx_big_normalize(a.log_base2k(), a.at_mut(0), &mut res_big, tmp_bytes);
279
280 // a[0] = AUTO([-b*AUTO(s, -p) + m + e], p) = [-AUTO(b, p)*s + AUTO(m, p) + AUTO(b, e)]
281 module.vec_znx_automorphism_inplace(b.p, a.at_mut(0));
282
283 // res_dft = IDFT(<DFT(a), DFT([A])>) = [b]
284 module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, b.value.at(1), tmp_bytes);
285 module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft);
286
287 (0..b_cols).for_each(|col_i| {
288 let raw: &[i64] = res_big.raw::<i64>(module);
289 println!("{:?}", &raw[col_i * module.n()..(col_i + 1) * module.n()])
290 });
291
292 // a[1] = b
293 module.vec_znx_big_normalize(a.log_base2k(), a.at_mut(1), &mut res_big, tmp_bytes);
294
295 // a[1] = AUTO(b, p)
296 module.vec_znx_automorphism_inplace(b.p, a.at_mut(1));
297}
298
299pub fn automorphism_big(
300 module: &Module,
301 c: &mut Ciphertext<VecZnxBig>,
302 a: &Ciphertext<VecZnx>,
303 b: &AutomorphismKey,
304 tmp_bytes: &mut [u8],
305) {
306 let cols = std::cmp::min(c.cols(), a.cols());
307
308 #[cfg(debug_assertions)]
309 {
310 assert!(
311 tmp_bytes.len()
312 >= automorphism_tmp_bytes(
313 module,
314 c.cols(),
315 a.cols(),
316 b.value.rows(),
317 b.value.cols()
318 )
319 );
320 assert_alignement(tmp_bytes.as_ptr());
321 }
322
323 let (tmp_bytes_b1_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
324 let (tmp_bytes_res_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
325
326 let mut a1_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_b1_dft);
327 let mut res_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_res_dft);
328
329 // a1_dft = DFT(a[1])
330 module.vec_znx_dft(&mut a1_dft, a.at(1));
331
332 // res_dft = IDFT(<DFT(a), DFT([-A*AUTO(s, -p) + 2^{-K*i}*s + E])>) = [-b*AUTO(s, -p) + a * s + e]
333 module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, b.value.at(0), tmp_bytes);
334 module.vec_znx_idft_tmp_a(c.at_mut(0), &mut res_dft);
335
336 // res_dft = [-b*AUTO(s, -p) + a * s + e] + [-a * s + m + e] = [-b*AUTO(s, -p) + m + e]
337 module.vec_znx_big_add_small_inplace(c.at_mut(0), a.at(0));
338
339 // c[0] = AUTO([-b*AUTO(s, -p) + m + e], p) = [-AUTO(b, p)*s + AUTO(m, p) + AUTO(b, e)]
340 module.vec_znx_big_automorphism_inplace(b.p, c.at_mut(0));
341
342 // res_dft = IDFT(<DFT(a), DFT([A])>) = [b]
343 module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, b.value.at(1), tmp_bytes);
344 module.vec_znx_idft_tmp_a(c.at_mut(1), &mut res_dft);
345
346 // c[1] = AUTO(b, p)
347 module.vec_znx_big_automorphism_inplace(b.p, c.at_mut(1));
348}
349
350#[cfg(test)]
351mod test {
352 use super::{AutomorphismKey, automorphism};
353 use crate::{
354 ciphertext::Ciphertext,
355 decryptor::decrypt_rlwe,
356 elem::ElemCommon,
357 encryptor::encrypt_rlwe_sk,
358 keys::SecretKey,
359 parameters::{Parameters, ParametersLiteral},
360 plaintext::Plaintext,
361 };
362 use base2k::{
363 BACKEND, Encoding, Module, SvpPPol, SvpPPolOps, VecZnx, VecZnxOps, alloc_aligned,
364 };
365 use sampling::source::{Source, new_seed};
366
367 #[test]
368 fn test_automorphism() {
369 let log_base2k: usize = 10;
370 let log_q: usize = 50;
371 let log_p: usize = 15;
372
373 // Basic parameters with enough limbs to test edge cases
374 let params_lit: ParametersLiteral = ParametersLiteral {
375 backend: BACKEND::FFT64,
376 log_n: 12,
377 log_q: log_q,
378 log_p: log_p,
379 log_base2k: log_base2k,
380 log_scale: 20,
381 xe: 3.2,
382 xs: 1 << 11,
383 };
384
385 let params: Parameters = Parameters::new(¶ms_lit);
386
387 let module: &Module = params.module();
388 let log_q: usize = params.log_q();
389 let log_qp: usize = params.log_qp();
390 let gct_rows: usize = params.cols_q();
391 let gct_cols: usize = params.cols_qp();
392
393 // scratch space
394 let mut tmp_bytes: Vec<u8> = alloc_aligned(
395 params.decrypt_rlwe_tmp_byte(log_q)
396 | params.encrypt_rlwe_sk_tmp_bytes(log_q)
397 | params.automorphism_key_new_tmp_bytes(gct_rows, log_qp)
398 | params.automorphism_tmp_bytes(log_q, log_q, log_qp),
399 );
400
401 // Samplers for public and private randomness
402 let mut source_xe: Source = Source::new(new_seed());
403 let mut source_xa: Source = Source::new(new_seed());
404 let mut source_xs: Source = Source::new(new_seed());
405
406 let mut sk: SecretKey = SecretKey::new(module);
407 sk.fill_ternary_hw(params.xs(), &mut source_xs);
408 let mut sk_svp_ppol: SvpPPol = module.new_svp_ppol();
409 module.svp_prepare(&mut sk_svp_ppol, &sk.0);
410
411 let p: i64 = -5;
412
413 let auto_key: AutomorphismKey = AutomorphismKey::new(
414 module,
415 p,
416 &sk,
417 log_base2k,
418 gct_rows,
419 log_qp,
420 &mut source_xa,
421 &mut source_xe,
422 params.xe(),
423 &mut tmp_bytes,
424 );
425
426 let mut data: Vec<i64> = vec![0i64; params.n()];
427
428 data.iter_mut().enumerate().for_each(|(i, x)| *x = i as i64);
429
430 let log_k: usize = 2 * log_base2k;
431
432 let mut ct: Ciphertext<VecZnx> = params.new_ciphertext(log_q);
433 let mut pt: Plaintext = params.new_plaintext(log_q);
434 let mut pt_auto: Plaintext = params.new_plaintext(log_q);
435
436 pt.at_mut(0).encode_vec_i64(log_base2k, log_k, &data, 32);
437 module.vec_znx_automorphism(p, pt_auto.at_mut(0), pt.at(0));
438
439 encrypt_rlwe_sk(
440 module,
441 &mut ct.elem_mut(),
442 Some(pt.at(0)),
443 &sk_svp_ppol,
444 &mut source_xa,
445 &mut source_xe,
446 params.xe(),
447 &mut tmp_bytes,
448 );
449
450 let mut ct_auto: Ciphertext<VecZnx> = params.new_ciphertext(log_q);
451
452 // ct <- AUTO(ct)
453 automorphism(
454 module,
455 &mut ct_auto,
456 &ct,
457 &auto_key,
458 gct_cols,
459 &mut tmp_bytes,
460 );
461
462 // pt = dec(auto(ct)) - auto(pt)
463 decrypt_rlwe(
464 module,
465 pt.elem_mut(),
466 ct_auto.elem(),
467 &sk_svp_ppol,
468 &mut tmp_bytes,
469 );
470
471 module.vec_znx_sub_ba_inplace(pt.at_mut(0), pt_auto.at(0));
472
473 //pt.at(0).print(pt.cols(), 16);
474
475 let noise_have: f64 = pt.at(0).std(log_base2k).log2();
476
477 let var_msg: f64 = (params.xs() as f64) / params.n() as f64;
478 let var_a_err: f64 = 1f64 / 12f64;
479
480 let noise_pred: f64 =
481 params.noise_grlwe_product(var_msg, var_a_err, ct_auto.log_q(), auto_key.value.log_q());
482
483 println!("noise_pred: {}", noise_pred);
484 println!("noise_have: {}", noise_have);
485
486 assert!(noise_have <= noise_pred + 1.0);
487 }
488}