rustytorch/crates/core/rtx-cubecl/src/runtime.rs

//! `CubeCL` Runtime abstraction
//!
//! Provides a unified interface to different `CubeCL` backends (WGPU, CUDA, HIP).

use crate::error::{CubeclError, Result};
use tracing::info;

/// Supported `CubeCL` runtime backends
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum RuntimeBackend {
    /// WebGPU backend (cross-platform)
    Wgpu,
    /// NVIDIA CUDA backend
    Cuda,
    /// AMD ROCm/HIP backend
    Hip,
    /// CPU backend (for testing)
    Cpu,
}

impl std::fmt::Display for RuntimeBackend {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            RuntimeBackend::Wgpu => write!(f, "WebGPU"),
            RuntimeBackend::Cuda => write!(f, "CUDA"),
            RuntimeBackend::Hip => write!(f, "HIP/ROCm"),
            RuntimeBackend::Cpu => write!(f, "CPU"),
        }
    }
}

/// `CubeCL` runtime wrapper
///
/// Manages the `CubeCL` runtime and provides kernel launching capabilities.
pub struct CubeclRuntime {
    backend: RuntimeBackend,
    device_index: usize,
}

impl CubeclRuntime {
    /// Create a new `CubeCL` runtime with the specified backend
    pub fn new(backend: RuntimeBackend, device_index: usize) -> Result<Self> {
        info!(
            "Initializing CubeCL runtime with {} backend on device {}",
            backend, device_index
        );

        // Validate backend availability
        Self::validate_backend(backend)?;

        Ok(Self {
            backend,
            device_index,
        })
    }

    /// Create a runtime with automatic backend detection
    pub fn auto() -> Result<Self> {
        let backend = Self::detect_best_backend()?;
        Self::new(backend, 0)
    }

    /// Detect the best available backend
    pub fn detect_best_backend() -> Result<RuntimeBackend> {
        // Priority: CUDA > HIP > WGPU > CPU

        #[cfg(feature = "cuda")]
        {
            if Self::is_cuda_available() {
                info!("CUDA backend available");
                return Ok(RuntimeBackend::Cuda);
            }
        }

        #[cfg(feature = "hip")]
        {
            if Self::is_hip_available() {
                info!("HIP/ROCm backend available");
                return Ok(RuntimeBackend::Hip);
            }
        }

        #[cfg(feature = "wgpu")]
        {
            if Self::is_wgpu_available() {
                info!("WebGPU backend available");
                return Ok(RuntimeBackend::Wgpu);
            }
        }

        #[cfg(feature = "cpu")]
        {
            info!("Falling back to CPU backend");
            return Ok(RuntimeBackend::Cpu);
        }

        #[allow(unreachable_code)]
        Err(CubeclError::RuntimeError(
            "No CubeCL backend available. Enable a feature: wgpu, cuda, hip, or cpu".to_string(),
        ))
    }

    /// Check if CUDA backend is available
    #[cfg(feature = "cuda")]
    pub fn is_cuda_available() -> bool {
        // TODO: Implement actual CUDA device detection
        true
    }

    #[cfg(not(feature = "cuda"))]
    pub fn is_cuda_available() -> bool {
        false
    }

    /// Check if HIP/ROCm backend is available
    #[cfg(feature = "hip")]
    pub fn is_hip_available() -> bool {
        // TODO: Implement actual HIP device detection
        true
    }

    #[cfg(not(feature = "hip"))]
    pub fn is_hip_available() -> bool {
        false
    }

    /// Check if WebGPU backend is available
    #[cfg(feature = "wgpu")]
    pub fn is_wgpu_available() -> bool {
        // TODO: Implement actual WebGPU adapter detection
        true
    }

    #[cfg(not(feature = "wgpu"))]
    pub fn is_wgpu_available() -> bool {
        false
    }

    /// Validate that the specified backend is available
    fn validate_backend(backend: RuntimeBackend) -> Result<()> {
        let available = match backend {
            RuntimeBackend::Cuda => Self::is_cuda_available(),
            RuntimeBackend::Hip => Self::is_hip_available(),
            RuntimeBackend::Wgpu => Self::is_wgpu_available(),
            RuntimeBackend::Cpu => true, // CPU is always available
        };

        if available {
            Ok(())
        } else {
            Err(CubeclError::DeviceError(format!(
                "{backend} backend not available. Enable the corresponding feature."
            )))
        }
    }

    /// Get the current backend
    pub fn backend(&self) -> RuntimeBackend {
        self.backend
    }

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

    /// Synchronize the device (wait for all operations to complete)
    pub fn synchronize(&self) -> Result<()> {
        // TODO: Implement synchronization for each backend
        Ok(())
    }
}

impl Default for CubeclRuntime {
    fn default() -> Self {
        Self::auto().expect("Failed to initialize default CubeCL runtime")
    }
}

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

    #[test]
    fn test_runtime_backend_display() {
        assert_eq!(format!("{}", RuntimeBackend::Wgpu), "WebGPU");
        assert_eq!(format!("{}", RuntimeBackend::Cuda), "CUDA");
        assert_eq!(format!("{}", RuntimeBackend::Hip), "HIP/ROCm");
        assert_eq!(format!("{}", RuntimeBackend::Cpu), "CPU");
    }
}