onnx-runtime-ep-cuda 0.1.0-dev.5

//! Shared CUDA runtime state: the driver context, its dedicated stream, and vendor
//! library backends. One [`CudaRuntime`] is created per
//! [`CudaExecutionProvider`] and shared (via `Arc`) into every kernel the
//! provider hands out, so the whole EP drives a single device + stream.

use std::collections::HashMap;
use std::ffi::{CStr, CString, c_void};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};

use cudarc::driver::sys::{CUdevice_attribute, CUdeviceptr, CUfunction_attribute_enum};
use cudarc::driver::{CudaContext, CudaFunction, CudaModule, CudaStream, LaunchConfig};

use onnx_runtime_ep_api::EpError;
use onnx_runtime_ep_api::Kernel;
use onnx_runtime_ep_api::Result;

use crate::blas::CublasLt;
use crate::cudnn::CudnnBackend;
use crate::dynamic_library::{CudaLibrary, require};
use crate::error::{driver_err, nvrtc_err};
use crate::graph::CudaGraphLifecycle;

/// Counts explicit device allocation/free calls made through a runtime.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct CudaAllocationCounts {
    pub allocations: u64,
    pub frees: u64,
}

fn nvrtc_include_paths() -> Vec<String> {
    let mut candidates = Vec::<PathBuf>::new();
    for variable in ["CUDA_HOME", "CUDA_PATH"] {
        if let Some(root) = std::env::var_os(variable) {
            candidates.push(PathBuf::from(root).join("include"));
        }
    }
    candidates.push(PathBuf::from("/usr/local/cuda/include"));

    if let Some(paths) = std::env::var_os("LD_LIBRARY_PATH") {
        for path in std::env::split_paths(&paths) {
            if path.ends_with(Path::new("nvidia/cuda_nvrtc/lib"))
                && let Some(nvidia) = path.parent().and_then(Path::parent)
            {
                candidates.push(nvidia.join("cuda_runtime/include"));
            }
        }
    }

    candidates.sort();
    candidates.dedup();
    candidates
        .into_iter()
        .filter(|path| path.join("cuda_fp16.h").is_file())
        .map(|path| path.to_string_lossy().into_owned())
        .collect()
}

fn ptx_arch_for(major: u32, minor: u32) -> String {
    format!("compute_{major}{minor}")
}

fn cubin_arch_for(major: u32, minor: u32) -> String {
    format!("sm_{major}{minor}")
}

const SAFE_MAX_THREADS_PER_BLOCK_FALLBACK: u32 = 256;
const SAFE_SHARED_MEMORY_PER_BLOCK_FALLBACK: u32 = 48 * 1024;

/// Hardware limits used to select portable CUDA launch configurations.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct CudaDeviceCapabilities {
    compute_capability: (u32, u32),
    max_threads_per_block: u32,
    max_shared_memory_per_block: u32,
    max_shared_memory_per_block_optin: u32,
    multiprocessor_count: u32,
}

impl CudaDeviceCapabilities {
    fn from_reported_limits(
        compute_capability: (u32, u32),
        max_threads_per_block: Option<u32>,
        max_shared_memory_per_block: Option<u32>,
        max_shared_memory_per_block_optin: Option<u32>,
        multiprocessor_count: Option<u32>,
    ) -> Self {
        let max_threads_per_block = max_threads_per_block
            .filter(|&value| value > 0)
            .unwrap_or(SAFE_MAX_THREADS_PER_BLOCK_FALLBACK);
        let max_shared_memory_per_block = max_shared_memory_per_block
            .filter(|&value| value > 0)
            .unwrap_or(SAFE_SHARED_MEMORY_PER_BLOCK_FALLBACK);
        let max_shared_memory_per_block_optin = max_shared_memory_per_block_optin
            .filter(|&value| value > 0)
            .unwrap_or(max_shared_memory_per_block)
            .max(max_shared_memory_per_block);
        let multiprocessor_count = multiprocessor_count.filter(|&value| value > 0).unwrap_or(1);
        Self {
            compute_capability,
            max_threads_per_block,
            max_shared_memory_per_block,
            max_shared_memory_per_block_optin,
            multiprocessor_count,
        }
    }

    pub fn compute_capability(self) -> (u32, u32) {
        self.compute_capability
    }

    pub fn max_shared_memory_per_block_optin(self) -> u32 {
        self.max_shared_memory_per_block_optin
    }

    pub fn max_threads_per_block(self) -> u32 {
        self.max_threads_per_block
    }

    pub fn multiprocessor_count(self) -> u32 {
        self.multiprocessor_count
    }
}

fn positive_attribute(context: &CudaContext, attribute: CUdevice_attribute) -> Option<u32> {
    context
        .attribute(attribute)
        .ok()
        .and_then(|value| u32::try_from(value).ok())
        .filter(|&value| value > 0)
}

fn reduction_launch_params(
    preferred_threads: u32,
    max_threads: u32,
    bytes_per_thread: u32,
    max_dynamic_shared_memory: u32,
) -> Option<(u32, u32)> {
    if preferred_threads == 0 || max_threads == 0 || bytes_per_thread == 0 {
        return None;
    }
    let threads_by_shared_memory = max_dynamic_shared_memory / bytes_per_thread;
    let thread_limit = preferred_threads
        .min(max_threads)
        .min(threads_by_shared_memory);
    if thread_limit == 0 {
        return None;
    }
    let threads = 1 << (31 - thread_limit.leading_zeros());
    Some((threads, threads * bytes_per_thread))
}

