Skip to content

Commit ee41b37

Browse files
SouchonTheoSouchonTheo
committed
chore(bench): migrate transciphering, bitonic shuffle and bin benchmarks into the spec
1 parent 2925fc7 commit ee41b37

24 files changed

Lines changed: 609 additions & 180 deletions

File tree

‎Makefile‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2527,7 +2527,7 @@ endif
25272527
RUSTFLAGS="$(RUSTFLAGS)" __TFHE_RS_PARAM_TYPE=$(BENCH_PARAM_TYPE) __TFHE_RS_BENCH_TYPE=$(BENCH_TYPE) \
25282528
cargo $(CARGO_RS_CHECK_TOOLCHAIN) bench \
25292529
--bench hlapi-bitonic-shuffle \
2530-
--features=integer,gpu,internal-keycache,pbs-stats -p tfhe-benchmark --profile release_lto_off -- 'hlapi::bitonic_shuffle_gpu::summary::'
2530+
--features=integer,gpu,internal-keycache,pbs-stats -p tfhe-benchmark --profile release_lto_off -- 'tfhe::hlapi::bitonic_shuffle::cuda::'
25312531

25322532
# Key-Value store (summary: N=64, u64 key, u64 value)
25332533
RUSTFLAGS="$(RUSTFLAGS)" __TFHE_RS_PARAM_TYPE=$(BENCH_PARAM_TYPE) __TFHE_RS_BENCH_TYPE=$(BENCH_TYPE) __TFHE_RS_BENCH_BIT_SIZES_SET=FAST \

‎tfhe-benchmark/benches/high_level_api/bitonic_shuffle.rs‎

