oxiproj-core 0.1.2

//! Thread-local CUDA runtime cache for OxiProj GPU kernels.
//!
//! A CUDA context is bound to one CPU thread at a time — `oxicuda::Context` is
//! `Send` but **not** `Sync` (the driver migrates a context across threads via
//! `cuCtxSetCurrent`). Rather than serialise the whole GPU behind a global mutex,
//! each thread that touches the GPU lazily builds and caches its own
//! [`GpuRuntime`]: device context, default stream, and JIT-compiled PTX modules.
//! This mirrors CUDA's native per-thread context model and needs no `unsafe`,
//! which keeps this crate's `#![forbid(unsafe_code)]` intact.

use std::cell::RefCell;
use std::collections::HashMap;
use std::sync::Arc;

use oxicuda::{Context, CudaError, CudaResult, Device, JitOptions, Kernel, Module, Stream};

use super::ptx_math::{ptx_target_for_cc, PtxTarget};

/// One thread's GPU resources: device context, default stream, the PTX target
/// for the active device, and a cache of JIT-compiled modules keyed by a stable
/// kernel identifier so each PTX program is compiled at most once per thread.
struct GpuRuntime {
    ctx: Arc<Context>,
    stream: Stream,
    target: PtxTarget,
    modules: HashMap<&'static str, Arc<Module>>,
}

impl GpuRuntime {
    fn new() -> CudaResult<Self> {
        oxicuda::init()?;
        let device = Device::get(0)?;
        let (cc_major, cc_minor) = device.compute_capability()?;
        // `Context::new` pushes the new context current on this (the creating)
        // thread, which is also the only thread that will use this runtime.
        let ctx = Arc::new(Context::new(&device)?);
        let stream = Stream::new(&ctx)?;
        Ok(Self {
            ctx,
            stream,
            target: ptx_target_for_cc(cc_major, cc_minor),
            modules: HashMap::new(),
        })
    }
}

/// Tri-state probe so a missing or broken GPU is detected exactly once per
/// thread instead of on every dispatch.
enum Probe {
    Unprobed,
    Unavailable,
    Ready(GpuRuntime),
}

thread_local! {
    static RUNTIME: RefCell<Probe> = const { RefCell::new(Probe::Unprobed) };
}

/// Ensure this thread's runtime exists, returning a mutable handle, or
/// [`CudaError::NoDevice`] when no CUDA device is reachable.
fn ensure(slot: &mut Probe) -> CudaResult<&mut GpuRuntime> {
    if matches!(slot, Probe::Unprobed) {
        *slot = match GpuRuntime::new() {
            Ok(rt) => Probe::Ready(rt),
            Err(_) => Probe::Unavailable,
        };
    }
    match slot {
        Probe::Ready(rt) => Ok(rt),
        _ => Err(CudaError::NoDevice),
    }
}

/// Returns `true` when a CUDA device is reachable from the current thread.
///
/// The probe result is cached per thread for the life of the process, so this is
/// cheap to call in hot paths as a GPU/CPU dispatch guard. When this returns
/// `false`, callers should run their CPU fallback.
#[must_use]
pub fn cuda_available() -> bool {
    RUNTIME.with(|cell| ensure(&mut cell.borrow_mut()).is_ok())
}

/// Returns the active GPU's name (e.g. `"NVIDIA RTX A4000"`), or `None` when no
/// CUDA device is reachable. Intended for diagnostics and CLI reporting.
#[must_use]
pub fn device_name() -> Option<String> {
    RUNTIME.with(|cell| {
        let mut slot = cell.borrow_mut();
        let rt = ensure(&mut slot).ok()?;
        rt.ctx.device().name().ok()
    })
}

/// Run `body` with a ready [`Kernel`] for `entry` and this thread's stream.
///
/// `build_ptx` is invoked with the active device's [`PtxTarget`] to produce the
/// kernel's PTX source. The PTX is JIT-compiled once and cached under `key`;
/// subsequent calls with the same `key` reuse the compiled module (only the
/// cheap function lookup is repeated). The context is made current on the calling
/// thread before `body` runs.
///
/// # Errors
///
/// Returns [`CudaError::NoDevice`] when no GPU is reachable, or any error raised
/// while compiling the module, resolving the entry point, or by `body` itself.
/// Callers are expected to fall back to a CPU path on `Err`.
pub fn with_kernel<B, F, R>(key: &'static str, entry: &str, build_ptx: B, body: F) -> CudaResult<R>
where
    B: FnOnce(&PtxTarget) -> String,
    F: FnOnce(&Kernel, &Stream) -> CudaResult<R>,
{
    RUNTIME.with(|cell| {
        let mut slot = cell.borrow_mut();
        let rt = ensure(&mut slot)?;
        rt.ctx.set_current()?;
        let module = match rt.modules.get(key) {
            Some(m) => Arc::clone(m),
            None => {
                let ptx = build_ptx(&rt.target);
                // Compile at -O4 so ptxas inlines the transcendental `.func`
                // calls (otherwise they execute through param/local memory,
                // which is ~100× slower than the inlined register-only form).
                let opts = JitOptions {
                    optimization_level: 4,
                    target_from_context: true,
                    ..Default::default()
                };
                let (compiled, log) = Module::from_ptx_with_options(&ptx, &opts)?;
                if log.has_errors() {
                    return Err(CudaError::InvalidImage);
                }
                let m = Arc::new(compiled);
                rt.modules.insert(key, Arc::clone(&m));
                m
            }
        };
        let kernel = Kernel::from_module(module, entry)?;
        body(&kernel, &rt.stream)
    })
}