/// Decide how a dynamic shared-memory request maps onto a device's per-block
/// budgets. `default_budget` is the non-opt-in ceiling (~48&nbsp;KB on every
/// architecture) and `optin_budget` the device-specific opt-in ceiling, both
/// already net of the kernel's static shared memory. Returns:
/// * `Err(())` — the request exceeds even the opt-in ceiling, so no launch on
///   this GPU can satisfy it and the caller must route to a portable fallback.
/// * `Ok(None)` — the request fits the default budget; launch as-is.
/// * `Ok(Some(bytes))` — the request needs the function opted into `bytes` of
///   dynamic shared memory before it can launch.
fn dynamic_shared_memory_optin(
    requested_bytes: u32,
    default_budget: u32,
    optin_budget: u32,
) -> std::result::Result<Option<u32>, ()> {
    if requested_bytes > optin_budget {
        return Err(());
    }
    if requested_bytes > default_budget {
        Ok(Some(requested_bytes))
    } else {
        Ok(None)
    }
}

/// Device context, stream, and vendor-library backends shared across the EP.
pub struct CudaRuntime {
    context: Arc<CudaContext>,
    stream: Arc<CudaStream>,
    graph: CudaGraphLifecycle,
    blas: CublasLt,
    cudnn: CudnnBackend,
    ordinal: u32,
    capabilities: CudaDeviceCapabilities,
    ptx_arch: String,
    cubin_arch: String,
    /// Cache of NVRTC-compiled modules, keyed by a stable module name, so each
    /// runtime compiles a given kernel (e.g. the fused attention softmax) at
    /// most once and reuses the loaded module for every kernel invocation.
    modules: Mutex<HashMap<&'static str, Arc<CudaModule>>>,
    /// Set after a driver rejects the toolkit's PTX ISA. Subsequent modules are
    /// compiled directly to the device's native SM CUBIN instead of repeating
    /// the failed load.
    nvrtc_cubin_fallback: AtomicBool,
    allocations: AtomicU64,
    frees: AtomicU64,
    /// Persistent four-byte device word into which capture-safe kernels latch an
    /// out-of-range bounds violation detected during CUDA-graph replay. It is set
    /// (via `atomicOr`) by device kernels and never auto-cleared on the device;
    /// only an explicit host [`CudaRuntime::reset_capture_error`] (invoked on
    /// graph reset / re-capture) returns it to zero. The host reads it at the
    /// existing per-step logits sync so a poisoned replay becomes a hard error
    /// before the produced token is consumed.
    capture_error: CUdeviceptr,
}

impl std::fmt::Debug for CudaRuntime {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CudaRuntime")
            .field("ordinal", &self.ordinal)
            .field("capabilities", &self.capabilities)
            .finish()
    }
}

impl CudaRuntime {
    /// Initialise the primary context on CUDA device `ordinal`, its dedicated
    /// stream, and a cuBLASLt handle. Returns an error (never panics) when no
    /// such device exists or the CUDA driver / cuBLASLt cannot be loaded.
    pub fn new(ordinal: u32) -> Result<Self> {
        require(CudaLibrary::Driver).map_err(|message| {
            EpError::KernelFailed(format!(
                "cuda_ep: {message}; CPU execution remains available"
            ))
        })?;
        require(CudaLibrary::CublasLt).map_err(|message| {
            EpError::KernelFailed(format!(
                "cuda_ep: {message}; CPU execution remains available"
            ))
        })?;
        let context =
            CudaContext::new(ordinal as usize).map_err(|e| driver_err("CudaContext::new", e))?;
        let major = context
            .attribute(CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR)
            .map_err(|e| driver_err("querying CUDA compute capability major", e))?;
        let minor = context
            .attribute(CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR)
            .map_err(|e| driver_err("querying CUDA compute capability minor", e))?;
        let major = u32::try_from(major).map_err(|_| {
            EpError::KernelFailed(format!(
                "cuda_ep: CUDA device {ordinal} reported invalid compute capability major {major}"
            ))
        })?;
        let minor = u32::try_from(minor).map_err(|_| {
            EpError::KernelFailed(format!(
                "cuda_ep: CUDA device {ordinal} reported invalid compute capability minor {minor}"
            ))
        })?;
        if major == 0 {
            return Err(EpError::KernelFailed(format!(
                "cuda_ep: CUDA device {ordinal} reported invalid compute capability {major}.{minor}"
            )));
        }
        let compute_capability = (major, minor);
        let capabilities = CudaDeviceCapabilities::from_reported_limits(
            compute_capability,
            positive_attribute(
                &context,
                CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK,
            ),
            positive_attribute(
                &context,
                CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK,
            ),
            positive_attribute(
                &context,
                CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN,
            ),
            positive_attribute(
                &context,
                CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT,
            ),
        );
        let ptx_arch = ptx_arch_for(major, minor);
        let cubin_arch = cubin_arch_for(major, minor);
        // A dedicated non-blocking stream (not the legacy NULL stream, which the
        // driver refuses to capture) so device-resident kernels are eligible for
        // CUDA-graph capture. The whole EP drives this single stream, so its
        // ordering is self-contained and host-blocking `*_sync` copies remain
        // correctly serialized against kernel launches.
        let stream = context
            .new_stream()
            .map_err(|e| driver_err("create compute stream", e))?;
        let blas = CublasLt::new()?;
        let cudnn = CudnnBackend::new(stream.clone());
        let graph = CudaGraphLifecycle::new(stream.clone());
        Self {
            context,
            stream,
            graph,
            blas,
            cudnn,
            ordinal,
            capabilities,
            ptx_arch,
            cubin_arch,
            modules: Mutex::new(HashMap::new()),
            nvrtc_cubin_fallback: AtomicBool::new(false),
            allocations: AtomicU64::new(0),
            frees: AtomicU64::new(0),
            capture_error: 0,
        }
        .with_capture_error_word()
    }

