f

tfhe/Cargo.toml

@@ -49,6 +49,7 @@ log = "0.4.19"
 cbindgen = { version = "0.26.0", optional = true }
 
 [dependencies]
+rand = "0.8.5"
 concrete-csprng = { version = "0.4.0", path = "../concrete-csprng", features = [
     "generator_fallback",
     "parallel",
@@ -183,12 +184,18 @@ rustdoc-args = ["--html-in-header", "katex-header.html"]
 # harness = false
 # required-features = ["boolean", "internal-keycache"]
 
-[[bench]]
-name = "shortint-bench"
-path = "benches/shortint/bench.rs"
-harness = false
+# [[bench]]
+# name = "shortint-bench"
+# path = "benches/shortint/bench.rs"
+# harness = false
+# required-features = ["shortint", "internal-keycache"]
+
+[[bin]]
+name = "pbs_bench"
+path = "src/pbs_bench.rs"
 required-features = ["shortint", "internal-keycache"]
 
+
 # [[bench]]
 # name = "integer-bench"
 # path = "benches/integer/bench.rs"

tfhe/src/pbs_bench.rs

@@ -0,0 +1,35 @@
+use lamellar::ActiveMessaging;
+use rand::Rng;
+use std::time::Instant;
+use tfhe::shortint::keycache::KEY_CACHE;
+use tfhe::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS_PBS;
+use tfhe::shortint::{Ciphertext, DISPATCHER};
+
+fn main() {
+    let count = 10000;
+
+    let keys = KEY_CACHE.get_from_param(PARAM_MESSAGE_2_CARRY_2_KS_PBS);
+    let (cks, sks) = (keys.client_key(), keys.server_key());
+
+    let mut rng = rand::thread_rng();
+
+    let modulus = cks.parameters.message_modulus().0 as u64;
+
+    let acc = sks.generate_lookup_table(|x| x);
+
+    let clear_0 = rng.gen::<u64>() % modulus;
+
+    let ctxt = cks.encrypt(clear_0);
+
+    let before = Instant::now();
+
+    let a: Vec<_> = (0..count)
+        .map(|_| sks.apply_lookup_table_future(ctxt.clone(), acc.clone()))
+        .collect();
+
+    for i in a {
+        let c: Ciphertext = DISPATCHER.world.block_on(i);
+    }
+
+    println!("{} PBS done in {}s", count, before.elapsed().as_secs());
+}