mesh_llm_gpu_bench/
runner.rs

1use crate::BenchmarkOutput;
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4pub enum BenchmarkBackend {
5    Metal,
6    Cuda,
7    Hip,
8    Intel,
9}
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub struct BenchmarkRunner {
13    pub backend: BenchmarkBackend,
14}
15
16pub fn runner_for(
17    os: &str,
18    gpu_count: u8,
19    gpu_name: Option<&str>,
20    is_soc: bool,
21) -> Option<BenchmarkRunner> {
22    if gpu_count == 0 {
23        tracing::debug!("no GPUs detected; skipping benchmark");
24        return None;
25    }
26
27    let gpu_upper = gpu_name.unwrap_or("").to_uppercase();
28
29    if os == "macos" && is_soc {
30        return Some(BenchmarkRunner {
31            backend: BenchmarkBackend::Metal,
32        });
33    }
34
35    if os == "linux" || os == "windows" {
36        if gpu_upper.contains("NVIDIA")
37            || gpu_upper.contains("ORIN")
38            || gpu_upper.contains("NVGPU")
39            || gpu_upper.contains("TEGRA")
40        {
41            return Some(BenchmarkRunner {
42                backend: BenchmarkBackend::Cuda,
43            });
44        }
45
46        if gpu_upper.contains("AMD") || gpu_upper.contains("RADEON") {
47            return Some(BenchmarkRunner {
48                backend: BenchmarkBackend::Hip,
49            });
50        }
51
52        if gpu_upper.contains("INTEL") || gpu_upper.contains("ARC") {
53            tracing::info!(
54                "Intel GPU benchmark is not supported in standard mesh-llm builds; skipping"
55            );
56            return None;
57        }
58
59        if os == "linux" && is_soc {
60            tracing::warn!("Jetson benchmark is unvalidated for ARM CUDA; attempting");
61            return Some(BenchmarkRunner {
62                backend: BenchmarkBackend::Cuda,
63            });
64        }
65    }
66
67    tracing::warn!("could not identify benchmark runner for GPU platform: {gpu_name:?}");
68    None
69}
70
71pub fn parse_benchmark_output(stdout: &[u8]) -> Option<Vec<BenchmarkOutput>> {
72    match serde_json::from_slice::<Vec<BenchmarkOutput>>(stdout) {
73        Ok(outputs) if !outputs.is_empty() => Some(outputs),
74        Ok(_) => {
75            tracing::debug!("benchmark returned empty device list");
76            None
77        }
78        Err(err) => {
79            let error_message = serde_json::from_slice::<serde_json::Value>(stdout)
80                .ok()
81                .and_then(|val| {
82                    val.get("error")
83                        .and_then(|v| v.as_str())
84                        .map(ToOwned::to_owned)
85                });
86            if let Some(msg) = error_message {
87                tracing::warn!("benchmark reported error: {msg}");
88                return None;
89            }
90            tracing::warn!("failed to parse benchmark output: {err}");
91            None
92        }
93    }
94}
mesh_llm_gpu_bench/runner.rs

mesh_llm_gpu_bench/
runner.rs
