foxhunt/crates/ml-core/src/gpu/capabilities.rs

//! GPU hardware capability detection.
//!
//! Detects GPU VRAM and device name via nvidia-smi.
//! Use `GpuCapabilities::detect()` at startup and cache the result.

use serde::{Deserialize, Serialize};
use std::process::Command;
use std::sync::OnceLock;
use tracing::{info, warn};

/// Detected GPU hardware capabilities.
///
/// Construct via `detect()` for real hardware or `cpu_fallback()` / `with_vram()` for testing.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct GpuCapabilities {
    /// GPU device name (e.g., "NVIDIA GeForce RTX 3050 Ti Laptop GPU")
    pub device_name: String,
    /// Total VRAM in MB (0.0 for CPU)
    pub total_vram_mb: f64,
    /// Free VRAM in MB at detection time (0.0 for CPU)
    pub free_vram_mb: f64,
    /// Whether a CUDA GPU was detected
    pub is_cuda: bool,
}

impl GpuCapabilities {
    /// Detect GPU capabilities by querying nvidia-smi.
    ///
    /// Falls back to CPU if nvidia-smi is unavailable or fails.
    /// Call this once at startup and cache the result.
    pub fn detect() -> Self {
        match Self::query_nvidia_smi() {
            Ok(caps) => {
                info!(
                    "GPU detected: {} (Total: {:.0} MB, Free: {:.0} MB)",
                    caps.device_name, caps.total_vram_mb, caps.free_vram_mb
                );
                caps
            }
            Err(reason) => {
                warn!("GPU detection failed ({}), using CPU fallback", reason);
                Self::cpu_fallback()
            }
        }
    }

    /// Explicit CPU-only capabilities (no GPU).
    pub fn cpu_fallback() -> Self {
        Self {
            device_name: "CPU".to_owned(),
            total_vram_mb: 0.0,
            free_vram_mb: 0.0,
            is_cuda: false,
        }
    }

    /// Construct with known VRAM values (for testing or manual override).
    pub fn with_vram<S: Into<String>>(device_name: S, total_mb: f64, free_mb: f64) -> Self {
        Self {
            device_name: device_name.into(),
            total_vram_mb: total_mb,
            free_vram_mb: free_mb,
            is_cuda: total_mb > 0.0,
        }
    }

    /// Query nvidia-smi for GPU memory and device name.
    fn query_nvidia_smi() -> Result<Self, String> {
        let output = Command::new("nvidia-smi")
            .args([
                "--query-gpu=memory.total,memory.free,name",
                "--format=csv,noheader,nounits",
            ])
            .output()
            .map_err(|e| format!("nvidia-smi not found: {}", e))?;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            return Err(format!("nvidia-smi failed: {}", stderr));
        }

        let stdout = String::from_utf8_lossy(&output.stdout);
        let parts: Vec<&str> = stdout.trim().split(',').collect();

        if parts.len() < 3 {
            return Err(format!("unexpected nvidia-smi format: {}", stdout));
        }

        let total_mb = parts
            .first()
            .ok_or("missing total")?
            .trim()
            .parse::<f64>()
            .map_err(|e| format!("parse total: {}", e))?;
        let free_mb = parts
            .get(1)
            .ok_or("missing free")?
            .trim()
            .parse::<f64>()
            .map_err(|e| format!("parse free: {}", e))?;
        let device_name = parts
            .get(2)
            .ok_or("missing name")?
            .trim()
            .to_string();

        Ok(Self {
            device_name,
            total_vram_mb: total_mb,
            free_vram_mb: free_mb,
            is_cuda: true,
        })
    }
}

/// Returns the maximum shared memory per SM in KB for the detected GPU.
///
/// On CUDA builds, queries the driver directly via
/// `CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_MULTIPROCESSOR` (device 0).
/// Falls back to name-based heuristic if the query fails or on non-CUDA builds.
///
/// Known heuristic values:
/// - H100/H200/GH200: 228 KB (compute capability 9.0)
/// - A100/L40/L4: 164 KB (compute capability 8.x)
/// - RTX 30xx/40xx/50xx: 100 KB (Ampere/Ada consumer)
/// - Older / unknown: 48 KB (conservative default, always safe)
pub fn max_shared_memory_kb(gpu_name: &str) -> usize {
    if let Some(kb) = query_max_shmem_per_sm_kb() {
        return kb;
    }
    max_shared_memory_kb_by_name(gpu_name)
}

/// Query `CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_MULTIPROCESSOR` for device 0.
/// Returns `None` if the query fails (no CUDA context, driver error, etc.).
#[allow(unsafe_code)]
fn query_max_shmem_per_sm_kb() -> Option<usize> {
    use cudarc::driver::sys::{CUdevice_attribute, CUresult};

    let mut bytes: i32 = 0;
    // SAFETY: cuDeviceGetAttribute is a read-only CUDA driver query that writes
    // to `bytes` through the provided pointer. Device ordinal 0 is always valid
    // when a CUDA driver is present.
    let result = unsafe {
        cudarc::driver::sys::cuDeviceGetAttribute(
            &mut bytes,
            CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_MULTIPROCESSOR,
            0, // device ordinal 0
        )
    };
    if result != CUresult::CUDA_SUCCESS || bytes <= 0 {
        return None;
    }
    Some(bytes as usize / 1024)
}

/// Name-based heuristic fallback for shared memory per SM.
fn max_shared_memory_kb_by_name(gpu_name: &str) -> usize {
    let name = gpu_name.to_uppercase();
    if name.contains("H100") || name.contains("H200") || name.contains("GH200") {
        228
    } else if name.contains("A100") || name.contains("L40") || name.contains("L4") {
        164
    } else if name.contains("RTX 40")
        || name.contains("RTX 50")
        || name.contains("RTX 30")
        || name.contains("RTX A")
    {
        100 // Ampere/Ada consumer
    } else {
        48 // Conservative default — fits any CUDA GPU
    }
}