Lines changed: 40 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,11 @@ use benchmark::params_aliases::BENCH_PARAM_GPU_MULTI_BIT_GROUP_4_MESSAGE_2_CARRY
33
use benchmark::params_aliases::BENCH_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
44
#[cfg(feature = "gpu")]
55
use benchmark::utilities::{get_param_type, ParamType};
6-
use benchmark::utilities::{write_to_json_unchecked, OperatorType};
7-
use benchmark_spec::{get_bench_type, BenchmarkType};
6+
use benchmark::utilities::{write_to_json, OperatorType};
7+
use benchmark_spec::{
8+
get_bench_type, Backend, BenchCrate, BenchmarkSpec, BenchmarkType, HlapiBench, OperandType,
9+
ShuffleConfig, TfheLayer,
10+
};
811
use criterion::{criterion_group, Criterion, Throughput};
912
use rand::prelude::*;
1013
use rayon::prelude::*;
@@ -131,7 +134,7 @@ fn bench_collision_probability(c: &mut Criterion, cks: &ClientKey, bench_name: &
131134
fn bench_shuffle_config<T>(
132135
group: &mut criterion::BenchmarkGroup<'_, criterion::measurement::WallTime>,
133136
cks: &ClientKey,
134-
bench_name: &str,
137+
backend: Backend,
135138
value_bits: u32,
136139
key_bits: u32,
137140
num_elements: usize,
@@ -142,18 +145,34 @@ fn bench_shuffle_config<T>(
142145
let key_size = BitonicShuffleKeySize::num_bits(key_bits);
143146
let params = cks.computation_parameters();
144147
let params_name = params.name();
145-
let stem =
146-
format!("{bench_name}::{value_bits}_bits::{num_elements}_elements::key_{key_bits}_bits");
148+
149+
// The CPU and GPU entry points both run in a GPU build, so the backend is
150+
// passed in rather than read from the enabled features.
151+
let spec = BenchmarkSpec::new(
152+
BenchCrate::Tfhe(TfheLayer::Hlapi(HlapiBench::BitonicShuffle)),
153+
backend,
154+
&params_name,
155+
OperandType::CipherText,
156+
Some(
157+
ShuffleConfig {
158+
value_bits,
159+
key_bits,
160+
}
161+
.into(),
162+
),
163+
get_bench_type(),
164+
Some(num_elements),
165+
);
166+
let bench_id = spec.to_string();
147167

148168
let encrypt_dataset = |rng: &mut ThreadRng| -> Vec<T> {
149169
(0..num_elements)
150170
.map(|_| T::encrypt(rng.gen::<u128>(), cks))
151171
.collect()
152172
};
153173

154-
let bench_id = match get_bench_type() {
174+
match get_bench_type() {
155175
BenchmarkType::Latency => {
156-
let bench_id = format!("{stem}::{params_name}");
157176
group.bench_function(&bench_id, |b| {
158177
b.iter_batched(
159178
|| (encrypt_dataset(&mut rng), Seed(rng.gen())),
@@ -165,7 +184,6 @@ fn bench_shuffle_config<T>(
165184
criterion::BatchSize::SmallInput,
166185
)
167186
});
168-
bench_id
169187
}
170188
BenchmarkType::Throughput => {
171189
let num_ops = {
@@ -196,7 +214,6 @@ fn bench_shuffle_config<T>(
196214
}
197215
};
198216

199-
let bench_id = format!("{stem}::throughput::{params_name}");
200217
group.throughput(Throughput::Elements(num_ops as u64));
201218
group.bench_function(&bench_id, |b| {
202219
b.iter_batched(
@@ -215,29 +232,27 @@ fn bench_shuffle_config<T>(
215232
criterion::BatchSize::SmallInput,
216233
)
217234
});
218-
bench_id
219235
}
220-
};
236+
}
221237

222-
write_to_json_unchecked(
223-
&bench_id,
224-
params_name,
238+
write_to_json(
239+
&spec,
225240
"bitonic_shuffle",
226241
&OperatorType::Atomic,
227242
value_bits,
228243
vec![],
229244
);
230245
}
231246

232-
fn bench_shuffle_configs(c: &mut Criterion, cks: &ClientKey, bench_name: &str) {
247+
fn bench_shuffle_configs(c: &mut Criterion, cks: &ClientKey, bench_name: &str, backend: Backend) {
233248
let mut group = c.benchmark_group(bench_name);
234249
group
235250
.sample_size(10)
236251
.measurement_time(std::time::Duration::from_secs(60));
237252

238-
bench_shuffle_config::<FheUint64>(&mut group, cks, bench_name, 64, 16, 16);
239-
bench_shuffle_config::<FheUint64>(&mut group, cks, bench_name, 64, 32, 16);
240-
bench_shuffle_config::<FheUint160>(&mut group, cks, bench_name, 160, 16, 16);
253+
bench_shuffle_config::<FheUint64>(&mut group, cks, backend, 64, 16, 16);
254+
bench_shuffle_config::<FheUint64>(&mut group, cks, backend, 64, 32, 16);
255+
bench_shuffle_config::<FheUint160>(&mut group, cks, backend, 160, 16, 16);
241256

242257
group.finish();
243258
}
@@ -363,7 +378,7 @@ pub fn bitonic_shuffle_cpu(c: &mut Criterion) {
363378
bench_non_pow2(c, &cks, "hlapi::bitonic_shuffle_cpu::non_pow2");
364379
bench_key_size_sweep(c, &cks, "hlapi::bitonic_shuffle_cpu::key_size_sweep");
365380
bench_collision_probability(c, &cks, "hlapi::bitonic_shuffle_cpu::collision_probability");
366-
bench_shuffle_configs(c, &cks, "hlapi::bitonic_shuffle_cpu::summary");
381+
bench_shuffle_configs(c, &cks, "hlapi::bitonic_shuffle_cpu::summary", Backend::Cpu);
367382

368383
let cpu_cks = tfhe::integer::ClientKey::new(param);
369384
bench_unchecked_with_keys_cpu_inner(
@@ -392,7 +407,12 @@ pub fn bitonic_shuffle_gpu(c: &mut Criterion) {
392407
bench_non_pow2(c, &cks, "hlapi::bitonic_shuffle_gpu::non_pow2");
393408
bench_key_size_sweep(c, &cks, "hlapi::bitonic_shuffle_gpu::key_size_sweep");
394409
bench_collision_probability(c, &cks, "hlapi::bitonic_shuffle_gpu::collision_probability");
395-
bench_shuffle_configs(c, &cks, "hlapi::bitonic_shuffle_gpu::summary");
410+
bench_shuffle_configs(
411+
c,
412+
&cks,
413+
"hlapi::bitonic_shuffle_gpu::summary",
414+
Backend::Cuda,
415+
);
396416

397417
let cpu_cks = tfhe::integer::ClientKey::new(param);
398418
bench_unchecked_with_keys_gpu_inner(

‎tfhe-benchmark/benches/integer/aes.rs‎

Lines changed: 33 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,9 @@ use criterion::Criterion;
33
#[cfg(feature = "gpu")]
44
pub mod cuda {
55
use benchmark::params_aliases::BENCH_PARAM_GPU_MULTI_BIT_GROUP_4_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
6-
use benchmark::utilities::{write_to_json_unchecked, OperatorType};
6+
use benchmark::utilities::{write_to_json, OperatorType};
7+
use benchmark_spec::tfhe::transciphering::aes::AesFlavor;
8+
use benchmark_spec::{BenchmarkMetric, BenchmarkSpec, PrecisionTag, TranscipheringBench};
79
use criterion::{criterion_group, Criterion};
810
use std::hint::black_box;
911
use tfhe::core_crypto::gpu::{check_valid_cuda_malloc, CudaStreams};
@@ -46,7 +48,14 @@ pub mod cuda {
4648
{
4749
const NUM_AES_INPUTS: usize = 1;
4850
const SBOX_PARALLELISM: usize = 16;
49-
let bench_id = format!("{bench_name}::{param_name}::{NUM_AES_INPUTS}_input_encryption");
51+
let spec = BenchmarkSpec::new_transciphering(
52+
TranscipheringBench::Aes(AesFlavor::Encryption),
53+
&param_name,
54+
Some(PrecisionTag::Bits(aes_op_bit_size as usize).into()),
55+
BenchmarkMetric::Latency,
56+
Some(NUM_AES_INPUTS),
57+
);
58+
let bench_id = spec.to_string();
5059

5160
let round_keys = sks.key_expansion(&d_key, &streams);
5261

@@ -63,9 +72,8 @@ pub mod cuda {
6372
})
6473
});
6574

66-
write_to_json_unchecked(
67-
&bench_id,
68-
param.name(),
75+
write_to_json(
76+
&spec,
6977
"aes_encryption",
7078
&OperatorType::Atomic,
7179
aes_op_bit_size,
@@ -74,16 +82,22 @@ pub mod cuda {
7482
}
7583

7684
{
77-
let bench_id = format!("{bench_name}::{param_name}::key_expansion");
85+
let spec = BenchmarkSpec::new_transciphering(
86+
TranscipheringBench::Aes(AesFlavor::KeyExpansion),
87+
&param_name,
88+
Some(PrecisionTag::Bits(aes_op_bit_size as usize).into()),
89+
BenchmarkMetric::Latency,
90+
None,
91+
);
92+
let bench_id = spec.to_string();
7893
bench_group.bench_function(&bench_id, |b| {
7994
b.iter(|| {
8095
black_box(sks.key_expansion(&d_key, &streams));
8196
})
8297
});
8398

84-
write_to_json_unchecked(
85-
&bench_id,
86-
param.name(),
99+
write_to_json(
100+
&spec,
87101
"aes_key_expansion",
88102
&OperatorType::Atomic,
89103
aes_op_bit_size,
@@ -94,8 +108,14 @@ pub mod cuda {
94108
{
95109
const NUM_AES_INPUTS: usize = 192;
96110
const SBOX_PARALLELISM: usize = 16;
97-
let bench_id =
98-
format!("{bench_name}::{param_name}::{NUM_AES_INPUTS}_inputs_encryption");
111+
let spec = BenchmarkSpec::new_transciphering(
112+
TranscipheringBench::Aes(AesFlavor::Encryption),
113+
&param_name,
114+
Some(PrecisionTag::Bits(aes_op_bit_size as usize).into()),
115+
BenchmarkMetric::Latency,
116+
Some(NUM_AES_INPUTS),
117+
);
118+
let bench_id = spec.to_string();
99119

100120
let streams = CudaStreams::new_multi_gpu();
101121
let (cpu_cks, _) = KEY_CACHE.get_from_params(atomic_param, IntegerKeyKind::Radix);
@@ -135,9 +155,8 @@ pub mod cuda {
135155
})
136156
});
137157

138-
write_to_json_unchecked(
139-
&bench_id,
140-
param.name(),
158+
write_to_json(
159+
&spec,
141160
"aes_encryption",
142161
&OperatorType::Atomic,
143162
aes_op_bit_size,

‎tfhe-benchmark/benches/integer/aes256.rs‎

Lines changed: 33 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,9 @@ use criterion::Criterion;
33
#[cfg(feature = "gpu")]
44
pub mod cuda {
55
use benchmark::params_aliases::BENCH_PARAM_GPU_MULTI_BIT_GROUP_4_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128;
6-
use benchmark::utilities::{write_to_json_unchecked, OperatorType};
6+
use benchmark::utilities::{write_to_json, OperatorType};
7+
use benchmark_spec::tfhe::transciphering::aes::AesFlavor;
8+
use benchmark_spec::{BenchmarkMetric, BenchmarkSpec, PrecisionTag, TranscipheringBench};
79
use criterion::{criterion_group, Criterion};
810
use std::hint::black_box;
911
use tfhe::core_crypto::gpu::{check_valid_cuda_malloc, CudaStreams};
@@ -50,7 +52,14 @@ pub mod cuda {
5052
{
5153
const NUM_AES_INPUTS: usize = 1;
5254
const SBOX_PARALLELISM: usize = 16;
53-
let bench_id = format!("{bench_name}::{param_name}::{NUM_AES_INPUTS}_input_encryption");
55+
let spec = BenchmarkSpec::new_transciphering(
56+
TranscipheringBench::Aes256(AesFlavor::Encryption),
57+
&param_name,
58+
Some(PrecisionTag::Bits(aes_block_op_bit_size as usize).into()),
59+
BenchmarkMetric::Latency,
60+
Some(NUM_AES_INPUTS),
61+
);
62+
let bench_id = spec.to_string();
5463

5564
let round_keys = sks.key_expansion_256(&d_key, &streams);
5665

@@ -67,9 +76,8 @@ pub mod cuda {
6776
})
6877
});
6978

70-
write_to_json_unchecked(
71-
&bench_id,
72-
param.name(),
79+
write_to_json(
80+
&spec,
7381
"aes_256_encryption",
7482
&OperatorType::Atomic,
7583
aes_block_op_bit_size,
@@ -78,16 +86,22 @@ pub mod cuda {
7886
}
7987

8088
{
81-
let bench_id = format!("{bench_name}::{param_name}::key_expansion");
89+
let spec = BenchmarkSpec::new_transciphering(
90+
TranscipheringBench::Aes256(AesFlavor::KeyExpansion),
91+
&param_name,
92+
Some(PrecisionTag::Bits(aes_key_op_bit_size as usize).into()),
93+
BenchmarkMetric::Latency,
94+
None,
95+
);
96+
let bench_id = spec.to_string();
8297
bench_group.bench_function(&bench_id, |b| {
8398
b.iter(|| {
8499
black_box(sks.key_expansion_256(&d_key, &streams));
85100
})
86101
});
87102

88-
write_to_json_unchecked(
89-
&bench_id,
90-
param.name(),
103+
write_to_json(
104+
&spec,
91105
"aes_256_key_expansion",
92106
&OperatorType::Atomic,
93107
aes_key_op_bit_size,
@@ -98,8 +112,14 @@ pub mod cuda {
98112
{
99113
const NUM_AES_INPUTS: usize = 192;
100114
const SBOX_PARALLELISM: usize = 16;
101-
let bench_id =
102-
format!("{bench_name}::{param_name}::{NUM_AES_INPUTS}_inputs_encryption");
115+
let spec = BenchmarkSpec::new_transciphering(
116+
TranscipheringBench::Aes256(AesFlavor::Encryption),
117+
&param_name,
118+
Some(PrecisionTag::Bits(aes_block_op_bit_size as usize).into()),
119+
BenchmarkMetric::Latency,
120+
Some(NUM_AES_INPUTS),
121+
);
122+
let bench_id = spec.to_string();
103123

104124
let streams = CudaStreams::new_multi_gpu();
105125
let (cpu_cks, _) = KEY_CACHE.get_from_params(atomic_param, IntegerKeyKind::Radix);
@@ -139,9 +159,8 @@ pub mod cuda {
139159
})
140160
});
141161

142-
write_to_json_unchecked(
143-
&bench_id,
144-
param.name(),
162+
write_to_json(
163+
&spec,
145164
"aes_256_encryption",
146165
&OperatorType::Atomic,
147166
aes_block_op_bit_size,

0 commit comments

Comments
 (0)