77 lines
2.3 KiB
Rust
77 lines
2.3 KiB
Rust
//! Criterion benchmark for tensor operations
|
|
|
|
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
|
|
use rtx_bench::{BenchmarkConfig, core::tensor_ops::TensorOperationBenchmarks};
|
|
use tokio::runtime::Runtime;
|
|
|
|
fn benchmark_tensor_creation(c: &mut Criterion) {
|
|
let rt = Runtime::new().unwrap();
|
|
let benchmarks = TensorOperationBenchmarks::new();
|
|
let config = BenchmarkConfig::default()
|
|
.with_measurement_iterations(10)
|
|
.with_warmup_iterations(3);
|
|
|
|
let mut group = c.benchmark_group("tensor_creation");
|
|
|
|
for size in [64, 256, 1024].iter() {
|
|
group.bench_with_input(BenchmarkId::new("zeros", size), size, |b, &size| {
|
|
b.to_async(&rt).iter(|| async {
|
|
benchmarks
|
|
.benchmark_tensor_zeros(&config, size, size)
|
|
.await
|
|
.unwrap();
|
|
});
|
|
});
|
|
|
|
group.bench_with_input(BenchmarkId::new("random", size), size, |b, &size| {
|
|
b.to_async(&rt).iter(|| async {
|
|
benchmarks
|
|
.benchmark_tensor_random(&config, size, size)
|
|
.await
|
|
.unwrap();
|
|
});
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn benchmark_tensor_arithmetic(c: &mut Criterion) {
|
|
let rt = Runtime::new().unwrap();
|
|
let benchmarks = TensorOperationBenchmarks::new();
|
|
let config = BenchmarkConfig::default()
|
|
.with_measurement_iterations(10)
|
|
.with_warmup_iterations(3);
|
|
|
|
let mut group = c.benchmark_group("tensor_arithmetic");
|
|
|
|
for size in [256, 512, 1024].iter() {
|
|
group.bench_with_input(BenchmarkId::new("add", size), size, |b, &size| {
|
|
b.to_async(&rt).iter(|| async {
|
|
benchmarks
|
|
.benchmark_tensor_add(&config, size, size)
|
|
.await
|
|
.unwrap();
|
|
});
|
|
});
|
|
|
|
group.bench_with_input(BenchmarkId::new("matmul", size), size, |b, &size| {
|
|
b.to_async(&rt).iter(|| async {
|
|
benchmarks
|
|
.benchmark_tensor_matmul(&config, size, size)
|
|
.await
|
|
.unwrap();
|
|
});
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
criterion_group!(
|
|
benches,
|
|
benchmark_tensor_creation,
|
|
benchmark_tensor_arithmetic
|
|
);
|
|
criterion_main!(benches);
|