chore(bench): add dedicated targets for aes cuda benchmarks

.github/workflows/benchmark_gpu.yml

@@ -31,6 +31,8 @@ on:
           - ks
           - ks_pbs
           - integer_zk
+          - integer_aes
+          - integer_aes256
           - hlapi_noise_squash
       op_flavor:
         description: "Operations set to run"

.github/workflows/benchmark_gpu_weekly.yml

@@ -54,14 +54,14 @@ jobs:
       SLAB_URL: ${{ secrets.SLAB_URL }}
       SLAB_BASE_URL: ${{ secrets.SLAB_BASE_URL }}
 
-  run-benchmarks-8-h100-sxm5-integer-zk:
-    name: benchmark_gpu_weekly/run-benchmarks-8-h100-sxm5-integer-zk
+  run-benchmarks-8-h100-sxm5-integer-zk-aes:
+    name: benchmark_gpu_weekly/run-benchmarks-8-h100-sxm5-integer-zk-aes
     if: github.repository == 'zama-ai/tfhe-rs'
     uses: ./.github/workflows/benchmark_gpu_common.yml
     with:
       profile: multi-h100-sxm5
       hardware_name: n3-H100-SXM5x8
-      command: integer_zk
+      command: integer_zk,integer_aes,integer_aes256
       op_flavor: default
       bench_type: both
       all_precisions: true

Makefile

@@ -1416,6 +1416,20 @@ bench_integer_zk_gpu: install_rs_check_toolchain
 	--bench integer-zk-pke \
 	--features=integer,internal-keycache,gpu,pbs-stats,zk-pok -p tfhe-benchmark --profile release_lto_off --
 
+.PHONY: bench_integer_aes_gpu # Run benchmarks for AES on GPU backend
+bench_integer_aes_gpu: install_rs_check_toolchain
+	RUSTFLAGS="$(RUSTFLAGS)" __TFHE_RS_BENCH_TYPE=$(BENCH_TYPE) \
+	cargo $(CARGO_RS_CHECK_TOOLCHAIN) bench \
+	--bench integer-aes \
+	--features=integer,internal-keycache,gpu, -p tfhe-benchmark --profile release_lto_off --
+
+.PHONY: bench_integer_aes256_gpu # Run benchmarks for AES256 on GPU backend
+bench_integer_aes256_gpu: install_rs_check_toolchain
+	RUSTFLAGS="$(RUSTFLAGS)" __TFHE_RS_BENCH_TYPE=$(BENCH_TYPE) \
+	cargo $(CARGO_RS_CHECK_TOOLCHAIN) bench \
+	--bench integer-aes256 \
+	--features=integer,internal-keycache,gpu, -p tfhe-benchmark --profile release_lto_off --
+
 .PHONY: bench_integer_multi_bit # Run benchmarks for unsigned integer using multi-bit parameters
 bench_integer_multi_bit: install_rs_check_toolchain
 	RUSTFLAGS="$(RUSTFLAGS)" __TFHE_RS_PARAM_TYPE=MULTI_BIT __TFHE_RS_BENCH_TYPE=$(BENCH_TYPE) \

tfhe-benchmark/Cargo.toml

@@ -126,6 +126,18 @@ path = "benches/integer/zk_pke.rs"
 harness = false
 required-features = ["integer", "zk-pok", "pbs-stats", "internal-keycache"]
 
+[[bench]]
+name = "integer-aes"
+path = "benches/integer/aes.rs"
+harness = false
+required-features = ["integer", "internal-keycache"]
+
+[[bench]]
+name = "integer-aes256"
+path = "benches/integer/aes256.rs"
+harness = false
+required-features = ["integer", "internal-keycache"]
+
 [[bench]]
 name = "core_crypto-ks"
 path = "benches/core_crypto/ks_bench.rs"

tfhe-benchmark/benches/integer/aes.rs

@@ -1,8 +1,10 @@
+use criterion::Criterion;
+
 #[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::{write_to_json, OperatorType};
-    use criterion::{black_box, Criterion};
+    use criterion::{black_box, criterion_group, Criterion};
     use tfhe::core_crypto::gpu::{check_valid_cuda_malloc, CudaStreams};
     use tfhe::integer::gpu::ciphertext::CudaUnsignedRadixCiphertext;
     use tfhe::integer::gpu::CudaServerKey;
@@ -150,4 +152,16 @@ pub mod cuda {
 
         bench_group.finish();
     }
+
+    criterion_group!(gpu_aes, cuda_aes);
+}
+
+#[cfg(feature = "gpu")]
+use cuda::gpu_aes;
+
+fn main() {
+    #[cfg(feature = "gpu")]
+    gpu_aes();
+
+    Criterion::default().configure_from_args().final_summary();
 }

tfhe-benchmark/benches/integer/aes256.rs

@@ -1,8 +1,10 @@
+use criterion::Criterion;
+
 #[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::{write_to_json, OperatorType};
-    use criterion::{black_box, Criterion};
+    use criterion::{black_box, criterion_group, Criterion};
     use tfhe::core_crypto::gpu::{check_valid_cuda_malloc, CudaStreams};
     use tfhe::integer::gpu::ciphertext::CudaUnsignedRadixCiphertext;
     use tfhe::integer::gpu::CudaServerKey;
@@ -154,4 +156,15 @@ pub mod cuda {
 
         bench_group.finish();
     }
+    criterion_group!(gpu_aes_256, cuda_aes_256);
+}
+
+#[cfg(feature = "gpu")]
+use cuda::gpu_aes_256;
+
+fn main() {
+    #[cfg(feature = "gpu")]
+    gpu_aes_256();
+
+    Criterion::default().configure_from_args().final_summary();
 }

tfhe-benchmark/benches/integer/bench.rs

@@ -2799,8 +2799,6 @@ mod cuda {
         cuda_trailing_ones,
         cuda_ilog2,
         oprf::cuda::cuda_unsigned_oprf,
-        aes::cuda::cuda_aes,
-        aes256::cuda::cuda_aes_256,
     );
 
     criterion_group!(
@@ -2829,8 +2827,6 @@ mod cuda {
         cuda_scalar_div,
         cuda_scalar_rem,
         oprf::cuda::cuda_unsigned_oprf,
-        aes::cuda::cuda_aes,
-        aes256::cuda::cuda_aes_256,
     );
 
     criterion_group!(