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