rustytorch/demos/rtx-slidescope/src/gpu/mod.rs

//! GPU abstraction layer for NMF computations
//!
//! Provides a unified interface for GPU-accelerated matrix operations
//! using the rtx-tensor backend for both CUDA (NVIDIA) and Metal (Apple).
//!
//! This module uses rtx-tensor's unified device abstraction, ensuring
//! consistency with other RustyTorch++ demos (hemodynamics, MRE, bioheat).

#[cfg(feature = "cuda")]
pub mod cuda;

#[cfg(feature = "metal")]
pub mod metal;

use anyhow::Result;
use ndarray::Array2;

/// GPU backend trait for NMF matrix operations
///
/// Implementations must provide efficient GPU-accelerated versions
/// of the core NMF operations: matrix multiplication, element-wise
/// operations, and transpose.
pub trait GpuBackend: Send + Sync {
    /// Get the name of this GPU backend
    fn name(&self) -> &'static str;

    /// Get device information
    fn device_info(&self) -> DeviceInfo;

    /// Check if the backend is available and functional
    fn is_available(&self) -> bool;

    /// Allocate a matrix on the GPU
    fn allocate(&self, rows: usize, cols: usize) -> Result<GpuMatrix>;

    /// Upload a CPU matrix to the GPU
    fn upload(&self, data: &Array2<f32>) -> Result<GpuMatrix>;

    /// Download a GPU matrix to the CPU
    fn download(&self, matrix: &GpuMatrix) -> Result<Array2<f32>>;

    /// Matrix multiplication: C = A @ B
    fn matmul(&self, a: &GpuMatrix, b: &GpuMatrix) -> Result<GpuMatrix>;

    /// Matrix transpose
    fn transpose(&self, a: &GpuMatrix) -> Result<GpuMatrix>;

    /// Element-wise multiplication: C = A .* B
    fn element_mul(&self, a: &GpuMatrix, b: &GpuMatrix) -> Result<GpuMatrix>;

    /// Element-wise division with epsilon: C = A ./ (B + eps)
    fn element_div_eps(&self, a: &GpuMatrix, b: &GpuMatrix, epsilon: f32) -> Result<GpuMatrix>;

    /// In-place element-wise update: A = A .* B ./ (C + eps)
    fn nmf_update_inplace(
        &self,
        target: &mut GpuMatrix,
        numerator: &GpuMatrix,
        denominator: &GpuMatrix,
        epsilon: f32,
    ) -> Result<()>;

    /// Compute Frobenius norm of (A - B)
    fn frobenius_diff(&self, a: &GpuMatrix, b: &GpuMatrix) -> Result<f32>;

    /// Synchronize GPU operations (wait for completion)
    fn synchronize(&self) -> Result<()>;
}

/// GPU matrix handle - opaque reference to GPU-allocated memory
#[derive(Debug)]
pub struct GpuMatrix {
    /// Unique identifier for this matrix
    pub id: u64,
    /// Number of rows
    pub rows: usize,
    /// Number of columns
    pub cols: usize,
    /// Backend-specific handle (raw pointer or index)
    #[allow(dead_code)]
    pub(crate) handle: GpuHandle,
}

impl GpuMatrix {
    /// Create a new GPU matrix reference
    pub fn new(id: u64, rows: usize, cols: usize, handle: GpuHandle) -> Self {
        Self {
            id,
            rows,
            cols,
            handle,
        }
    }

    /// Get the shape as (rows, cols)
    pub fn shape(&self) -> (usize, usize) {
        (self.rows, self.cols)
    }

    /// Get total number of elements
    pub fn len(&self) -> usize {
        self.rows * self.cols
    }

    /// Check if matrix is empty
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

/// Backend-specific handle for GPU memory
#[derive(Debug)]
pub enum GpuHandle {
    /// Raw device pointer (CUDA)
    DevicePtr(u64),
    /// Buffer index (Metal/wgpu)
    BufferIndex(u32),
    /// Null/invalid handle
    None,
}

/// Information about a GPU device
#[derive(Debug, Clone)]
pub struct DeviceInfo {
    /// Device name
    pub name: String,
    /// Total memory in bytes
    pub total_memory: u64,
    /// Available memory in bytes (if known)
    pub available_memory: Option<u64>,
    /// Compute capability or feature level
    pub compute_capability: String,
    /// Backend type
    pub backend: BackendType,
}

/// Type of GPU backend
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BackendType {
    /// NVIDIA CUDA
    Cuda,
    /// Apple Metal
    Metal,
    /// CPU fallback
    Cpu,
}

impl std::fmt::Display for BackendType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            BackendType::Cuda => write!(f, "CUDA"),
            BackendType::Metal => write!(f, "Metal"),
            BackendType::Cpu => write!(f, "CPU"),
        }
    }
}

