1115 lines
38 KiB
Rust
1115 lines
38 KiB
Rust
//! Criterion benchmarks for PINN MRE Helmholtz solver
|
|
//!
|
|
//! Benchmarks:
|
|
//! - Forward pass (network evaluation)
|
|
//! - PDE residual computation
|
|
//! - Single training step
|
|
//! - Full training epoch
|
|
//!
|
|
//! Run with: cargo bench --bench pinn_benchmark
|
|
|
|
use criterion::{criterion_group, criterion_main, Criterion, BenchmarkId, Throughput};
|
|
use pinn_mre_helmholtz::{Config, Mre1DPinnSolver, LffnUNet1D, ForwardWorkspace, GpuLossAccumulator, synthesize_displacement, calculate_k};
|
|
use rtx_tensor::{Tensor, Device};
|
|
|
|
/// Benchmark forward pass (network evaluation only)
|
|
fn benchmark_forward_pass(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("forward_pass");
|
|
group.sample_size(50);
|
|
|
|
for n_points in [200, 1000, 10000] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 1,
|
|
..Config::default()
|
|
};
|
|
|
|
// Setup solver to get network and input tensor
|
|
let solver = Mre1DPinnSolver::new(cfg.clone()).expect("Failed to create solver");
|
|
|
|
group.throughput(Throughput::Elements(n_points as u64));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("lffn_mlp", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
solver.u_net.forward(&solver.x_data).expect("Forward pass failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark optimized forward pass with workspace (Phase 2 optimizations)
|
|
fn benchmark_forward_pass_optimized(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("forward_pass_optimized");
|
|
group.sample_size(50);
|
|
|
|
// Include 100000 to test rayon parallelism (threshold is 100k)
|
|
for n_points in [200, 1000, 10000, 100000] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 1,
|
|
..Config::default()
|
|
};
|
|
|
|
#[cfg(feature = "cuda")]
|
|
let device = Device::Cuda(0);
|
|
#[cfg(not(feature = "cuda"))]
|
|
let device = Device::Cpu;
|
|
|
|
// Setup solver and workspace
|
|
let solver = Mre1DPinnSolver::new(cfg.clone()).expect("Failed to create solver");
|
|
let mut ws = ForwardWorkspace::new(n_points, &cfg, &device).expect("Failed to create workspace");
|
|
|
|
group.throughput(Throughput::Elements(n_points as u64));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("lffn_mlp_workspace", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
solver.u_net.forward_with_workspace(&solver.x_data, &mut ws)
|
|
.expect("Forward pass failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark PDE residual computation (analytical derivatives)
|
|
fn benchmark_pde_residual_analytical(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("pde_residual_analytical");
|
|
group.sample_size(50);
|
|
|
|
for n_points in [200, 1000, 10000] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 1,
|
|
..Config::default()
|
|
};
|
|
|
|
let solver = Mre1DPinnSolver::new(cfg.clone()).expect("Failed to create solver");
|
|
|
|
group.throughput(Throughput::Elements(n_points as u64));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("helmholtz_cpu", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
solver.compute_pde_residual_analytical()
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark PDE residual computation (tensor operations - GPU accelerable)
|
|
fn benchmark_pde_residual_tensor(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("pde_residual_tensor");
|
|
group.sample_size(50);
|
|
|
|
for n_points in [200, 1000, 10000] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 1,
|
|
..Config::default()
|
|
};
|
|
|
|
let solver = Mre1DPinnSolver::new(cfg.clone()).expect("Failed to create solver");
|
|
|
|
group.throughput(Throughput::Elements(n_points as u64));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("helmholtz_tensor", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
solver.compute_pde_residual_tensor().expect("PDE residual failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark single training step (forward + loss + scheduler)
|
|
fn benchmark_training_step(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("training_step");
|
|
group.sample_size(30);
|
|
|
|
for n_points in [200, 1000] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 1,
|
|
..Config::default()
|
|
};
|
|
|
|
// Need mutable solver for training step
|
|
let mut solver = Mre1DPinnSolver::new(cfg.clone()).expect("Failed to create solver");
|
|
|
|
group.throughput(Throughput::Elements(n_points as u64));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("single_step", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
solver.training_step().expect("Training step failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark optimized training step with workspace (Phase 2 optimizations)
|
|
fn benchmark_training_step_optimized(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("training_step_optimized");
|
|
group.sample_size(30);
|
|
|
|
for n_points in [200, 1000] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 1,
|
|
..Config::default()
|
|
};
|
|
|
|
// Need mutable solver and workspace for training step
|
|
let mut solver = Mre1DPinnSolver::new(cfg.clone()).expect("Failed to create solver");
|
|
let mut ws = solver.create_workspace().expect("Failed to create workspace");
|
|
|
|
group.throughput(Throughput::Elements(n_points as u64));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("single_step_workspace", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
solver.training_step_with_workspace(&mut ws).expect("Training step failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark data generation (synthetic displacement)
|
|
fn benchmark_data_generation(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("data_generation");
|
|
group.sample_size(100);
|
|
|
|
for n_points in [200, 1000, 10000, 100000] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
..Config::default()
|
|
};
|
|
|
|
group.throughput(Throughput::Elements(n_points as u64));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("synthesize_displacement", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
synthesize_displacement(&cfg)
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark wave number calculation
|
|
fn benchmark_wave_number(c: &mut Criterion) {
|
|
let cfg = Config::default();
|
|
|
|
c.bench_function("calculate_k", |b| {
|
|
b.iter(|| calculate_k(&cfg))
|
|
});
|
|
}
|
|
|
|
/// Benchmark MSE loss computation
|
|
fn benchmark_mse_loss(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("mse_loss");
|
|
group.sample_size(50);
|
|
|
|
#[cfg(feature = "cuda")]
|
|
let device = Device::Cuda(0);
|
|
#[cfg(not(feature = "cuda"))]
|
|
let device = Device::Cpu;
|
|
|
|
for n_points in [200, 1000, 10000] {
|
|
// Create random tensors for MSE computation
|
|
let data1: Vec<f32> = (0..n_points * 2).map(|_| fastrand::f32()).collect();
|
|
let data2: Vec<f32> = (0..n_points * 2).map(|_| fastrand::f32()).collect();
|
|
|
|
let tensor1 = Tensor::from_slice(&data1, &[n_points, 2], &device)
|
|
.expect("Failed to create tensor");
|
|
let tensor2 = Tensor::from_slice(&data2, &[n_points, 2], &device)
|
|
.expect("Failed to create tensor");
|
|
|
|
group.throughput(Throughput::Elements(n_points as u64 * 2));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("mse_computation", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
Mre1DPinnSolver::mse_loss(&tensor1, &tensor2).expect("MSE failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark full training (100 epochs) - UNOPTIMIZED baseline
|
|
fn benchmark_training_100_epochs(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("training_100_epochs");
|
|
group.sample_size(10);
|
|
|
|
for n_points in [200] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 100,
|
|
print_every: 1000, // Disable printing
|
|
..Config::default()
|
|
};
|
|
|
|
group.throughput(Throughput::Elements(100));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("train_100", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
let mut solver = Mre1DPinnSolver::new(cfg.clone())
|
|
.expect("Failed to create solver");
|
|
solver.train().expect("Training failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark full training (100 epochs) - OPTIMIZED with cached PDE tensors
|
|
/// This benchmark demonstrates the Phase 1 optimization that eliminates ~89% overhead
|
|
/// by pre-computing PDE tensors at solver construction.
|
|
fn benchmark_training_100_epochs_cached(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("training_100_epochs_cached");
|
|
group.sample_size(10);
|
|
|
|
for n_points in [200] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 100,
|
|
print_every: 1000, // Disable printing
|
|
..Config::default()
|
|
};
|
|
|
|
group.throughput(Throughput::Elements(100));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("train_100_cached", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
let mut solver = Mre1DPinnSolver::new(cfg.clone())
|
|
.expect("Failed to create solver");
|
|
solver.train_cached().expect("Training failed") // Use cached version!
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark single training step with CACHED PDE tensors
|
|
fn benchmark_training_step_cached(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("training_step_cached");
|
|
group.sample_size(30);
|
|
|
|
for n_points in [200, 1000] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 1,
|
|
..Config::default()
|
|
};
|
|
|
|
let mut solver = Mre1DPinnSolver::new(cfg.clone()).expect("Failed to create solver");
|
|
|
|
group.throughput(Throughput::Elements(n_points as u64));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("single_step_cached", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
solver.training_step_cached().expect("Training step failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark FULLY OPTIMIZED training step: workspace + cached PDE
|
|
fn benchmark_training_step_fully_optimized(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("training_step_fully_optimized");
|
|
group.sample_size(30);
|
|
|
|
for n_points in [200, 1000] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 1,
|
|
..Config::default()
|
|
};
|
|
|
|
let mut solver = Mre1DPinnSolver::new(cfg.clone()).expect("Failed to create solver");
|
|
let mut ws = solver.create_workspace().expect("Failed to create workspace");
|
|
|
|
group.throughput(Throughput::Elements(n_points as u64));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("single_step_fully_opt", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
solver.training_step_fully_optimized(&mut ws).expect("Training step failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark FULLY OPTIMIZED 100 epochs: workspace + cached PDE
|
|
fn benchmark_training_100_epochs_fully_optimized(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("training_100_epochs_fully_optimized");
|
|
group.sample_size(10);
|
|
|
|
for n_points in [200] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 100,
|
|
print_every: 1000,
|
|
..Config::default()
|
|
};
|
|
|
|
group.throughput(Throughput::Elements(100));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("train_100_fully_opt", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
let mut solver = Mre1DPinnSolver::new(cfg.clone())
|
|
.expect("Failed to create solver");
|
|
solver.train_fully_optimized().expect("Training failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark forward pass with CUDA graph acceleration
|
|
#[cfg(feature = "cuda")]
|
|
fn benchmark_forward_pass_cuda_graph(c: &mut Criterion) {
|
|
use pinn_mre_helmholtz::CapturedForwardGraph;
|
|
|
|
let mut group = c.benchmark_group("forward_pass_cuda_graph");
|
|
group.sample_size(50);
|
|
|
|
for n_points in [200, 1000, 10000] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 1,
|
|
..Config::default()
|
|
};
|
|
|
|
let device = Device::Cuda(0);
|
|
let solver = Mre1DPinnSolver::new(cfg.clone()).expect("Failed to create solver");
|
|
let mut ws = ForwardWorkspace::new(n_points, &cfg, &device).expect("Failed to create workspace");
|
|
let mut graph = CapturedForwardGraph::new();
|
|
|
|
// Warm up: capture the graph on first call
|
|
solver.u_net.forward_with_graph(&solver.x_data, &mut ws, &mut graph)
|
|
.expect("Initial graph capture failed");
|
|
|
|
group.throughput(Throughput::Elements(n_points as u64));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("lffn_mlp_cuda_graph", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
solver.u_net.forward_with_graph(&solver.x_data, &mut ws, &mut graph)
|
|
.expect("Graph launch failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark forward pass with uber-kernel (single launch, forward mode AD)
|
|
#[cfg(feature = "cuda")]
|
|
fn benchmark_forward_pass_uber_kernel(c: &mut Criterion) {
|
|
use pinn_mre_helmholtz::uber_kernel::{forward_uber_value_only, forward_uber_with_gradients};
|
|
|
|
let mut group = c.benchmark_group("forward_pass_uber_kernel");
|
|
group.sample_size(100);
|
|
|
|
for n_points in [200, 1000, 10000] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 1,
|
|
..Config::default()
|
|
};
|
|
|
|
let solver = Mre1DPinnSolver::new(cfg.clone()).expect("Failed to create solver");
|
|
|
|
// Extract weights using new helper method
|
|
let weights = solver.u_net.to_uber_kernel_weights(&cfg);
|
|
|
|
let scale = 2.0 * std::f64::consts::PI as f32;
|
|
|
|
// Benchmark value-only (no gradients)
|
|
group.throughput(Throughput::Elements(n_points as u64));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("uber_value_only", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
forward_uber_value_only(&solver.x_data, &weights, scale)
|
|
.expect("Uber-kernel forward failed")
|
|
})
|
|
},
|
|
);
|
|
|
|
// Benchmark with gradients (forward mode AD)
|
|
group.bench_with_input(
|
|
BenchmarkId::new("uber_with_gradients", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
forward_uber_with_gradients(&solver.x_data, &weights, scale)
|
|
.expect("Uber-kernel forward with gradients failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark forward pass with MATMULS-ONLY CUDA graph
|
|
///
|
|
/// Strategy: Pre/Post kernel execution with matmul graph
|
|
/// 1. Fourier features: Safe API (before graph)
|
|
/// 2. Matmuls (5x): CUDA Graph (captured and replayed)
|
|
/// 3. Bias+activation: Safe API (after graph)
|
|
#[cfg(feature = "cuda")]
|
|
fn benchmark_forward_pass_matmuls_graph(c: &mut Criterion) {
|
|
use pinn_mre_helmholtz::{CachedGpuPtrs, UnsafeGraph};
|
|
|
|
let mut group = c.benchmark_group("forward_pass_matmuls_graph");
|
|
group.sample_size(100);
|
|
|
|
for n_points in [200, 1000, 10000] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 1,
|
|
..Config::default()
|
|
};
|
|
|
|
// Create solver with ALL tensors on capture stream
|
|
let (solver, mut ws, ctx) = Mre1DPinnSolver::new_for_graph_capture(cfg.clone())
|
|
.expect("Failed to create graph-ready solver");
|
|
|
|
let stream = ctx.stream().clone();
|
|
|
|
// Force sync before capture
|
|
ctx.device_synchronize().expect("Device sync failed");
|
|
ctx.join_with_default_stream().expect("Stream join failed");
|
|
|
|
// Pre-cache GPU pointers
|
|
let ptrs = CachedGpuPtrs::from_forward_pass(
|
|
&solver.x_data,
|
|
solver.u_net.b_learnable(),
|
|
solver.u_net.layers(),
|
|
&ws,
|
|
&stream,
|
|
).expect("Failed to cache pointers");
|
|
|
|
// Get raw cuBLAS handle
|
|
let cublas_handle = *ctx.cublas().handle();
|
|
|
|
// Capture the matmuls graph
|
|
ctx.synchronize().expect("Pre-capture sync failed");
|
|
let mut graph = UnsafeGraph::capture(stream.clone(), || {
|
|
unsafe { ptrs.forward_matmuls_only(cublas_handle) }
|
|
}).expect("Graph capture failed");
|
|
|
|
// Warm up
|
|
for _ in 0..10 {
|
|
ctx.fourier_features_out(&solver.x_data, solver.u_net.b_learnable(),
|
|
2.0 * std::f32::consts::PI, &mut ws.features).unwrap();
|
|
graph.launch().unwrap();
|
|
for i in 0..solver.u_net.layers().len() - 1 {
|
|
if let Some(bias) = solver.u_net.layers()[i].bias() {
|
|
ctx.bias_add_tanh_(&mut ws.hidden_layers[i], bias).unwrap();
|
|
}
|
|
}
|
|
let last_idx = solver.u_net.layers().len() - 1;
|
|
if let Some(bias) = solver.u_net.layers()[last_idx].bias() {
|
|
ctx.add_bias_(&mut ws.output, bias).unwrap();
|
|
}
|
|
}
|
|
ctx.synchronize().expect("Warm-up sync failed");
|
|
|
|
group.throughput(Throughput::Elements(n_points as u64));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("hybrid_graph", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
// Pre-graph: Fourier features
|
|
ctx.fourier_features_out(&solver.x_data, solver.u_net.b_learnable(),
|
|
2.0 * std::f32::consts::PI, &mut ws.features).unwrap();
|
|
|
|
// Graph: 5 matmuls
|
|
graph.launch().unwrap();
|
|
|
|
// Post-graph: bias+activation
|
|
for i in 0..solver.u_net.layers().len() - 1 {
|
|
if let Some(bias) = solver.u_net.layers()[i].bias() {
|
|
ctx.bias_add_tanh_(&mut ws.hidden_layers[i], bias).unwrap();
|
|
}
|
|
}
|
|
let last_idx = solver.u_net.layers().len() - 1;
|
|
if let Some(bias) = solver.u_net.layers()[last_idx].bias() {
|
|
ctx.add_bias_(&mut ws.output, bias).unwrap();
|
|
}
|
|
|
|
// Sync to measure actual completion
|
|
ctx.synchronize().unwrap();
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark forward pass with uber-kernel ZERO-ALLOCATION path
|
|
/// This is the fastest possible inference: pre-allocated output, cached weights
|
|
#[cfg(feature = "cuda")]
|
|
fn benchmark_forward_pass_uber_kernel_zero_alloc(c: &mut Criterion) {
|
|
use pinn_mre_helmholtz::uber_kernel::forward_uber_value_only_out;
|
|
|
|
let mut group = c.benchmark_group("forward_pass_uber_kernel_zero_alloc");
|
|
group.sample_size(200); // More samples for µs-level timing
|
|
|
|
for n_points in [200, 1000, 10000] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 1,
|
|
..Config::default()
|
|
};
|
|
|
|
let solver = Mre1DPinnSolver::new(cfg.clone()).expect("Failed to create solver");
|
|
|
|
// Pre-compute weights ONCE (outside benchmark loop)
|
|
let weights = solver.u_net.to_uber_kernel_weights(&cfg);
|
|
let scale = 2.0 * std::f64::consts::PI as f32;
|
|
|
|
// Pre-allocate output tensor ONCE
|
|
let device = Device::Cuda(0);
|
|
let mut output = Tensor::zeros([n_points, 2], &device)
|
|
.expect("Failed to allocate output");
|
|
|
|
group.throughput(Throughput::Elements(n_points as u64));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("uber_zero_alloc", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
forward_uber_value_only_out(&solver.x_data, &weights, scale, &mut output)
|
|
.expect("Uber-kernel failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark FUSED-LITE kernel with shared memory weight caching
|
|
/// This is the target: single-digit µs inference
|
|
#[cfg(feature = "cuda")]
|
|
fn benchmark_forward_pass_fused_lite(c: &mut Criterion) {
|
|
use pinn_mre_helmholtz::uber_kernel::forward_fused_lite;
|
|
|
|
let mut group = c.benchmark_group("forward_pass_fused_lite");
|
|
group.sample_size(200);
|
|
|
|
for n_points in [200, 1000, 10000] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 1,
|
|
..Config::default()
|
|
};
|
|
|
|
let solver = Mre1DPinnSolver::new(cfg.clone()).expect("Failed to create solver");
|
|
|
|
// Pre-compute weights ONCE
|
|
let weights = solver.u_net.to_uber_kernel_weights(&cfg);
|
|
let scale = 2.0 * std::f64::consts::PI as f32;
|
|
|
|
// Pre-allocate output tensor
|
|
let device = Device::Cuda(0);
|
|
let mut output = Tensor::zeros([n_points, 2], &device)
|
|
.expect("Failed to allocate output");
|
|
|
|
group.throughput(Throughput::Elements(n_points as u64));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("fused_lite", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
forward_fused_lite(&solver.x_data, &weights, scale, &mut output)
|
|
.expect("Fused-lite kernel failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark training step with CUDA graph acceleration
|
|
#[cfg(feature = "cuda")]
|
|
fn benchmark_training_step_cuda_graph(c: &mut Criterion) {
|
|
use pinn_mre_helmholtz::CapturedForwardGraph;
|
|
|
|
let mut group = c.benchmark_group("training_step_cuda_graph");
|
|
group.sample_size(30);
|
|
|
|
for n_points in [200, 1000] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 1,
|
|
..Config::default()
|
|
};
|
|
|
|
let mut solver = Mre1DPinnSolver::new(cfg.clone()).expect("Failed to create solver");
|
|
let mut ws = solver.create_workspace().expect("Failed to create workspace");
|
|
let mut graph = CapturedForwardGraph::new();
|
|
|
|
// Warm up: capture the graph on first call
|
|
solver.training_step_with_graph(&mut ws, &mut graph)
|
|
.expect("Initial graph capture failed");
|
|
|
|
group.throughput(Throughput::Elements(n_points as u64));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("single_step_cuda_graph", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
solver.training_step_with_graph(&mut ws, &mut graph)
|
|
.expect("Training step failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark 100 epochs with CUDA graph acceleration
|
|
#[cfg(feature = "cuda")]
|
|
fn benchmark_training_100_epochs_cuda_graph(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("training_100_epochs_cuda_graph");
|
|
group.sample_size(10);
|
|
|
|
for n_points in [200] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 100,
|
|
print_every: 1000,
|
|
..Config::default()
|
|
};
|
|
|
|
group.throughput(Throughput::Elements(100));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("train_100_cuda_graph", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
let mut solver = Mre1DPinnSolver::new(cfg.clone())
|
|
.expect("Failed to create solver");
|
|
solver.train_with_graph().expect("Training failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
// =============================================================================
|
|
// ZERO-SYNC TRAINING BENCHMARKS (Eliminates GPU-CPU sync overhead)
|
|
// =============================================================================
|
|
|
|
/// Benchmark DEFERRED-LOSS training step: only computes loss when needed
|
|
/// This should be faster than fully_optimized when loss is skipped
|
|
fn benchmark_training_step_deferred_loss(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("training_step_deferred_loss");
|
|
group.sample_size(30);
|
|
|
|
for n_points in [200, 1000] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 1,
|
|
..Config::default()
|
|
};
|
|
|
|
let mut solver = Mre1DPinnSolver::new(cfg.clone()).expect("Failed to create solver");
|
|
let mut ws = solver.create_workspace().expect("Failed to create workspace");
|
|
|
|
group.throughput(Throughput::Elements(n_points as u64));
|
|
|
|
// Benchmark step WITHOUT loss computation (most steps in training)
|
|
group.bench_with_input(
|
|
BenchmarkId::new("no_loss_step", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
let mut step = 1; // Not divisible by 100, so no loss computed
|
|
b.iter(|| {
|
|
solver.training_step_deferred_loss(&mut ws, step, 100)
|
|
.expect("Deferred loss step failed");
|
|
step += 1;
|
|
if step % 100 == 0 { step += 1; } // Skip sync steps
|
|
})
|
|
},
|
|
);
|
|
|
|
// Benchmark step WITH loss computation (every sync_interval steps)
|
|
group.bench_with_input(
|
|
BenchmarkId::new("with_loss_step", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
// Step 100 will compute loss
|
|
solver.training_step_deferred_loss(&mut ws, 100, 100)
|
|
.expect("Deferred loss step failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark ZERO-SYNC training step: no GPU-CPU sync in hot path
|
|
/// NOTE: This was found to be SLOWER than fully_optimized due to accumulator overhead
|
|
fn benchmark_training_step_zero_sync(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("training_step_zero_sync");
|
|
group.sample_size(30);
|
|
|
|
for n_points in [200, 1000] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 1,
|
|
..Config::default()
|
|
};
|
|
|
|
let device = Device::Cuda(0);
|
|
let mut solver = Mre1DPinnSolver::new(cfg.clone()).expect("Failed to create solver");
|
|
let mut ws = solver.create_workspace().expect("Failed to create workspace");
|
|
let mut accumulator = GpuLossAccumulator::new(&device, 100).expect("Failed to create accumulator");
|
|
|
|
group.throughput(Throughput::Elements(n_points as u64));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("single_step_zero_sync", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
solver.training_step_zero_sync(&mut ws, &mut accumulator)
|
|
.expect("Zero-sync training step failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark DEFERRED-LOSS 100 epochs: skips loss on 99% of steps
|
|
fn benchmark_training_100_epochs_deferred_loss(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("training_100_epochs_deferred_loss");
|
|
group.sample_size(10);
|
|
|
|
for n_points in [200] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 100,
|
|
print_every: 1000,
|
|
..Config::default()
|
|
};
|
|
|
|
group.throughput(Throughput::Elements(100));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("train_100_deferred_loss", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
let mut solver = Mre1DPinnSolver::new(cfg.clone())
|
|
.expect("Failed to create solver");
|
|
solver.train_deferred_loss(100).expect("Training failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark ZERO-SYNC 100 epochs: GPU accumulation (slower due to accumulator overhead)
|
|
fn benchmark_training_100_epochs_zero_sync(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("training_100_epochs_zero_sync");
|
|
group.sample_size(10);
|
|
|
|
for n_points in [200] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 100,
|
|
print_every: 1000,
|
|
..Config::default()
|
|
};
|
|
|
|
group.throughput(Throughput::Elements(100));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("train_100_zero_sync", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
let mut solver = Mre1DPinnSolver::new(cfg.clone())
|
|
.expect("Failed to create solver");
|
|
solver.train_zero_sync(100).expect("Training failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark FP16 Tensor Core matmul vs FP32 SGEMM
|
|
/// Tests raw matmul performance for typical PINN layer sizes
|
|
#[cfg(feature = "cuda")]
|
|
fn benchmark_matmul_fp16_tensor_cores(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("matmul_tensor_cores");
|
|
group.sample_size(100);
|
|
|
|
let device = Device::Cuda(0);
|
|
|
|
// Test typical PINN layer sizes
|
|
for (batch, _m, k, n) in [
|
|
(200, 200, 128, 64), // Fourier features → first hidden
|
|
(200, 200, 64, 64), // Hidden → hidden
|
|
(1000, 1000, 128, 64), // Larger batch
|
|
(1000, 1000, 64, 64), // Larger batch hidden
|
|
(10000, 10000, 128, 64), // Very large batch
|
|
] {
|
|
let label = format!("{}x{}x{}", batch, k, n);
|
|
|
|
// Create FP32 matrices (randn creates F32 by default)
|
|
let a_f32 = Tensor::randn(&[batch, k], &device)
|
|
.expect("Failed to create tensor A");
|
|
let b_f32 = Tensor::randn(&[k, n], &device)
|
|
.expect("Failed to create tensor B");
|
|
|
|
// FP32 benchmark (SGEMM)
|
|
group.throughput(Throughput::Elements((2 * batch * k * n) as u64)); // FLOPs
|
|
group.bench_with_input(
|
|
BenchmarkId::new("fp32_sgemm", &label),
|
|
&batch,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
a_f32.matmul(&b_f32).expect("FP32 matmul failed")
|
|
})
|
|
},
|
|
);
|
|
|
|
// Create FP16 matrices
|
|
let a_f16 = a_f32.to_half().expect("Failed to convert A to FP16");
|
|
let b_f16 = b_f32.to_half().expect("Failed to convert B to FP16");
|
|
|
|
// FP16 benchmark (HGEMM / Tensor Cores)
|
|
group.bench_with_input(
|
|
BenchmarkId::new("fp16_tensor_cores", &label),
|
|
&batch,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
a_f16.matmul(&b_f16).expect("FP16 matmul failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// =============================================================================
|
|
// PHASE 8: ANALYTICAL BACKPROP TRAINING BENCHMARKS
|
|
// =============================================================================
|
|
|
|
/// Benchmark ANALYTICAL training step with proper GPU synchronization.
|
|
///
|
|
/// This measures the TRUE latency of a complete training step:
|
|
/// - Forward pass (cache activations)
|
|
/// - Backward pass (analytical gradients via chain rule)
|
|
/// - Optimizer step (GPU-native Adam)
|
|
///
|
|
/// CRITICAL: Includes device synchronization to measure actual GPU execution time,
|
|
/// not just CPU dispatch speed.
|
|
#[cfg(feature = "cuda")]
|
|
fn benchmark_training_step_analytical(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("training_step_analytical");
|
|
group.sample_size(200); // More samples for µs-level timing
|
|
group.measurement_time(std::time::Duration::from_secs(5));
|
|
|
|
for n_points in [200, 1000] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 1,
|
|
..Config::default()
|
|
};
|
|
|
|
let mut solver = Mre1DPinnSolver::new(cfg.clone()).expect("Failed to create solver");
|
|
|
|
// Get CUDA context for synchronization
|
|
let ctx = rtx_tensor::storage::cuda_manager::get_or_create_context(0)
|
|
.expect("Failed to get CUDA context");
|
|
|
|
// Warm up
|
|
for _ in 0..10 {
|
|
solver.train_step_data_only().expect("Warm-up step failed");
|
|
}
|
|
ctx.synchronize().expect("Warm-up sync failed");
|
|
|
|
group.throughput(Throughput::Elements(n_points as u64));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("analytical_backprop", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
// 1. Dispatch the entire training step (Forward + Backward + Optimizer)
|
|
solver.train_step_data_only().expect("Training step failed");
|
|
|
|
// 2. Block CPU until GPU finishes
|
|
// This measures the TRUE latency of the step
|
|
ctx.synchronize().expect("Sync failed");
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark 100 epochs with analytical backprop.
|
|
/// Uses the convenience train_analytical() method.
|
|
#[cfg(feature = "cuda")]
|
|
fn benchmark_training_100_epochs_analytical(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("training_100_epochs_analytical");
|
|
group.sample_size(10);
|
|
|
|
for n_points in [200] {
|
|
let cfg = Config {
|
|
n_data: n_points,
|
|
n_pde: n_points,
|
|
epochs: 100,
|
|
print_every: 1000, // Disable printing
|
|
..Config::default()
|
|
};
|
|
|
|
group.throughput(Throughput::Elements(100));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("train_100_analytical", n_points),
|
|
&n_points,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
let mut solver = Mre1DPinnSolver::new(cfg.clone())
|
|
.expect("Failed to create solver");
|
|
solver.train_analytical(100, 1000).expect("Training failed")
|
|
})
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
#[cfg(feature = "cuda")]
|
|
criterion_group!(
|
|
benches,
|
|
benchmark_wave_number,
|
|
benchmark_data_generation,
|
|
benchmark_forward_pass,
|
|
benchmark_forward_pass_optimized, // workspace-based forward
|
|
benchmark_forward_pass_cuda_graph, // CUDA graph-accelerated forward (OLD - may fail)
|
|
benchmark_forward_pass_matmuls_graph, // MATMULS-ONLY CUDA graph (NEW - hybrid approach)
|
|
benchmark_forward_pass_uber_kernel, // UBER-KERNEL: single launch forward + AD
|
|
benchmark_forward_pass_uber_kernel_zero_alloc, // UBER-KERNEL: ZERO-ALLOCATION path
|
|
benchmark_forward_pass_fused_lite, // FUSED-LITE: shared memory optimized
|
|
benchmark_matmul_fp16_tensor_cores, // FP16 Tensor Core vs FP32 comparison
|
|
benchmark_pde_residual_analytical,
|
|
benchmark_pde_residual_tensor,
|
|
benchmark_mse_loss,
|
|
benchmark_training_step,
|
|
benchmark_training_step_optimized, // workspace-based training step
|
|
benchmark_training_step_cached, // cached PDE tensors only
|
|
benchmark_training_step_fully_optimized, // workspace + cached PDE
|
|
benchmark_training_step_deferred_loss, // DEFERRED-LOSS training step (skips loss on most steps)
|
|
benchmark_training_step_zero_sync, // ZERO-SYNC training step (GPU accumulation - slower)
|
|
benchmark_training_step_cuda_graph, // CUDA graph-accelerated training step
|
|
benchmark_training_100_epochs,
|
|
benchmark_training_100_epochs_cached, // cached PDE only
|
|
benchmark_training_100_epochs_fully_optimized, // workspace + cached PDE
|
|
benchmark_training_100_epochs_deferred_loss, // DEFERRED-LOSS training (skips loss on 99% of steps)
|
|
benchmark_training_100_epochs_zero_sync, // ZERO-SYNC training (GPU accumulation - slower)
|
|
benchmark_training_100_epochs_cuda_graph, // CUDA graph-accelerated training
|
|
benchmark_training_step_analytical, // PHASE 8: Analytical backprop training step
|
|
benchmark_training_100_epochs_analytical, // PHASE 8: Analytical backprop 100 epochs
|
|
);
|
|
|
|
#[cfg(not(feature = "cuda"))]
|
|
criterion_group!(
|
|
benches,
|
|
benchmark_wave_number,
|
|
benchmark_data_generation,
|
|
benchmark_forward_pass,
|
|
benchmark_forward_pass_optimized, // workspace-based forward
|
|
benchmark_pde_residual_analytical,
|
|
benchmark_pde_residual_tensor,
|
|
benchmark_mse_loss,
|
|
benchmark_training_step,
|
|
benchmark_training_step_optimized, // workspace-based training step
|
|
benchmark_training_step_cached, // cached PDE tensors only
|
|
benchmark_training_step_fully_optimized, // workspace + cached PDE
|
|
benchmark_training_100_epochs,
|
|
benchmark_training_100_epochs_cached, // cached PDE only
|
|
benchmark_training_100_epochs_fully_optimized, // workspace + cached PDE
|
|
);
|
|
|
|
criterion_main!(benches);
|