rustytorch/crates/core/rtx-backend-cuda/src/lib.rs

//! # RustyTorch++ CUDA Backend
//!
//! NVIDIA GPU backend implementation using hand-optimized CUDA kernels.
//!
//! ## Features
//!
//! - **cuBLAS Integration**: High-performance matrix operations via cuBLAS
//! - **cuDNN Integration**: Optimized convolutions and normalization
//! - **Flash Attention**: Hand-optimized FlashAttention-3 kernels
//! - **Tensor Core Support**: FP16/BF16/FP8 with Tensor Core acceleration
//!
//! ## Architecture
//!
//! ```text
//! CudaBackend
//! ├── CudaTensorPrimitive - GPU memory handle with cudarc
//! ├── CudaDevice - Device context and stream management
//! └── Ops
//!     ├── Basic - Add, mul, etc. (cuBLAS)
//!     ├── GEMM - Matrix multiply (cuBLAS LT)
//!     └── Attention - Flash Attention (hand-optimized)
//! ```
//!
//! ## Performance Notes
//!
//! - Uses RTX 5090 (sm_90) optimizations by default
//! - Tensor Core scheduling for BF16 and FP8 operations
//! - CUDA Graph capture for reduced kernel launch overhead
//!
//! ## Example
//!
//! ```rust,ignore
//! use rtx_backend_cuda::{CudaBackend, CudaDevice};
//! use rtx_backend::Backend;
//!
//! let device = CudaDevice::new(0)?;
//! let a = CudaBackend::zeros([1024, 1024], &device);
//! let b = CudaBackend::randn([1024, 1024], &device);
//! let c = CudaBackend::matmul(&a, &b);
//! ```

#![warn(missing_docs)]

// Note: When the cuda feature is disabled, stub implementations are provided.
// Enable the cuda feature for actual CUDA functionality.

mod device;
mod error;
mod tensor;

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

#[cfg(feature = "cuda")]
/// Operations module containing GPU-accelerated tensor operations.
pub mod ops;

#[cfg(not(feature = "cuda"))]
/// Operations module (stub - requires cuda feature).
pub mod ops {
    //! Stub operations module - requires cuda feature.

    /// Tensor creation operations (stub).
    pub mod creation {}
    /// Basic arithmetic operations (stub).
    pub mod basic {}
    /// Unary mathematical operations (stub).
    pub mod unary {}
    /// General matrix multiplication operations (stub).
    pub mod gemm {}
    /// Reduction operations like sum and mean (stub).
    pub mod reduction {}
    /// Shape manipulation operations (stub).
    pub mod shape {}
    /// Activation functions (stub).
    pub mod activation {}
    /// Normalization operations (stub).
    pub mod normalization {}
    /// Attention mechanism operations (stub).
    pub mod attention {}
    /// Device management operations (stub).
    pub mod device {}

    use parking_lot::Mutex;

    static RNG: Mutex<Option<rand::rngs::StdRng>> = Mutex::new(None);

    /// Seeds the random number generator with the given seed (stub).
    pub fn seed_rng(_seed: u64) {}

    pub(crate) fn get_rng() -> rand::rngs::StdRng {
        use rand::SeedableRng;
        rand::rngs::StdRng::from_entropy()
    }
}

/// Windows-specific CUDA support.
#[cfg(target_os = "windows")]
pub mod windows;
#[cfg(target_os = "windows")]
pub use windows::{WindowsCudaConfig, WindowsGpuInfo};

pub use device::CudaDevice;
pub use error::{CudaError, CudaResult};
pub use tensor::CudaTensorPrimitive;

#[cfg(feature = "cuda")]
use rtx_backend::{Backend, BoolU8, DeviceId, DeviceOps};
use std::fmt::Debug;

/// CUDA backend for RustyTorch++.
///
/// This backend uses NVIDIA GPUs via cudarc and provides:
/// - cuBLAS for matrix operations
/// - cuDNN for convolutions
/// - Hand-optimized Flash Attention kernels
/// - Tensor Core acceleration for FP16/BF16/FP8
#[derive(Clone, Debug, Default)]
pub struct CudaBackend;

#[cfg(feature = "cuda")]
impl Backend for CudaBackend {
    type TensorPrimitive<const D: usize> = CudaTensorPrimitive<D>;
    type Device = CudaDevice;
    type FloatElem = f32; // Default to f32, with FP16/BF16 support via type parameter
    type IntElem = i32;
    type BoolElem = BoolU8;

