ringkernel_cuda/

lib.rs

//! CUDA Backend for RingKernel
//!
//! This crate provides NVIDIA CUDA GPU support for RingKernel using cudarc.
//!
//! # Features
//!
//! - Persistent kernel execution (cooperative groups)
//! - Lock-free message queues in GPU global memory
//! - PTX compilation via NVRTC
//! - Multi-GPU support
//!
//! # Requirements
//!
//! - NVIDIA GPU with Compute Capability 6.0+ (Pascal and newer) for the core
//!   persistent-actor/cooperative-groups path. Some features have higher
//!   floors (e.g. Thread Block Clusters/DSMEM/TMA/Green Contexts require
//!   Hopper, CC 9.0+) — see the top-level README's "Feature to minimum
//!   compute capability" table for the full breakdown.
//! - CUDA Toolkit 11.0+
//! - Native Linux (persistent kernels) or WSL2 (event-driven fallback)
//!
//! # Example
//!
//! ```ignore
//! use ringkernel_cuda::CudaRuntime;
//! use ringkernel_core::runtime::RingKernelRuntime;
//!
//! #[tokio::main]
//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
//!     let runtime = CudaRuntime::new().await?;
//!     let kernel = runtime.launch("vector_add", Default::default()).await?;
//!     kernel.activate().await?;
//!     Ok(())
//! }
//! ```

#![warn(missing_docs)]
#![warn(clippy::unwrap_used)]

#[cfg(feature = "ptx-cache")]
pub mod compile;
#[cfg(feature = "cooperative")]
pub mod cooperative;
#[cfg(feature = "cuda")]
mod device;
#[cfg(feature = "cuda")]
pub mod driver_api;
#[cfg(feature = "cuda")]
pub mod hopper;
#[cfg(feature = "cuda")]
pub mod k2k_gpu;
#[cfg(feature = "cuda")]
mod kernel;
#[cfg(feature = "cuda")]
pub mod launch_config;
#[cfg(feature = "cuda")]
mod memory;
#[cfg(feature = "cuda")]
pub mod memory_pool;
#[cfg(feature = "cuda")]
pub mod multi_gpu;
#[cfg(feature = "cuda")]
pub mod persistent;
#[cfg(feature = "cuda")]
pub mod phases;
#[cfg(feature = "profiling")]
pub mod profiling;
#[cfg(feature = "cuda")]
pub mod reduction;
#[cfg(feature = "cuda")]
mod runtime;
#[cfg(feature = "cuda")]
mod stencil;
#[cfg(feature = "cuda")]
pub mod stream;

#[cfg(feature = "cuda")]
pub use device::CudaDevice;
#[cfg(feature = "cuda")]
pub use kernel::CudaKernel;
#[cfg(feature = "cuda")]
pub use memory::{CudaBuffer, CudaControlBlock, CudaMemoryPool, CudaMessageQueue};
#[cfg(feature = "cuda")]
pub use persistent::CudaMappedBuffer;
#[cfg(feature = "cuda")]
pub use phases::{
    InterPhaseReduction, KernelPhase, MultiPhaseConfig, MultiPhaseExecutor, PhaseExecutionStats,
    SyncMode,
};
#[cfg(feature = "cuda")]
pub use reduction::{
    generate_block_reduce_code, generate_grid_reduce_code, generate_reduce_and_broadcast_code,
    CacheKey, CacheStats, CachedReductionBuffer, ReductionBuffer, ReductionBufferBuilder,
    ReductionBufferCache,
};
#[cfg(feature = "cuda")]
pub use runtime::CudaRuntime;
#[cfg(feature = "cuda")]
pub use stencil::{CompiledStencilKernel, LaunchConfig, StencilKernelLoader};

// Profiling re-exports
#[cfg(feature = "profiling")]
pub use profiling::{
    CudaEvent, CudaEventFlags, CudaMemoryKind, CudaMemoryTracker, CudaNvtxProfiler,
    GpuChromeTraceBuilder, GpuEventArgs, GpuTimer, GpuTimerPool, GpuTraceEvent, KernelMetrics,
    ProfilingSession, TrackedAllocation, TransferDirection, TransferMetrics,
};