    /// Allocate and zero the persistent capture-error latch word. Split out of
    /// [`CudaRuntime::new`] so it can use the runtime's own bound-context
    /// alloc/copy helpers.
    fn with_capture_error_word(mut self) -> Result<Self> {
        let ptr = self.alloc_raw(std::mem::size_of::<u32>())?;
        // SAFETY: `ptr` is a fresh four-byte device allocation owned by this
        // runtime; zeroing it establishes the un-latched initial state.
        unsafe { self.htod(&0_u32.to_ne_bytes(), ptr) }?;
        self.capture_error = ptr;
        Ok(self)
    }

    /// The CUDA device ordinal this runtime drives.
    pub fn ordinal(&self) -> u32 {
        self.ordinal
    }

    /// Hardware capabilities reported by the selected CUDA device.
    pub fn capabilities(&self) -> CudaDeviceCapabilities {
        self.capabilities
    }

    /// The cuBLASLt handle.
    pub fn blas(&self) -> &CublasLt {
        &self.blas
    }

    /// The lazily initialized cuDNN backend bound to this runtime's stream.
    pub fn cudnn(&self) -> &CudnnBackend {
        &self.cudnn
    }

    /// The raw CUDA stream the EP submits work on.
    pub fn stream_ptr(&self) -> cudarc::driver::sys::CUstream {
        self.stream.cu_stream()
    }

    /// The EP's compute stream (for `launch_builder`-based kernel launches).
    pub fn stream(&self) -> &Arc<CudaStream> {
        &self.stream
    }

    /// Begin capture on the EP stream after auditing the complete kernel sequence.
    pub fn begin_graph_capture(&self, kernels: &[&dyn Kernel]) -> Result<()> {
        crate::capture::require_subgraph_graph_capturable(kernels)?;
        self.graph.begin()
    }

    /// End stream capture and install the instantiated graph executable.
    pub fn end_graph_capture(&self) -> Result<()> {
        self.graph.end()
    }

    /// Abort an in-progress stream capture, discarding any half-recorded graph
    /// and returning the lifecycle to idle so a subsequent [`reset_graph`]
    /// succeeds. Used on the error path of segmented capture.
    pub fn abort_graph_capture(&self) -> Result<()> {
        self.graph.abort()
    }

    /// Launch the installed graph executable on the same EP stream.
    ///
    /// Replays every installed segment in capture order (one graph for a
    /// whole-subgraph capture).
    pub fn replay_graph(&self) -> Result<()> {
        self.graph.replay()
    }

    /// Launch one installed segment by its zero-based capture-order index.
    pub fn replay_graph_segment(&self, index: usize) -> Result<()> {
        self.graph.replay_segment(index)
    }

    /// Number of installed captured segments (1 for a whole-subgraph capture).
    pub fn graph_segment_count(&self) -> Result<usize> {
        self.graph.segment_count()
    }

    /// Destroy the installed graph and graph-exec handles.
    ///
    /// Returns whether an executable was invalidated. Reset is rejected while a
    /// capture is active; callers must end the capture first.
    pub fn reset_graph(&self) -> Result<bool> {
        self.graph.reset()
    }

    /// Whether this runtime currently owns an instantiated graph executable.
    pub fn has_graph_executable(&self) -> Result<bool> {
        self.graph.has_executable()
    }

    /// Raw device pointer to the persistent capture-error latch word, passed to
    /// capture-safe kernels so they can `atomicOr` a bounds-violation code into
    /// it (and read it back to propagate the poison to later kernels/replays).
    pub fn capture_error_ptr(&self) -> CUdeviceptr {
        self.capture_error
    }

    /// Read the latching capture-error word device → host, returning the raw
    /// violation bitmask (zero when no capture-safe kernel has tripped).
    ///
    /// This does not clear the latch: once set, every subsequent captured replay
    /// stays poisoned until [`CudaRuntime::reset_capture_error`]. Callers invoke
    /// this at the per-step logits D2H sync, so the trailing stream synchronize
    /// is already satisfied and no additional serialization is introduced.
    pub fn check_capture_error(&self) -> Result<u32> {
        let mut bytes = [0_u8; std::mem::size_of::<u32>()];
        // SAFETY: `capture_error` is a live four-byte device allocation owned by
        // this runtime for its whole lifetime.
        unsafe { self.dtoh(&mut bytes, self.capture_error) }?;
        Ok(u32::from_ne_bytes(bytes))
    }

    /// Clear the latching capture-error word back to the un-poisoned state.
    /// Invoked on graph reset / re-capture so a fresh generation starts clean.
    pub fn reset_capture_error(&self) -> Result<()> {
        // SAFETY: `capture_error` is a live four-byte device allocation owned by
        // this runtime for its whole lifetime.
        unsafe { self.htod(&0_u32.to_ne_bytes(), self.capture_error) }
    }

    /// Driver-reported capture status for the EP stream.
    pub fn graph_capture_status(&self) -> Result<cudarc::driver::sys::CUstreamCaptureStatus> {
        self.graph.capture_status()
    }

    /// Snapshot explicit device allocation/free calls made through this runtime.
    pub fn allocation_counts(&self) -> CudaAllocationCounts {
        CudaAllocationCounts {
            allocations: self.allocations.load(Ordering::Relaxed),
            frees: self.frees.load(Ordering::Relaxed),
        }
    }

