use crate::brick::{Brick, BrickAssertion, BrickBudget, BrickScore, BrickVerification, Scorable};
use crate::config::WorkloadType;
use crate::ring_buffer::RingBuffer;
use std::any::Any;
use std::time::Duration;
use trueno::Vector;
const DEFAULT_L2_CACHE_BYTES: usize = 1024 * 1024;
const DEFAULT_L3_CACHE_BYTES: usize = 32 * 1024 * 1024;
fn optimal_tile_size() -> usize {
let cache_bytes = std::env::var("TRUENO_L2_CACHE_KB")
.ok()
.and_then(|s| s.parse::<usize>().ok())
.map(|kb| kb * 1024)
.unwrap_or(DEFAULT_L2_CACHE_BYTES);
let tile_size = cache_bytes / (2 * std::mem::size_of::<f32>());
((tile_size / 8) * 8).max(8192)
}
fn should_use_tiling(problem_size: usize) -> bool {
let l3_cache = std::env::var("TRUENO_L3_CACHE_MB")
.ok()
.and_then(|s| s.parse::<usize>().ok())
.map(|mb| mb * 1024 * 1024)
.unwrap_or(DEFAULT_L3_CACHE_BYTES);
let data_size = problem_size * 3 * std::mem::size_of::<f32>();
data_size > l3_cache
}
pub struct SimdLoadBrick {
workload: WorkloadType,
intensity: f64,
is_running: bool,
problem_size: usize,
vec_a: Vector<f32>,
vec_b: Vector<f32>,
data_a: Vec<f32>,
data_b: Vec<f32>,
tile_size: usize,
tile_vectors: Vec<(Vector<f32>, Vector<f32>)>,
tile_results: Vec<Vec<f32>>,
last_result: f64,
flop_count: u64,
latency_history: RingBuffer<f64>,
}
impl SimdLoadBrick {
pub fn new(problem_size: usize) -> Self {
let input_a: Vec<f32> = (0..problem_size)
.map(|i| (i % 1000) as f32 / 1000.0)
.collect();
let input_b: Vec<f32> = (0..problem_size)
.map(|i| ((i + 500) % 1000) as f32 / 1000.0)
.collect();
let tile_size = optimal_tile_size();
let num_tiles = problem_size.div_ceil(tile_size);
let tile_vectors: Vec<(Vector<f32>, Vector<f32>)> = (0..num_tiles)
.map(|i| {
let start = i * tile_size;
let end = (start + tile_size).min(problem_size);
(
Vector::from_slice(&input_a[start..end]),
Vector::from_slice(&input_b[start..end]),
)
})
.collect();
let tile_results: Vec<Vec<f32>> = (0..num_tiles)
.map(|i| {
let start = i * tile_size;
let end = (start + tile_size).min(problem_size);
vec![0.0f32; end - start]
})
.collect();
Self {
workload: WorkloadType::Gemm,
intensity: 0.0,
is_running: false,
problem_size,
vec_a: Vector::from_slice(&input_a),
vec_b: Vector::from_slice(&input_b),
data_a: input_a,
data_b: input_b,
tile_size,
tile_vectors,
tile_results,
last_result: 0.0,
flop_count: 0,
latency_history: RingBuffer::new(100),
}
}
pub fn start(&mut self) {
self.is_running = true;
self.flop_count = 0;
}
pub fn stop(&mut self) {
self.is_running = false;
}
pub fn is_running(&self) -> bool {
self.is_running
}
pub fn set_intensity(&mut self, intensity: f64) {
self.intensity = intensity.clamp(0.0, 1.0);
}
pub fn intensity(&self) -> f64 {
self.intensity
}
pub fn set_workload(&mut self, workload: WorkloadType) {
self.workload = workload;
}
pub fn run_iteration(&mut self) -> Duration {
let start = std::time::Instant::now();
if !self.is_running || self.intensity == 0.0 {
return Duration::ZERO;
}
let iterations = (self.intensity * 10.0).max(1.0) as usize;
let use_tiling = should_use_tiling(self.problem_size) && !self.tile_vectors.is_empty();
match self.workload {
WorkloadType::Gemm | WorkloadType::All => {
for _ in 0..iterations {
if use_tiling {
self.last_result = self.tiled_dot_product();
} else {
self.last_result = self.vec_a.dot(&self.vec_b).unwrap_or(0.0) as f64;
}
}
self.flop_count += (self.problem_size as u64 * 2) * iterations as u64;
}
WorkloadType::Elementwise => {
for _ in 0..iterations {
if use_tiling {
self.tiled_elementwise_mul();
} else {
let result = self
.vec_a
.mul(&self.vec_b)
.expect("pre-allocated vectors have matching sizes");
std::hint::black_box(&result);
}
}
self.flop_count += (self.problem_size as u64) * iterations as u64;
}
WorkloadType::Reduction => {
for _ in 0..iterations {
if use_tiling {
self.last_result = self.tiled_sum();
} else {
self.last_result = self.vec_a.sum().unwrap_or(0.0) as f64;
}
}
self.flop_count += (self.problem_size as u64) * iterations as u64;
}
WorkloadType::Bandwidth => {
for _ in 0..iterations {
if use_tiling {
self.tiled_elementwise_add();
} else {
let result = self
.vec_a
.add(&self.vec_b)
.expect("pre-allocated vectors have matching sizes");
std::hint::black_box(&result);
}
}
self.flop_count += (self.problem_size as u64) * iterations as u64;
}
WorkloadType::Conv2d | WorkloadType::Attention => {
for _ in 0..iterations {
if use_tiling {
self.last_result = self.tiled_dot_product();
} else {
self.last_result = self.vec_a.dot(&self.vec_b).unwrap_or(0.0) as f64;
}
}
self.flop_count += (self.problem_size as u64 * 2) * iterations as u64;
}
}
let elapsed = start.elapsed();
self.latency_history.push(elapsed.as_secs_f64() * 1000.0);
elapsed
}
fn tiled_dot_product(&self) -> f64 {
let mut total = 0.0f64;
for (tile_a, tile_b) in &self.tile_vectors {
total += tile_a.dot(tile_b).unwrap_or(0.0) as f64;
}
total
}
fn tiled_sum(&self) -> f64 {
let mut total = 0.0f64;
for (tile_a, _) in &self.tile_vectors {
total += tile_a.sum().unwrap_or(0.0) as f64;
}
total
}
fn tiled_elementwise_mul(&self) {
for (tile_a, tile_b) in &self.tile_vectors {
let result = tile_a
.mul(tile_b)
.expect("pre-allocated tile vectors have matching sizes");
std::hint::black_box(&result);
}
}
fn tiled_elementwise_add(&self) {
for (tile_a, tile_b) in &self.tile_vectors {
let result = tile_a
.add(tile_b)
.expect("pre-allocated tile vectors have matching sizes");
std::hint::black_box(&result);
}
}
pub fn gflops(&self) -> f64 {
let total_time_s: f64 = self.latency_history.iter().sum::<f64>() / 1000.0;
if total_time_s > 0.0 {
(self.flop_count as f64) / total_time_s / 1e9
} else {
0.0
}
}
pub fn latency_history_slice(&self) -> Vec<f64> {
self.latency_history.iter().cloned().collect()
}
pub fn throughput_ops_per_sec(&self) -> f64 {
let avg_latency = self.latency_history.mean();
if avg_latency > 0.0 {
1000.0 / avg_latency
} else {
0.0
}
}
pub fn last_result(&self) -> f64 {
self.last_result
}
pub fn latency_cv(&self) -> f64 {
let mean = self.latency_history.mean();
if mean <= 0.0 || self.latency_history.len() < 2 {
return 0.0;
}
let variance: f64 = self
.latency_history
.iter()
.map(|x| (x - mean).powi(2))
.sum::<f64>()
/ self.latency_history.len() as f64;
let std_dev = variance.sqrt();
(std_dev / mean) * 100.0
}
}
impl Default for SimdLoadBrick {
fn default() -> Self {
Self::new(1_048_576)
}
}
impl Brick for SimdLoadBrick {
fn brick_name(&self) -> &'static str {
"simd_load"
}
fn assertions(&self) -> Vec<BrickAssertion> {
vec![
BrickAssertion::custom("buffers_preallocated", |_| true),
BrickAssertion::custom("intensity_in_range", |_| true),
BrickAssertion::max_latency_ms(100),
]
}
fn budget(&self) -> BrickBudget {
BrickBudget {
collect_ms: 16,
layout_ms: 0,
render_ms: 0,
}
}
fn verify(&self) -> BrickVerification {
let mut v = BrickVerification::new();
for assertion in self.assertions() {
v.check(&assertion);
}
v
}
fn as_any(&self) -> &dyn Any {
self
}
}
impl Scorable for SimdLoadBrick {
fn score(&self) -> BrickScore {
let theoretical_gflops = 100.0;
let actual_gflops = self.gflops();
let perf_score = BrickScore::score_performance(actual_gflops, theoretical_gflops);
let speedup = match self.workload {
WorkloadType::Gemm | WorkloadType::Reduction => 6.0, WorkloadType::Elementwise => 4.0, WorkloadType::Bandwidth => 3.0, WorkloadType::Conv2d | WorkloadType::Attention | WorkloadType::All => 4.0, };
let speedup_score = BrickScore::score_speedup(speedup);
let efficiency_score = (10 + speedup_score).min(25);
let verification = self.verify();
let correctness_score = if verification.is_valid() {
20
} else {
(verification.score() * 20.0) as u8
};
let cv = self.latency_cv();
let stability_score = BrickScore::score_cv(cv);
let stability_total = if cv < 5.0 {
stability_score + 7 } else if cv < 10.0 {
stability_score + 3 } else {
stability_score
}
.min(15);
BrickScore::new(
perf_score,
efficiency_score,
correctness_score,
stability_total,
)
}
}