507 lines
16 KiB
Rust
507 lines
16 KiB
Rust
//! Benchmarks for single residual calculation
|
|
//!
|
|
//! ADR-014 Performance Target: < 1us per residual calculation
|
|
//!
|
|
//! Residual is the core primitive: r_e = rho_u(x_u) - rho_v(x_v)
|
|
//! This measures the local constraint violation at each edge.
|
|
|
|
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
|
|
|
|
// ============================================================================
|
|
// Restriction Map Types (Simulated for benchmarking)
|
|
// ============================================================================
|
|
|
|
/// Linear restriction map: y = Ax + b
|
|
/// Maps node state to shared constraint space
|
|
#[derive(Clone)]
|
|
pub struct RestrictionMap {
|
|
/// Linear transformation matrix (row-major, output_dim x input_dim)
|
|
pub matrix: Vec<f32>,
|
|
/// Bias vector
|
|
pub bias: Vec<f32>,
|
|
/// Input dimension
|
|
pub input_dim: usize,
|
|
/// Output dimension
|
|
pub output_dim: usize,
|
|
}
|
|
|
|
impl RestrictionMap {
|
|
/// Create identity restriction map
|
|
pub fn identity(dim: usize) -> Self {
|
|
let mut matrix = vec![0.0f32; dim * dim];
|
|
for i in 0..dim {
|
|
matrix[i * dim + i] = 1.0;
|
|
}
|
|
Self {
|
|
matrix,
|
|
bias: vec![0.0; dim],
|
|
input_dim: dim,
|
|
output_dim: dim,
|
|
}
|
|
}
|
|
|
|
/// Create random restriction map for testing
|
|
pub fn random(input_dim: usize, output_dim: usize, seed: u64) -> Self {
|
|
use std::collections::hash_map::DefaultHasher;
|
|
use std::hash::{Hash, Hasher};
|
|
|
|
let mut matrix = Vec::with_capacity(output_dim * input_dim);
|
|
let mut bias = Vec::with_capacity(output_dim);
|
|
|
|
for i in 0..(output_dim * input_dim) {
|
|
let mut hasher = DefaultHasher::new();
|
|
(seed, i).hash(&mut hasher);
|
|
let val = (hasher.finish() % 1000) as f32 / 1000.0 - 0.5;
|
|
matrix.push(val);
|
|
}
|
|
|
|
for i in 0..output_dim {
|
|
let mut hasher = DefaultHasher::new();
|
|
(seed, i, "bias").hash(&mut hasher);
|
|
let val = (hasher.finish() % 1000) as f32 / 1000.0 - 0.5;
|
|
bias.push(val);
|
|
}
|
|
|
|
Self {
|
|
matrix,
|
|
bias,
|
|
input_dim,
|
|
output_dim,
|
|
}
|
|
}
|
|
|
|
/// Apply restriction map: y = Ax + b
|
|
#[inline]
|
|
pub fn apply(&self, input: &[f32]) -> Vec<f32> {
|
|
debug_assert_eq!(input.len(), self.input_dim);
|
|
let mut output = self.bias.clone();
|
|
|
|
for i in 0..self.output_dim {
|
|
let row_start = i * self.input_dim;
|
|
for j in 0..self.input_dim {
|
|
output[i] += self.matrix[row_start + j] * input[j];
|
|
}
|
|
}
|
|
|
|
output
|
|
}
|
|
|
|
/// Apply restriction map with SIMD-friendly layout (output buffer provided)
|
|
#[inline]
|
|
pub fn apply_into(&self, input: &[f32], output: &mut [f32]) {
|
|
debug_assert_eq!(input.len(), self.input_dim);
|
|
debug_assert_eq!(output.len(), self.output_dim);
|
|
|
|
// Copy bias first
|
|
output.copy_from_slice(&self.bias);
|
|
|
|
// Matrix-vector multiply
|
|
for i in 0..self.output_dim {
|
|
let row_start = i * self.input_dim;
|
|
for j in 0..self.input_dim {
|
|
output[i] += self.matrix[row_start + j] * input[j];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Edge with restriction maps
|
|
pub struct SheafEdge {
|
|
pub source: u64,
|
|
pub target: u64,
|
|
pub weight: f32,
|
|
pub rho_source: RestrictionMap,
|
|
pub rho_target: RestrictionMap,
|
|
}
|
|
|
|
impl SheafEdge {
|
|
/// Calculate the edge residual (local mismatch)
|
|
/// r_e = rho_u(x_u) - rho_v(x_v)
|
|
#[inline]
|
|
pub fn residual(&self, source_state: &[f32], target_state: &[f32]) -> Vec<f32> {
|
|
let projected_source = self.rho_source.apply(source_state);
|
|
let projected_target = self.rho_target.apply(target_state);
|
|
|
|
projected_source
|
|
.iter()
|
|
.zip(projected_target.iter())
|
|
.map(|(a, b)| a - b)
|
|
.collect()
|
|
}
|
|
|
|
/// Calculate residual with pre-allocated buffers (zero allocation)
|
|
#[inline]
|
|
pub fn residual_into(
|
|
&self,
|
|
source_state: &[f32],
|
|
target_state: &[f32],
|
|
source_buf: &mut [f32],
|
|
target_buf: &mut [f32],
|
|
residual: &mut [f32],
|
|
) {
|
|
self.rho_source.apply_into(source_state, source_buf);
|
|
self.rho_target.apply_into(target_state, target_buf);
|
|
|
|
for i in 0..residual.len() {
|
|
residual[i] = source_buf[i] - target_buf[i];
|
|
}
|
|
}
|
|
|
|
/// Calculate weighted residual norm squared: w_e * |r_e|^2
|
|
#[inline]
|
|
pub fn weighted_residual_energy(&self, source: &[f32], target: &[f32]) -> f32 {
|
|
let r = self.residual(source, target);
|
|
let norm_sq: f32 = r.iter().map(|x| x * x).sum();
|
|
self.weight * norm_sq
|
|
}
|
|
|
|
/// Weighted residual energy with pre-allocated buffers
|
|
#[inline]
|
|
pub fn weighted_residual_energy_into(
|
|
&self,
|
|
source: &[f32],
|
|
target: &[f32],
|
|
source_buf: &mut [f32],
|
|
target_buf: &mut [f32],
|
|
) -> f32 {
|
|
self.rho_source.apply_into(source, source_buf);
|
|
self.rho_target.apply_into(target, target_buf);
|
|
|
|
let mut norm_sq = 0.0f32;
|
|
for i in 0..source_buf.len() {
|
|
let diff = source_buf[i] - target_buf[i];
|
|
norm_sq += diff * diff;
|
|
}
|
|
|
|
self.weight * norm_sq
|
|
}
|
|
}
|
|
|
|
// ============================================================================
|
|
// Benchmarks
|
|
// ============================================================================
|
|
|
|
fn generate_state(dim: usize, seed: u64) -> Vec<f32> {
|
|
use std::collections::hash_map::DefaultHasher;
|
|
use std::hash::{Hash, Hasher};
|
|
|
|
(0..dim)
|
|
.map(|i| {
|
|
let mut hasher = DefaultHasher::new();
|
|
(seed, i).hash(&mut hasher);
|
|
(hasher.finish() % 1000) as f32 / 1000.0 - 0.5
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
/// Benchmark single residual calculation at various dimensions
|
|
fn bench_single_residual(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("residual_single");
|
|
group.throughput(Throughput::Elements(1));
|
|
|
|
// Test dimensions relevant for coherence engine:
|
|
// 8: Minimal state
|
|
// 32: Compact embedding
|
|
// 64: Standard embedding
|
|
// 128: Rich state
|
|
// 256: Large state
|
|
for dim in [8, 32, 64, 128, 256] {
|
|
let rho_source = RestrictionMap::identity(dim);
|
|
let rho_target = RestrictionMap::identity(dim);
|
|
let source_state = generate_state(dim, 42);
|
|
let target_state = generate_state(dim, 123);
|
|
|
|
let edge = SheafEdge {
|
|
source: 0,
|
|
target: 1,
|
|
weight: 1.0,
|
|
rho_source,
|
|
rho_target,
|
|
};
|
|
|
|
group.bench_with_input(BenchmarkId::new("identity_map", dim), &dim, |b, _| {
|
|
b.iter(|| edge.residual(black_box(&source_state), black_box(&target_state)))
|
|
});
|
|
}
|
|
|
|
// Test with projection (non-identity maps)
|
|
for (input_dim, output_dim) in [(64, 32), (128, 64), (256, 128)] {
|
|
let rho_source = RestrictionMap::random(input_dim, output_dim, 42);
|
|
let rho_target = RestrictionMap::random(input_dim, output_dim, 123);
|
|
let source_state = generate_state(input_dim, 42);
|
|
let target_state = generate_state(input_dim, 123);
|
|
|
|
let edge = SheafEdge {
|
|
source: 0,
|
|
target: 1,
|
|
weight: 1.0,
|
|
rho_source,
|
|
rho_target,
|
|
};
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("projection_map", format!("{}to{}", input_dim, output_dim)),
|
|
&(input_dim, output_dim),
|
|
|b, _| b.iter(|| edge.residual(black_box(&source_state), black_box(&target_state))),
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark residual calculation with pre-allocated buffers (zero allocation)
|
|
fn bench_residual_zero_alloc(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("residual_zero_alloc");
|
|
group.throughput(Throughput::Elements(1));
|
|
|
|
for dim in [32, 64, 128, 256] {
|
|
let rho_source = RestrictionMap::identity(dim);
|
|
let rho_target = RestrictionMap::identity(dim);
|
|
let source_state = generate_state(dim, 42);
|
|
let target_state = generate_state(dim, 123);
|
|
|
|
let edge = SheafEdge {
|
|
source: 0,
|
|
target: 1,
|
|
weight: 1.0,
|
|
rho_source,
|
|
rho_target,
|
|
};
|
|
|
|
// Pre-allocate buffers
|
|
let mut source_buf = vec![0.0f32; dim];
|
|
let mut target_buf = vec![0.0f32; dim];
|
|
let mut residual = vec![0.0f32; dim];
|
|
|
|
group.bench_with_input(BenchmarkId::new("dim", dim), &dim, |b, _| {
|
|
b.iter(|| {
|
|
edge.residual_into(
|
|
black_box(&source_state),
|
|
black_box(&target_state),
|
|
black_box(&mut source_buf),
|
|
black_box(&mut target_buf),
|
|
black_box(&mut residual),
|
|
)
|
|
})
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark weighted residual energy computation
|
|
fn bench_weighted_energy(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("residual_weighted_energy");
|
|
group.throughput(Throughput::Elements(1));
|
|
|
|
for dim in [32, 64, 128, 256] {
|
|
let rho_source = RestrictionMap::identity(dim);
|
|
let rho_target = RestrictionMap::identity(dim);
|
|
let source_state = generate_state(dim, 42);
|
|
let target_state = generate_state(dim, 123);
|
|
|
|
let edge = SheafEdge {
|
|
source: 0,
|
|
target: 1,
|
|
weight: 1.5,
|
|
rho_source,
|
|
rho_target,
|
|
};
|
|
|
|
group.bench_with_input(BenchmarkId::new("allocating", dim), &dim, |b, _| {
|
|
b.iter(|| {
|
|
edge.weighted_residual_energy(black_box(&source_state), black_box(&target_state))
|
|
})
|
|
});
|
|
|
|
// Pre-allocate buffers for zero-alloc version
|
|
let mut source_buf = vec![0.0f32; dim];
|
|
let mut target_buf = vec![0.0f32; dim];
|
|
|
|
group.bench_with_input(BenchmarkId::new("zero_alloc", dim), &dim, |b, _| {
|
|
b.iter(|| {
|
|
edge.weighted_residual_energy_into(
|
|
black_box(&source_state),
|
|
black_box(&target_state),
|
|
black_box(&mut source_buf),
|
|
black_box(&mut target_buf),
|
|
)
|
|
})
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark batch residual computation (for parallel evaluation)
|
|
fn bench_batch_residual(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("residual_batch");
|
|
|
|
for batch_size in [10, 100, 1000] {
|
|
let dim = 64;
|
|
|
|
// Create batch of edges
|
|
let edges: Vec<SheafEdge> = (0..batch_size)
|
|
.map(|i| SheafEdge {
|
|
source: i as u64,
|
|
target: (i + 1) as u64,
|
|
weight: 1.0,
|
|
rho_source: RestrictionMap::identity(dim),
|
|
rho_target: RestrictionMap::identity(dim),
|
|
})
|
|
.collect();
|
|
|
|
let states: Vec<Vec<f32>> = (0..batch_size + 1)
|
|
.map(|i| generate_state(dim, i as u64))
|
|
.collect();
|
|
|
|
group.throughput(Throughput::Elements(batch_size as u64));
|
|
|
|
// Sequential computation
|
|
group.bench_with_input(
|
|
BenchmarkId::new("sequential", batch_size),
|
|
&batch_size,
|
|
|b, _| {
|
|
b.iter(|| {
|
|
let mut total_energy = 0.0f32;
|
|
for (i, edge) in edges.iter().enumerate() {
|
|
total_energy += edge.weighted_residual_energy(
|
|
black_box(&states[i]),
|
|
black_box(&states[i + 1]),
|
|
);
|
|
}
|
|
black_box(total_energy)
|
|
})
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark restriction map application alone
|
|
fn bench_restriction_map(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("restriction_map");
|
|
group.throughput(Throughput::Elements(1));
|
|
|
|
// Identity maps
|
|
for dim in [32, 64, 128, 256] {
|
|
let rho = RestrictionMap::identity(dim);
|
|
let input = generate_state(dim, 42);
|
|
let mut output = vec![0.0f32; dim];
|
|
|
|
group.bench_with_input(BenchmarkId::new("identity_apply", dim), &dim, |b, _| {
|
|
b.iter(|| rho.apply(black_box(&input)))
|
|
});
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("identity_apply_into", dim),
|
|
&dim,
|
|
|b, _| b.iter(|| rho.apply_into(black_box(&input), black_box(&mut output))),
|
|
);
|
|
}
|
|
|
|
// Projection maps (dense matrix multiply)
|
|
for (input_dim, output_dim) in [(64, 32), (128, 64), (256, 128), (512, 256)] {
|
|
let rho = RestrictionMap::random(input_dim, output_dim, 42);
|
|
let input = generate_state(input_dim, 42);
|
|
let mut output = vec![0.0f32; output_dim];
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("projection_apply", format!("{}x{}", input_dim, output_dim)),
|
|
&(input_dim, output_dim),
|
|
|b, _| b.iter(|| rho.apply(black_box(&input))),
|
|
);
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new(
|
|
"projection_apply_into",
|
|
format!("{}x{}", input_dim, output_dim),
|
|
),
|
|
&(input_dim, output_dim),
|
|
|b, _| b.iter(|| rho.apply_into(black_box(&input), black_box(&mut output))),
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark SIMD-optimized residual patterns
|
|
fn bench_simd_patterns(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("residual_simd_patterns");
|
|
group.throughput(Throughput::Elements(1));
|
|
|
|
// Aligned dimensions for SIMD (multiples of 8 for AVX2, 16 for AVX-512)
|
|
for dim in [32, 64, 128, 256, 512] {
|
|
let a = generate_state(dim, 42);
|
|
let b = generate_state(dim, 123);
|
|
|
|
// Scalar subtraction and norm
|
|
group.bench_with_input(
|
|
BenchmarkId::new("scalar_diff_norm", dim),
|
|
&dim,
|
|
|b_iter, _| {
|
|
b_iter.iter(|| {
|
|
let mut norm_sq = 0.0f32;
|
|
for i in 0..dim {
|
|
let diff = a[i] - b[i];
|
|
norm_sq += diff * diff;
|
|
}
|
|
black_box(norm_sq)
|
|
})
|
|
},
|
|
);
|
|
|
|
// Iterator-based (auto-vectorization friendly)
|
|
group.bench_with_input(
|
|
BenchmarkId::new("iter_diff_norm", dim),
|
|
&dim,
|
|
|b_iter, _| {
|
|
b_iter.iter(|| {
|
|
let norm_sq: f32 = a
|
|
.iter()
|
|
.zip(b.iter())
|
|
.map(|(x, y)| {
|
|
let d = x - y;
|
|
d * d
|
|
})
|
|
.sum();
|
|
black_box(norm_sq)
|
|
})
|
|
},
|
|
);
|
|
|
|
// Chunked for explicit SIMD opportunity
|
|
group.bench_with_input(
|
|
BenchmarkId::new("chunked_diff_norm", dim),
|
|
&dim,
|
|
|b_iter, _| {
|
|
b_iter.iter(|| {
|
|
let mut accum = [0.0f32; 8];
|
|
for (chunk_a, chunk_b) in a.chunks(8).zip(b.chunks(8)) {
|
|
for i in 0..chunk_a.len() {
|
|
let d = chunk_a[i] - chunk_b[i];
|
|
accum[i] += d * d;
|
|
}
|
|
}
|
|
black_box(accum.iter().sum::<f32>())
|
|
})
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
criterion_group!(
|
|
benches,
|
|
bench_single_residual,
|
|
bench_residual_zero_alloc,
|
|
bench_weighted_energy,
|
|
bench_batch_residual,
|
|
bench_restriction_map,
|
|
bench_simd_patterns,
|
|
);
|
|
|
|
criterion_main!(benches);
|