feat(gpu): add benchmark for CUDA-accelerated compression

tfhe/benches/integer/glwe_packing_compression.rs

@@ -1,10 +1,25 @@
 use criterion::{black_box, criterion_group, Criterion};
+#[cfg(feature = "gpu")]
+use tfhe::core_crypto::gpu::CudaStreams;
 use tfhe::integer::ciphertext::CompressedCiphertextListBuilder;
 use tfhe::integer::{ClientKey, RadixCiphertext};
+
+#[cfg(feature = "gpu")]
+use tfhe::integer::gpu::ciphertext::compressed_ciphertext_list::CudaCompressedCiphertextListBuilder;
+
+#[cfg(feature = "gpu")]
+use tfhe::integer::gpu::ciphertext::CudaRadixCiphertext;
+use tfhe::integer::gpu::ciphertext::CudaUnsignedRadixCiphertext;
+
+#[cfg(feature = "gpu")]
+use tfhe::integer::gpu::gen_keys_radix_gpu;
 use tfhe::shortint::parameters::list_compression::COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
 use tfhe::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
 
-fn glwe_packing(c: &mut Criterion) {
+#[cfg(feature = "gpu")]
+use tfhe::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS_PBS;
+
+fn cpu_glwe_packing(c: &mut Criterion) {
     let param = PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
 
     let comp_param = COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
@@ -73,9 +88,90 @@ fn glwe_packing(c: &mut Criterion) {
     }
 }
 
-criterion_group!(glwe_packing2, glwe_packing);
+#[cfg(feature = "gpu")]
+fn gpu_glwe_packing(c: &mut Criterion) {
+    let param = PARAM_MESSAGE_2_CARRY_2_KS_PBS;
+
+    let comp_param = COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
+
+    let bench_name = "integer_cuda_packing_compression";
+
+    let mut bench_group = c.benchmark_group(bench_name);
+
+    let gpu_index = 0;
+    let streams = CudaStreams::new_single_gpu(gpu_index);
+
+    let log_message_modulus = param.message_modulus.0.ilog2() as usize;
+
+    for num_bits in [
+        8,
+        16,
+        32,
+        64,
+        128,
+        256,
+        256,
+        comp_param.lwe_per_glwe.0 * log_message_modulus,
+    ] {
+        assert_eq!(num_bits % log_message_modulus, 0);
+        let num_blocks = num_bits / log_message_modulus;
+
+        let (cks, _) = gen_keys_radix_gpu(param, num_blocks, &streams);
+
+        let private_compression_key = cks.new_compression_private_key(comp_param);
+
+        let (cuda_compression_key, cuda_decompression_key) =
+            cks.new_cuda_compression_decompression_keys(&private_compression_key, &streams);
+
+        let ct = cks.encrypt(0_u32);
+
+        let d_ct = CudaUnsignedRadixCiphertext::from_radix_ciphertext(&ct, &streams);
+
+        let mut builder = CudaCompressedCiphertextListBuilder::new();
+
+        builder.push(d_ct, &streams);
+
+        bench_group.bench_function(format!("pack_u{num_bits}"), |b| {
+            b.iter(|| {
+                let compressed = builder.build(&cuda_compression_key, &streams, gpu_index);
+
+                _ = black_box(compressed);
+            })
+        });
+
+        let compressed = builder.build(&cuda_compression_key, &streams, gpu_index);
+
+        bench_group.bench_function(format!("unpack_u{num_bits}"), |b| {
+            b.iter(|| {
+                let unpacked: CudaRadixCiphertext =
+                    compressed.get(0, &cuda_decompression_key, &streams, gpu_index);
+
+                _ = black_box(unpacked);
+            })
+        });
+
+        bench_group.bench_function(format!("pack_unpack_u{num_bits}"), |b| {
+            b.iter(|| {
+                let compressed = builder.build(&cuda_compression_key, &streams, gpu_index);
+
+                let unpacked: CudaRadixCiphertext =
+                    compressed.get(0, &cuda_decompression_key, &streams, gpu_index);
+
+                _ = black_box(unpacked);
+            })
+        });
+    }
+}
+
+criterion_group!(cpu_glwe_packing2, cpu_glwe_packing);
+#[cfg(feature = "gpu")]
+criterion_group!(gpu_glwe_packing2, gpu_glwe_packing);
 
 fn main() {
-    glwe_packing2();
+    #[cfg(feature = "gpu")]
+    gpu_glwe_packing2();
+    #[cfg(not(feature = "gpu"))]
+    cpu_glwe_packing2();
+
     Criterion::default().configure_from_args().final_summary();
 }

tfhe/src/integer/gpu/ciphertext/compressed_ciphertext_list.rs

@@ -178,23 +178,23 @@ mod tests {
             .build(&cuda_compression_key, &streams, gpu_index);
 
         let d_decompressed1 = CudaUnsignedRadixCiphertext {
-            ciphertext: cuda_compressed.get(0, &cuda_decompression_key, &streams, 0),
+            ciphertext: cuda_compressed.get(0, &cuda_decompression_key, &streams, gpu_index),
         };
 
         let decompressed1 = d_decompressed1.to_radix_ciphertext(&streams);
         let decrypted: u32 = cks.decrypt(&decompressed1);
 
         assert_eq!(decrypted, 3_u32);
         let d_decompressed2 = CudaUnsignedRadixCiphertext {
-            ciphertext: cuda_compressed.get(1, &cuda_decompression_key, &streams, 0),
+            ciphertext: cuda_compressed.get(1, &cuda_decompression_key, &streams, gpu_index),
         };
 
         let decompressed2 = d_decompressed2.to_radix_ciphertext(&streams);
         let decrypted: u32 = cks.decrypt(&decompressed2);
 
         assert_eq!(decrypted, 2_u32);
         let d_decompressed3 = CudaSignedRadixCiphertext {
-            ciphertext: cuda_compressed.get(2, &cuda_decompression_key, &streams, 0),
+            ciphertext: cuda_compressed.get(2, &cuda_decompression_key, &streams, gpu_index),
         };
 
         let decompressed3 = d_decompressed3.to_signed_radix_ciphertext(&streams);