poulpy/rlwe/src/gadget_product.rs
rlwe/src/gadget_product.rsRaw
1use crate::{ciphertext::Ciphertext, elem::ElemCommon, parameters::Parameters};
2use base2k::{Module, VecZnx, VecZnxBig, VecZnxBigOps, VecZnxDft, VecZnxDftOps, VmpPMat, VmpPMatOps};
3use std::cmp::min;
4
5pub fn gadget_product_core_tmp_bytes(
6 module: &Module,
7 log_base2k: usize,
8 res_log_q: usize,
9 in_log_q: usize,
10 gct_rows: usize,
11 gct_log_q: usize,
12) -> usize {
13 let gct_cols: usize = (gct_log_q + log_base2k - 1) / log_base2k;
14 let in_cols: usize = (in_log_q + log_base2k - 1) / log_base2k;
15 let out_cols: usize = (res_log_q + log_base2k - 1) / log_base2k;
16 module.vmp_apply_dft_to_dft_tmp_bytes(out_cols, in_cols, gct_rows, gct_cols)
17}
18
19impl Parameters {
20 pub fn gadget_product_tmp_bytes(&self, res_log_q: usize, in_log_q: usize, gct_rows: usize, gct_log_q: usize) -> usize {
21 gadget_product_core_tmp_bytes(
22 self.module(),
23 self.log_base2k(),
24 res_log_q,
25 in_log_q,
26 gct_rows,
27 gct_log_q,
28 )
29 }
30}
31
32pub fn gadget_product_core(
33 module: &Module,
34 res_dft_0: &mut VecZnxDft,
35 res_dft_1: &mut VecZnxDft,
36 a: &VecZnx,
37 b: &Ciphertext<VmpPMat>,
38 b_cols: usize,
39 tmp_bytes: &mut [u8],
40) {
41 assert!(b_cols <= b.cols());
42 module.vec_znx_dft(res_dft_1, a);
43 module.vmp_apply_dft_to_dft(res_dft_0, res_dft_1, b.at(0), tmp_bytes);
44 module.vmp_apply_dft_to_dft_inplace(res_dft_1, b.at(1), tmp_bytes);
45}
46
47pub fn gadget_product_big_tmp_bytes(module: &Module, c_cols: usize, a_cols: usize, b_rows: usize, b_cols: usize) -> usize {
48 return module.vmp_apply_dft_to_dft_tmp_bytes(c_cols, a_cols, b_rows, b_cols)
49 + 2 * module.bytes_of_vec_znx_dft(min(c_cols, a_cols));
50}
51
52/// Evaluates the gadget product: c.at(i) = IDFT(<DFT(a.at(i)), b.at(i)>)
53///
54/// # Arguments
55///
56/// * `module`: backend support for operations mod (X^N + 1).
57/// * `c`: a [Ciphertext<VecZnxBig>] with cols_c cols.
58/// * `a`: a [Ciphertext<VecZnx>] with cols_a cols.
59/// * `b`: a [Ciphertext<VmpPMat>] with at least min(cols_c, cols_a) rows.
60pub fn gadget_product_big(
61 module: &Module,
62 c: &mut Ciphertext<VecZnxBig>,
63 a: &Ciphertext<VecZnx>,
64 b: &Ciphertext<VmpPMat>,
65 tmp_bytes: &mut [u8],
66) {
67 let cols: usize = min(c.cols(), a.cols());
68
69 let (tmp_bytes_b1_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
70 let (tmp_bytes_res_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
71
72 let mut a1_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_b1_dft);
73 let mut res_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_res_dft);
74
75 // a1_dft = DFT(a[1])
76 module.vec_znx_dft(&mut a1_dft, a.at(1));
77
78 // c[i] = IDFT(DFT(a[1]) * b[i])
79 (0..2).for_each(|i| {
80 module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, b.at(i), tmp_bytes);
81 module.vec_znx_idft_tmp_a(c.at_mut(i), &mut res_dft);
82 })
83}
84
85/// Evaluates the gadget product: c.at(i) = NORMALIZE(IDFT(<DFT(a.at(i)), b.at(i)>)
86///
87/// # Arguments
88///
89/// * `module`: backend support for operations mod (X^N + 1).
90/// * `c`: a [Ciphertext<VecZnx>] with cols_c cols.
91/// * `a`: a [Ciphertext<VecZnx>] with cols_a cols.
92/// * `b`: a [Ciphertext<VmpPMat>] with at least min(cols_c, cols_a) rows.
93pub fn gadget_product(
94 module: &Module,
95 c: &mut Ciphertext<VecZnx>,
96 a: &Ciphertext<VecZnx>,
97 b: &Ciphertext<VmpPMat>,
98 tmp_bytes: &mut [u8],
99) {
100 let cols: usize = min(c.cols(), a.cols());
101
102 let (tmp_bytes_b1_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
103 let (tmp_bytes_res_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
104
105 let mut a1_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_b1_dft);
106 let mut res_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_res_dft);
107 let mut res_big: VecZnxBig = res_dft.as_vec_znx_big();
108
109 // a1_dft = DFT(a[1])
110 module.vec_znx_dft(&mut a1_dft, a.at(1));
111
112 // c[i] = IDFT(DFT(a[1]) * b[i])
113 (0..2).for_each(|i| {
114 module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, b.at(i), tmp_bytes);
115 module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft);
116 module.vec_znx_big_normalize(c.log_base2k(), c.at_mut(i), &mut res_big, tmp_bytes);
117 })
118}
119
120#[cfg(test)]
121mod test {
122 use crate::{
123 ciphertext::{Ciphertext, new_gadget_ciphertext},
124 decryptor::decrypt_rlwe,
125 elem::{Elem, ElemCommon, ElemVecZnx},
126 encryptor::encrypt_grlwe_sk,
127 gadget_product::gadget_product_core,
128 keys::SecretKey,
129 parameters::{Parameters, ParametersLiteral},
130 plaintext::Plaintext,
131 };
132 use base2k::{
133 BACKEND, Infos, Sampling, SvpPPolOps, VecZnx, VecZnxBig, VecZnxBigOps, VecZnxDft, VecZnxDftOps, VecZnxOps, VmpPMat,
134 alloc_aligned_u8,
135 };
136 use sampling::source::{Source, new_seed};
137
138 #[test]
139 fn test_gadget_product_core() {
140 let log_base2k: usize = 10;
141 let q_cols: usize = 7;
142 let p_cols: usize = 1;
143
144 // Basic parameters with enough limbs to test edge cases
145 let params_lit: ParametersLiteral = ParametersLiteral {
146 backend: BACKEND::FFT64,
147 log_n: 12,
148 log_q: q_cols * log_base2k,
149 log_p: p_cols * log_base2k,
150 log_base2k: log_base2k,
151 log_scale: 20,
152 xe: 3.2,
153 xs: 1 << 11,
154 };
155
156 let params: Parameters = Parameters::new(¶ms_lit);
157
158 // scratch space
159 let mut tmp_bytes: Vec<u8> = alloc_aligned_u8(
160 params.decrypt_rlwe_tmp_byte(params.log_qp())
161 | params.gadget_product_tmp_bytes(
162 params.log_qp(),
163 params.log_qp(),
164 params.cols_qp(),
165 params.log_qp(),
166 )
167 | params.encrypt_grlwe_sk_tmp_bytes(params.cols_qp(), params.log_qp()),
168 );
169
170 // Samplers for public and private randomness
171 let mut source_xe: Source = Source::new(new_seed());
172 let mut source_xa: Source = Source::new(new_seed());
173 let mut source_xs: Source = Source::new(new_seed());
174
175 // Two secret keys
176 let mut sk0: SecretKey = SecretKey::new(params.module());
177 sk0.fill_ternary_hw(params.xs(), &mut source_xs);
178 let mut sk0_svp_ppol: base2k::SvpPPol = params.module().new_svp_ppol();
179 params.module().svp_prepare(&mut sk0_svp_ppol, &sk0.0);
180
181 let mut sk1: SecretKey = SecretKey::new(params.module());
182 sk1.fill_ternary_hw(params.xs(), &mut source_xs);
183 let mut sk1_svp_ppol: base2k::SvpPPol = params.module().new_svp_ppol();
184 params.module().svp_prepare(&mut sk1_svp_ppol, &sk1.0);
185
186 // The gadget ciphertext
187 let mut gadget_ct: Ciphertext<VmpPMat> = new_gadget_ciphertext(
188 params.module(),
189 log_base2k,
190 params.cols_qp(),
191 params.log_qp(),
192 );
193
194 // gct = [-b*sk1 + g(sk0) + e, b]
195 encrypt_grlwe_sk(
196 params.module(),
197 &mut gadget_ct,
198 &sk0.0,
199 &sk1_svp_ppol,
200 &mut source_xa,
201 &mut source_xe,
202 params.xe(),
203 &mut tmp_bytes,
204 );
205
206 // Intermediate buffers
207
208 // Input polynopmial, uniformly distributed
209 let mut a: VecZnx = params.module().new_vec_znx(params.cols_q());
210 params
211 .module()
212 .fill_uniform(log_base2k, &mut a, params.cols_q(), &mut source_xa);
213
214 // res = g^-1(a) * gct
215 let mut elem_res: Elem<VecZnx> = Elem::<VecZnx>::new(params.module(), log_base2k, params.log_qp(), 2);
216
217 // Ideal output = a * s
218 let mut a_dft: VecZnxDft = params.module().new_vec_znx_dft(a.cols());
219 let mut a_big: VecZnxBig = a_dft.as_vec_znx_big();
220 let mut a_times_s: VecZnx = params.module().new_vec_znx(a.cols());
221
222 // a * sk0
223 params.module().svp_apply_dft(&mut a_dft, &sk0_svp_ppol, &a);
224 params.module().vec_znx_idft_tmp_a(&mut a_big, &mut a_dft);
225 params
226 .module()
227 .vec_znx_big_normalize(params.log_base2k(), &mut a_times_s, &a_big, &mut tmp_bytes);
228
229 // Plaintext for decrypted output of gadget product
230 let mut pt: Plaintext = Plaintext::new(params.module(), params.log_base2k(), params.log_qp());
231
232 // Iterates over all possible cols values for input/output polynomials and gadget ciphertext.
233
234 (1..a.cols() + 1).for_each(|a_cols| {
235 let mut a_trunc: VecZnx = params.module().new_vec_znx(a_cols);
236 a_trunc.copy_from(&a);
237
238 (1..gadget_ct.cols() + 1).for_each(|b_cols| {
239 let mut res_dft_0: VecZnxDft = params.module().new_vec_znx_dft(b_cols);
240 let mut res_dft_1: VecZnxDft = params.module().new_vec_znx_dft(b_cols);
241 let mut res_big_0: VecZnxBig = res_dft_0.as_vec_znx_big();
242 let mut res_big_1: VecZnxBig = res_dft_1.as_vec_znx_big();
243
244 pt.elem_mut().zero();
245 elem_res.zero();
246
247 // let b_cols: usize = min(a_cols+1, gadget_ct.cols());
248
249 println!("a_cols: {} b_cols: {}", a_cols, b_cols);
250
251 // res_dft_0 = DFT(gct_[0] * ct[1] = a * (-bs' + s + e) = -cs' + as + e')
252 // res_dft_1 = DFT(gct_[1] * ct[1] = a * b = c)
253 gadget_product_core(
254 params.module(),
255 &mut res_dft_0,
256 &mut res_dft_1,
257 &a_trunc,
258 &gadget_ct,
259 b_cols,
260 &mut tmp_bytes,
261 );
262
263 // res_big_0 = IDFT(res_dft_0)
264 params
265 .module()
266 .vec_znx_idft_tmp_a(&mut res_big_0, &mut res_dft_0);
267 // res_big_1 = IDFT(res_dft_1);
268 params
269 .module()
270 .vec_znx_idft_tmp_a(&mut res_big_1, &mut res_dft_1);
271
272 // res_big_0 = normalize(res_big_0)
273 params
274 .module()
275 .vec_znx_big_normalize(log_base2k, elem_res.at_mut(0), &res_big_0, &mut tmp_bytes);
276
277 // res_big_1 = normalize(res_big_1)
278 params
279 .module()
280 .vec_znx_big_normalize(log_base2k, elem_res.at_mut(1), &res_big_1, &mut tmp_bytes);
281
282 // <(-c*sk1 + a*sk0 + e, a), (1, sk1)> = a*sk0 + e
283 decrypt_rlwe(
284 params.module(),
285 pt.elem_mut(),
286 &elem_res,
287 &sk1_svp_ppol,
288 &mut tmp_bytes,
289 );
290
291 // a * sk0 + e - a*sk0 = e
292 params
293 .module()
294 .vec_znx_sub_ab_inplace(pt.at_mut(0), &mut a_times_s);
295 pt.at_mut(0).normalize(log_base2k, &mut tmp_bytes);
296
297 // pt.at(0).print(pt.elem().cols(), 16);
298
299 let noise_have: f64 = pt.at(0).std(log_base2k).log2();
300
301 let var_a_err: f64;
302
303 if a_cols < a.cols() {
304 var_a_err = 1f64 / 12f64;
305 } else {
306 var_a_err = 0f64;
307 }
308
309 let a_logq: usize = a_cols * log_base2k;
310 let b_logq: usize = b_cols * log_base2k;
311 let var_msg: f64 = (params.xs() as f64) / params.n() as f64;
312
313 println!("{} {} {} {}", var_msg, var_a_err, a_logq, b_logq);
314
315 let noise_pred: f64 = params.noise_grlwe_product(var_msg, var_a_err, a_logq, b_logq);
316
317 println!("noise_pred: {}", noise_pred);
318 println!("noise_have: {}", noise_have);
319
320 // assert!(noise_have <= noise_pred + 1.0);
321 });
322 });
323 }
324}
325
326impl Parameters {
327 pub fn noise_grlwe_product(&self, var_msg: f64, var_a_err: f64, a_logq: usize, b_logq: usize) -> f64 {
328 let n: f64 = self.n() as f64;
329 let var_xs: f64 = self.xs() as f64;
330
331 let var_gct_err_lhs: f64;
332 let var_gct_err_rhs: f64;
333 if b_logq < self.log_qp() {
334 let var_round: f64 = 1f64 / 12f64;
335 var_gct_err_lhs = var_round;
336 var_gct_err_rhs = var_round;
337 } else {
338 var_gct_err_lhs = self.xe() * self.xe();
339 var_gct_err_rhs = 0f64;
340 }
341
342 noise_grlwe_product(
343 n,
344 self.log_base2k(),
345 var_xs,
346 var_msg,
347 var_a_err,
348 var_gct_err_lhs,
349 var_gct_err_rhs,
350 a_logq,
351 b_logq,
352 )
353 }
354}
355
356pub fn noise_grlwe_product(
357 n: f64,
358 log_base2k: usize,
359 var_xs: f64,
360 var_msg: f64,
361 var_a_err: f64,
362 var_gct_err_lhs: f64,
363 var_gct_err_rhs: f64,
364 a_logq: usize,
365 b_logq: usize,
366) -> f64 {
367 let a_logq: usize = min(a_logq, b_logq);
368 let a_cols: usize = (a_logq + log_base2k - 1) / log_base2k;
369
370 let b_scale = 2.0f64.powi(b_logq as i32);
371 let a_scale: f64 = 2.0f64.powi((b_logq - a_logq) as i32);
372
373 let base: f64 = (1 << (log_base2k)) as f64;
374 let var_base: f64 = base * base / 12f64;
375
376 // lhs = a_cols * n * (var_base * var_gct_err_lhs + var_e_a * var_msg * p^2)
377 // rhs = a_cols * n * var_base * var_gct_err_rhs * var_xs
378 let mut noise: f64 = (a_cols as f64) * n * var_base * (var_gct_err_lhs + var_xs * var_gct_err_rhs);
379 noise += var_msg * var_a_err * a_scale * a_scale * n;
380 noise = noise.sqrt();
381 noise /= b_scale;
382 noise.log2().min(-1.0) // max noise is [-2^{-1}, 2^{-1}]
383}