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