    /// Validate a requested dynamic shared-memory allocation against the device
    /// limits and, when it exceeds the default (non-opt-in) per-block budget,
    /// opt the function into the larger dynamic size the hardware supports.
    ///
    /// Every architecture caps *non-opt-in* dynamic shared memory at roughly
    /// 48&nbsp;KB, while the opt-in ceiling is device specific (for example
    /// ~100&nbsp;KB on sm_86/sm_89 consumer cards, ~163&nbsp;KB on sm_80, and
    /// ~227&nbsp;KB on sm_90). A kernel that requests more than 48&nbsp;KB without
    /// setting `CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES` fails to launch
    /// on *any* GPU, and one that requests more than the device's opt-in ceiling
    /// fails on that specific GPU — a request sized for an H200 can therefore
    /// crash a consumer card outright. This helper returns a loud error (never
    /// launching) when even the opt-in maximum cannot satisfy the request, so the
    /// caller can route to a portable fallback instead of hitting a hard launch
    /// failure. The static shared memory the function already reserves is
    /// subtracted from both budgets.
    pub fn configure_dynamic_shared_memory(
        &self,
        function: &CudaFunction,
        requested_bytes: u32,
    ) -> Result<()> {
        let static_shared_memory = function
            .shared_size_bytes()
            .map_err(|error| driver_err("querying CUDA function static shared memory", error))?;
        let static_shared_memory = u32::try_from(static_shared_memory).map_err(|_| {
            EpError::KernelFailed(format!(
                "cuda_ep: CUDA function reported invalid static shared memory {static_shared_memory}"
            ))
        })?;
        let optin_budget = self
            .capabilities
            .max_shared_memory_per_block_optin
            .saturating_sub(static_shared_memory);
        let default_budget = self
            .capabilities
            .max_shared_memory_per_block
            .saturating_sub(static_shared_memory);
        match dynamic_shared_memory_optin(requested_bytes, default_budget, optin_budget) {
            Err(()) => Err(EpError::KernelFailed(format!(
                "cuda_ep: kernel requests {requested_bytes} dynamic shared-memory bytes, but \
                 device SM {}.{} allows at most {optin_budget} opt-in bytes \
                 ({static_shared_memory} already reserved statically); route this shape to a \
                 portable kernel instead of launching",
                self.capabilities.compute_capability.0, self.capabilities.compute_capability.1,
            ))),
            Ok(None) => Ok(()),
            Ok(Some(bytes)) => {
                let bytes_i32 = i32::try_from(bytes).map_err(|_| {
                    EpError::KernelFailed(format!(
                        "cuda_ep: dynamic shared-memory request {bytes} exceeds i32"
                    ))
                })?;
                function
                    .set_attribute(
                        CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
                        bytes_i32,
                    )
                    .map_err(|error| {
                        driver_err("opting CUDA function into dynamic shared memory", error)
                    })
            }
        }
    }

    /// Build a power-of-two reduction launch that fits both the function and
    /// device thread/shared-memory limits. If the launch exceeds the legacy
    /// shared-memory limit, opt the function into the required dynamic size.
    pub fn reduction_launch_config(
        &self,
        function: &CudaFunction,
        grid_x: u32,
        preferred_threads: u32,
        bytes_per_thread: u32,
    ) -> Result<LaunchConfig> {
        let function_max_threads = function
            .max_threads_per_block()
            .map_err(|error| driver_err("querying CUDA function max threads", error))?;
        let function_max_threads = u32::try_from(function_max_threads).map_err(|_| {
            EpError::KernelFailed(format!(
                "cuda_ep: CUDA function reported invalid max threads {function_max_threads}"
            ))
        })?;
        let static_shared_memory = function
            .shared_size_bytes()
            .map_err(|error| driver_err("querying CUDA function static shared memory", error))?;
        let static_shared_memory = u32::try_from(static_shared_memory).map_err(|_| {
            EpError::KernelFailed(format!(
                "cuda_ep: CUDA function reported invalid static shared memory {static_shared_memory}"
            ))
        })?;
        let max_dynamic_shared_memory = self
            .capabilities
            .max_shared_memory_per_block_optin
            .saturating_sub(static_shared_memory);
        let max_threads = self
            .capabilities
            .max_threads_per_block
            .min(function_max_threads);
        let (threads, shared_mem_bytes) = reduction_launch_params(
            preferred_threads,
            max_threads,
            bytes_per_thread,
            max_dynamic_shared_memory,
        )
        .ok_or_else(|| {
            EpError::KernelFailed(format!(
                "cuda_ep: reduction launch needs {bytes_per_thread} shared-memory bytes per \
                 thread, but device SM {}.{} allows {max_dynamic_shared_memory} dynamic bytes",
                self.capabilities.compute_capability.0, self.capabilities.compute_capability.1,
            ))
        })?;

        let default_dynamic_shared_memory = self
            .capabilities
            .max_shared_memory_per_block
            .saturating_sub(static_shared_memory);
        if shared_mem_bytes > default_dynamic_shared_memory {
            let shared_mem_bytes_i32 = i32::try_from(shared_mem_bytes).map_err(|_| {
                EpError::KernelFailed(format!(
                    "cuda_ep: dynamic shared-memory request {shared_mem_bytes} exceeds i32"
                ))
            })?;
            function
                .set_attribute(
                    CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
                    shared_mem_bytes_i32,
                )
                .map_err(|error| {
                    driver_err("opting CUDA function into dynamic shared memory", error)
                })?;
        }

        Ok(LaunchConfig {
            grid_dim: (grid_x, 1, 1),
            block_dim: (threads, 1, 1),
            shared_mem_bytes,
        })
    }

