rustytorch/crates/integration/rtx-burn/src/backend.rs

//! Burn backend configuration and detection

use serde::{Deserialize, Serialize};
use tracing::info;

/// Available Burn backends
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum BurnBackend {
    /// `NdArray` backend (CPU, always available)
    NdArray,
    /// WebGPU backend (cross-platform GPU)
    Wgpu,
    /// CUDA backend (NVIDIA GPUs)
    Cuda,
    /// `PyTorch` backend (via tch-rs)
    Tch,
}

impl Default for BurnBackend {
    fn default() -> Self {
        detect_best_backend()
    }
}

impl BurnBackend {
    /// Get the name of the backend
    pub fn name(&self) -> &'static str {
        match self {
            BurnBackend::NdArray => "ndarray",
            BurnBackend::Wgpu => "wgpu",
            BurnBackend::Cuda => "cuda",
            BurnBackend::Tch => "tch",
        }
    }

    /// Check if this backend supports GPU acceleration
    pub fn is_gpu(&self) -> bool {
        matches!(
            self,
            BurnBackend::Wgpu | BurnBackend::Cuda | BurnBackend::Tch
        )
    }

    /// Check if this backend is available
    pub fn is_available(&self) -> bool {
        match self {
            BurnBackend::NdArray => cfg!(feature = "ndarray"),
            BurnBackend::Wgpu => cfg!(feature = "wgpu"),
            BurnBackend::Cuda => cfg!(feature = "cuda") && is_cuda_runtime_available(),
            BurnBackend::Tch => cfg!(feature = "tch"),
        }
    }
}

/// Detect the best available backend based on features and hardware
pub fn detect_best_backend() -> BurnBackend {
    // Prefer CUDA if available
    #[cfg(feature = "cuda")]
    {
        if is_cuda_runtime_available() {
            info!("Burn: Using CUDA backend");
            return BurnBackend::Cuda;
        }
    }

    // Then WGPU for cross-platform GPU
    #[cfg(feature = "wgpu")]
    {
        info!("Burn: Using WebGPU backend");
        return BurnBackend::Wgpu;
    }

    // Then tch if available
    #[cfg(feature = "tch")]
    {
        info!("Burn: Using PyTorch (tch) backend");
        return BurnBackend::Tch;
    }

    // Fall back to ndarray
    info!("Burn: Using NdArray (CPU) backend");
    BurnBackend::NdArray
}

/// Check if CUDA runtime is available
fn is_cuda_runtime_available() -> bool {
    #[cfg(feature = "cuda")]
    {
        // Check for CUDA environment
        std::env::var("CUDA_VISIBLE_DEVICES").is_ok()
            || std::path::Path::new("/usr/local/cuda").exists()
    }
    #[cfg(not(feature = "cuda"))]
    {
        false
    }
}

/// Backend configuration options
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BackendConfig {
    /// Selected backend
    pub backend: BurnBackend,
    /// Device index (for multi-GPU)
    pub device_index: usize,
    /// Enable memory optimizations
    pub memory_efficient: bool,
    /// Seed for reproducibility
    pub seed: Option<u64>,
}

impl Default for BackendConfig {
    fn default() -> Self {
        Self {
            backend: detect_best_backend(),
            device_index: 0,
            memory_efficient: true,
            seed: None,
        }
    }
}

impl BackendConfig {
    /// Create config with specific backend
    pub fn with_backend(mut self, backend: BurnBackend) -> Self {
        self.backend = backend;
        self
    }

    /// Set device index
    pub fn with_device(mut self, index: usize) -> Self {
        self.device_index = index;
        self
    }

    /// Set seed for reproducibility
    pub fn with_seed(mut self, seed: u64) -> Self {
        self.seed = Some(seed);
        self
    }
}

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

    #[test]
    fn test_backend_name() {
        assert_eq!(BurnBackend::NdArray.name(), "ndarray");
        assert_eq!(BurnBackend::Wgpu.name(), "wgpu");
        assert_eq!(BurnBackend::Cuda.name(), "cuda");
        assert_eq!(BurnBackend::Tch.name(), "tch");
    }

    #[test]
    fn test_backend_is_gpu() {
        assert!(!BurnBackend::NdArray.is_gpu());
        assert!(BurnBackend::Wgpu.is_gpu());
        assert!(BurnBackend::Cuda.is_gpu());
        assert!(BurnBackend::Tch.is_gpu());
    }

    #[test]
    fn test_detect_backend() {
        let backend = detect_best_backend();
        // Should always succeed
        assert!(backend.name().len() > 0);
    }

    #[test]
    fn test_backend_config() {
        let config = BackendConfig::default()
            .with_backend(BurnBackend::NdArray)
            .with_device(1)
            .with_seed(42);

        assert_eq!(config.backend, BurnBackend::NdArray);
        assert_eq!(config.device_index, 1);
        assert_eq!(config.seed, Some(42));
    }
}