/// Select the best available GPU backend
///
/// Uses rtx-tensor's unified device detection to find available GPUs.
/// Priority: CUDA > Metal > CPU
pub fn select_backend() -> Result<Box<dyn GpuBackend>> {
    use rtx_tensor::Device;

    let devices = Device::available_devices();

    // Try CUDA first (highest priority for compute)
    #[cfg(feature = "cuda")]
    {
        if devices.iter().any(|d| d.is_cuda()) {
            if let Ok(backend) = cuda::CudaBackend::new() {
                if backend.is_available() {
                    tracing::info!("Using CUDA GPU backend: {}", backend.device_info().name);
                    return Ok(Box::new(backend));
                }
            }
        }
    }

    // Try Metal on macOS
    #[cfg(feature = "metal")]
    {
        if devices.iter().any(|d| d.is_metal())
            && let Ok(backend) = metal::MetalBackend::new()
            && backend.is_available()
        {
            tracing::info!(
                "Using Metal GPU backend (rtx-tensor): {}",
                backend.device_info().name
            );
            return Ok(Box::new(backend));
        }
    }

    // Return CPU fallback
    tracing::info!("Using CPU fallback backend");
    Ok(Box::new(CpuFallbackBackend::new()))
}

/// Check if any GPU backend is available
///
/// Uses rtx-tensor's device detection for consistency with other demos.
pub fn gpu_available() -> bool {
    use rtx_tensor::Device;

    let devices = Device::available_devices();

    #[cfg(feature = "cuda")]
    {
        if devices.iter().any(|d| d.is_cuda()) {
            return true;
        }
    }

    #[cfg(feature = "metal")]
    {
        if devices.iter().any(|d| d.is_metal()) {
            return true;
        }
    }

    false
}

/// CPU fallback backend for systems without GPU
pub struct CpuFallbackBackend {
    next_id: std::sync::atomic::AtomicU64,
    matrices: std::sync::RwLock<std::collections::HashMap<u64, Array2<f32>>>,
}

impl CpuFallbackBackend {
    /// Create a new CPU fallback backend
    pub fn new() -> Self {
        Self {
            next_id: std::sync::atomic::AtomicU64::new(1),
            matrices: std::sync::RwLock::new(std::collections::HashMap::new()),
        }
    }

    fn next_id(&self) -> u64 {
        self.next_id
            .fetch_add(1, std::sync::atomic::Ordering::SeqCst)
    }

    fn get_matrix(&self, id: u64) -> Result<Array2<f32>> {
        let matrices = self.matrices.read().unwrap();
        matrices
            .get(&id)
            .cloned()
            .ok_or_else(|| anyhow::anyhow!("Matrix {} not found", id))
    }

    fn store_matrix(&self, data: Array2<f32>) -> GpuMatrix {
        let id = self.next_id();
        let (rows, cols) = data.dim();
        self.matrices.write().unwrap().insert(id, data);
        GpuMatrix::new(id, rows, cols, GpuHandle::BufferIndex(id as u32))
    }
}

impl Default for CpuFallbackBackend {
    fn default() -> Self {
        Self::new()
    }
}

