rustytorch/crates/core/rtx-backend-rocm/src/device.rs

//! ROCm device wrapper for Backend trait.
//!
//! Provides AMD GPU device management via HIP runtime.
//!
//! Note: This module requires HIP runtime to be installed on the system.
//! Testing is done via CI on AMD GPU systems.

use rtx_backend::{DeviceId, DeviceOps};
use std::sync::Arc;

use crate::{RocmBackend, RocmBackendError, RocmBackendResult};

/// Device information for ROCm/HIP devices.
#[derive(Clone, Debug)]
pub struct RocmDeviceInfo {
    /// Device name (e.g., "AMD Radeon RX 7900 XTX")
    pub name: String,
    /// Total global memory in bytes
    pub total_memory: usize,
    /// Compute units
    pub compute_units: u32,
    /// Max clock frequency in MHz
    pub max_clock_mhz: u32,
    /// GCN/CDNA architecture name
    pub arch_name: String,
}

impl Default for RocmDeviceInfo {
    fn default() -> Self {
        Self {
            name: "Unknown ROCm Device".to_string(),
            total_memory: 0,
            compute_units: 0,
            max_clock_mhz: 0,
            arch_name: "unknown".to_string(),
        }
    }
}

/// ROCm device wrapper implementing DeviceOps.
///
/// Wraps HIP runtime device handle and provides backend device trait.
#[derive(Clone)]
pub struct RocmDevice {
    /// Device index
    index: usize,
    /// Device info (cached)
    info: Arc<RocmDeviceInfo>,
}

impl RocmDevice {
    /// Create a new ROCm device for the default device (device 0).
    pub fn new(index: usize) -> RocmBackendResult<Self> {
        // Check if ROCm is available
        if !is_available() {
            return Err(RocmBackendError::NotAvailable);
        }

        // Get device info
        let info = get_device_info(index)?;

        Ok(Self {
            index,
            info: Arc::new(info),
        })
    }

    /// Get the device index.
    pub fn device_index(&self) -> usize {
        self.index
    }

    /// Get device name.
    pub fn name(&self) -> &str {
        &self.info.name
    }

    /// Get device info.
    pub fn info(&self) -> &RocmDeviceInfo {
        &self.info
    }

    /// Get the number of available ROCm devices.
    pub fn device_count() -> usize {
        get_device_count()
    }

    /// Synchronize all pending operations on this device.
    pub fn synchronize(&self) -> RocmBackendResult<()> {
        synchronize_device(self.index)
    }
}

impl Default for RocmDevice {
    fn default() -> Self {
        Self::new(0).expect("No ROCm device available")
    }
}

impl std::fmt::Debug for RocmDevice {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("RocmDevice")
            .field("index", &self.index)
            .field("name", &self.info.name)
            .field("arch", &self.info.arch_name)
            .finish()
    }
}

impl PartialEq for RocmDevice {
    fn eq(&self, other: &Self) -> bool {
        self.index == other.index
    }
}

impl Eq for RocmDevice {}

impl std::hash::Hash for RocmDevice {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.index.hash(state);
    }
}

impl DeviceOps<RocmBackend> for RocmDevice {
    fn id(&self) -> DeviceId {
        DeviceId::Rocm(self.index)
    }

    fn memory_capacity(&self) -> usize {
        self.info.total_memory
    }

    fn memory_available(&self) -> usize {
        // Try to get actual available memory via HIP runtime
        get_memory_available().unwrap_or_else(|_| {
            // Fallback: estimate 80% available
            self.info.total_memory * 8 / 10
        })
    }

    fn compute_capability(&self) -> Option<(u32, u32)> {
        // Return GCN/CDNA version as compute capability equivalent
        // Map architecture to version numbers
        match self.info.arch_name.as_str() {
            s if s.starts_with("gfx90") => Some((9, 0)), // MI100/MI200
            s if s.starts_with("gfx94") => Some((9, 4)), // MI300
            s if s.starts_with("gfx110") => Some((11, 0)), // RDNA 3
            s if s.starts_with("gfx103") => Some((10, 3)), // RDNA 2
            _ => Some((9, 0)),
        }
    }

    fn synchronize(&self) {
        let _ = self.synchronize();
    }

    fn is_available(&self) -> bool {
        is_available()
    }
}

// ============== HIP Runtime Integration ==============
// These functions wrap HIP runtime calls using the hip_ffi module.

#[cfg(feature = "hip-runtime")]
use crate::hip_ffi::HipRuntime;

