feat(gpu): aes 128

tfhe-benchmark/benches/integer/aes.rs

@@ -0,0 +1,154 @@
+#[cfg(feature = "gpu")]
+pub mod cuda {
+    use benchmark::params_aliases::BENCH_PARAM_GPU_MULTI_BIT_GROUP_4_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
+    use benchmark::utilities::{get_bench_type, write_to_json, BenchmarkType, OperatorType};
+    use criterion::{black_box, Criterion, Throughput};
+    use tfhe::core_crypto::gpu::CudaStreams;
+    use tfhe::integer::gpu::ciphertext::CudaUnsignedRadixCiphertext;
+    use tfhe::integer::gpu::CudaServerKey;
+    use tfhe::integer::keycache::KEY_CACHE;
+    use tfhe::integer::{IntegerKeyKind, RadixClientKey};
+    use tfhe::keycache::NamedParam;
+    use tfhe::shortint::AtomicPatternParameters;
+
+    pub fn cuda_aes(c: &mut Criterion) {
+        let bench_name = "integer::cuda::aes";
+
+        let mut bench_group = c.benchmark_group(bench_name);
+        bench_group
+            .sample_size(15)
+            .measurement_time(std::time::Duration::from_secs(60))
+            .warm_up_time(std::time::Duration::from_secs(60));
+
+        let param = BENCH_PARAM_GPU_MULTI_BIT_GROUP_4_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
+        let atomic_param: AtomicPatternParameters = param.into();
+
+        let key: u128 = 0x2b7e151628aed2a6abf7158809cf4f3c;
+        let iv: u128 = 0xf0f1f2f3f4f5f6f7f8f9fafbfcfdfeff;
+        let aes_op_bit_size = 128;
+
+        let param_name = param.name();
+
+        match get_bench_type() {
+            BenchmarkType::Latency => {
+                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 cks = RadixClientKey::from((cpu_cks, 1));
+
+                let ct_key = cks.encrypt_u128_for_aes_ctr(key);
+                let ct_iv = cks.encrypt_u128_for_aes_ctr(iv);
+
+                let d_key = CudaUnsignedRadixCiphertext::from_radix_ciphertext(&ct_key, &streams);
+                let d_iv = CudaUnsignedRadixCiphertext::from_radix_ciphertext(&ct_iv, &streams);
+
+                {
+                    const NUM_AES_INPUTS: usize = 1;
+                    const SBOX_PARALLELISM: usize = 16;
+                    let bench_id = format!("{param_name}::{NUM_AES_INPUTS}_input_encryption");
+
+                    let round_keys = unsafe { sks.key_expansion_async(&d_key, &streams) };
+                    streams.synchronize();
+
+                    bench_group.bench_function(&bench_id, |b| {
+                        b.iter(|| {
+                            unsafe {
+                                black_box(sks.aes_encrypt_async(
+                                    &d_iv,
+                                    &round_keys,
+                                    0,
+                                    NUM_AES_INPUTS,
+                                    SBOX_PARALLELISM,
+                                    &streams,
+                                ));
+                            }
+                            streams.synchronize();
+                        })
+                    });
+
+                    write_to_json::<u64, _>(
+                        &bench_id,
+                        atomic_param,
+                        param.name(),
+                        "aes_encryption",
+                        &OperatorType::Atomic,
+                        aes_op_bit_size,
+                        vec![atomic_param.message_modulus().0.ilog2(); aes_op_bit_size as usize],
+                    );
+                }
+
+                {
+                    let bench_id = format!("{param_name}::key_expansion");
+
+                    bench_group.bench_function(&bench_id, |b| {
+                        b.iter(|| {
+                            unsafe {
+                                black_box(sks.key_expansion_async(&d_key, &streams));
+                            }
+                            streams.synchronize();
+                        })
+                    });
+
+                    write_to_json::<u64, _>(
+                        &bench_id,
+                        atomic_param,
+                        param.name(),
+                        "aes_key_expansion",
+                        &OperatorType::Atomic,
+                        aes_op_bit_size,
+                        vec![atomic_param.message_modulus().0.ilog2(); aes_op_bit_size as usize],
+                    );
+                }
+            }
+            BenchmarkType::Throughput => {
+                const NUM_AES_INPUTS: usize = 192;
+                const SBOX_PARALLELISM: usize = 16;
+                let bench_id = format!("throughput::{param_name}::{NUM_AES_INPUTS}_inputs");
+
+                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 cks = RadixClientKey::from((cpu_cks, 1));
+
+                bench_group.throughput(Throughput::Elements(NUM_AES_INPUTS as u64));
+
+                let ct_key = cks.encrypt_u128_for_aes_ctr(key);
+                let ct_iv = cks.encrypt_u128_for_aes_ctr(iv);
+
+                let d_key = CudaUnsignedRadixCiphertext::from_radix_ciphertext(&ct_key, &streams);
+                let d_iv = CudaUnsignedRadixCiphertext::from_radix_ciphertext(&ct_iv, &streams);
+
+                let round_keys = unsafe { sks.key_expansion_async(&d_key, &streams) };
+                streams.synchronize();
+
+                bench_group.bench_function(&bench_id, |b| {
+                    b.iter(|| {
+                        unsafe {
+                            black_box(sks.aes_encrypt_async(
+                                &d_iv,
+                                &round_keys,
+                                0,
+                                NUM_AES_INPUTS,
+                                SBOX_PARALLELISM,
+                                &streams,
+                            ));
+                        }
+                        streams.synchronize();
+                    })
+                });
+
+                write_to_json::<u64, _>(
+                    &bench_id,
+                    atomic_param,
+                    param.name(),
+                    "aes_encryption",
+                    &OperatorType::Atomic,
+                    aes_op_bit_size,
+                    vec![atomic_param.message_modulus().0.ilog2(); aes_op_bit_size as usize],
+                );
+            }
+        };
+
+        bench_group.finish();
+    }
+}

tfhe-benchmark/benches/integer/bench.rs

@@ -1,5 +1,6 @@
 #![allow(dead_code)]
 
+mod aes;
 mod oprf;
 
 use benchmark::params::ParamsAndNumBlocksIter;
@@ -2795,6 +2796,7 @@ mod cuda {
         cuda_trailing_ones,
         cuda_ilog2,
         oprf::cuda::cuda_unsigned_oprf,
+        aes::cuda::cuda_aes,
     );
 
     criterion_group!(
@@ -2823,6 +2825,7 @@ mod cuda {
         cuda_scalar_div,
         cuda_scalar_rem,
         oprf::cuda::cuda_unsigned_oprf,
+        aes::cuda::cuda_aes,
     );
 
     criterion_group!(