    /// Get a [`CudaFunction`] for entry point `entry` in the NVRTC module named
    /// `module_key`, compiling `src` to PTX and loading it on first use and
    /// reusing the cached module thereafter.
    ///
    /// The compile targets the device's detected virtual compute architecture.
    /// If the installed NVRTC emits a PTX ISA newer than the driver accepts,
    /// compilation is retried for the matching real SM architecture and the
    /// resulting CUBIN is loaded instead. An NVRTC failure surfaces the compiler
    /// log via [`nvrtc_err`] (RULES.md #1).
    pub fn nvrtc_function(
        &self,
        module_key: &'static str,
        src: &str,
        entry: &str,
    ) -> Result<CudaFunction> {
        require(CudaLibrary::Nvrtc).map_err(|message| {
            EpError::KernelFailed(format!(
                "cuda_ep: {message}; CPU execution remains available"
            ))
        })?;
        self.bind()?;
        let module = {
            let mut cache = self.modules.lock().expect("cuda_ep module cache poisoned");
            if let Some(m) = cache.get(module_key) {
                m.clone()
            } else {
                let include_paths = nvrtc_include_paths();
                let m = if self.nvrtc_cubin_fallback.load(Ordering::Relaxed) {
                    self.load_nvrtc_cubin(module_key, src, &include_paths)?
                } else {
                    let opts = cudarc::nvrtc::CompileOptions {
                        include_paths: include_paths.clone(),
                        options: vec![format!("--gpu-architecture={}", self.ptx_arch)],
                        ..Default::default()
                    };
                    let ptx = cudarc::nvrtc::compile_ptx_with_opts(src, opts).map_err(|e| {
                        nvrtc_err(&format!("compiling NVRTC module '{module_key}'"), e)
                    })?;
                    match self.context.load_module(ptx) {
                        Ok(module) => module,
                        Err(error)
                            if error.0
                                == cudarc::driver::sys::CUresult::CUDA_ERROR_UNSUPPORTED_PTX_VERSION =>
                        {
                            self.nvrtc_cubin_fallback.store(true, Ordering::Relaxed);
                            self.load_nvrtc_cubin(module_key, src, &include_paths)?
                        }
                        Err(error) => {
                            return Err(driver_err(
                                &format!("loading NVRTC module '{module_key}'"),
                                error,
                            ));
                        }
                    }
                };
                cache.insert(module_key, m.clone());
                m
            }
        };
        module
            .load_function(entry)
            .map_err(|e| driver_err(&format!("loading NVRTC function '{entry}'"), e))
    }

    fn load_nvrtc_cubin(
        &self,
        module_key: &'static str,
        src: &str,
        include_paths: &[String],
    ) -> Result<Arc<CudaModule>> {
        let source = CString::new(src).map_err(|_| {
            EpError::KernelFailed(format!(
                "cuda_ep: compiling NVRTC module '{module_key}': source contains a NUL byte"
            ))
        })?;
        let name = CString::new(module_key).expect("static module key cannot contain a NUL byte");
        let program =
            cudarc::nvrtc::result::create_program(source.as_c_str(), Some(name.as_c_str()))
                .map_err(|error| {
                    EpError::KernelFailed(format!(
                        "cuda_ep: creating NVRTC CUBIN module '{module_key}': {error:?}"
                    ))
                })?;
        let mut options = include_paths
            .iter()
            .map(|path| format!("--include-path={path}"))
            .collect::<Vec<_>>();
        options.push(format!("--gpu-architecture={}", self.cubin_arch));

        // SAFETY: `program` is live until the matching destroy call below.
        let compile_result = unsafe { cudarc::nvrtc::result::compile_program(program, &options) };
        if let Err(error) = compile_result {
            // SAFETY: compilation may fail, but the live program still owns its log.
            let log = unsafe { cudarc::nvrtc::result::get_program_log(program) }
                .ok()
                .map(|bytes| {
                    // SAFETY: NVRTC returns a NUL-terminated compiler log.
                    unsafe { CStr::from_ptr(bytes.as_ptr()) }
                        .to_string_lossy()
                        .into_owned()
                })
                .unwrap_or_else(|| "<compiler log unavailable>".into());
            // SAFETY: this is the single destroy for the live program.
            let _ = unsafe { cudarc::nvrtc::result::destroy_program(program) };
            return Err(EpError::KernelFailed(format!(
                "cuda_ep: compiling NVRTC CUBIN module '{module_key}' failed ({error:?}); compiler log:\n{log}"
            )));
        }

        let cubin: Result<Vec<u8>> = (|| {
            let mut size = 0usize;
            // SAFETY: `program` compiled successfully and `size` is writable.
            unsafe { cudarc::nvrtc::sys::nvrtcGetCUBINSize(program, &mut size) }
                .result()
                .map_err(|error| {
                    EpError::KernelFailed(format!(
                        "cuda_ep: getting NVRTC CUBIN size for '{module_key}': {error:?}"
                    ))
                })?;
            let mut image = vec![0u8; size];
            // SAFETY: `image` has the exact size reported by NVRTC.
            unsafe { cudarc::nvrtc::sys::nvrtcGetCUBIN(program, image.as_mut_ptr().cast()) }
                .result()
                .map_err(|error| {
                    EpError::KernelFailed(format!(
                        "cuda_ep: getting NVRTC CUBIN for '{module_key}': {error:?}"
                    ))
                })?;
            Ok(image)
        })();
        // SAFETY: this is the single destroy for the live program.
        let destroy_result = unsafe { cudarc::nvrtc::result::destroy_program(program) };
        let image = cubin?;
        destroy_result.map_err(|error| {
            EpError::KernelFailed(format!(
                "cuda_ep: destroying NVRTC CUBIN program '{module_key}': {error:?}"
            ))
        })?;
        self.context
            .load_module(cudarc::nvrtc::Ptx::from_binary(image))
            .map_err(|error| {
                driver_err(
                    &format!("loading NVRTC CUBIN fallback module '{module_key}'"),
                    error,
                )
            })
    }

