aprender-compute 0.64.0

//! Hardware Capability Detection (PMAT-447)
//!
//! Detects CPU SIMD capabilities, GPU presence, and calculates
//! theoretical peak performance for roofline analysis.
//!
//! Integrates with `pmat brick-score` for hardware-aware profiling.

use serde::{Deserialize, Serialize};
use std::fs;
use std::path::Path;

/// Get hostname (native only, returns "wasm" on WASM targets)
#[cfg(not(target_arch = "wasm32"))]
fn get_hostname() -> String {
    hostname::get().map(|h| h.to_string_lossy().to_string()).unwrap_or_else(|e| {
        eprintln!("warning: failed to get hostname: {e}");
        "unknown".to_string()
    })
}

/// Get hostname (WASM fallback)
#[cfg(target_arch = "wasm32")]
fn get_hostname() -> String {
    "wasm".to_string()
}

/// SIMD instruction set width
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum SimdWidth {
    /// No SIMD (scalar)
    Scalar,
    /// ARM NEON (128-bit, 4×f32)
    Neon128,
    /// SSE2 (128-bit, 4×f32)
    Sse2,
    /// AVX2 (256-bit, 8×f32)
    Avx2,
    /// AVX-512 (512-bit, 16×f32)
    Avx512,
    /// WebAssembly SIMD (128-bit, 4×f32)
    WasmSimd128,
}

impl SimdWidth {
    /// Number of f32 lanes
    pub fn lanes(&self) -> usize {
        match self {
            SimdWidth::Scalar => 1,
            SimdWidth::Neon128 | SimdWidth::Sse2 | SimdWidth::WasmSimd128 => 4,
            SimdWidth::Avx2 => 8,
            SimdWidth::Avx512 => 16,
        }
    }

    /// Bit width
    pub fn bits(&self) -> usize {
        self.lanes() * 32
    }

    /// Typical speedup factor for compute-bound operations
    pub fn compute_speedup(&self) -> f64 {
        match self {
            SimdWidth::Scalar => 1.0,
            SimdWidth::Neon128 | SimdWidth::Sse2 | SimdWidth::WasmSimd128 => 4.0,
            SimdWidth::Avx2 => 10.0,   // 8-12x measured in trueno-zram
            SimdWidth::Avx512 => 12.0, // 8-13x measured
        }
    }
}

/// GPU compute backend
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum GpuBackend {
    /// No GPU available
    None,
    /// NVIDIA CUDA
    Cuda,
    /// WebGPU (cross-platform)
    Wgpu,
    /// Apple Metal
    Metal,
    /// Vulkan compute
    Vulkan,
}

/// CPU capabilities
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CpuCapability {
    /// CPU vendor (Intel, AMD, Apple, etc.)
    pub vendor: String,
    /// CPU model name
    pub model: String,
    /// Number of physical cores
    pub cores: usize,
    /// Number of logical threads
    pub threads: usize,
    /// Best available SIMD width
    pub simd: SimdWidth,
    /// Base frequency in GHz
    pub base_freq_ghz: f64,
    /// Theoretical peak GFLOP/s (FMA)
    pub peak_gflops: f64,
    /// Memory bandwidth in GB/s (estimated)
    pub memory_bw_gbps: f64,
}

/// GPU capabilities
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GpuCapability {
    /// GPU vendor
    pub vendor: String,
    /// GPU model name
    pub model: String,
    /// Compute backend
    pub backend: GpuBackend,
    /// CUDA compute capability (e.g., "8.9" for RTX 4090)
    pub compute_capability: Option<String>,
    /// Peak FP32 TFLOP/s
    pub peak_tflops_fp32: f64,
    /// Peak Tensor Core TFLOP/s (NVIDIA only)
    pub peak_tflops_tensor: Option<f64>,
    /// Memory bandwidth in GB/s
    pub memory_bw_gbps: f64,
    /// VRAM in GB
    pub vram_gb: f64,
}

/// Complete hardware capability profile
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HardwareCapability {
    /// Detection timestamp
    pub timestamp: String,
    /// Hostname
    pub hostname: String,
    /// CPU capabilities
    pub cpu: CpuCapability,
    /// GPU capabilities (if present)
    pub gpu: Option<GpuCapability>,
    /// Roofline model parameters
    pub roofline: RooflineParams,
    /// PMAT-452: Byte budget configuration for compression/I/O workloads
    #[serde(default)]
    pub byte_budget: Option<crate::brick::ByteBudget>,
}

