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;
13
14/// Stores DFT([-A*AUTO(s, -p) + 2^{-K*i}*s + E, A]) where AUTO(X, p): X^{i} -> X^{i*p}
15pub struct AutomorphismKey {
16 pub value: Ciphertext<VmpPMat>,
17 pub p: i64,
18}
19
20pub fn automorphis_key_new_tmp_bytes(
21 module: &Module,
22 log_base2k: usize,
23 rows: usize,
24 log_q: usize,
25) -> usize {
26 module.bytes_of_scalar()
27 + module.bytes_of_svp_ppol()
28 + encrypt_grlwe_sk_tmp_bytes(module, log_base2k, rows, log_q)
29}
30
31impl Parameters {
32 pub fn automorphism_key_new_tmp_bytes(&self, rows: usize, log_q: usize) -> usize {
33 automorphis_key_new_tmp_bytes(self.module(), self.log_base2k(), rows, log_q)
34 }
35}
36
37impl AutomorphismKey {
38 pub fn new(
39 module: &Module,
40 p: i64,
41 sk: &SecretKey,
42 log_base2k: usize,
43 rows: usize,
44 log_q: usize,
45 source_xa: &mut Source,
46 source_xe: &mut Source,
47 sigma: f64,
48 tmp_bytes: &mut [u8],
49 ) -> Self {
50 let (sk_auto_bytes, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_scalar());
51 let (sk_out_bytes, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_svp_ppol());
52
53 let sk_auto: Scalar = module.new_scalar_from_bytes_borrow(sk_auto_bytes);
54 let mut sk_out: SvpPPol = module.new_svp_ppol_from_bytes_borrow(sk_out_bytes);
55 let mut value: Ciphertext<VmpPMat> = new_gadget_ciphertext(module, log_base2k, rows, log_q);
56
57 let p_inv: i64 = module.galois_element_inv(p);
58 module.vec_znx_automorphism(p_inv, &mut sk_auto.as_vec_znx(), &sk.0.as_vec_znx());
59 module.svp_prepare(&mut sk_out, &sk_auto);
60 encrypt_grlwe_sk(
61 module, &mut value, &sk.0, &sk_out, source_xa, source_xe, sigma, tmp_bytes,
62 );
63
64 Self { value: value, p: p }
65 }
66}
67
68pub fn automorphism_tmp_bytes(
69 module: &Module,
70 c_cols: usize,
71 a_cols: usize,
72 b_rows: usize,
73 b_cols: usize,
74) -> usize {
75 return module.vmp_apply_dft_to_dft_tmp_bytes(c_cols, a_cols, b_rows, b_cols)
76 + 2 * module.bytes_of_vec_znx_dft(std::cmp::min(c_cols, a_cols));
77}
78
79pub fn automorphism(
80 module: &Module,
81 c: &mut Ciphertext<VecZnx>,
82 a: &Ciphertext<VecZnx>,
83 b: &AutomorphismKey,
84 tmp_bytes: &mut [u8],
85) {
86 let cols = std::cmp::min(c.cols(), a.cols());
87
88 #[cfg(debug_assertions)]
89 {
90 assert!(
91 tmp_bytes.len()
92 >= automorphism_tmp_bytes(
93 module,
94 c.cols(),
95 a.cols(),
96 b.value.rows(),
97 b.value.cols()
98 )
99 );
100 assert_alignement(tmp_bytes.as_ptr());
101 }
102
103 let (tmp_bytes_b1_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
104 let (tmp_bytes_res_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
105
106 let mut a1_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_b1_dft);
107 let mut res_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_res_dft);
108 let mut res_big: VecZnxBig = res_dft.as_vec_znx_big();
109
110 // a1_dft = DFT(a[1])
111 module.vec_znx_dft(&mut a1_dft, a.at(1));
112
113 // res_dft = IDFT(<DFT(a), DFT([-A*AUTO(s, -p) + 2^{-K*i}*s + E])>) = [-b*AUTO(s, -p) + a * s + e]
114 module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, b.value.at(0), tmp_bytes);
115 module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft);
116
117 // res_dft = [-b*AUTO(s, -p) + a * s + e] + [-a * s + m + e] = [-b*AUTO(s, -p) + m + e]
118 module.vec_znx_big_add_small_inplace(&mut res_big, a.at(0));
119
120 // c[0] = NORMALIZE([-b*AUTO(s, -p) + m + e])
121 module.vec_znx_big_normalize(c.log_base2k(), c.at_mut(0), &mut res_big, tmp_bytes);
122
123 // c[0] = AUTO([-b*AUTO(s, -p) + m + e], p) = [-AUTO(b, p)*s + AUTO(m, p) + AUTO(b, e)]
124 module.vec_znx_automorphism_inplace(b.p, c.at_mut(0));
125
126 // res_dft = IDFT(<DFT(a), DFT([A])>) = [b]
127 module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, b.value.at(1), tmp_bytes);
128 module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft);
129
130 // c[1] = b
131 module.vec_znx_big_normalize(c.log_base2k(), c.at_mut(1), &mut res_big, tmp_bytes);
132
133 // c[1] = AUTO(b, p)
134 module.vec_znx_automorphism_inplace(b.p, c.at_mut(1));
135}
136
137pub fn automorphism_big(
138 module: &Module,
139 c: &mut Ciphertext<VecZnxBig>,
140 a: &Ciphertext<VecZnx>,
141 b: &AutomorphismKey,
142 tmp_bytes: &mut [u8],
143) {
144 let cols = std::cmp::min(c.cols(), a.cols());
145
146 #[cfg(debug_assertions)]
147 {
148 assert!(
149 tmp_bytes.len()
150 >= automorphism_tmp_bytes(
151 module,
152 c.cols(),
153 a.cols(),
154 b.value.rows(),
155 b.value.cols()
156 )
157 );
158 assert_alignement(tmp_bytes.as_ptr());
159 }
160
161 let (tmp_bytes_b1_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
162 let (tmp_bytes_res_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
163
164 let mut a1_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_b1_dft);
165 let mut res_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_res_dft);
166
167 // a1_dft = DFT(a[1])
168 module.vec_znx_dft(&mut a1_dft, a.at(1));
169
170 // res_dft = IDFT(<DFT(a), DFT([-A*AUTO(s, -p) + 2^{-K*i}*s + E])>) = [-b*AUTO(s, -p) + a * s + e]
171 module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, b.value.at(0), tmp_bytes);
172 module.vec_znx_idft_tmp_a(c.at_mut(0), &mut res_dft);
173
174 // res_dft = [-b*AUTO(s, -p) + a * s + e] + [-a * s + m + e] = [-b*AUTO(s, -p) + m + e]
175 module.vec_znx_big_add_small_inplace(c.at_mut(0), a.at(0));
176
177 // c[0] = AUTO([-b*AUTO(s, -p) + m + e], p) = [-AUTO(b, p)*s + AUTO(m, p) + AUTO(b, e)]
178 module.vec_znx_big_automorphism_inplace(b.p, c.at_mut(0));
179
180 // res_dft = IDFT(<DFT(a), DFT([A])>) = [b]
181 module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, b.value.at(1), tmp_bytes);
182 module.vec_znx_idft_tmp_a(c.at_mut(1), &mut res_dft);
183
184 // c[1] = AUTO(b, p)
185 module.vec_znx_big_automorphism_inplace(b.p, c.at_mut(1));
186}
187
188#[cfg(test)]
189mod test {
190 use crate::{
191 ciphertext::{new_gadget_ciphertext, Ciphertext}, decryptor::decrypt_rlwe, elem::{Elem, ElemCommon, ElemVecZnx}, encryptor::encrypt_rlwe_sk, keys::SecretKey, parameters::{Parameters, ParametersLiteral}, plaintext::Plaintext
192 };
193 use base2k::{
194 alloc_aligned, Encoding, Infos, Sampling, SvpPPolOps, VecZnx, VecZnxBig, VecZnxBigOps, VecZnxDft, VecZnxDftOps, VecZnxOps, VmpPMat, BACKEND
195 };
196 use sampling::source::{Source, new_seed};
197
198 use super::{AutomorphismKey, automorphis_key_new_tmp_bytes};
199
200 #[test]
201 fn test_automorphism() {
202 let log_base2k: usize = 10;
203 let q_cols: usize = 4;
204 let p_cols: usize = 1;
205
206 // Basic parameters with enough limbs to test edge cases
207 let params_lit: ParametersLiteral = ParametersLiteral {
208 backend: BACKEND::FFT64,
209 log_n: 12,
210 log_q: q_cols * log_base2k,
211 log_p: p_cols * log_base2k,
212 log_base2k: log_base2k,
213 log_scale: 20,
214 xe: 3.2,
215 xs: 1 << 11,
216 };
217
218 let params: Parameters = Parameters::new(¶ms_lit);
219
220 let rows: usize = params.cols_q();
221
222 // scratch space
223 let mut tmp_bytes: Vec<u8> = alloc_aligned(
224 params.decrypt_rlwe_tmp_byte(params.log_q())
225 | params.encrypt_rlwe_sk_tmp_bytes(params.log_q())
226 | params.gadget_product_tmp_bytes(
227 params.log_qp(),
228 params.log_qp(),
229 params.cols_qp(),
230 params.log_qp(),
231 )
232 | params.encrypt_grlwe_sk_tmp_bytes(rows, params.log_qp())
233 | params.automorphism_key_new_tmp_bytes(rows, params.log_qp()),
234 );
235
236 // Samplers for public and private randomness
237 let mut source_xe: Source = Source::new(new_seed());
238 let mut source_xa: Source = Source::new(new_seed());
239 let mut source_xs: Source = Source::new(new_seed());
240
241 // Two secret keys
242 let mut sk: SecretKey = SecretKey::new(params.module());
243 sk.fill_ternary_hw(params.xs(), &mut source_xs);
244 let mut sk_svp_ppol: base2k::SvpPPol = params.module().new_svp_ppol();
245 params.module().svp_prepare(&mut sk_svp_ppol, &sk.0);
246
247 let p: i64 = -5;
248
249 let auto_key: AutomorphismKey = AutomorphismKey::new(
250 params.module(),
251 p,
252 &sk,
253 params.log_base2k(),
254 rows,
255 params.log_qp(),
256 &mut source_xa,
257 &mut source_xe,
258 params.xe(),
259 &mut tmp_bytes,
260 );
261
262 let data: Vec<i64> = vec![0i64; params.n()];
263
264 let mut ct: Ciphertext<VecZnx> = Ciphertext::new(params.module(), params.log_base2k(), params.log_q(), 2);
265 let mut pt: Plaintext = Plaintext::new(params.module(), params.log_base2k(), params.log_q());
266
267 pt.at_mut(0).encode_vec_i64(params.log_base2k(), 2*params.log_base2k(), &data, 32);
268
269 encrypt_rlwe_sk(params.module(), &mut ct.elem_mut(), Some(&pt.elem()), &sk_svp_ppol, &mut source_xa, &mut source_xe, params.xe(), &mut tmp_bytes);
270
271 }
272}