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;
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
38impl AutomorphismKey {
39 pub fn new(
40 module: &Module,
41 p: i64,
42 sk: &SecretKey,
43 log_base2k: usize,
44 rows: usize,
45 log_q: usize,
46 source_xa: &mut Source,
47 source_xe: &mut Source,
48 sigma: f64,
49 tmp_bytes: &mut [u8],
50 ) -> Self {
51 let (sk_auto_bytes, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_scalar());
52 let (sk_out_bytes, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_svp_ppol());
53
54 let sk_auto: Scalar = module.new_scalar_from_bytes_borrow(sk_auto_bytes);
55 let mut sk_out: SvpPPol = module.new_svp_ppol_from_bytes_borrow(sk_out_bytes);
56 let mut value: Ciphertext<VmpPMat> = new_gadget_ciphertext(module, log_base2k, rows, log_q);
57
58 let p_inv: i64 = module.galois_element_inv(p);
59
60 module.vec_znx_automorphism(p_inv, &mut sk_auto.as_vec_znx(), &sk.0.as_vec_znx());
61 module.svp_prepare(&mut sk_out, &sk_auto);
62 encrypt_grlwe_sk(
63 module, &mut value, &sk.0, &sk_out, source_xa, source_xe, sigma, tmp_bytes,
64 );
65
66 Self { value: value, p: p }
67 }
68}
69
70pub fn automorphism_tmp_bytes(
71 module: &Module,
72 c_cols: usize,
73 a_cols: usize,
74 b_rows: usize,
75 b_cols: usize,
76) -> usize {
77 return module.vmp_apply_dft_to_dft_tmp_bytes(c_cols, a_cols, b_rows, b_cols)
78 + 2 * module.bytes_of_vec_znx_dft(std::cmp::min(c_cols, a_cols));
79}
80
81pub fn automorphism(
82 module: &Module,
83 c: &mut Ciphertext<VecZnx>,
84 a: &Ciphertext<VecZnx>,
85 b: &AutomorphismKey,
86 tmp_bytes: &mut [u8],
87) {
88 let cols: usize = min(min(c.cols(), a.cols()), b.value.rows());
89
90 #[cfg(debug_assertions)]
91 {
92 assert!(
93 tmp_bytes.len()
94 >= automorphism_tmp_bytes(
95 module,
96 c.cols(),
97 a.cols(),
98 b.value.rows(),
99 b.value.cols()
100 )
101 );
102 assert_alignement(tmp_bytes.as_ptr());
103 }
104
105 let (tmp_bytes_b1_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
106 let (tmp_bytes_res_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
107
108 let mut a1_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_b1_dft);
109 let mut res_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_res_dft);
110 let mut res_big: VecZnxBig = res_dft.as_vec_znx_big();
111
112 // a1_dft = DFT(a[1])
113 module.vec_znx_dft(&mut a1_dft, a.at(1));
114
115 // res_dft = IDFT(<DFT(a), DFT([-A*AUTO(s, -p) + 2^{-K*i}*s + E])>) = [-b*AUTO(s, -p) + a * s + e]
116 module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, b.value.at(0), tmp_bytes);
117 module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft);
118
119 // res_dft = [-b*AUTO(s, -p) + a * s + e] + [-a * s + m + e] = [-b*AUTO(s, -p) + m + e]
120 module.vec_znx_big_add_small_inplace(&mut res_big, a.at(0));
121
122 // c[0] = NORMALIZE([-b*AUTO(s, -p) + m + e])
123 module.vec_znx_big_normalize(c.log_base2k(), c.at_mut(0), &mut res_big, tmp_bytes);
124
125 // c[0] = AUTO([-b*AUTO(s, -p) + m + e], p) = [-AUTO(b, p)*s + AUTO(m, p) + AUTO(b, e)]
126 module.vec_znx_automorphism_inplace(b.p, c.at_mut(0));
127
128 // res_dft = IDFT(<DFT(a), DFT([A])>) = [b]
129 module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, b.value.at(1), tmp_bytes);
130 module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft);
131
132 // c[1] = b
133 module.vec_znx_big_normalize(c.log_base2k(), c.at_mut(1), &mut res_big, tmp_bytes);
134
135 // c[1] = AUTO(b, p)
136 module.vec_znx_automorphism_inplace(b.p, c.at_mut(1));
137}
138
139pub fn automorphism_inplace_tmp_bytes(
140 module: &Module,
141 c_cols: usize,
142 a_cols: usize,
143 b_rows: usize,
144 b_cols: usize,
145) -> usize {
146 return module.vmp_apply_dft_to_dft_tmp_bytes(c_cols, a_cols, b_rows, b_cols)
147 + 2 * module.bytes_of_vec_znx_dft(std::cmp::min(c_cols, a_cols));
148}
149
150pub fn automorphism_inplace(
151 module: &Module,
152 a: &mut Ciphertext<VecZnx>,
153 b: &AutomorphismKey,
154 tmp_bytes: &mut [u8],
155) {
156 let cols: usize = min(a.cols(), b.value.rows());
157
158 #[cfg(debug_assertions)]
159 {
160 assert!(
161 tmp_bytes.len()
162 >= automorphism_inplace_tmp_bytes(
163 module,
164 a.cols(),
165 a.cols(),
166 b.value.rows(),
167 b.value.cols()
168 )
169 );
170 assert_alignement(tmp_bytes.as_ptr());
171 }
172
173 let (tmp_bytes_b1_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
174 let (tmp_bytes_res_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
175
176 let mut a1_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_b1_dft);
177 let mut res_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_res_dft);
178 let mut res_big: VecZnxBig = res_dft.as_vec_znx_big();
179
180 // a1_dft = DFT(a[1])
181 module.vec_znx_dft(&mut a1_dft, a.at(1));
182
183 // res_dft = IDFT(<DFT(a), DFT([-A*AUTO(s, -p) + 2^{-K*i}*s + E])>) = [-b*AUTO(s, -p) + a * s + e]
184 module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, b.value.at(0), tmp_bytes);
185 module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft);
186
187 // res_dft = [-b*AUTO(s, -p) + a * s + e] + [-a * s + m + e] = [-b*AUTO(s, -p) + m + e]
188 module.vec_znx_big_add_small_inplace(&mut res_big, a.at(0));
189
190 // a[0] = NORMALIZE([-b*AUTO(s, -p) + m + e])
191 module.vec_znx_big_normalize(a.log_base2k(), a.at_mut(0), &mut res_big, tmp_bytes);
192
193 // a[0] = AUTO([-b*AUTO(s, -p) + m + e], p) = [-AUTO(b, p)*s + AUTO(m, p) + AUTO(b, e)]
194 module.vec_znx_automorphism_inplace(b.p, a.at_mut(0));
195
196 // res_dft = IDFT(<DFT(a), DFT([A])>) = [b]
197 module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, b.value.at(1), tmp_bytes);
198 module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft);
199
200 // a[1] = b
201 module.vec_znx_big_normalize(a.log_base2k(), a.at_mut(1), &mut res_big, tmp_bytes);
202
203 // a[1] = AUTO(b, p)
204 module.vec_znx_automorphism_inplace(b.p, a.at_mut(1));
205}
206
207pub fn automorphism_big(
208 module: &Module,
209 c: &mut Ciphertext<VecZnxBig>,
210 a: &Ciphertext<VecZnx>,
211 b: &AutomorphismKey,
212 tmp_bytes: &mut [u8],
213) {
214 let cols = std::cmp::min(c.cols(), a.cols());
215
216 #[cfg(debug_assertions)]
217 {
218 assert!(
219 tmp_bytes.len()
220 >= automorphism_tmp_bytes(
221 module,
222 c.cols(),
223 a.cols(),
224 b.value.rows(),
225 b.value.cols()
226 )
227 );
228 assert_alignement(tmp_bytes.as_ptr());
229 }
230
231 let (tmp_bytes_b1_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
232 let (tmp_bytes_res_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
233
234 let mut a1_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_b1_dft);
235 let mut res_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_res_dft);
236
237 // a1_dft = DFT(a[1])
238 module.vec_znx_dft(&mut a1_dft, a.at(1));
239
240 // res_dft = IDFT(<DFT(a), DFT([-A*AUTO(s, -p) + 2^{-K*i}*s + E])>) = [-b*AUTO(s, -p) + a * s + e]
241 module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, b.value.at(0), tmp_bytes);
242 module.vec_znx_idft_tmp_a(c.at_mut(0), &mut res_dft);
243
244 // res_dft = [-b*AUTO(s, -p) + a * s + e] + [-a * s + m + e] = [-b*AUTO(s, -p) + m + e]
245 module.vec_znx_big_add_small_inplace(c.at_mut(0), a.at(0));
246
247 // c[0] = AUTO([-b*AUTO(s, -p) + m + e], p) = [-AUTO(b, p)*s + AUTO(m, p) + AUTO(b, e)]
248 module.vec_znx_big_automorphism_inplace(b.p, c.at_mut(0));
249
250 // res_dft = IDFT(<DFT(a), DFT([A])>) = [b]
251 module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, b.value.at(1), tmp_bytes);
252 module.vec_znx_idft_tmp_a(c.at_mut(1), &mut res_dft);
253
254 // c[1] = AUTO(b, p)
255 module.vec_znx_big_automorphism_inplace(b.p, c.at_mut(1));
256}
257
258#[cfg(test)]
259mod test {
260 use crate::{
261 ciphertext::{new_gadget_ciphertext, Ciphertext}, decryptor::decrypt_rlwe, elem::{Elem, ElemCommon, ElemVecZnx}, encryptor::encrypt_rlwe_sk, keys::SecretKey, parameters::{Parameters, ParametersLiteral}, plaintext::Plaintext
262 };
263 use base2k::{
264 alloc_aligned, Encoding, Infos, Sampling, SvpPPolOps, VecZnx, VecZnxBig, VecZnxBigOps, VecZnxDft, VecZnxDftOps, VecZnxOps, VmpPMat, BACKEND
265 };
266 use sampling::source::{Source, new_seed};
267
268 use super::{automorphis_key_new_tmp_bytes, automorphism, AutomorphismKey};
269
270 #[test]
271 fn test_automorphism() {
272 let log_base2k: usize = 10;
273 let q_cols: usize = 4;
274 let p_cols: usize = 1;
275
276 // Basic parameters with enough limbs to test edge cases
277 let params_lit: ParametersLiteral = ParametersLiteral {
278 backend: BACKEND::FFT64,
279 log_n: 12,
280 log_q: q_cols * log_base2k,
281 log_p: p_cols * log_base2k,
282 log_base2k: log_base2k,
283 log_scale: 20,
284 xe: 3.2,
285 xs: 1 << 11,
286 };
287
288 let params: Parameters = Parameters::new(¶ms_lit);
289
290 let module: &base2k::Module = params.module();
291 let log_q: usize = params.log_q();
292 let log_qp: usize = params.log_qp();
293 let rows: usize = params.cols_q();
294
295 // scratch space
296 let mut tmp_bytes: Vec<u8> = alloc_aligned(
297 params.decrypt_rlwe_tmp_byte(log_q)
298 | params.encrypt_rlwe_sk_tmp_bytes(log_q)
299 | params.gadget_product_tmp_bytes(
300 log_qp,
301 log_qp,
302 rows,
303 log_qp,
304 )
305 | params.encrypt_grlwe_sk_tmp_bytes(rows, log_qp)
306 | params.automorphism_key_new_tmp_bytes(rows, log_qp),
307 );
308
309 // Samplers for public and private randomness
310 let mut source_xe: Source = Source::new(new_seed());
311 let mut source_xa: Source = Source::new(new_seed());
312 let mut source_xs: Source = Source::new(new_seed());
313
314 // Two secret keys
315 let mut sk: SecretKey = SecretKey::new(module);
316 sk.fill_ternary_hw(params.xs(), &mut source_xs);
317 let mut sk_svp_ppol: base2k::SvpPPol = module.new_svp_ppol();
318 module.svp_prepare(&mut sk_svp_ppol, &sk.0);
319
320 let p: i64 = -5;
321
322 let auto_key: AutomorphismKey = AutomorphismKey::new(
323 module,
324 p,
325 &sk,
326 log_base2k,
327 rows,
328 log_qp,
329 &mut source_xa,
330 &mut source_xe,
331 params.xe(),
332 &mut tmp_bytes,
333 );
334
335 let mut data: Vec<i64> = vec![0i64; params.n()];
336
337 data.iter_mut().enumerate().for_each(|(i, x)|{
338 *x = i as i64
339 });
340
341 let log_k: usize = 2*log_base2k;
342
343 let mut ct: Ciphertext<VecZnx> = Ciphertext::new(module, log_base2k, log_q, 2);
344 let mut pt: Plaintext = Plaintext::new(module, log_base2k, log_q);
345
346 pt.at_mut(0).encode_vec_i64(log_base2k, log_k, &data, 32);
347
348 encrypt_rlwe_sk(module, &mut ct.elem_mut(), Some(pt.elem()), &sk_svp_ppol, &mut source_xa, &mut source_xe, params.xe(), &mut tmp_bytes);
349
350 module.vec_znx_automorphism_inplace(p, pt.at_mut(0));
351
352 let mut ct_auto: Ciphertext<VecZnx> = Ciphertext::new(module, log_base2k, log_q, 2);
353
354 automorphism(module, &mut ct_auto, &ct, &auto_key, &mut tmp_bytes);
355
356 module.vec_znx_sub_inplace(ct_auto.at_mut(0), pt.at(0));
357 ct_auto.at_mut(0).normalize(log_base2k, &mut tmp_bytes);
358
359 decrypt_rlwe(module, pt.elem_mut(), ct_auto.elem(), &sk_svp_ppol, &mut tmp_bytes);
360
361 let noise_have: f64 = pt.at(0).std(log_base2k).log2();
362
363 let var_a_err: f64;
364 if ct_auto.cols() < ct.cols() {
365 var_a_err = 1f64 / 12f64;
366 } else {
367 var_a_err = 0f64;
368 }
369
370 let var_msg: f64 = (params.xs() as f64) / params.n() as f64;
371
372 let noise_pred: f64 =
373 params.noise_grlwe_product(var_msg, var_a_err, ct_auto.log_q(), auto_key.value.log_q());
374
375 println!("noise_pred: {}", noise_have);
376 println!("noise_have: {}", noise_pred);
377
378 assert!(noise_have <= noise_pred + 1.0);
379
380
381
382 }
383}