impl GpuBackend for CpuFallbackBackend {
    fn name(&self) -> &'static str {
        "CPU Fallback"
    }

    fn device_info(&self) -> DeviceInfo {
        DeviceInfo {
            name: "CPU".to_string(),
            total_memory: 0,
            available_memory: None,
            compute_capability: "N/A".to_string(),
            backend: BackendType::Cpu,
        }
    }

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

    fn allocate(&self, rows: usize, cols: usize) -> Result<GpuMatrix> {
        let data = Array2::zeros((rows, cols));
        Ok(self.store_matrix(data))
    }

    fn upload(&self, data: &Array2<f32>) -> Result<GpuMatrix> {
        Ok(self.store_matrix(data.clone()))
    }

    fn download(&self, matrix: &GpuMatrix) -> Result<Array2<f32>> {
        self.get_matrix(matrix.id)
    }

    fn matmul(&self, a: &GpuMatrix, b: &GpuMatrix) -> Result<GpuMatrix> {
        let a_data = self.get_matrix(a.id)?;
        let b_data = self.get_matrix(b.id)?;
        let result = a_data.dot(&b_data);
        Ok(self.store_matrix(result))
    }

    fn transpose(&self, a: &GpuMatrix) -> Result<GpuMatrix> {
        let data = self.get_matrix(a.id)?;
        let result = data.t().to_owned();
        Ok(self.store_matrix(result))
    }

    fn element_mul(&self, a: &GpuMatrix, b: &GpuMatrix) -> Result<GpuMatrix> {
        let a_data = self.get_matrix(a.id)?;
        let b_data = self.get_matrix(b.id)?;
        let result = &a_data * &b_data;
        Ok(self.store_matrix(result))
    }

    fn element_div_eps(&self, a: &GpuMatrix, b: &GpuMatrix, epsilon: f32) -> Result<GpuMatrix> {
        let a_data = self.get_matrix(a.id)?;
        let b_data = self.get_matrix(b.id)?;
        let result = &a_data / &(&b_data + epsilon);
        Ok(self.store_matrix(result))
    }

    fn nmf_update_inplace(
        &self,
        target: &mut GpuMatrix,
        numerator: &GpuMatrix,
        denominator: &GpuMatrix,
        epsilon: f32,
    ) -> Result<()> {
        let mut target_data = self.get_matrix(target.id)?;
        let num_data = self.get_matrix(numerator.id)?;
        let denom_data = self.get_matrix(denominator.id)?;

        for ((t, n), d) in target_data
            .iter_mut()
            .zip(num_data.iter())
            .zip(denom_data.iter())
        {
            *t = *t * n / (d + epsilon);
        }

        self.matrices
            .write()
            .unwrap()
            .insert(target.id, target_data);
        Ok(())
    }

    fn frobenius_diff(&self, a: &GpuMatrix, b: &GpuMatrix) -> Result<f32> {
        let a_data = self.get_matrix(a.id)?;
        let b_data = self.get_matrix(b.id)?;
        let diff = &a_data - &b_data;
        let norm = diff.iter().map(|x| x * x).sum::<f32>().sqrt();
        Ok(norm)
    }

    fn synchronize(&self) -> Result<()> {
        // No-op for CPU
        Ok(())
    }
}

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

    #[test]
    fn test_cpu_fallback_matmul() {
        let backend = CpuFallbackBackend::new();

        let a = Array2::from_shape_vec((2, 3), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let b = Array2::from_shape_vec((3, 2), vec![7.0, 8.0, 9.0, 10.0, 11.0, 12.0]).unwrap();

        let gpu_a = backend.upload(&a).unwrap();
        let gpu_b = backend.upload(&b).unwrap();
        let gpu_c = backend.matmul(&gpu_a, &gpu_b).unwrap();
        let c = backend.download(&gpu_c).unwrap();

        let expected = a.dot(&b);
        assert_eq!(c.shape(), expected.shape());
        for (actual, exp) in c.iter().zip(expected.iter()) {
            assert!((actual - exp).abs() < 1e-5);
        }
    }

    #[test]
    fn test_cpu_fallback_nmf_update() {
        let backend = CpuFallbackBackend::new();

        let target = Array2::from_elem((3, 3), 1.0_f32);
        let numerator = Array2::from_elem((3, 3), 2.0_f32);
        let denominator = Array2::from_elem((3, 3), 1.0_f32);

        let mut gpu_target = backend.upload(&target).unwrap();
        let gpu_num = backend.upload(&numerator).unwrap();
        let gpu_denom = backend.upload(&denominator).unwrap();

        backend
            .nmf_update_inplace(&mut gpu_target, &gpu_num, &gpu_denom, 1e-10)
            .unwrap();

        let result = backend.download(&gpu_target).unwrap();
        // target * numerator / denominator = 1.0 * 2.0 / 1.0 = 2.0
        for val in result.iter() {
            assert!((val - 2.0).abs() < 1e-5);
        }
    }

    #[test]
    fn test_gpu_matrix_shape() {
        let matrix = GpuMatrix::new(1, 10, 20, GpuHandle::None);
        assert_eq!(matrix.shape(), (10, 20));
        assert_eq!(matrix.len(), 200);
        assert!(!matrix.is_empty());
    }

    #[test]
    fn test_backend_type_display() {
        assert_eq!(format!("{}", BackendType::Cuda), "CUDA");
        assert_eq!(format!("{}", BackendType::Metal), "Metal");
        assert_eq!(format!("{}", BackendType::Cpu), "CPU");
    }
}