    pub fn require_nvrtc_half_headers(&self, op: &str) -> Result<()> {
        if nvrtc_include_paths().is_empty() {
            return Err(EpError::KernelFailed(format!(
                "cuda_ep {op}: f16/bf16 NVRTC kernels require cuda_fp16.h and cuda_bf16.h. \
                 Install the CUDA runtime headers (for pip CUDA 13: `pip install \
                 nvidia-cuda-runtime`; alternatively set CUDA_HOME/CUDA_PATH)."
            )));
        }
        Ok(())
    }

    /// Bind this runtime's context to the calling thread. Required before any
    /// driver call (`malloc`, `memcpy`, cuBLASLt) that targets the current
    /// context.
    pub fn bind(&self) -> Result<()> {
        self.context
            .bind_to_thread()
            .map_err(|e| driver_err("bind_to_thread", e))
    }

    /// Block until all submitted work on the EP's dedicated stream completes.
    pub fn synchronize(&self) -> Result<()> {
        self.stream
            .synchronize()
            .map_err(|e| driver_err("stream synchronize", e))
    }

    /// Whether the EP's compute stream is currently capturing into a CUDA graph.
    /// A stream synchronize is illegal during capture, so device-resident kernels
    /// use this to skip the trailing sync while a graph is being recorded.
    pub fn is_capturing(&self) -> Result<bool> {
        Ok(self.graph_capture_status()?
            != cudarc::driver::sys::CUstreamCaptureStatus::CU_STREAM_CAPTURE_STATUS_NONE)
    }

    /// Allocate `bytes` (>= 1) of device memory, returning the raw device
    /// pointer. Binds the context first.
    pub fn alloc_raw(&self, bytes: usize) -> Result<CUdeviceptr> {
        self.bind()?;
        // SAFETY: `malloc_sync` returns a fresh device allocation on the current
        // (bound) context; we own it and free it exactly once via `free_raw`.
        let ptr = unsafe { cudarc::driver::result::malloc_sync(bytes.max(1)) }
            .map_err(|e| driver_err("cuMemAlloc", e))?;
        self.allocations.fetch_add(1, Ordering::Relaxed);
        Ok(ptr)
    }

    /// Free a device pointer previously returned by [`CudaRuntime::alloc_raw`].
    ///
    /// # Safety
    /// `ptr` must have come from this runtime's `alloc_raw` and not been freed.
    pub unsafe fn free_raw(&self, ptr: CUdeviceptr) -> Result<()> {
        self.bind()?;
        // SAFETY: caller upholds the single-free contract.
        unsafe { cudarc::driver::result::free_sync(ptr) }
            .map_err(|e| driver_err("cuMemFree", e))?;
        self.frees.fetch_add(1, Ordering::Relaxed);
        Ok(())
    }

    /// Copy `bytes` host → device (H2D). `dst` must be large enough.
    ///
    /// # Safety
    /// `dst` is a live device allocation of at least `src.len()` bytes.
    pub unsafe fn htod(&self, src: &[u8], dst: CUdeviceptr) -> Result<()> {
        self.bind()?;
        // SAFETY: bound context; `dst` covers `src.len()` bytes per the contract.
        unsafe { cudarc::driver::result::memcpy_htod_sync(dst, src) }
            .map_err(|e| driver_err("cuMemcpyHtoD", e))
    }

    /// Copy `dst.len()` bytes device → host (D2H). `src` must be large enough.
    ///
    /// # Safety
    /// `src` is a live device allocation of at least `dst.len()` bytes.
    pub unsafe fn dtoh(&self, dst: &mut [u8], src: CUdeviceptr) -> Result<()> {
        self.bind()?;
        // Kernels enqueue work on the EP's dedicated non-default stream. Wait
        // before issuing the synchronous driver copy so the host never observes
        // bytes that were still being produced on that stream.
        self.synchronize()?;
        // SAFETY: bound context; `src` covers `dst.len()` bytes per the contract.
        unsafe { cudarc::driver::result::memcpy_dtoh_sync(dst, src) }
            .map_err(|e| driver_err("cuMemcpyDtoH", e))
    }

    /// Copy `bytes` device → device (D2D).
    ///
    /// # Safety
    /// Both pointers are live allocations of at least `bytes` bytes.
    pub unsafe fn dtod(&self, src: CUdeviceptr, dst: CUdeviceptr, bytes: usize) -> Result<()> {
        self.bind()?;
        // Kernels enqueue their writes on the EP's dedicated non-default stream,
        // but `cuMemcpyDtoD` issues on the legacy default stream. On a
        // non-blocking compute stream the two are not implicitly ordered, so the
        // copy can race a producer kernel that is still writing `src` (or a
        // consumer that already queued a read of `dst`). Drain the EP stream
        // first so the synchronous copy always sees fully-produced bytes. This
        // mirrors `dtoh`, which synchronizes for the same reason.
        self.synchronize()?;
        // SAFETY: bound context; both endpoints cover `bytes` per the contract.
        unsafe { cudarc::driver::result::memcpy_dtod_sync(dst, src, bytes) }
            .map_err(|e| driver_err("cuMemcpyDtoD", e))
    }

    /// Enqueue a device → device copy on the EP stream.
    ///
    /// # Safety
    /// Both pointers are live allocations of at least `bytes` bytes and remain
    /// live until the stream has completed the copy.
    pub unsafe fn dtod_async(
        &self,
        src: CUdeviceptr,
        dst: CUdeviceptr,
        bytes: usize,
    ) -> Result<()> {
        self.bind()?;
        // SAFETY: bound context; both endpoints cover `bytes` and the runtime
        // owns the stream on which the copy is ordered.
        unsafe {
            cudarc::driver::result::memcpy_dtod_async(dst, src, bytes, self.stream.cu_stream())
        }
        .map_err(|e| driver_err("cuMemcpyDtoDAsync", e))
    }
}

