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::{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(1, 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(1, cols));
 70    let (tmp_bytes_res_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(1, cols));
 71
 72    let mut a1_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(1, cols, tmp_bytes_b1_dft);
 73    let mut res_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(1, 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(1, cols));
103    let (tmp_bytes_res_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(1, cols));
104
105    let mut a1_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(1, cols, tmp_bytes_b1_dft);
106    let mut res_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(1, 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(&params_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(1, 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(1, 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(1, 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(1, 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(1, b_cols);
240                let mut res_dft_1: VecZnxDft = params.module().new_vec_znx_dft(1, 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(0, 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}