/// Check if ROCm/HIP is available on this system.
pub fn is_available() -> bool {
    // First check if HIP runtime feature is enabled and working
    #[cfg(feature = "hip-runtime")]
    {
        HipRuntime::is_available()
    }

    // Fallback: check for ROCm tools on Linux
    #[cfg(all(target_os = "linux", not(feature = "hip-runtime")))]
    {
        std::process::Command::new("rocminfo")
            .output()
            .map(|o| o.status.success())
            .unwrap_or(false)
    }

    #[cfg(all(not(target_os = "linux"), not(feature = "hip-runtime")))]
    {
        false
    }
}

/// Get the number of available HIP devices.
fn get_device_count() -> usize {
    #[cfg(feature = "hip-runtime")]
    {
        HipRuntime::device_count()
    }

    #[cfg(not(feature = "hip-runtime"))]
    {
        // Try to parse rocminfo output as fallback
        if !is_available() {
            return 0;
        }

        std::process::Command::new("rocminfo")
            .output()
            .ok()
            .and_then(|output| {
                let stdout = String::from_utf8_lossy(&output.stdout);
                // Count lines containing "Agent" and "GPU"
                Some(
                    stdout
                        .lines()
                        .filter(|line| line.contains("Agent") && line.contains("GPU"))
                        .count(),
                )
            })
            .unwrap_or(0)
    }
}

/// Get device information.
fn get_device_info(index: usize) -> RocmBackendResult<RocmDeviceInfo> {
    if !is_available() {
        return Err(RocmBackendError::NotAvailable);
    }

    #[cfg(feature = "hip-runtime")]
    {
        let props = HipRuntime::get_device_properties(index as i32).map_err(|e| {
            RocmBackendError::DeviceInit(format!("Failed to get device properties: {:?}", e))
        })?;

        Ok(RocmDeviceInfo {
            name: props.name,
            total_memory: props.total_global_mem,
            compute_units: props.multi_processor_count as u32,
            max_clock_mhz: props.clock_rate_mhz as u32,
            arch_name: props.gcn_arch_name,
        })
    }

    #[cfg(not(feature = "hip-runtime"))]
    {
        // Parse rocminfo output as fallback
        let output = std::process::Command::new("rocminfo")
            .output()
            .map_err(|e| RocmBackendError::DeviceInit(format!("Failed to run rocminfo: {}", e)))?;

        let stdout = String::from_utf8_lossy(&output.stdout);
        let mut current_device = 0usize;
        let mut name = String::new();
        let mut arch_name = String::new();
        let mut compute_units = 0u32;

        for line in stdout.lines() {
            if line.contains("Marketing Name:") {
                if current_device == index {
                    name = line
                        .split(':')
                        .nth(1)
                        .map(|s| s.trim().to_string())
                        .unwrap_or_default();
                }
            } else if line.contains("Name:") && line.contains("gfx") {
                if current_device == index {
                    arch_name = line
                        .split(':')
                        .nth(1)
                        .map(|s| s.trim().to_string())
                        .unwrap_or_default();
                }
            } else if line.contains("Compute Unit:") {
                if current_device == index {
                    compute_units = line
                        .split(':')
                        .nth(1)
                        .and_then(|s| s.trim().parse().ok())
                        .unwrap_or(0);
                }
            } else if line.contains("Agent") && line.contains("GPU") {
                current_device += 1;
            }
        }

        if name.is_empty() {
            return Err(RocmBackendError::DeviceInit(format!(
                "Device {} not found",
                index
            )));
        }

        Ok(RocmDeviceInfo {
            name,
            total_memory: 0, // Not available without HIP runtime
            compute_units,
            max_clock_mhz: 0, // Not available without HIP runtime
            arch_name,
        })
    }
}

/// Get available memory on the current device.
#[allow(dead_code)]
pub fn get_memory_available() -> RocmBackendResult<usize> {
    #[cfg(feature = "hip-runtime")]
    {
        let (free, _total) = HipRuntime::get_memory_info().map_err(|e| {
            RocmBackendError::DeviceInit(format!("Failed to get memory info: {:?}", e))
        })?;
        Ok(free)
    }

    #[cfg(not(feature = "hip-runtime"))]
    {
        Err(RocmBackendError::NotAvailable)
    }
}

/// Synchronize device.
fn synchronize_device(_index: usize) -> RocmBackendResult<()> {
    if !is_available() {
        return Err(RocmBackendError::NotAvailable);
    }

    #[cfg(feature = "hip-runtime")]
    {
        HipRuntime::synchronize()
            .map_err(|e| RocmBackendError::Synchronization(format!("Device sync failed: {:?}", e)))
    }

    #[cfg(not(feature = "hip-runtime"))]
    {
        Ok(()) // No-op without HIP runtime
    }
}

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

    #[test]
    fn test_rocm_availability() {
        let available = is_available();
        println!("ROCm available: {}", available);
        // This should pass regardless of ROCm availability
    }

    #[test]
    fn test_device_count() {
        let count = RocmDevice::device_count();
        println!("ROCm device count: {}", count);
    }
}