From c465256b6356408c86cadde8befa27156ce57989 Mon Sep 17 00:00:00 2001
From: "Mayeul@Zama" <69792125+mayeul-zama@users.noreply.github.com>
Date: Thu, 5 Dec 2024 12:45:41 +0100
Subject: [PATCH] chore(bench): add modulus switch noise reduction bench

---
 tfhe/Cargo.toml                               |  7 ++
 .../modulus_switch_noise_reduction.rs         | 89 +++++++++++++++++++
 2 files changed, 96 insertions(+)
 create mode 100644 tfhe/benches/core_crypto/modulus_switch_noise_reduction.rs

diff --git a/tfhe/Cargo.toml b/tfhe/Cargo.toml
index 1d01e4c211..f0acbf0966 100644
--- a/tfhe/Cargo.toml
+++ b/tfhe/Cargo.toml
@@ -165,6 +165,13 @@ path = "benches/core_crypto/dev_bench.rs"
 harness = false
 required-features = ["internal-keycache"]
 
+[[bench]]
+name = "modulus_switch_noise_reduction"
+path = "benches/core_crypto/modulus_switch_noise_reduction.rs"
+harness = false
+required-features = ["shortint"]
+
+
 [[bench]]
 name = "pbs128-bench"
 path = "benches/core_crypto/pbs128_bench.rs"
diff --git a/tfhe/benches/core_crypto/modulus_switch_noise_reduction.rs b/tfhe/benches/core_crypto/modulus_switch_noise_reduction.rs
new file mode 100644
index 0000000000..fd71e85957
--- /dev/null
+++ b/tfhe/benches/core_crypto/modulus_switch_noise_reduction.rs
@@ -0,0 +1,89 @@
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+use modulus_switch_noise_reduction::improve_lwe_ciphertext_modulus_switch_noise_for_binary_key;
+use tfhe::core_crypto::prelude::modulus_switch_noise_reduction::{
+    NoiseEstimationMeasureBound, RSigmaFactor,
+};
+use tfhe::core_crypto::prelude::*;
+
+fn modulus_switch_noise_reduction(c: &mut Criterion) {
+    // TODO: use shortint params
+    let lwe_dimension = LweDimension(918);
+    let noise_distribution = DynamicDistribution::new_t_uniform(46);
+    let ciphertext_modulus = CiphertextModulus::new_native();
+    let bound = NoiseEstimationMeasureBound((1_u64 << (64 - 1 - 4 - 1)) as f64);
+    let r_sigma_factor = RSigmaFactor(14.658999256586121);
+    let log_modulus = PolynomialSize(2048).to_blind_rotation_input_modulus_log();
+
+    for count in [10, 50, 100, 1_000, 10_000, 100_000] {
+        let mut boxed_seeder = new_seeder();
+        let seeder = boxed_seeder.as_mut();
+
+        let mut secret_generator =
+            SecretRandomGenerator::<DefaultRandomGenerator>::new(seeder.seed());
+
+        let mut encryption_generator =
+            EncryptionRandomGenerator::<DefaultRandomGenerator>::new(seeder.seed(), seeder);
+
+        let sk =
+            allocate_and_generate_new_binary_lwe_secret_key(lwe_dimension, &mut secret_generator);
+
+        let clean_lwe = allocate_and_encrypt_new_lwe_ciphertext(
+            &sk,
+            Plaintext(0),
+            noise_distribution,
+            ciphertext_modulus,
+            &mut encryption_generator,
+        );
+
+        let mut encryptions_of_zero = LweCiphertextList::new(
+            0,
+            lwe_dimension.to_lwe_size(),
+            LweCiphertextCount(count),
+            ciphertext_modulus,
+        );
+
+        let plaintext_list = PlaintextList::new(0, PlaintextCount(count));
+
+        encrypt_lwe_ciphertext_list(
+            &sk,
+            &mut encryptions_of_zero,
+            &plaintext_list,
+            noise_distribution,
+            &mut encryption_generator,
+        );
+
+        let mut lwe =
+            LweCiphertext::new(0_u64, sk.lwe_dimension().to_lwe_size(), ciphertext_modulus);
+
+        let bench_name = "modulus_switch_noise_reduction";
+
+        let mut bench_group = c.benchmark_group(bench_name);
+        bench_group
+            .sample_size(15)
+            .measurement_time(std::time::Duration::from_secs(5));
+
+        let bench_name = format!("modulus_switch_noise_reduction_{count}");
+
+        bench_group.bench_function(&bench_name, |b| {
+            b.iter(|| {
+                lwe.as_mut().copy_from_slice(clean_lwe.as_ref());
+
+                improve_lwe_ciphertext_modulus_switch_noise_for_binary_key(
+                    &mut lwe,
+                    &encryptions_of_zero,
+                    r_sigma_factor,
+                    bound,
+                    log_modulus,
+                );
+
+                black_box(&lwe);
+            });
+        });
+    }
+}
+
+criterion_group!(
+    modulus_switch_noise_reduction2,
+    modulus_switch_noise_reduction
+);
+criterion_main!(modulus_switch_noise_reduction2);