/// Roofline model parameters
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RooflineParams {
    /// CPU arithmetic intensity threshold (GFLOP/s ÷ GB/s)
    pub cpu_arithmetic_intensity: f64,
    /// GPU arithmetic intensity threshold
    pub gpu_arithmetic_intensity: Option<f64>,
}

impl HardwareCapability {
    /// Detect hardware capabilities at runtime
    pub fn detect() -> Self {
        let cpu = detect_cpu();
        let gpu = detect_gpu();

        let cpu_ai = cpu.peak_gflops / cpu.memory_bw_gbps;
        let gpu_ai = gpu.as_ref().map(|g| g.peak_tflops_fp32 * 1000.0 / g.memory_bw_gbps);
        // PMAT-452: Extract memory bandwidth before moving cpu
        let byte_budget_throughput = cpu.memory_bw_gbps.min(25.0);

        HardwareCapability {
            timestamp: chrono::Utc::now().to_rfc3339(),
            hostname: get_hostname(),
            cpu,
            gpu,
            roofline: RooflineParams {
                cpu_arithmetic_intensity: cpu_ai,
                gpu_arithmetic_intensity: gpu_ai,
            },
            // PMAT-452: Default byte budget based on memory bandwidth
            byte_budget: Some(crate::brick::ByteBudget::from_throughput(byte_budget_throughput)),
        }
    }

    /// Load from TOML file or detect if missing
    pub fn load_or_detect(path: &Path) -> Self {
        if path.exists() {
            if let Ok(content) = fs::read_to_string(path) {
                if let Ok(cap) = toml::from_str(&content) {
                    return cap;
                }
            }
        }
        let cap = Self::detect();
        // Try to cache it
        let _ = cap.save(path);
        cap
    }

    /// Save to TOML file
    pub fn save(&self, path: &Path) -> std::io::Result<()> {
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent)?;
        }
        let content = toml::to_string_pretty(self)
            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
        fs::write(path, content)
    }

    /// Get the best available backend for a workload
    pub fn best_backend(&self) -> GpuBackend {
        self.gpu.as_ref().map(|g| g.backend).unwrap_or(GpuBackend::None)
    }

    /// Calculate expected throughput for a brick given its arithmetic intensity
    pub fn expected_throughput_gflops(&self, arithmetic_intensity: f64, use_gpu: bool) -> f64 {
        if use_gpu {
            if let Some(gpu) = &self.gpu {
                let memory_bound = gpu.memory_bw_gbps * arithmetic_intensity;
                let compute_bound = gpu.peak_tflops_fp32 * 1000.0;
                memory_bound.min(compute_bound)
            } else {
                self.cpu_expected_throughput(arithmetic_intensity)
            }
        } else {
            self.cpu_expected_throughput(arithmetic_intensity)
        }
    }

    fn cpu_expected_throughput(&self, arithmetic_intensity: f64) -> f64 {
        let memory_bound = self.cpu.memory_bw_gbps * arithmetic_intensity;
        let compute_bound = self.cpu.peak_gflops;
        memory_bound.min(compute_bound)
    }

    /// Determine if workload is memory-bound or compute-bound
    pub fn bottleneck(&self, arithmetic_intensity: f64, use_gpu: bool) -> Bottleneck {
        let threshold = if use_gpu {
            self.roofline.gpu_arithmetic_intensity.unwrap_or(f64::MAX)
        } else {
            self.roofline.cpu_arithmetic_intensity
        };

        if arithmetic_intensity < threshold {
            Bottleneck::Memory
        } else {
            Bottleneck::Compute
        }
    }
}

/// Workload bottleneck classification
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum Bottleneck {
    /// Limited by memory bandwidth
    Memory,
    /// Limited by compute throughput
    Compute,
}

/// Map a raw x86 `vendor_id` to a human-readable vendor name.
fn normalize_cpu_vendor(raw: &str) -> String {
    match raw.trim() {
        "GenuineIntel" => "Intel".to_string(),
        "AuthenticAMD" => "AMD".to_string(),
        "" => "Unknown".to_string(),
        other => other.to_string(),
    }
}

