fixed gadget product & vec_znx_big_add

base2k/src/vec_znx_big.rs

@@ -75,9 +75,9 @@ impl Module {
                 c.0,
                 c.limbs() as u64,
                 b.0,
-                a.limbs() as u64,
-                a.as_ptr(),
                 b.limbs() as u64,
+                a.as_ptr(),
+                a.limbs() as u64,
                 a.n() as u64,
             )
         }
@@ -91,7 +91,7 @@ impl Module {
                 b.0,
                 b.limbs() as u64,
                 b.0,
-                a.limbs() as u64,
+                b.limbs() as u64,
                 a.as_ptr(),
                 a.limbs() as u64,
                 a.n() as u64,

rlwe/examples/gadget_product.rs

@@ -1,4 +1,6 @@
-use base2k::{FFT64, Infos, Sampling, Scalar, SvpPPolOps, VecZnx, VecZnxBig, VecZnxDft, VecZnxOps};
+use base2k::{
+    Encoding, FFT64, Infos, Sampling, Scalar, SvpPPolOps, VecZnx, VecZnxBig, VecZnxDft, VecZnxOps,
+};
 use rlwe::{
     ciphertext::{Ciphertext, GadgetCiphertext},
     decryptor::{Decryptor, decrypt_rlwe_thread_safe, decrypt_rlwe_thread_safe_tmp_byte},
@@ -7,7 +9,7 @@ use rlwe::{
         EncryptorSk, encrypt_grlwe_sk_thread_safe, encrypt_grlwe_sk_tmp_bytes,
         encrypt_rlwe_sk_tmp_bytes,
     },
-    evaluator::{gadget_product_thread_safe, gadget_product_tmp_bytes},
+    evaluator::{gadget_product_inplace_thread_safe, gadget_product_tmp_bytes},
     key_generator::{gen_switching_key_thread_safe, gen_switching_key_thread_safe_tmp_bytes},
     keys::{SecretKey, SwitchingKey},
     parameters::{Parameters, ParametersLiteral},
@@ -56,9 +58,11 @@ fn main() {
 
     let mut source: Source = Source::new([3; 32]);
 
-    let mut sk: SecretKey = SecretKey::new(params.module());
+    let mut sk0: SecretKey = SecretKey::new(params.module());
+    let mut sk1: SecretKey = SecretKey::new(params.module());
 
-    sk.fill_ternary_hw(params.xs(), &mut source);
+    sk0.fill_ternary_hw(params.xs(), &mut source);
+    sk1.fill_ternary_hw(params.xs(), &mut source);
 
     let mut want = vec![i64::default(); params.n()];
 
@@ -71,8 +75,11 @@ fn main() {
     let mut source_xe: Source = Source::new([4; 32]);
     let mut source_xa: Source = Source::new([5; 32]);
 
-    let mut sk_svp_ppol: base2k::SvpPPol = params.module().new_svp_ppol();
+    let mut sk0_svp_ppol: base2k::SvpPPol = params.module().new_svp_ppol();
-    params.module().svp_prepare(&mut sk_svp_ppol, &sk.0);
+    params.module().svp_prepare(&mut sk0_svp_ppol, &sk0.0);
+
+    let mut sk1_svp_ppol: base2k::SvpPPol = params.module().new_svp_ppol();
+    params.module().svp_prepare(&mut sk1_svp_ppol, &sk1.0);
 
     let mut gadget_ct: GadgetCiphertext = GadgetCiphertext::new(
         params.module(),
@@ -81,67 +88,68 @@ fn main() {
         params.log_qp(),
     );
 
-    let mut m: Scalar = Scalar::new(params.n());
-    m.fill_ternary_prob(0.5, &mut source_xa);
-
     encrypt_grlwe_sk_thread_safe(
         params.module(),
         &mut gadget_ct,
-        &m,
+        &sk0.0,
-        &sk_svp_ppol,
+        &sk1_svp_ppol,
         &mut source_xa,
         &mut source_xe,
         params.xe(),
         &mut tmp_bytes,
     );
 
-    let mut res: Elem = Elem::new(params.module(), log_base2k, params.log_q(), 1, 0);
+    let mut pt: Plaintext = Plaintext::new(
-    let mut a: VecZnx = VecZnx::new(params.module().n(), params.limbs_q());
-    a.fill_uniform(params.log_base2k(), a.limbs(), &mut source_xa);
-    gadget_product_thread_safe(params.module(), &mut res, &a, &gadget_ct, &mut tmp_bytes);
-
-    println!("a.limbs()={}", a.limbs());
-    println!("gadget_ct.rows()={}", gadget_ct.rows());
-    println!("gadget_ct.cols()={}", gadget_ct.cols());
-    println!("res.limbs()={}", res.limbs());
-    println!();
-
-    println!("a:");
-    a.print_limbs(a.limbs(), 16);
-    println!();
-
-    println!("m:");
-    println!("{:?}", &m.0[..16]);
-    println!();
-
-    let mut a_res: Elem = Elem::new(params.module(), params.log_base2k(), params.log_q(), 0, 0);
-
-    decrypt_rlwe_thread_safe(
         params.module(),
-        &mut a_res,
+        params.log_base2k(),
-        &res,
+        params.log_q(),
-        &sk_svp_ppol,
+        params.log_scale(),
+    );
+
+    let mut want = vec![i64::default(); params.n()];
+    want.iter_mut().enumerate().for_each(|(i, x)| *x = i as i64);
+    pt.0.value[0].encode_vec_i64(log_base2k, log_k, &want, 32);
+    pt.0.value[0].normalize(log_base2k, &mut tmp_bytes);
+
+    let mut ct: Ciphertext = params.new_ciphertext(params.log_q());
+
+    params.encrypt_rlwe_sk_thread_safe(
+        &mut ct,
+        Some(&pt),
+        &sk0_svp_ppol,
+        &mut source_xa,
+        &mut source_xe,
         &mut tmp_bytes,
     );
 
-    let mut m_svp_ppol = params.module().new_svp_ppol();
+    gadget_product_inplace_thread_safe::<true>(
-    params.module().svp_prepare(&mut m_svp_ppol, &m);
+        params.module(),
+        &mut ct.0,
+        &gadget_ct,
+        &mut tmp_bytes,
+    );
 
-    let mut a_dft: VecZnxDft = params.module().new_vec_znx_dft(a.limbs());
+    println!("ct.limbs()={}", ct.limbs());
-    let mut a_big: VecZnxBig = a_dft.as_vec_znx_big();
+    println!("gadget_ct.rows()={}", gadget_ct.rows());
+    println!("gadget_ct.cols()={}", gadget_ct.cols());
+    println!("res.limbs()={}", ct.limbs());
+    println!();
 
-    params
+    decrypt_rlwe_thread_safe(
-        .module()
+        params.module(),
-        .svp_apply_dft(&mut a_dft, &m_svp_ppol, &a, a.limbs());
+        &mut pt.0,
-    params
+        &ct.0,
-        .module()
+        &sk1_svp_ppol,
-        .vec_znx_idft_tmp_a(&mut a_big, &mut a_dft, a.limbs());
+        &mut tmp_bytes,
-    params
+    );
-        .module()
-        .vec_znx_big_normalize(params.log_base2k(), &mut a, &a_big, &mut tmp_bytes);
 
-    params.module().vec_znx_sub_inplace(&mut a, &a_res.value[0]);
+    pt.0.value[0].print_limbs(pt.limbs(), 16);
 
-    println!("a*m - dec(a * GRLWE(m))");
+    let mut have = vec![i64::default(); params.n()];
-    a.print_limbs(a.limbs(), 16);
+
+    println!("pt: {}", log_k);
+    pt.0.value[0].decode_vec_i64(pt.log_base2k(), log_k, &mut have);
+
+    println!("want: {:?}", &want[..16]);
+    println!("have: {:?}", &have[..16]);
 }

rlwe/src/evaluator.rs

@@ -1,8 +1,9 @@
 use crate::{
     ciphertext::{Ciphertext, GadgetCiphertext},
     elem::Elem,
+    keys::SwitchingKey,
 };
-use base2k::{Infos, Module, VecZnx, VecZnxBig, VecZnxDft, VmpPMatOps};
+use base2k::{Infos, Module, VecZnx, VecZnxBig, VecZnxDft, VecZnxOps, VmpPMatOps};
 
 pub fn gadget_product_tmp_bytes(
     module: &Module,
@@ -19,11 +20,23 @@ pub fn gadget_product_tmp_bytes(
         + 2 * module.bytes_of_vec_znx_dft(gct_cols)
 }
 
+pub fn gadget_product_inplace_thread_safe<const OVERWRITE: bool>(
+    module: &Module,
+    res: &mut Elem,
+    b: &GadgetCiphertext,
+    tmp_bytes: &mut [u8],
+) {
+    unsafe {
+        let a_ptr: *const VecZnx = res.at(1) as *const VecZnx;
+        gadget_product_thread_safe::<OVERWRITE>(module, res, &*a_ptr, b, tmp_bytes);
+    }
+}
+
 /// Evaluates the gadget product res <- a x b.
 ///
 /// # Arguments
 ///
-/// * `module`: backend support for operations mod (X^N + 1)
+/// * `module`: backend support for operations mod (X^N + 1).
 /// * `res`: an [Elem] to store (-cs + m * a + e, c) with res_ncols limbs.
 /// * `a`: a [VecZnx] of a_ncols limbs.
 /// * `b`: a [GadgetCiphertext] as a vector of (-Bs + m * 2^{-k} + E, B)
@@ -33,7 +46,7 @@ pub fn gadget_product_tmp_bytes(
 ///
 /// res = sum[min(a_ncols, b_nrows)] decomp(a, i) * (-B[i]s + m * 2^{-k*i} + E[i], B[i])
 ///     = (cs + m * a + e, c) with min(res_limbs, b_cols) limbs.
-pub fn gadget_product_thread_safe(
+pub fn gadget_product_thread_safe<const OVERWRITE: bool>(
     module: &Module,
     res: &mut Elem,
     a: &VecZnx,
@@ -52,24 +65,35 @@ pub fn gadget_product_thread_safe(
     let mut res_dft: VecZnxDft = module.new_vec_znx_from_bytes(cols, tmp_bytes_res_dft);
     let mut res_big: VecZnxBig = res_dft.as_vec_znx_big();
 
-    // c1_dft <- DFT(c1) [cols]
+    // a_dft <- DFT(a) [cols]
     module.vec_znx_dft(&mut c1_dft, a, a.limbs());
 
-    // res_dft <- sum[rows] DFT(c1)[cols] x GadgetCiphertext[0][cols]
+    // >>>>>>>> RES[0]
+
+    // res_dft <- sum[rows] DFT(a)[cols] x GadgetCiphertext[0][cols]
     module.vmp_apply_dft_to_dft(&mut res_dft, &c1_dft, &b.value[0], tmp_bytes_vmp_apply_dft);
 
-    // res_big <- IDFT(DFT(c1) x GadgetCiphertext[0])
+    // res_big <- IDFT(DFT(a) x GadgetCiphertext[0])
     module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft, cols);
 
-    // res[0] = normalize(c0 + c1_dft x GadgetCiphertext[0])
+    // res_big <- res[0] + a_dft x GadgetCiphertext[0]
+    module.vec_znx_big_add_small_inplace(&mut res_big, res.at(0));
     module.vec_znx_big_normalize(log_base2k, res.at_mut(0), &res_big, tmp_bytes_vmp_apply_dft);
 
+    // >>>>>>>> RES[1]
+
     // res_dft <- DFT(c1) x GadgetCiphertext[1]
     module.vmp_apply_dft_to_dft(&mut res_dft, &c1_dft, &b.value[1], tmp_bytes_vmp_apply_dft);
 
     // res_big <- IDFT(DFT(c1) x GadgetCiphertext[1])
     module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft, cols);
 
-    // res[1] = normalize(c1_dft x GadgetCiphertext[1])
+    if OVERWRITE {
+        // res[1] = normalize(a_dft x GadgetCiphertext[1])
         module.vec_znx_big_normalize(log_base2k, res.at_mut(1), &res_big, tmp_bytes_vmp_apply_dft);
+    } else {
+        // res[1] = normalize(a_dft x GadgetCiphertext[1] + res[1])
+        module.vec_znx_big_add_small_inplace(&mut res_big, res.at(0));
+        module.vec_znx_big_normalize(log_base2k, res.at_mut(0), &res_big, tmp_bytes_vmp_apply_dft);
+    }
 }