feat(trivium): add bench for casting and packing

apps/trivium/Cargo.toml

@@ -21,4 +21,4 @@ criterion = { version = "0.4", features = [ "html_reports" ]}
 
 [[bench]]
 name = "trivium"
-harness = false
+harness = false

apps/trivium/benches/kreyvium_shortint.rs

@@ -1,6 +1,6 @@
 use tfhe::prelude::*;
 use tfhe::shortint::prelude::*;
-use tfhe::shortint::CastingKey;
+use tfhe::shortint::KeySwitchingKey;
 use tfhe::{generate_keys, ConfigBuilder, FheUint64};
 
 use tfhe_trivium::{KreyviumStreamShortint, TransCiphering};
@@ -12,8 +12,16 @@ pub fn kreyvium_shortint_warmup(c: &mut Criterion) {
         .enable_default_integers()
         .build();
     let (hl_client_key, hl_server_key) = generate_keys(config);
+    let underlying_ck: tfhe::shortint::ClientKey = (*hl_client_key.as_ref()).clone().into();
+    let underlying_sk: tfhe::shortint::ServerKey = (*hl_server_key.as_ref()).clone().into();
+
     let (client_key, server_key): (ClientKey, ServerKey) = gen_keys(PARAM_MESSAGE_1_CARRY_1);
-    let ksk = CastingKey::new((&client_key, &server_key), (&hl_client_key, &hl_server_key));
+
+    let ksk = KeySwitchingKey::new(
+        (&client_key, &server_key),
+        (&underlying_ck, &underlying_sk),
+        PARAM_KEYSWITCH_1_1_TO_2_2,
+    );
 
     let key_string = "0053A6F94C9FF24598EB000000000000".to_string();
     let mut key = [0; 128];
@@ -40,8 +48,13 @@ pub fn kreyvium_shortint_warmup(c: &mut Criterion) {
     c.bench_function("kreyvium 1_1 warmup", |b| {
         b.iter(|| {
             let cipher_key = key.map(|x| client_key.encrypt(x));
-            let _kreyvium =
-                KreyviumStreamShortint::new(cipher_key, iv, &server_key, &ksk, &hl_server_key);
+            let _kreyvium = KreyviumStreamShortint::new(
+                cipher_key,
+                iv,
+                server_key.clone(),
+                ksk.clone(),
+                hl_server_key.clone(),
+            );
         })
     });
 }
@@ -51,8 +64,16 @@ pub fn kreyvium_shortint_gen(c: &mut Criterion) {
         .enable_default_integers()
         .build();
     let (hl_client_key, hl_server_key) = generate_keys(config);
+    let underlying_ck: tfhe::shortint::ClientKey = (*hl_client_key.as_ref()).clone().into();
+    let underlying_sk: tfhe::shortint::ServerKey = (*hl_server_key.as_ref()).clone().into();
+
     let (client_key, server_key): (ClientKey, ServerKey) = gen_keys(PARAM_MESSAGE_1_CARRY_1);
-    let ksk = CastingKey::new((&client_key, &server_key), (&hl_client_key, &hl_server_key));
+
+    let ksk = KeySwitchingKey::new(
+        (&client_key, &server_key),
+        (&underlying_ck, &underlying_sk),
+        PARAM_KEYSWITCH_1_1_TO_2_2,
+    );
 
     let key_string = "0053A6F94C9FF24598EB000000000000".to_string();
     let mut key = [0; 128];
@@ -78,8 +99,7 @@ pub fn kreyvium_shortint_gen(c: &mut Criterion) {
 
     let cipher_key = key.map(|x| client_key.encrypt(x));
 
-    let mut kreyvium =
-        KreyviumStreamShortint::new(cipher_key, iv, &server_key, &ksk, &hl_server_key);
+    let mut kreyvium = KreyviumStreamShortint::new(cipher_key, iv, server_key, ksk, hl_server_key);
 
     c.bench_function("kreyvium 1_1 generate 64 bits", |b| {
         b.iter(|| kreyvium.next_64())
@@ -91,8 +111,16 @@ pub fn kreyvium_shortint_trans(c: &mut Criterion) {
         .enable_default_integers()
         .build();
     let (hl_client_key, hl_server_key) = generate_keys(config);
+    let underlying_ck: tfhe::shortint::ClientKey = (*hl_client_key.as_ref()).clone().into();
+    let underlying_sk: tfhe::shortint::ServerKey = (*hl_server_key.as_ref()).clone().into();
+
     let (client_key, server_key): (ClientKey, ServerKey) = gen_keys(PARAM_MESSAGE_1_CARRY_1);
-    let ksk = CastingKey::new((&client_key, &server_key), (&hl_client_key, &hl_server_key));
+
+    let ksk = KeySwitchingKey::new(
+        (&client_key, &server_key),
+        (&underlying_ck, &underlying_sk),
+        PARAM_KEYSWITCH_1_1_TO_2_2,
+    );
 
     let key_string = "0053A6F94C9FF24598EB000000000000".to_string();
     let mut key = [0; 128];
@@ -119,8 +147,7 @@ pub fn kreyvium_shortint_trans(c: &mut Criterion) {
     let cipher_key = key.map(|x| client_key.encrypt(x));
 
     let ciphered_message = FheUint64::try_encrypt(0u64, &hl_client_key).unwrap();
-    let mut kreyvium =
-        KreyviumStreamShortint::new(cipher_key, iv, &server_key, &ksk, &hl_server_key);
+    let mut kreyvium = KreyviumStreamShortint::new(cipher_key, iv, server_key, ksk, hl_server_key);
 
     c.bench_function("kreyvium 1_1 transencrypt 64 bits", |b| {
         b.iter(|| kreyvium.trans_encrypt_64(ciphered_message.clone()))

apps/trivium/benches/trivium_shortint.rs

@@ -1,6 +1,6 @@
 use tfhe::prelude::*;
 use tfhe::shortint::prelude::*;
-use tfhe::shortint::CastingKey;
+use tfhe::shortint::KeySwitchingKey;
 use tfhe::{generate_keys, ConfigBuilder, FheUint64};
 
 use tfhe_trivium::{TransCiphering, TriviumStreamShortint};
@@ -12,8 +12,16 @@ pub fn trivium_shortint_warmup(c: &mut Criterion) {
         .enable_default_integers()
         .build();
     let (hl_client_key, hl_server_key) = generate_keys(config);
+    let underlying_ck: tfhe::shortint::ClientKey = (*hl_client_key.as_ref()).clone().into();
+    let underlying_sk: tfhe::shortint::ServerKey = (*hl_server_key.as_ref()).clone().into();
+
     let (client_key, server_key): (ClientKey, ServerKey) = gen_keys(PARAM_MESSAGE_1_CARRY_1);
-    let ksk = CastingKey::new((&client_key, &server_key), (&hl_client_key, &hl_server_key));
+
+    let ksk = KeySwitchingKey::new(
+        (&client_key, &server_key),
+        (&underlying_ck, &underlying_sk),
+        PARAM_KEYSWITCH_1_1_TO_2_2,
+    );
 
     let key_string = "0053A6F94C9FF24598EB".to_string();
     let mut key = [0; 80];
@@ -40,8 +48,13 @@ pub fn trivium_shortint_warmup(c: &mut Criterion) {
     c.bench_function("trivium 1_1 warmup", |b| {
         b.iter(|| {
             let cipher_key = key.map(|x| client_key.encrypt(x));
-            let _trivium =
-                TriviumStreamShortint::new(cipher_key, iv, &server_key, &ksk, &hl_server_key);
+            let _trivium = TriviumStreamShortint::new(
+                cipher_key,
+                iv,
+                server_key.clone(),
+                ksk.clone(),
+                hl_server_key.clone(),
+            );
         })
     });
 }
@@ -51,8 +64,16 @@ pub fn trivium_shortint_gen(c: &mut Criterion) {
         .enable_default_integers()
         .build();
     let (hl_client_key, hl_server_key) = generate_keys(config);
+    let underlying_ck: tfhe::shortint::ClientKey = (*hl_client_key.as_ref()).clone().into();
+    let underlying_sk: tfhe::shortint::ServerKey = (*hl_server_key.as_ref()).clone().into();
+
     let (client_key, server_key): (ClientKey, ServerKey) = gen_keys(PARAM_MESSAGE_1_CARRY_1);
-    let ksk = CastingKey::new((&client_key, &server_key), (&hl_client_key, &hl_server_key));
+
+    let ksk = KeySwitchingKey::new(
+        (&client_key, &server_key),
+        (&underlying_ck, &underlying_sk),
+        PARAM_KEYSWITCH_1_1_TO_2_2,
+    );
 
     let key_string = "0053A6F94C9FF24598EB".to_string();
     let mut key = [0; 80];
@@ -78,7 +99,7 @@ pub fn trivium_shortint_gen(c: &mut Criterion) {
 
     let cipher_key = key.map(|x| client_key.encrypt(x));
 
-    let mut trivium = TriviumStreamShortint::new(cipher_key, iv, &server_key, &ksk, &hl_server_key);
+    let mut trivium = TriviumStreamShortint::new(cipher_key, iv, server_key, ksk, hl_server_key);
 
     c.bench_function("trivium 1_1 generate 64 bits", |b| {
         b.iter(|| trivium.next_64())
@@ -90,8 +111,16 @@ pub fn trivium_shortint_trans(c: &mut Criterion) {
         .enable_default_integers()
         .build();
     let (hl_client_key, hl_server_key) = generate_keys(config);
+    let underlying_ck: tfhe::shortint::ClientKey = (*hl_client_key.as_ref()).clone().into();
+    let underlying_sk: tfhe::shortint::ServerKey = (*hl_server_key.as_ref()).clone().into();
+
     let (client_key, server_key): (ClientKey, ServerKey) = gen_keys(PARAM_MESSAGE_1_CARRY_1);
-    let ksk = CastingKey::new((&client_key, &server_key), (&hl_client_key, &hl_server_key));
+
+    let ksk = KeySwitchingKey::new(
+        (&client_key, &server_key),
+        (&underlying_ck, &underlying_sk),
+        PARAM_KEYSWITCH_1_1_TO_2_2,
+    );
 
     let key_string = "0053A6F94C9FF24598EB".to_string();
     let mut key = [0; 80];
@@ -118,7 +147,7 @@ pub fn trivium_shortint_trans(c: &mut Criterion) {
     let cipher_key = key.map(|x| client_key.encrypt(x));
 
     let ciphered_message = FheUint64::try_encrypt(0u64, &hl_client_key).unwrap();
-    let mut trivium = TriviumStreamShortint::new(cipher_key, iv, &server_key, &ksk, &hl_server_key);
+    let mut trivium = TriviumStreamShortint::new(cipher_key, iv, server_key, ksk, hl_server_key);
 
     c.bench_function("trivium 1_1 transencrypt 64 bits", |b| {
         b.iter(|| trivium.trans_encrypt_64(ciphered_message.clone()))

tfhe/benches/shortint/bench.rs

@@ -703,9 +703,18 @@ criterion_group!(
     scalar_not_equal
 );
 
+mod casting;
+criterion_group!(
+    casting,
+    casting::pack_cast_64,
+    casting::pack_cast,
+    casting::cast
+);
+
 criterion_main!(
     // arithmetic_operation,
     // arithmetic_scalar_operation,
+    casting,
     default_ops,
     default_scalar_ops,
 );

tfhe/benches/shortint/casting.rs

@@ -0,0 +1,78 @@
+use tfhe::shortint::prelude::*;
+
+use rayon::prelude::*;
+
+use criterion::Criterion;
+
+pub fn pack_cast_64(c: &mut Criterion) {
+    let (client_key_1, server_key_1): (ClientKey, ServerKey) = gen_keys(PARAM_MESSAGE_1_CARRY_1);
+    let (client_key_2, server_key_2): (ClientKey, ServerKey) = gen_keys(PARAM_MESSAGE_2_CARRY_2);
+
+    let ksk = KeySwitchingKey::new(
+        (&client_key_1, &server_key_1),
+        (&client_key_2, &server_key_2),
+        PARAM_KEYSWITCH_1_1_TO_2_2,
+    );
+
+    let vec_ct = vec![client_key_1.encrypt(1); 64];
+
+    c.bench_function("pack_cast_64", |b| {
+        b.iter(|| {
+            let _ = (0..32)
+                .into_par_iter()
+                .map(|i| {
+                    let byte_idx = 7 - i / 4;
+                    let pair_idx = i % 4;
+
+                    let b0 = &vec_ct[8 * byte_idx + 2 * pair_idx];
+                    let b1 = &vec_ct[8 * byte_idx + 2 * pair_idx + 1];
+
+                    ksk.cast(
+                        &server_key_1.unchecked_add(b0, &server_key_1.unchecked_scalar_mul(b1, 2)),
+                    )
+                })
+                .collect::<Vec<_>>();
+        });
+    });
+}
+
+pub fn pack_cast(c: &mut Criterion) {
+    let (client_key_1, server_key_1): (ClientKey, ServerKey) = gen_keys(PARAM_MESSAGE_1_CARRY_1);
+    let (client_key_2, server_key_2): (ClientKey, ServerKey) = gen_keys(PARAM_MESSAGE_2_CARRY_2);
+
+    let ksk = KeySwitchingKey::new(
+        (&client_key_1, &server_key_1),
+        (&client_key_2, &server_key_2),
+        PARAM_KEYSWITCH_1_1_TO_2_2,
+    );
+
+    let ct_1 = client_key_1.encrypt(1);
+    let ct_2 = client_key_1.encrypt(1);
+
+    c.bench_function("pack_cast", |b| {
+        b.iter(|| {
+            let _ = ksk.cast(
+                &server_key_1.unchecked_add(&ct_1, &server_key_1.unchecked_scalar_mul(&ct_2, 2)),
+            );
+        });
+    });
+}
+
+pub fn cast(c: &mut Criterion) {
+    let (client_key_1, server_key_1): (ClientKey, ServerKey) = gen_keys(PARAM_MESSAGE_1_CARRY_1);
+    let (client_key_2, server_key_2): (ClientKey, ServerKey) = gen_keys(PARAM_MESSAGE_2_CARRY_2);
+
+    let ksk = KeySwitchingKey::new(
+        (&client_key_1, &server_key_1),
+        (&client_key_2, &server_key_2),
+        PARAM_KEYSWITCH_1_1_TO_2_2,
+    );
+
+    let ct = client_key_1.encrypt(1);
+
+    c.bench_function("cast", |b| {
+        b.iter(|| {
+            let _ = ksk.cast(&ct);
+        });
+    });
+}