// PTX cache re-exports
#[cfg(feature = "ptx-cache")]
pub use compile::{PtxCache, PtxCacheError, PtxCacheResult, PtxCacheStats, CACHE_VERSION};

// GPU memory pool re-exports
#[cfg(feature = "cuda")]
pub use memory_pool::{
    GpuBucketStats, GpuPoolConfig, GpuPoolDiagnostics, GpuSizeClass, GpuStratifiedPool,
};

// Stream manager re-exports
#[cfg(feature = "cuda")]
pub use stream::{
    OverlapMetrics, StreamConfig, StreamConfigBuilder, StreamError, StreamId, StreamManager,
    StreamPool, StreamPoolStats, StreamResult,
};

/// Re-export memory module for advanced usage.
#[cfg(feature = "cuda")]
pub mod memory_exports {
    pub use super::memory::{CudaBuffer, CudaControlBlock, CudaMemoryPool, CudaMessageQueue};
}

// Placeholder implementations when CUDA is not available
#[cfg(not(feature = "cuda"))]
mod stub {
    ringkernel_core::unavailable_backend!(
        CudaRuntime,
        ringkernel_core::runtime::Backend::Cuda,
        "CUDA"
    );
}

#[cfg(not(feature = "cuda"))]
pub use stub::CudaRuntime;

/// Check if CUDA is available at runtime.
///
/// This function returns false if:
/// - CUDA feature is not enabled
/// - CUDA libraries are not installed on the system
/// - No CUDA devices are present
///
/// It safely catches panics from cudarc when CUDA is not installed.
pub fn is_cuda_available() -> bool {
    #[cfg(feature = "cuda")]
    {
        // cudarc panics if CUDA libraries are not found, so we catch that
        std::panic::catch_unwind(|| {
            cudarc::driver::CudaContext::device_count()
                .map(|c| c > 0)
                .unwrap_or(false)
        })
        .unwrap_or(false)
    }
    #[cfg(not(feature = "cuda"))]
    {
        false
    }
}

/// Get CUDA device count.
///
/// Returns 0 if CUDA is not available or libraries are not installed.
pub fn cuda_device_count() -> usize {
    #[cfg(feature = "cuda")]
    {
        // cudarc panics if CUDA libraries are not found, so we catch that
        std::panic::catch_unwind(|| {
            cudarc::driver::CudaContext::device_count().unwrap_or(0) as usize
        })
        .unwrap_or(0)
    }
    #[cfg(not(feature = "cuda"))]
    {
        0
    }
}

/// Compile CUDA C source code to PTX using NVRTC.
///
/// This wraps `cudarc::nvrtc::compile_ptx` to provide PTX compilation
/// without requiring downstream crates to depend on cudarc directly.
///
/// # Arguments
///
/// * `cuda_source` - CUDA C source code string
///
/// # Returns
///
/// PTX assembly as a string, or an error if compilation fails.
///
/// # Example
///
/// ```ignore
/// use ringkernel_cuda::compile_ptx;
///
/// let cuda_source = r#"
///     extern "C" __global__ void add(float* a, float* b, float* c, int n) {
///         int i = blockIdx.x * blockDim.x + threadIdx.x;
///         if (i < n) c[i] = a[i] + b[i];
///     }
/// "#;
///
/// let ptx = compile_ptx(cuda_source)?;
/// ```
#[cfg(feature = "cuda")]
pub fn compile_ptx(cuda_source: &str) -> ringkernel_core::error::Result<String> {
    use ringkernel_core::error::RingKernelError;

    let ptx = cudarc::nvrtc::compile_ptx(cuda_source).map_err(|e| {
        RingKernelError::CompilationError(format!("NVRTC compilation failed: {}", e))
    })?;

    Ok(ptx.to_src().to_string())
}

