refactor(gpu): creating benchmarks for match_value

Author: enzodimaria

tfhe-benchmark/benches/integer/bench.rs

@@ -3,6 +3,7 @@
 mod aes;
 mod aes256;
 mod oprf;
+mod vector_find;
 
 mod rerand;
 
@@ -2802,6 +2803,7 @@ mod cuda {
         cuda_trailing_ones,
         cuda_ilog2,
         oprf::cuda::cuda_unsigned_oprf,
+        vector_find::cuda::cuda_match_value,
     );
 
     criterion_group!(
@@ -2830,6 +2832,7 @@ mod cuda {
         cuda_scalar_div,
         cuda_scalar_rem,
         oprf::cuda::cuda_unsigned_oprf,
+        vector_find::cuda::cuda_match_value,
     );
 
     criterion_group!(
@@ -3742,6 +3745,8 @@ criterion_group!(misc, full_propagate, full_propagate_parallelized);
 
 criterion_group!(oprf, oprf::unsigned_oprf);
 
+criterion_group!(vector_find, vector_find::match_value);
+
 #[cfg(feature = "gpu")]
 fn go_through_gpu_bench_groups(val: &str) {
     match val.to_lowercase().as_str() {
@@ -3789,10 +3794,12 @@ fn go_through_cpu_bench_groups(val: &str) {
             default_scalar_parallelized_ops();
             default_scalar_parallelized_ops_comp();
             cast_ops();
-            oprf()
+            oprf();
+            vector_find();
         }
         "fast_default" => {
             default_dedup_ops();
+            vector_find();
         }
         "smart" => {
             smart_ops();
@@ -3831,7 +3838,8 @@ fn main() {
             default_scalar_parallelized_ops();
             default_scalar_parallelized_ops_comp();
             cast_ops();
-            oprf()
+            oprf();
+            vector_find();
         }
     };
 

tfhe-benchmark/benches/integer/vector_find.rs

@@ -0,0 +1,109 @@
+use benchmark::params_aliases::BENCH_PARAM_GPU_MULTI_BIT_GROUP_4_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
+use benchmark::utilities::{write_to_json, OperatorType};
+use criterion::{black_box, Criterion};
+use tfhe::integer::keycache::KEY_CACHE;
+use tfhe::integer::{IntegerKeyKind, RadixClientKey};
+use tfhe::keycache::NamedParam;
+use tfhe::shortint::AtomicPatternParameters;
+use tfhe::MatchValues;
+
+pub fn match_value(c: &mut Criterion) {
+    let bench_name = "integer::match_value";
+
+    let param = BENCH_PARAM_GPU_MULTI_BIT_GROUP_4_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
+    let atomic_param: AtomicPatternParameters = param.into();
+    let param_name = param.name();
+
+    let (cpu_cks, cpu_sks) = KEY_CACHE.get_from_params(atomic_param, IntegerKeyKind::Radix);
+
+    let scenarios = vec![(64, 10, 32), (8, 256, 4)];
+
+    for (bits, num_elements, num_blocks) in scenarios {
+        let bench_id = format!("{bench_name}::{param_name}::{bits}bit_{num_elements}elements");
+
+        let cks = RadixClientKey::from((cpu_cks.clone(), num_blocks));
+
+        let mapping_data: Vec<(u64, u64)> = (0..num_elements as u64).map(|i| (i, i)).collect();
+        let match_values = MatchValues::new(mapping_data).expect("Invalid match values");
+
+        let input_val = 1u64;
+        let ct_input = cks.encrypt(input_val);
+
+        let mut group = c.benchmark_group(&bench_id);
+        group.sample_size(15);
+
+        group.bench_function(&bench_id, |b| {
+            b.iter(|| {
+                black_box(cpu_sks.match_value_parallelized(&ct_input, &match_values));
+            })
+        });
+        group.finish();
+
+        write_to_json::<u64, _>(
+            &bench_id,
+            atomic_param,
+            param.name(),
+            "match_value_latency",
+            &OperatorType::Atomic,
+            bits as u32,
+            vec![atomic_param.message_modulus().0.ilog2(); bits],
+        );
+    }
+}
+
+#[cfg(feature = "gpu")]
+pub mod cuda {
+    use super::*;
+    use tfhe::core_crypto::gpu::CudaStreams;
+    use tfhe::integer::gpu::ciphertext::CudaUnsignedRadixCiphertext;
+    use tfhe::integer::gpu::CudaServerKey;
+
+    pub fn cuda_match_value(c: &mut Criterion) {
+        let bench_name = "integer::cuda::match_value";
+
+        let param = BENCH_PARAM_GPU_MULTI_BIT_GROUP_4_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
+        let atomic_param: AtomicPatternParameters = param.into();
+        let param_name = param.name();
+
+        let streams = CudaStreams::new_multi_gpu();
+
+        let (cpu_cks, _) = KEY_CACHE.get_from_params(atomic_param, IntegerKeyKind::Radix);
+        let sks = CudaServerKey::new(&cpu_cks, &streams);
+
+        let scenarios = vec![(64, 10, 32), (8, 256, 4)];
+
+        for (bits, num_elements, num_blocks) in scenarios {
+            let bench_id = format!("{bench_name}::{param_name}::{bits}bit_{num_elements}elements");
+
+            let cks = RadixClientKey::from((cpu_cks.clone(), num_blocks));
+
+            let mapping_data: Vec<(u64, u64)> = (0..num_elements as u64).map(|i| (i, i)).collect();
+            let match_values = MatchValues::new(mapping_data).expect("Invalid match values");
+
+            let input_val = 1u64;
+            let ct_input = cks.encrypt(input_val);
+            let d_ct_input =
+                CudaUnsignedRadixCiphertext::from_radix_ciphertext(&ct_input, &streams);
+
+            let mut group = c.benchmark_group(&bench_id);
+            group.sample_size(15);
+
+            group.bench_function(&bench_id, |b| {
+                b.iter(|| {
+                    black_box(sks.match_value(&d_ct_input, &match_values, &streams));
+                })
+            });
+            group.finish();
+
+            write_to_json::<u64, _>(
+                &bench_id,
+                atomic_param,
+                param.name(),
+                "match_value_latency",
+                &OperatorType::Atomic,
+                bits as u32,
+                vec![atomic_param.message_modulus().0.ilog2(); bits],
+            );
+        }
+    }
+}