/// Cached GPU capabilities -- detected once, reused across all callers.
///
/// Use this instead of `GpuCapabilities::detect()` when you don't need fresh data.
/// The first call runs nvidia-smi; subsequent calls return the cached result.
pub fn cached_capabilities() -> &'static GpuCapabilities {
    static CAPS: OnceLock<GpuCapabilities> = OnceLock::new();
    CAPS.get_or_init(GpuCapabilities::detect)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_cpu_fallback() {
        let caps = GpuCapabilities::cpu_fallback();
        assert_eq!(caps.device_name, "CPU");
        assert_eq!(caps.total_vram_mb, 0.0);
        assert_eq!(caps.free_vram_mb, 0.0);
        assert!(!caps.is_cuda);
    }

    #[test]
    fn test_with_vram_rtx_3050_ti() {
        let caps = GpuCapabilities::with_vram("RTX 3050 Ti", 4096.0, 3700.0);
        assert_eq!(caps.device_name, "RTX 3050 Ti");
        assert_eq!(caps.total_vram_mb, 4096.0);
        assert_eq!(caps.free_vram_mb, 3700.0);
        assert!(caps.is_cuda);
    }

    #[test]
    fn test_with_vram_a100() {
        let caps = GpuCapabilities::with_vram("A100-SXM4-80GB", 81920.0, 80000.0);
        assert_eq!(caps.total_vram_mb, 81920.0);
        assert!(caps.is_cuda);
    }

    #[test]
    fn test_with_vram_zero_is_cpu() {
        let caps = GpuCapabilities::with_vram("none", 0.0, 0.0);
        assert!(!caps.is_cuda);
    }

    #[test]
    fn test_detect_does_not_panic() {
        // detect() should never panic -- falls back to CPU gracefully
        let caps = GpuCapabilities::detect();
        assert!(!caps.device_name.is_empty());
    }

    #[test]
    fn test_serde_roundtrip() {
        let caps = GpuCapabilities::with_vram("RTX 4090", 24576.0, 23000.0);
        let json = serde_json::to_string(&caps).unwrap();
        let deserialized: GpuCapabilities = serde_json::from_str(&json).unwrap();
        assert_eq!(caps, deserialized);
    }

    #[test]
    fn test_cached_capabilities_returns_same_instance() {
        let caps1 = cached_capabilities();
        let caps2 = cached_capabilities();
        assert_eq!(caps1, caps2);
    }

    #[test]
    fn test_max_shared_memory_kb_h100() {
        // Use name-based lookup directly so the test is not affected by whichever
        // GPU the CUDA driver reports on the current machine.
        assert_eq!(max_shared_memory_kb_by_name("NVIDIA H100 80GB HBM3"), 228);
        assert_eq!(max_shared_memory_kb_by_name("NVIDIA H200"), 228);
        assert_eq!(max_shared_memory_kb_by_name("GH200"), 228);
    }

    #[test]
    fn test_max_shared_memory_kb_queries_real_device() {
        // max_shared_memory_kb() queries the CUDA driver first, falling back to
        // name-based heuristic.  On any GPU the result must be at least 48 KB
        // (conservative default) and at most 256 KB (no current GPU exceeds this).
        let kb = max_shared_memory_kb("irrelevant-name");
        assert!(kb >= 48, "shared memory too low: {kb} KB");
        assert!(kb <= 256, "shared memory unexpectedly high: {kb} KB");
    }

    #[test]
    fn test_max_shared_memory_kb_a100() {
        // Use name-based lookup directly — max_shared_memory_kb() returns the
        // CUDA driver value on any GPU with a driver, ignoring the name arg.
        assert_eq!(max_shared_memory_kb_by_name("NVIDIA A100-SXM4-80GB"), 164);
        assert_eq!(max_shared_memory_kb_by_name("NVIDIA L40S"), 164);
        assert_eq!(max_shared_memory_kb_by_name("NVIDIA L4"), 164);
    }

    #[test]
    fn test_max_shared_memory_kb_consumer() {
        assert_eq!(max_shared_memory_kb_by_name("NVIDIA GeForce RTX 4090"), 100);
        assert_eq!(max_shared_memory_kb_by_name("NVIDIA GeForce RTX 3050 Ti Laptop GPU"), 100);
        assert_eq!(max_shared_memory_kb_by_name("NVIDIA RTX A6000"), 100);
    }

    #[test]
    fn test_max_shared_memory_kb_default() {
        assert_eq!(max_shared_memory_kb_by_name("Tesla V100"), 48);
        assert_eq!(max_shared_memory_kb_by_name("Tesla T4"), 48);
        assert_eq!(max_shared_memory_kb_by_name("CPU"), 48);
        assert_eq!(max_shared_memory_kb_by_name(""), 48);
    }

    #[test]
    fn test_max_shared_memory_kb_by_name_fallback() {
        // Verify the name-based heuristic works independently of CUDA driver
        assert_eq!(max_shared_memory_kb_by_name("NVIDIA H100 80GB HBM3"), 228);
        assert_eq!(max_shared_memory_kb_by_name("NVIDIA A100-SXM4-80GB"), 164);
        assert_eq!(max_shared_memory_kb_by_name("NVIDIA GeForce RTX 4090"), 100);
        assert_eq!(max_shared_memory_kb_by_name("Tesla V100"), 48);
    }
}