/// Stub compile_ptx when CUDA is not available.
#[cfg(not(feature = "cuda"))]
pub fn compile_ptx(_cuda_source: &str) -> ringkernel_core::error::Result<String> {
    Err(ringkernel_core::error::RingKernelError::BackendUnavailable(
        "CUDA feature not enabled".to_string(),
    ))
}

/// PTX kernel source template for persistent ring kernel, generic over the
/// target compute capability (major, minor).
///
/// This is a minimal kernel that immediately marks itself as terminated.
///
/// PTX `.target` directives are only forward-compatible (PTX built for
/// `sm_X` runs on `sm_X` and newer, never older) — so this must be
/// generated per-device rather than hardcoded to a single value. Maps the
/// device's actual compute capability to the corresponding PTX ISA
/// version, matching the minimum ISA version that introduced each `sm_XX`
/// target (per NVIDIA's PTX ISA documentation).
pub fn ring_kernel_ptx_template_for(major: u32, minor: u32) -> String {
    let ptx_version = match (major, minor) {
        // Pascal
        (6, 0) | (6, 1) | (6, 2) => "5.0",
        // Volta / Turing
        (7, 0) | (7, 2) => "6.0",
        (7, 5) => "6.3",
        // Ampere
        (8, 0) => "7.0",
        (8, 6) | (8, 7) => "7.1",
        (8, 9) => "8.0",
        // Hopper
        (9, 0) => "8.0",
        // Blackwell and newer — fall through to the newest known ISA
        (major, _) if major >= 10 => "8.5",
        // Below Pascal (Maxwell/Kepler) or anything unrecognized: fall back
        // to the oldest ISA version this template's instructions need.
        _ => "5.0",
    };
    format!(
        r#"
.version {ptx_version}
.target sm_{major}{minor}
.address_size 64

.visible .entry ring_kernel_main(
    .param .u64 control_block_ptr,
    .param .u64 input_queue_ptr,
    .param .u64 output_queue_ptr,
    .param .u64 shared_state_ptr
) {{
    .reg .u64 %cb_ptr;
    .reg .u32 %one;

    // Load control block pointer
    ld.param.u64 %cb_ptr, [control_block_ptr];

    // Mark as terminated immediately (offset 8)
    mov.u32 %one, 1;
    st.global.u32 [%cb_ptr + 8], %one;

    ret;
}}
"#
    )
}

/// PTX kernel source template for persistent ring kernel — **deprecated**
/// fixed-`sm_75` fallback, kept only for API compatibility with existing
/// callers that don't have a device handle available.
///
/// PTX built for `sm_75` will fail to load (`CUDA_ERROR_INVALID_PTX`) on any
/// GPU with a compute capability below 7.5 (e.g. Pascal, sm_61) — PTX
/// forward-compatibility only extends to equal-or-newer architectures.
/// Prefer [`ring_kernel_ptx_template_for`] with the actual target device's
/// compute capability (`CudaDevice::compute_capability`, not part of this
/// crate's public API surface).
#[deprecated(
    since = "1.1.1",
    note = "hardcodes sm_75, breaking on Pascal and older GPUs — use ring_kernel_ptx_template_for(major, minor) with the real device's compute capability instead"
)]
pub const RING_KERNEL_PTX_TEMPLATE: &str = r#"
.version 8.0
.target sm_75
.address_size 64

.visible .entry ring_kernel_main(
    .param .u64 control_block_ptr,
    .param .u64 input_queue_ptr,
    .param .u64 output_queue_ptr,
    .param .u64 shared_state_ptr
) {
    .reg .u64 %cb_ptr;
    .reg .u32 %one;

    // Load control block pointer
    ld.param.u64 %cb_ptr, [control_block_ptr];

    // Mark as terminated immediately (offset 8)
    mov.u32 %one, 1;
    st.global.u32 [%cb_ptr + 8], %one;

    ret;
}
"#;