445 lines
13 KiB
Rust
445 lines
13 KiB
Rust
//! Core Benchmark Types
|
|
//!
|
|
//! Defines the main benchmark structures and implementation wrappers.
|
|
|
|
use std::fmt;
|
|
use std::sync::Arc;
|
|
use std::time::{Duration, Instant};
|
|
|
|
use crate::{Backend, BenchmarkConfig, BenchmarkError, Metrics, OperationType, Result};
|
|
|
|
/// A function that executes a kernel implementation
|
|
pub type KernelFn = Arc<dyn Fn(&[Vec<f32>], &[usize]) -> Vec<f32> + Send + Sync>;
|
|
|
|
/// A named implementation of a kernel
|
|
#[derive(Clone)]
|
|
pub struct Implementation {
|
|
/// Name of this implementation (e.g., "cubecl", "handcrafted")
|
|
pub name: String,
|
|
/// Description of the implementation
|
|
pub description: String,
|
|
/// The kernel function
|
|
pub kernel: KernelFn,
|
|
/// Backend this implementation targets
|
|
pub backend: Backend,
|
|
}
|
|
|
|
impl fmt::Debug for Implementation {
|
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
f.debug_struct("Implementation")
|
|
.field("name", &self.name)
|
|
.field("description", &self.description)
|
|
.field("backend", &self.backend)
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
impl Implementation {
|
|
/// Create a new implementation
|
|
pub fn new<F>(name: impl Into<String>, backend: Backend, kernel: F) -> Self
|
|
where
|
|
F: Fn(&[Vec<f32>], &[usize]) -> Vec<f32> + Send + Sync + 'static,
|
|
{
|
|
Self {
|
|
name: name.into(),
|
|
description: String::new(),
|
|
kernel: Arc::new(kernel),
|
|
backend,
|
|
}
|
|
}
|
|
|
|
/// Set the description
|
|
pub fn with_description(mut self, desc: impl Into<String>) -> Self {
|
|
self.description = desc.into();
|
|
self
|
|
}
|
|
|
|
/// Execute the kernel
|
|
pub fn execute(&self, inputs: &[Vec<f32>], shape: &[usize]) -> Vec<f32> {
|
|
(self.kernel)(inputs, shape)
|
|
}
|
|
}
|
|
|
|
/// Result of a single benchmark measurement
|
|
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
|
pub struct Measurement {
|
|
/// Duration of this measurement
|
|
pub duration: Duration,
|
|
/// Memory used (bytes) if available
|
|
pub memory_bytes: Option<u64>,
|
|
/// Throughput (elements/second) if applicable
|
|
pub throughput: Option<f64>,
|
|
}
|
|
|
|
/// Result of running a benchmark
|
|
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
|
pub struct BenchmarkResult {
|
|
/// Name of the benchmark
|
|
pub name: String,
|
|
/// Operation type
|
|
pub operation: OperationType,
|
|
/// Implementation name
|
|
pub implementation: String,
|
|
/// Backend used
|
|
pub backend: Backend,
|
|
/// Input shape
|
|
pub shape: Vec<usize>,
|
|
/// Total number of elements
|
|
pub num_elements: usize,
|
|
/// All measurements
|
|
pub measurements: Vec<Measurement>,
|
|
/// Computed metrics
|
|
pub metrics: Metrics,
|
|
/// Timestamp when benchmark was run
|
|
pub timestamp: chrono::DateTime<chrono::Utc>,
|
|
}
|
|
|
|
impl BenchmarkResult {
|
|
/// Create a new benchmark result from measurements
|
|
pub fn from_measurements(
|
|
name: String,
|
|
operation: OperationType,
|
|
implementation: String,
|
|
backend: Backend,
|
|
shape: Vec<usize>,
|
|
measurements: Vec<Measurement>,
|
|
) -> Self {
|
|
let num_elements = shape.iter().product();
|
|
let metrics = Metrics::from_measurements(&measurements);
|
|
|
|
Self {
|
|
name,
|
|
operation,
|
|
implementation,
|
|
backend,
|
|
shape,
|
|
num_elements,
|
|
measurements,
|
|
metrics,
|
|
timestamp: chrono::Utc::now(),
|
|
}
|
|
}
|
|
|
|
/// Get the mean duration
|
|
pub fn mean_duration(&self) -> Duration {
|
|
self.metrics.timing.mean
|
|
}
|
|
|
|
/// Get throughput in GFLOPS (assuming 2 FLOPs per element for most ops)
|
|
pub fn gflops(&self) -> f64 {
|
|
let flops = self.num_elements as f64 * 2.0; // Approximate
|
|
let seconds = self.metrics.timing.mean.as_secs_f64();
|
|
flops / seconds / 1e9
|
|
}
|
|
|
|
/// Get throughput in GB/s
|
|
pub fn bandwidth_gbps(&self) -> f64 {
|
|
let bytes = self.num_elements as f64 * 4.0 * 2.0; // Read + write, f32
|
|
let seconds = self.metrics.timing.mean.as_secs_f64();
|
|
bytes / seconds / 1e9
|
|
}
|
|
}
|
|
|
|
/// Builder for kernel benchmarks
|
|
#[derive(Debug)]
|
|
pub struct KernelBenchmark {
|
|
/// Name of the benchmark
|
|
name: String,
|
|
/// Operation being benchmarked
|
|
operation: OperationType,
|
|
/// Shapes to benchmark
|
|
shapes: Vec<Vec<usize>>,
|
|
/// Implementation A (typically CubeCL)
|
|
impl_a: Option<Implementation>,
|
|
/// Implementation B (typically hand-crafted)
|
|
impl_b: Option<Implementation>,
|
|
/// Additional implementations
|
|
other_impls: Vec<Implementation>,
|
|
}
|
|
|
|
impl KernelBenchmark {
|
|
/// Create a new kernel benchmark
|
|
pub fn new(name: impl Into<String>, operation: OperationType) -> Self {
|
|
Self {
|
|
name: name.into(),
|
|
operation,
|
|
shapes: Vec::new(),
|
|
impl_a: None,
|
|
impl_b: None,
|
|
other_impls: Vec::new(),
|
|
}
|
|
}
|
|
|
|
/// Add shapes to benchmark
|
|
pub fn with_shapes(mut self, shapes: Vec<Vec<usize>>) -> Self {
|
|
self.shapes = shapes;
|
|
self
|
|
}
|
|
|
|
/// Add a single shape
|
|
pub fn add_shape(mut self, shape: Vec<usize>) -> Self {
|
|
self.shapes.push(shape);
|
|
self
|
|
}
|
|
|
|
/// Set implementation A (typically CubeCL or reference)
|
|
pub fn implementation_a(mut self, impl_a: Implementation) -> Self {
|
|
self.impl_a = Some(impl_a);
|
|
self
|
|
}
|
|
|
|
/// Set implementation B (typically hand-crafted or optimized)
|
|
pub fn implementation_b(mut self, impl_b: Implementation) -> Self {
|
|
self.impl_b = Some(impl_b);
|
|
self
|
|
}
|
|
|
|
/// Add additional implementation for comparison
|
|
pub fn add_implementation(mut self, impl_: Implementation) -> Self {
|
|
self.other_impls.push(impl_);
|
|
self
|
|
}
|
|
|
|
/// Get all implementations
|
|
pub fn implementations(&self) -> Vec<&Implementation> {
|
|
let mut impls = Vec::new();
|
|
if let Some(ref a) = self.impl_a {
|
|
impls.push(a);
|
|
}
|
|
if let Some(ref b) = self.impl_b {
|
|
impls.push(b);
|
|
}
|
|
impls.extend(self.other_impls.iter());
|
|
impls
|
|
}
|
|
|
|
/// Get the benchmark name
|
|
pub fn name(&self) -> &str {
|
|
&self.name
|
|
}
|
|
|
|
/// Get the operation type
|
|
pub fn operation(&self) -> &OperationType {
|
|
&self.operation
|
|
}
|
|
|
|
/// Run the benchmark with the given configuration
|
|
pub fn run(&self, config: &BenchmarkConfig) -> Result<Vec<BenchmarkResult>> {
|
|
let mut results = Vec::new();
|
|
|
|
// Validate
|
|
if self.shapes.is_empty() {
|
|
return Err(BenchmarkError::ConfigError("No shapes specified".into()));
|
|
}
|
|
|
|
let impls = self.implementations();
|
|
if impls.is_empty() {
|
|
return Err(BenchmarkError::ConfigError(
|
|
"No implementations specified".into(),
|
|
));
|
|
}
|
|
|
|
// Run benchmarks for each shape and implementation
|
|
for shape in &self.shapes {
|
|
// Generate test inputs
|
|
let inputs = self.generate_inputs(shape);
|
|
|
|
for impl_ in &impls {
|
|
if !impl_.backend.is_available() && !config.skip_unavailable_backends {
|
|
return Err(BenchmarkError::BackendNotAvailable(
|
|
impl_.backend.name().to_string(),
|
|
));
|
|
}
|
|
|
|
if !impl_.backend.is_available() {
|
|
continue;
|
|
}
|
|
|
|
let measurements = self.run_implementation(impl_, &inputs, shape, config)?;
|
|
|
|
results.push(BenchmarkResult::from_measurements(
|
|
self.name.clone(),
|
|
self.operation.clone(),
|
|
impl_.name.clone(),
|
|
impl_.backend,
|
|
shape.clone(),
|
|
measurements,
|
|
));
|
|
}
|
|
}
|
|
|
|
Ok(results)
|
|
}
|
|
|
|
/// Generate random inputs for a given shape
|
|
fn generate_inputs(&self, shape: &[usize]) -> Vec<Vec<f32>> {
|
|
let num_elements: usize = shape.iter().product();
|
|
|
|
// Generate 1-2 input tensors depending on operation type
|
|
let num_inputs = if self.operation.is_elementwise() {
|
|
2
|
|
} else {
|
|
1
|
|
};
|
|
|
|
(0..num_inputs)
|
|
.map(|_| {
|
|
(0..num_elements)
|
|
.map(|i| ((i as f32 * 0.001) % 1.0) - 0.5)
|
|
.collect()
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
/// Run a single implementation
|
|
fn run_implementation(
|
|
&self,
|
|
impl_: &Implementation,
|
|
inputs: &[Vec<f32>],
|
|
shape: &[usize],
|
|
config: &BenchmarkConfig,
|
|
) -> Result<Vec<Measurement>> {
|
|
use crate::kernels::memory::MemorySnapshot;
|
|
|
|
// Warmup
|
|
for _ in 0..config.warmup.iterations {
|
|
let _ = impl_.execute(inputs, shape);
|
|
}
|
|
|
|
// Actual measurements
|
|
let mut measurements = Vec::with_capacity(config.iterations);
|
|
|
|
for _ in 0..config.iterations {
|
|
// Take memory snapshot before if tracking enabled
|
|
let mem_before = if config.track_memory {
|
|
Some(MemorySnapshot::now())
|
|
} else {
|
|
None
|
|
};
|
|
|
|
let start = Instant::now();
|
|
let _ = impl_.execute(inputs, shape);
|
|
let duration = start.elapsed();
|
|
|
|
// Take memory snapshot after and compute delta
|
|
let memory_bytes = if config.track_memory {
|
|
let mem_after = MemorySnapshot::now();
|
|
match (mem_before, mem_after.total()) {
|
|
(Some(_before), Some(after)) => {
|
|
// Report current memory usage (not delta, which could be negative)
|
|
Some(after)
|
|
}
|
|
(None, Some(after)) => Some(after),
|
|
_ => None,
|
|
}
|
|
} else {
|
|
None
|
|
};
|
|
|
|
measurements.push(Measurement {
|
|
duration,
|
|
memory_bytes,
|
|
throughput: None, // Computed later in Metrics
|
|
});
|
|
}
|
|
|
|
Ok(measurements)
|
|
}
|
|
}
|
|
|
|
/// Standard benchmark shapes for common operations
|
|
pub mod shapes {
|
|
/// Small shapes for quick testing
|
|
pub fn small() -> Vec<Vec<usize>> {
|
|
vec![vec![64, 64], vec![128, 128], vec![256, 256]]
|
|
}
|
|
|
|
/// Medium shapes for typical workloads
|
|
pub fn medium() -> Vec<Vec<usize>> {
|
|
vec![vec![512, 512], vec![1024, 1024], vec![2048, 2048]]
|
|
}
|
|
|
|
/// Large shapes for stress testing
|
|
pub fn large() -> Vec<Vec<usize>> {
|
|
vec![vec![4096, 4096], vec![8192, 8192]]
|
|
}
|
|
|
|
/// LLM-specific shapes (batch, seq_len, hidden_dim)
|
|
pub fn llm() -> Vec<Vec<usize>> {
|
|
vec![
|
|
vec![1, 512, 4096], // Single sequence
|
|
vec![8, 512, 4096], // Small batch
|
|
vec![32, 512, 4096], // Medium batch
|
|
vec![1, 2048, 4096], // Long sequence
|
|
vec![8, 2048, 4096], // Long batch
|
|
]
|
|
}
|
|
|
|
/// Attention shapes (batch, heads, seq_len, head_dim)
|
|
pub fn attention() -> Vec<Vec<usize>> {
|
|
vec![
|
|
vec![1, 32, 512, 128],
|
|
vec![8, 32, 512, 128],
|
|
vec![1, 32, 2048, 128],
|
|
vec![8, 32, 2048, 128],
|
|
]
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
fn dummy_kernel(inputs: &[Vec<f32>], shape: &[usize]) -> Vec<f32> {
|
|
let size: usize = shape.iter().product();
|
|
vec![0.0; size]
|
|
}
|
|
|
|
#[test]
|
|
fn test_implementation_creation() {
|
|
let impl_ = Implementation::new("test", Backend::Cpu, dummy_kernel)
|
|
.with_description("A test implementation");
|
|
|
|
assert_eq!(impl_.name, "test");
|
|
assert_eq!(impl_.backend, Backend::Cpu);
|
|
}
|
|
|
|
#[test]
|
|
fn test_benchmark_builder() {
|
|
let bench = KernelBenchmark::new("test_add", OperationType::Add)
|
|
.with_shapes(vec![vec![64, 64], vec![128, 128]])
|
|
.implementation_a(Implementation::new("impl_a", Backend::Cpu, dummy_kernel))
|
|
.implementation_b(Implementation::new("impl_b", Backend::Cpu, dummy_kernel));
|
|
|
|
assert_eq!(bench.implementations().len(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_run_benchmark() {
|
|
let bench = KernelBenchmark::new("test_add", OperationType::Add)
|
|
.with_shapes(vec![vec![64, 64]])
|
|
.implementation_a(Implementation::new("cpu", Backend::Cpu, dummy_kernel));
|
|
|
|
let config = BenchmarkConfig {
|
|
iterations: 3,
|
|
warmup: crate::WarmupConfig {
|
|
iterations: 1,
|
|
..Default::default()
|
|
},
|
|
skip_unavailable_backends: true,
|
|
..Default::default()
|
|
};
|
|
|
|
let results = bench.run(&config).unwrap();
|
|
assert_eq!(results.len(), 1);
|
|
assert_eq!(results[0].measurements.len(), 3);
|
|
}
|
|
|
|
#[test]
|
|
fn test_shape_generators() {
|
|
assert!(!shapes::small().is_empty());
|
|
assert!(!shapes::medium().is_empty());
|
|
assert!(!shapes::llm().is_empty());
|
|
}
|
|
}
|