/// Parse `(vendor, model)` out of `/proc/cpuinfo` contents.
///
/// GH-2395: `detect_cpu` hardcoded `vendor: "Unknown", model: "Unknown"`, so the
/// roofline block of `apr profile` reported `Hardware: Unknown Unknown (24 cores, …)`
/// while the peak GFLOPS and bandwidth printed beside it — the basis for the
/// MEMORY BOUND verdict — claimed to describe that machine.
///
/// Split out from the filesystem read so it is testable without `/proc`. Handles
/// the x86 layout (`vendor_id` / `model name`) and the aarch64 layout
/// (`CPU implementer` / `Model`, with no `model name` on many kernels).
pub fn parse_cpuinfo(contents: &str) -> (Option<String>, Option<String>) {
    let mut vendor = None;
    let mut model = None;
    for line in contents.lines() {
        let Some((key, value)) = line.split_once(':') else {
            continue;
        };
        let key = key.trim();
        let value = value.trim();
        if value.is_empty() {
            continue;
        }
        match key {
            "vendor_id" if vendor.is_none() => vendor = Some(normalize_cpu_vendor(value)),
            "model name" | "Model" if model.is_none() => model = Some(value.to_string()),
            // aarch64 has no vendor_id; implementer 0x41 is ARM Ltd.
            "CPU implementer" if vendor.is_none() => {
                vendor = Some(match value {
                    "0x41" => "ARM".to_string(),
                    "0x51" => "Qualcomm".to_string(),
                    "0x4e" => "NVIDIA".to_string(),
                    other => other.to_string(),
                });
            }
            _ => {}
        }
    }
    (vendor, model)
}

/// Read CPU vendor/model from the OS, falling back to `"Unknown"`.
fn detect_cpu_identity() -> (String, String) {
    #[cfg(target_os = "linux")]
    {
        if let Ok(contents) = fs::read_to_string("/proc/cpuinfo") {
            let (vendor, model) = parse_cpuinfo(&contents);
            return (
                vendor.unwrap_or_else(|| "Unknown".to_string()),
                model.unwrap_or_else(|| "Unknown".to_string()),
            );
        }
    }
    ("Unknown".to_string(), "Unknown".to_string())
}

/// Detect CPU capabilities
fn detect_cpu() -> CpuCapability {
    let simd = detect_simd();
    let cores = num_cpus::get_physical();
    let threads = num_cpus::get();
    let (vendor, model) = detect_cpu_identity();

    // Estimate frequency (fallback to 3.0 GHz if unknown)
    let base_freq_ghz = 3.0;

    // Calculate peak GFLOP/s: cores × lanes × 2 (FMA) × freq
    let peak_gflops = (cores as f64) * (simd.lanes() as f64) * 2.0 * base_freq_ghz;

    // Estimate memory bandwidth (DDR5-5600 dual channel ≈ 89.6 GB/s)
    let memory_bw_gbps = 80.0; // Conservative estimate

    CpuCapability {
        vendor,
        model,
        cores,
        threads,
        simd,
        base_freq_ghz,
        peak_gflops,
        memory_bw_gbps,
    }
}

/// Detect best available SIMD width
fn detect_simd() -> SimdWidth {
    #[cfg(target_arch = "x86_64")]
    {
        if is_x86_feature_detected!("avx512f") {
            return SimdWidth::Avx512;
        }
        if is_x86_feature_detected!("avx2") {
            return SimdWidth::Avx2;
        }
        if is_x86_feature_detected!("sse2") {
            return SimdWidth::Sse2;
        }
    }

    #[cfg(target_arch = "aarch64")]
    {
        // NEON is always available on aarch64
        return SimdWidth::Neon128;
    }

    #[cfg(target_arch = "wasm32")]
    {
        return SimdWidth::WasmSimd128;
    }

    SimdWidth::Scalar
}

/// Detect GPU capabilities
fn detect_gpu() -> Option<GpuCapability> {
    // Check for CUDA first (highest performance)
    #[cfg(feature = "cuda")]
    {
        if let Some(gpu) = detect_cuda_gpu() {
            return Some(gpu);
        }
    }

    // Fallback: no GPU detected
    None
}

#[cfg(feature = "cuda")]
fn detect_cuda_gpu() -> Option<GpuCapability> {
    // This would use cuDeviceGetAttribute in a real implementation
    // For now, return None and let the caller provide GPU info
    None
}

/// Default hardware.toml path
pub fn default_hardware_path() -> std::path::PathBuf {
    #[cfg(feature = "hardware-detect")]
    {
        dirs::home_dir()
            .unwrap_or_else(|| std::path::PathBuf::from("."))
            .join(".pmat")
            .join("hardware.toml")
    }
    #[cfg(not(feature = "hardware-detect"))]
    {
        std::path::PathBuf::from(".pmat").join("hardware.toml")
    }
}

#[cfg(test)]
mod tests;