remove 1 NTT from galois_auto AND add NTT benches

src/bool/noise.rs

@@ -28,7 +28,7 @@ mod test {
             ModularOpsU64<CiphertextModulus<u64>>,
         >::new(SMALL_MP_BOOL_PARAMS);
 
-        let parties = 16;
+        let parties = 2;
 
         let mut rng = DefaultSecureRng::new();
         let mut pk_cr_seed = [0u8; 32];

src/bool/parameters.rs

@@ -353,7 +353,7 @@ mod tests {
     #[test]
     fn find_prime() {
         let bits = 55;
-        let ring_size = 1 << 11;
+        let ring_size = 1 << 15;
         let prime = generate_prime(bits, ring_size * 2, 1 << bits).unwrap();
         dbg!(prime);
     }

src/lib.rs

@@ -20,6 +20,8 @@ mod rgsw;
 mod shortint;
 mod utils;
 
+pub use ntt::{Ntt, NttBackendU64, NttInit};
+
 pub trait Matrix: AsRef<[Self::R]> {
     type MatElement;
     type R: Row<Element = Self::MatElement>;

src/pbs.rs

@@ -180,6 +180,7 @@ pub(crate) fn pbs<
             });
     }
 
+    // let now = std::time::Instant::now();
     // blind rotate
     blind_rotation(
         &mut trivial_rlwe_test_poly,
@@ -195,6 +196,7 @@ pub(crate) fn pbs<
         pbs_info,
         pbs_key,
     );
+    // println!("Blind rotation time: {:?}", now.elapsed());
 
     // sample extract
     sample_extract(lwe_in, &trivial_rlwe_test_poly, pbs_info.modop_rlweq(), 0);
@@ -238,6 +240,7 @@ fn blind_rotation<
         let s_indices = &gk_to_si[q_by_4 + i];
 
         s_indices.iter().for_each(|s_index| {
+            let new = std::time::Instant::now();
             rlwe_by_rgsw(
                 trivial_rlwe_test_poly,
                 pbs_key.rgsw_ct_lwe_si(*s_index),
@@ -246,11 +249,14 @@ fn blind_rotation<
                 ntt_op,
                 mod_op,
             );
+            println!("Rlwe x Rgsw time: {:?}", new.elapsed());
         });
         v += 1;
 
         if gk_to_si[q_by_4 + i - 1].len() != 0 || v == w || i == 1 {
             let (auto_map_index, auto_map_sign) = parameters.rlwe_auto_map(v);
+
+            let now = std::time::Instant::now();
             galois_auto(
                 trivial_rlwe_test_poly,
                 pbs_key.galois_key_for_auto(v),
@@ -261,6 +267,8 @@ fn blind_rotation<
                 ntt_op,
                 auto_decomposer,
             );
+            println!("Auto time: {:?}", now.elapsed());
+
             count += 1;
             v = 0;
         }

src/rgsw.rs

@@ -556,22 +556,11 @@ pub(crate) fn galois_auto<
 
     let (scratch_matrix_d_ring, tmp_rlwe_out) = scratch_matrix.split_at_row_mut(d);
 
-    // send b(X) -> b(X^k)
-    izip!(
-        rlwe_in.get_row(1),
-        auto_map_index.iter(),
-        auto_map_sign.iter()
-    )
-    .for_each(|(el_in, to_index, sign)| {
-        if !*sign {
-            tmp_rlwe_out[1].as_mut()[*to_index] = mod_op.neg(el_in);
-        } else {
-            tmp_rlwe_out[1].as_mut()[*to_index] = *el_in;
-            // scratch_matrix_dplus2_ring.set(d + 1, *to_index, *el_in);
-        }
-    });
-
     if !rlwe_in.is_trivial() {
+        tmp_rlwe_out.iter_mut().for_each(|r| {
+            r.as_mut().fill(Mmut::MatElement::zero());
+        });
+
         // send a(X) -> a(X^k) and decompose a(X^k)
         izip!(
             rlwe_in.get_row(0),
@@ -595,7 +584,6 @@ pub(crate) fn galois_auto<
         // RLWE(m^k) = a', b'; RLWE(m) = a, b
         // key switch: (a * RLWE'(s(X^k)))
         let (ksk_a, ksk_b) = ksk.split_at_row(d);
-        tmp_rlwe_out[0].as_mut().fill(Mmut::MatElement::zero());
         // a' = decomp<a> * RLWE'_A(s(X^k))
         routine(
             tmp_rlwe_out[0].as_mut(),
@@ -603,9 +591,7 @@ pub(crate) fn galois_auto<
             ksk_a,
             mod_op,
         );
-        // send b(X^k) to evaluation domain
-        ntt_op.forward(tmp_rlwe_out[1].as_mut());
-        // b' = b(X^k)
+
         // b' += decomp<a(X^k)> * RLWE'_B(s(X^k))
         routine(
             tmp_rlwe_out[1].as_mut(),
@@ -619,11 +605,43 @@ pub(crate) fn galois_auto<
             .iter_mut()
             .for_each(|r| ntt_op.backward(r.as_mut()));
 
+        // send b(X) -> b(X^k) and then b'(X) += b(X^k)
+        izip!(
+            rlwe_in.get_row(1),
+            auto_map_index.iter(),
+            auto_map_sign.iter()
+        )
+        .for_each(|(el_in, to_index, sign)| {
+            let row = tmp_rlwe_out[1].as_mut();
+            if !*sign {
+                row[*to_index] = mod_op.sub(&row[*to_index], el_in);
+            } else {
+                row[*to_index] = mod_op.add(&row[*to_index], el_in);
+            }
+        });
+
+        // copy over A; Leave B for later
         rlwe_in
             .get_row_mut(0)
             .copy_from_slice(tmp_rlwe_out[0].as_ref());
+    } else {
+        // RLWE is trivial, a(X) is 0.
+        // send b(X) -> b(X^k)
+        izip!(
+            rlwe_in.get_row(1),
+            auto_map_index.iter(),
+            auto_map_sign.iter()
+        )
+        .for_each(|(el_in, to_index, sign)| {
+            if !*sign {
+                tmp_rlwe_out[1].as_mut()[*to_index] = mod_op.neg(el_in);
+            } else {
+                tmp_rlwe_out[1].as_mut()[*to_index] = *el_in;
+            }
+        });
     }
 
+    // Copy over B
     rlwe_in
         .get_row_mut(1)
         .copy_from_slice(tmp_rlwe_out[1].as_ref());