Initial commit

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

@@ -0,0 +1,425 @@
+//! 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");
+    }
+}