#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum WorkloadCategory {
Gemm,
Bandwidth,
Attention,
Conv2d,
Elementwise,
Reduction,
Unknown,
}
impl WorkloadCategory {
pub fn name(&self) -> &'static str {
match self {
Self::Gemm => "gemm",
Self::Bandwidth => "bandwidth",
Self::Attention => "attention",
Self::Conv2d => "conv2d",
Self::Elementwise => "elementwise",
Self::Reduction => "reduction",
Self::Unknown => "unknown",
}
}
pub fn is_compute_bound(&self) -> bool {
matches!(self, Self::Gemm | Self::Conv2d)
}
pub fn is_memory_bound(&self) -> bool {
matches!(self, Self::Bandwidth | Self::Elementwise)
}
pub fn typical_intensity_range(&self) -> (f64, f64) {
match self {
Self::Gemm => (10.0, 100.0),
Self::Bandwidth => (0.1, 1.0),
Self::Attention => (1.0, 20.0),
Self::Conv2d => (5.0, 50.0),
Self::Elementwise => (0.1, 0.5),
Self::Reduction => (0.5, 5.0),
Self::Unknown => (0.0, 100.0),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RecommendedBackend {
CpuSimd,
Gpu,
Either,
}
impl RecommendedBackend {
pub fn name(&self) -> &'static str {
match self {
Self::CpuSimd => "cpu_simd",
Self::Gpu => "gpu",
Self::Either => "either",
}
}
}
#[derive(Debug, Clone)]
pub struct WorkloadFeatures {
pub arithmetic_intensity: f64,
pub memory_footprint: usize,
pub working_set: usize,
pub access_pattern: f64,
pub compute_density: f64,
pub branch_rate: f64,
pub data_reuse: f64,
}
impl Default for WorkloadFeatures {
fn default() -> Self {
Self {
arithmetic_intensity: 1.0,
memory_footprint: 0,
working_set: 0,
access_pattern: 0.5,
compute_density: 1.0,
branch_rate: 0.0,
data_reuse: 1.0,
}
}
}
impl WorkloadFeatures {
pub fn new() -> Self {
Self::default()
}
pub fn with_intensity(mut self, intensity: f64) -> Self {
self.arithmetic_intensity = intensity.max(0.0);
self
}
pub fn with_memory(mut self, footprint: usize, working_set: usize) -> Self {
self.memory_footprint = footprint;
self.working_set = working_set;
self
}
pub fn with_access_pattern(mut self, pattern: f64) -> Self {
self.access_pattern = pattern.clamp(0.0, 1.0);
self
}
pub fn with_compute_density(mut self, density: f64) -> Self {
self.compute_density = density.max(0.0);
self
}
pub fn with_branch_rate(mut self, rate: f64) -> Self {
self.branch_rate = rate.clamp(0.0, 1.0);
self
}
pub fn with_data_reuse(mut self, reuse: f64) -> Self {
self.data_reuse = reuse.max(1.0);
self
}
pub fn normalize(&self, means: &[f64], stds: &[f64]) -> Vec<f64> {
let features = self.to_vec();
features
.iter()
.enumerate()
.map(|(i, &v)| {
if stds[i] > 1e-10 {
(v - means[i]) / stds[i]
} else {
0.0
}
})
.collect()
}
pub fn to_vec(&self) -> Vec<f64> {
vec![
self.arithmetic_intensity,
self.memory_footprint as f64,
self.working_set as f64,
self.access_pattern,
self.compute_density,
self.branch_rate,
self.data_reuse,
]
}
pub fn distance(&self, other: &Self) -> f64 {
let a = self.to_vec();
let b = other.to_vec();
a.iter()
.zip(b.iter())
.map(|(x, y)| (x - y).powi(2))
.sum::<f64>()
.sqrt()
}
pub fn cosine_similarity(&self, other: &Self) -> f64 {
let a = self.to_vec();
let b = other.to_vec();
let dot: f64 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
let norm_a: f64 = a.iter().map(|x| x.powi(2)).sum::<f64>().sqrt();
let norm_b: f64 = b.iter().map(|x| x.powi(2)).sum::<f64>().sqrt();
if norm_a < 1e-10 || norm_b < 1e-10 {
return 0.0;
}
(dot / (norm_a * norm_b)).clamp(-1.0, 1.0)
}
}
#[derive(Debug, Clone)]
pub struct ClassificationResult {
pub category: WorkloadCategory,
pub confidence: f64,
pub distance: f64,
pub recommended_backend: RecommendedBackend,
pub gpu_crossover_size: Option<usize>,
}
impl ClassificationResult {
pub fn is_confident(&self) -> bool {
self.confidence > 0.7
}
}
#[derive(Debug)]
pub struct WorkloadCharacterizer {
prototypes: Vec<(WorkloadCategory, WorkloadFeatures)>,
gpu_thresholds: Vec<(WorkloadCategory, usize)>,
}
impl Default for WorkloadCharacterizer {
fn default() -> Self {
Self::new()
}
}
impl WorkloadCharacterizer {
pub fn new() -> Self {
let prototypes = vec![
(
WorkloadCategory::Gemm,
WorkloadFeatures::new()
.with_intensity(50.0)
.with_compute_density(8.0)
.with_access_pattern(0.8)
.with_data_reuse(32.0)
.with_branch_rate(0.01),
),
(
WorkloadCategory::Bandwidth,
WorkloadFeatures::new()
.with_intensity(0.25)
.with_compute_density(0.5)
.with_access_pattern(1.0)
.with_data_reuse(1.0)
.with_branch_rate(0.0),
),
(
WorkloadCategory::Attention,
WorkloadFeatures::new()
.with_intensity(5.0)
.with_compute_density(4.0)
.with_access_pattern(0.6)
.with_data_reuse(4.0)
.with_branch_rate(0.05),
),
(
WorkloadCategory::Conv2d,
WorkloadFeatures::new()
.with_intensity(20.0)
.with_compute_density(6.0)
.with_access_pattern(0.7)
.with_data_reuse(9.0)
.with_branch_rate(0.02),
),
(
WorkloadCategory::Elementwise,
WorkloadFeatures::new()
.with_intensity(0.125)
.with_compute_density(1.0)
.with_access_pattern(1.0)
.with_data_reuse(1.0)
.with_branch_rate(0.0),
),
(
WorkloadCategory::Reduction,
WorkloadFeatures::new()
.with_intensity(1.0)
.with_compute_density(2.0)
.with_access_pattern(0.5)
.with_data_reuse(2.0)
.with_branch_rate(0.1),
),
];
let gpu_thresholds = vec![
(WorkloadCategory::Gemm, 10_000), (WorkloadCategory::Bandwidth, 1_000_000), (WorkloadCategory::Attention, 50_000), (WorkloadCategory::Conv2d, 100_000), (WorkloadCategory::Elementwise, 500_000), (WorkloadCategory::Reduction, 100_000), ];
Self {
prototypes,
gpu_thresholds,
}
}
pub fn extract_features(
&self,
flops: f64,
bytes_accessed: f64,
memory_footprint: usize,
working_set: usize,
) -> WorkloadFeatures {
let intensity = if bytes_accessed > 0.0 {
flops / bytes_accessed
} else {
0.0
};
WorkloadFeatures::new()
.with_intensity(intensity)
.with_memory(memory_footprint, working_set)
}
pub fn classify(&self, features: &WorkloadFeatures) -> ClassificationResult {
let mut best_category = WorkloadCategory::Unknown;
let mut best_distance = f64::MAX;
let mut second_best_distance = f64::MAX;
for (category, prototype) in &self.prototypes {
let distance = features.distance(prototype);
if distance < best_distance {
second_best_distance = best_distance;
best_distance = distance;
best_category = *category;
} else if distance < second_best_distance {
second_best_distance = distance;
}
}
let confidence = if second_best_distance > 1e-10 {
(1.0 - best_distance / second_best_distance).clamp(0.0, 1.0)
} else {
1.0
};
let recommended_backend = self.recommend_backend(best_category, features.memory_footprint);
let gpu_crossover_size = self
.gpu_thresholds
.iter()
.find(|(c, _)| *c == best_category)
.map(|(_, t)| *t);
ClassificationResult {
category: best_category,
confidence,
distance: best_distance,
recommended_backend,
gpu_crossover_size,
}
}
pub fn workload_similarity(&self, a: &WorkloadFeatures, b: &WorkloadFeatures) -> f64 {
(a.cosine_similarity(b) + 1.0) / 2.0
}
pub fn recommend_backend(&self, category: WorkloadCategory, size: usize) -> RecommendedBackend {
let threshold = self
.gpu_thresholds
.iter()
.find(|(c, _)| *c == category)
.map(|(_, t)| *t)
.unwrap_or(100_000);
if size < threshold / 2 {
RecommendedBackend::CpuSimd
} else if size > threshold * 2 {
RecommendedBackend::Gpu
} else {
RecommendedBackend::Either
}
}
pub fn predict_crossover(&self, category: WorkloadCategory) -> Option<usize> {
self.gpu_thresholds
.iter()
.find(|(c, _)| *c == category)
.map(|(_, t)| *t)
}
pub fn add_prototype(&mut self, category: WorkloadCategory, features: WorkloadFeatures) {
self.prototypes.push((category, features));
}
pub fn get_prototypes(&self) -> &[(WorkloadCategory, WorkloadFeatures)] {
&self.prototypes
}
}
#[cfg(test)]
mod tests;