impl Drop for CudaRuntime {
    fn drop(&mut self) {
        if self.capture_error != 0 {
            // SAFETY: `capture_error` was allocated by this runtime's `alloc_raw`
            // in `with_capture_error_word` and is freed exactly once here.
            let _ = unsafe { self.free_raw(self.capture_error) };
            self.capture_error = 0;
        }
    }
}

/// Reinterpret an EP [`onnx_runtime_ep_api::DeviceBuffer`] raw pointer (or a
/// [`onnx_runtime_ep_api::TensorView`] data pointer) as a CUDA device pointer.
/// CUDA device pointers are integer addresses; the EP stores them in the opaque
/// pointer slot, so this is a value reinterpretation, never a host deref.
#[inline]
pub fn cuptr(raw: *const c_void) -> CUdeviceptr {
    raw as usize as CUdeviceptr
}

/// Inverse of [`cuptr`]: pack a CUDA device pointer into the opaque pointer slot
/// used by [`onnx_runtime_ep_api::DeviceBuffer`].
#[inline]
pub fn raw_ptr(dptr: CUdeviceptr) -> *mut c_void {
    dptr as usize as *mut c_void
}

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

    #[test]
    fn derives_ptx_arch_from_compute_capability() {
        for (major, minor, expected) in [
            (6, 0, "compute_60"),
            (7, 5, "compute_75"),
            (8, 0, "compute_80"),
            (8, 6, "compute_86"),
            (8, 9, "compute_89"),
            (9, 0, "compute_90"),
            (10, 0, "compute_100"),
            (12, 0, "compute_120"),
        ] {
            assert_eq!(ptx_arch_for(major, minor), expected);
        }
    }

    #[test]
    fn derives_cubin_arch_from_compute_capability() {
        for (major, minor, expected) in [
            (6, 0, "sm_60"),
            (7, 5, "sm_75"),
            (8, 0, "sm_80"),
            (8, 6, "sm_86"),
            (8, 9, "sm_89"),
            (9, 0, "sm_90"),
            (10, 0, "sm_100"),
            (12, 0, "sm_120"),
        ] {
            assert_eq!(cubin_arch_for(major, minor), expected);
        }
    }

    #[test]
    fn capability_limits_use_conservative_fallbacks() {
        let capabilities =
            CudaDeviceCapabilities::from_reported_limits((7, 0), None, None, None, None);
        assert_eq!(capabilities.compute_capability(), (7, 0));
        assert_eq!(capabilities.max_threads_per_block, 256);
        assert_eq!(
            capabilities.max_shared_memory_per_block,
            SAFE_SHARED_MEMORY_PER_BLOCK_FALLBACK
        );
        assert_eq!(
            capabilities.max_shared_memory_per_block_optin(),
            SAFE_SHARED_MEMORY_PER_BLOCK_FALLBACK
        );
        assert_eq!(capabilities.multiprocessor_count(), 1);
    }

    #[test]
    fn capability_limits_never_reduce_optin_below_default() {
        let capabilities = CudaDeviceCapabilities::from_reported_limits(
            (12, 0),
            Some(1024),
            Some(64 * 1024),
            Some(48 * 1024),
            Some(200),
        );
        assert_eq!(capabilities.max_shared_memory_per_block_optin(), 64 * 1024);
        assert_eq!(capabilities.multiprocessor_count(), 200);
    }

    #[test]
    fn reduction_launch_is_clamped_to_device_limits() {
        assert_eq!(
            reduction_launch_params(256, 1024, 4, 227 * 1024),
            Some((256, 1024))
        );
        assert_eq!(
            reduction_launch_params(256, 128, 4, 227 * 1024),
            Some((128, 512))
        );
        assert_eq!(reduction_launch_params(256, 1024, 4, 768), Some((128, 512)));
        assert_eq!(reduction_launch_params(256, 1024, 8, 0), None);
    }

    #[test]
    fn dynamic_shared_memory_optin_respects_device_budgets() {
        let default_budget = 48 * 1024;
        // Fits the 48 KB non-opt-in budget: launch as-is, no attribute change.
        assert_eq!(
            dynamic_shared_memory_optin(32 * 1024, default_budget, 227 * 1024),
            Ok(None)
        );
        // Boundary: exactly the default budget still needs no opt-in.
        assert_eq!(
            dynamic_shared_memory_optin(default_budget, default_budget, 100 * 1024),
            Ok(None)
        );
        // Over 48 KB but within a consumer (sm_86/sm_89) ~100 KB opt-in ceiling:
        // must opt the function into the exact request.
        assert_eq!(
            dynamic_shared_memory_optin(64 * 1024, default_budget, 100 * 1024),
            Ok(Some(64 * 1024))
        );
        // Sized for an H200 (227 KB) but launched on a 100 KB consumer card:
        // reject loudly rather than crash at launch.
        assert_eq!(
            dynamic_shared_memory_optin(160 * 1024, default_budget, 100 * 1024),
            Err(())
        );
    }

    #[test]
    fn nvrtc_include_paths_only_returns_cuda_header_dirs() {
        assert!(
            nvrtc_include_paths()
                .iter()
                .all(|path| Path::new(path).join("cuda_fp16.h").is_file())
        );
    }

    fn maybe_runtime() -> Option<Arc<CudaRuntime>> {
        std::panic::catch_unwind(|| CudaRuntime::new(0).ok().map(Arc::new))
            .ok()
            .flatten()
    }

    // Regression guard for the DeepSeek-V2-Lite garbage-decode race: kernels run
    // on the EP's non-default stream, but `cuMemcpyDtoD` issues on the legacy
    // default stream, which is NOT implicitly ordered against a non-blocking
    // compute stream. `dtod` must therefore drain the EP stream before copying,
    // so it never reads bytes a producer kernel is still writing. Without the
    // synchronize this test observes the pre-launch poison values; with it the
    // copy always sees the fully-produced sentinel.
    #[test]
    fn dtod_waits_for_pending_stream_writes() {
        const MODULE: &str = "runtime_dtod_race_test";
        // Each thread spins on the GPU clock (~a few ms) before storing its
        // sentinel, guaranteeing the store is still in flight when the racing
        // default-stream copy would otherwise run.
        const SOURCE: &str = r#"
extern "C" __global__ void slow_fill(float* out, unsigned long long n, long long spin) {
    unsigned long long i = (unsigned long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    long long start = clock64();
    while (clock64() - start < spin) { }
    out[i] = 1.0f + (float)(i % 7);
}
"#;
        let Some(runtime) = maybe_runtime() else {
            eprintln!("skipping dtod race regression: CUDA runtime unavailable");
            return;
        };
        let function = runtime.nvrtc_function(MODULE, SOURCE, "slow_fill").unwrap();
        let n = 4096usize;
        let bytes = n * std::mem::size_of::<f32>();
        let src = runtime.alloc_raw(bytes).unwrap();
        let dst = runtime.alloc_raw(bytes).unwrap();

        // Poison the source so a premature (racing) copy is detectable.
        let poison = vec![-999.0f32; n];
        let poison_bytes =
            unsafe { std::slice::from_raw_parts(poison.as_ptr().cast::<u8>(), bytes) };
        // Run several iterations; a race is probabilistic per launch, but the
        // fix must make every iteration correct.
        for _ in 0..8 {
            unsafe { runtime.htod(poison_bytes, src) }.unwrap();
            runtime.synchronize().unwrap();

            // Enqueue the slow producer on the EP stream, then immediately copy
            // WITHOUT an explicit synchronize — `dtod` must order this itself.
            let spin: i64 = 8_000_000;
            let mut builder = runtime.stream().launch_builder(&function);
            let n_u64 = n as u64;
            builder.arg(&src).arg(&n_u64).arg(&spin);
            unsafe {
                builder
                    .launch(LaunchConfig::for_num_elems(n as u32))
                    .unwrap();
            }
            unsafe { runtime.dtod(src, dst, bytes) }.unwrap();

            let mut out = vec![0.0f32; n];
            let out_bytes =
                unsafe { std::slice::from_raw_parts_mut(out.as_mut_ptr().cast::<u8>(), bytes) };
            unsafe { runtime.dtoh(out_bytes, dst) }.unwrap();

            for (i, value) in out.iter().enumerate() {
                let expected = 1.0f32 + (i % 7) as f32;
                assert_eq!(
                    *value, expected,
                    "dtod observed unsynchronized/poison data at index {i}: \
                     got {value}, expected {expected} (stream-ordering race)"
                );
            }
        }

        unsafe { runtime.free_raw(src) }.unwrap();
        unsafe { runtime.free_raw(dst) }.unwrap();
    }

    // Companion to the sync-`dtod` guard: a stream-ordered `dtod_async` issued on
    // the EP compute stream (as `copy_reshape` uses for Reshape/Squeeze) must be
    // implicitly ordered after a producer kernel on the same stream, WITHOUT any
    // host synchronize. A later `dtoh` (which drains the stream) must then read
    // the fully-produced sentinel, never the pre-launch poison.
    #[test]
    fn dtod_async_is_ordered_after_same_stream_producer() {
        const MODULE: &str = "runtime_dtod_async_order_test";
        const SOURCE: &str = r#"
extern "C" __global__ void slow_fill(float* out, unsigned long long n, long long spin) {
    unsigned long long i = (unsigned long long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    long long start = clock64();
    while (clock64() - start < spin) { }
    out[i] = 1.0f + (float)(i % 7);
}
"#;
        let Some(runtime) = maybe_runtime() else {
            eprintln!("skipping dtod_async ordering test: CUDA runtime unavailable");
            return;
        };
        let function = runtime.nvrtc_function(MODULE, SOURCE, "slow_fill").unwrap();
        let n = 4096usize;
        let bytes = n * std::mem::size_of::<f32>();
        let src = runtime.alloc_raw(bytes).unwrap();
        let dst = runtime.alloc_raw(bytes).unwrap();
        let poison = vec![-999.0f32; n];
        let poison_bytes =
            unsafe { std::slice::from_raw_parts(poison.as_ptr().cast::<u8>(), bytes) };

        for _ in 0..8 {
            unsafe { runtime.htod(poison_bytes, src) }.unwrap();
            runtime.synchronize().unwrap();

            let spin: i64 = 8_000_000;
            let mut builder = runtime.stream().launch_builder(&function);
            let n_u64 = n as u64;
            builder.arg(&src).arg(&n_u64).arg(&spin);
            unsafe {
                builder
                    .launch(LaunchConfig::for_num_elems(n as u32))
                    .unwrap();
            }
            // Stream-ordered copy: no explicit synchronize, ordering is by stream.
            unsafe { runtime.dtod_async(src, dst, bytes) }.unwrap();

            let mut out = vec![0.0f32; n];
            let out_bytes =
                unsafe { std::slice::from_raw_parts_mut(out.as_mut_ptr().cast::<u8>(), bytes) };
            unsafe { runtime.dtoh(out_bytes, dst) }.unwrap();

            for (i, value) in out.iter().enumerate() {
                let expected = 1.0f32 + (i % 7) as f32;
                assert_eq!(
                    *value, expected,
                    "dtod_async observed poison at index {i}: got {value}, expected {expected} \
                     (same-stream ordering violated)"
                );
            }
        }

        unsafe { runtime.free_raw(src) }.unwrap();
        unsafe { runtime.free_raw(dst) }.unwrap();
    }
}