    fn name() -> &'static str {
        "cuda"
    }

    fn seed(seed: u64) {
        ops::seed_rng(seed);
    }

    // ==================== Tensor Creation ====================

    fn zeros<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D> {
        ops::creation::zeros(shape, device)
    }

    fn ones<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D> {
        ops::creation::ones(shape, device)
    }

    fn full<const D: usize>(
        shape: [usize; D],
        fill_value: Self::FloatElem,
        device: &Self::Device,
    ) -> Self::TensorPrimitive<D> {
        ops::creation::full(shape, fill_value, device)
    }

    fn rand<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D> {
        ops::creation::rand(shape, device)
    }

    fn randn<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D> {
        ops::creation::randn(shape, device)
    }

    fn from_data<const D: usize>(
        data: &[Self::FloatElem],
        shape: [usize; D],
        device: &Self::Device,
    ) -> Self::TensorPrimitive<D> {
        ops::creation::from_data(data, shape, device)
    }

    // ==================== Basic Operations ====================

    fn add<const D: usize>(
        lhs: Self::TensorPrimitive<D>,
        rhs: Self::TensorPrimitive<D>,
    ) -> Self::TensorPrimitive<D> {
        ops::basic::add(&lhs, &rhs)
    }

    fn sub<const D: usize>(
        lhs: Self::TensorPrimitive<D>,
        rhs: Self::TensorPrimitive<D>,
    ) -> Self::TensorPrimitive<D> {
        ops::basic::sub(&lhs, &rhs)
    }

    fn mul<const D: usize>(
        lhs: Self::TensorPrimitive<D>,
        rhs: Self::TensorPrimitive<D>,
    ) -> Self::TensorPrimitive<D> {
        ops::basic::mul(&lhs, &rhs)
    }

    fn div<const D: usize>(
        lhs: Self::TensorPrimitive<D>,
        rhs: Self::TensorPrimitive<D>,
    ) -> Self::TensorPrimitive<D> {
        ops::basic::div(&lhs, &rhs)
    }

    fn matmul(
        lhs: Self::TensorPrimitive<2>,
        rhs: Self::TensorPrimitive<2>,
    ) -> Self::TensorPrimitive<2> {
        ops::gemm::matmul(&lhs, &rhs)
    }

    fn bmm(
        lhs: Self::TensorPrimitive<3>,
        rhs: Self::TensorPrimitive<3>,
    ) -> Self::TensorPrimitive<3> {
        ops::gemm::bmm(&lhs, &rhs)
    }

    // ==================== Unary Operations ====================

    fn neg<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::unary::neg(&tensor)
    }

    fn exp<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::unary::exp(&tensor)
    }

    fn log<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::unary::log(&tensor)
    }

    fn sqrt<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::unary::sqrt(&tensor)
    }

    fn abs<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::unary::abs(&tensor)
    }

    fn sin<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::unary::sin(&tensor)
    }

    fn cos<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::unary::cos(&tensor)
    }

    fn pow<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        exp: Self::FloatElem,
    ) -> Self::TensorPrimitive<D> {
        ops::unary::pow(&tensor, exp)
    }

    fn clamp<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        min: Self::FloatElem,
        max: Self::FloatElem,
    ) -> Self::TensorPrimitive<D> {
        ops::unary::clamp(&tensor, min, max)
    }

    // ==================== Activation Functions ====================

    // ==================== Index Operations ====================
    // Native kernels: the trait defaults round-trip through host memory.

    fn index_select<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        indices: &[usize],
    ) -> Self::TensorPrimitive<D> {
        ops::index::index_select(&tensor, indices)
    }

    fn index_add<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        indices: &[usize],
        num_rows: usize,
    ) -> Self::TensorPrimitive<D> {
        ops::index::index_add(&tensor, indices, num_rows)
    }

    fn relu<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::activation::relu(&tensor)
    }

    fn sigmoid<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::activation::sigmoid(&tensor)
    }

    fn tanh<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::unary::tanh(&tensor)
    }

    // ==================== Reduction Operations ====================

    fn sum<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
        ops::reduction::sum(&tensor)
    }

    fn sum_dim<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        dim: usize,
    ) -> Self::TensorPrimitive<D> {
        ops::reduction::sum_dim(&tensor, dim)
    }

    fn mean<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
        ops::reduction::mean(&tensor)
    }

    fn mean_dim<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        dim: usize,
    ) -> Self::TensorPrimitive<D> {
        ops::reduction::mean_dim(&tensor, dim)
    }

    fn var<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
        ops::reduction::var(&tensor)
    }

    fn var_dim<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        dim: usize,
    ) -> Self::TensorPrimitive<D> {
        ops::reduction::var_dim(&tensor, dim)
    }

    fn max<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
        ops::reduction::max(&tensor)
    }

    fn min<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
        ops::reduction::min(&tensor)
    }

    // ==================== Shape Operations ====================

    fn shape<const D: usize>(tensor: &Self::TensorPrimitive<D>) -> [usize; D] {
        tensor.shape
    }

    fn reshape<const D1: usize, const D2: usize>(
        tensor: Self::TensorPrimitive<D1>,
        shape: [usize; D2],
    ) -> Self::TensorPrimitive<D2> {
        ops::shape::reshape(tensor, shape)
    }

    fn transpose<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::shape::transpose(&tensor)
    }

    fn swap_dims<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        dim1: usize,
        dim2: usize,
    ) -> Self::TensorPrimitive<D> {
        ops::shape::swap_dims(&tensor, dim1, dim2)
    }

    // ==================== LLM-Specific Operations ====================

    fn flash_attention(
        query: Self::TensorPrimitive<4>,
        key: Self::TensorPrimitive<4>,
        value: Self::TensorPrimitive<4>,
        mask: Option<&Self::TensorPrimitive<4>>,
        scale: Self::FloatElem,
        causal: bool,
    ) -> Self::TensorPrimitive<4> {
        ops::attention::flash_attention(&query, &key, &value, mask, scale, causal)
    }

    fn softmax<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        dim: usize,
    ) -> Self::TensorPrimitive<D> {
        ops::activation::softmax(&tensor, dim)
    }

    fn layer_norm<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        weight: &Self::TensorPrimitive<1>,
        bias: Option<&Self::TensorPrimitive<1>>,
        eps: Self::FloatElem,
    ) -> Self::TensorPrimitive<D> {
        ops::normalization::layer_norm(&tensor, weight, bias, eps)
    }

    fn rms_norm<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        weight: &Self::TensorPrimitive<1>,
        eps: Self::FloatElem,
    ) -> Self::TensorPrimitive<D> {
        ops::normalization::rms_norm(&tensor, weight, eps)
    }

    fn rope<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        cos: &Self::TensorPrimitive<2>,
        sin: &Self::TensorPrimitive<2>,
    ) -> Self::TensorPrimitive<D> {
        ops::attention::rope(&tensor, cos, sin)
    }

    fn gelu<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::activation::gelu(&tensor)
    }

    fn silu<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
        ops::activation::silu(&tensor)
    }

    fn leaky_relu<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        negative_slope: Self::FloatElem,
    ) -> Self::TensorPrimitive<D> {
        ops::activation::leaky_relu(&tensor, negative_slope)
    }

    fn elu<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        alpha: Self::FloatElem,
    ) -> Self::TensorPrimitive<D> {
        ops::activation::elu(&tensor, alpha)
    }

    // ==================== Comparison Operations ====================

    fn gt_scalar<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        value: Self::FloatElem,
    ) -> Self::TensorPrimitive<D> {
        ops::unary::gt_scalar(&tensor, value)
    }

    // ==================== Convolution Operations ====================

    fn conv2d(
        input: Self::TensorPrimitive<4>,
        weight: &Self::TensorPrimitive<4>,
        bias: Option<&Self::TensorPrimitive<1>>,
        stride: [usize; 2],
        padding: [usize; 2],
        dilation: [usize; 2],
        groups: usize,
    ) -> Self::TensorPrimitive<4> {
        ops::conv::conv2d(input, weight, bias, stride, padding, dilation, groups)
    }

    // ==================== Pooling Operations ====================

    fn max_pool2d(
        input: Self::TensorPrimitive<4>,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
    ) -> Self::TensorPrimitive<4> {
        ops::conv::max_pool2d(input, kernel_size, stride, padding)
    }

    fn avg_pool2d(
        input: Self::TensorPrimitive<4>,
        kernel_size: [usize; 2],
        stride: [usize; 2],
        padding: [usize; 2],
        count_include_pad: bool,
    ) -> Self::TensorPrimitive<4> {
        ops::conv::avg_pool2d(input, kernel_size, stride, padding, count_include_pad)
    }

    // ==================== Device Management ====================

    fn device<const D: usize>(tensor: &Self::TensorPrimitive<D>) -> Self::Device {
        tensor.device.clone()
    }

    fn to_device<const D: usize>(
        tensor: Self::TensorPrimitive<D>,
        device: &Self::Device,
    ) -> Self::TensorPrimitive<D> {
        if tensor.device == *device {
            tensor
        } else {
            ops::device::copy_to_device(tensor, device)
        }
    }

    fn to_data<const D: usize>(tensor: &Self::TensorPrimitive<D>) -> Vec<Self::FloatElem> {
        ops::device::copy_to_host(tensor)
    }

    fn sync(device: &Self::Device) {
        device.synchronize();
    }
}

/// Type alias for training with CUDA + autodiff.
pub type CudaTraining = CudaBackend; // Will become Autodiff<CudaBackend>

/// Type alias for inference with CUDA (no autodiff overhead).
pub type CudaInference = CudaBackend;