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