Skip to main content

tenferro_gpu/cubecl/raw/
nvrtc.rs

1//! NVRTC compilation surface (issue #1597).
2//!
3//! [`NvrtcOptions`] is the narrow, typed option set forwarded to the NVRTC
4//! compiler through `cudarc`. [`compile_nvrtc`] compiles CUDA source on the
5//! host and returns the resulting PTX image without touching the GPU.
6
7/// Narrow, typed NVRTC compile options.
8///
9/// Only a conservative subset of NVRTC flags is exposed. The `arch` option is
10/// forwarded as `--gpu-architecture=...` (e.g. `"compute_80"`).
11#[derive(Clone, Debug, Default, PartialEq)]
12pub struct NvrtcOptions {
13    /// Pass `--gpu-architecture=<arch>` when set (e.g. `Some("compute_80")`).
14    pub arch: Option<String>,
15    /// Pass `--std=<std>` when set (e.g. `Some("c++17")`).
16    pub std: Option<String>,
17    /// Extra raw flags forwarded verbatim.
18    pub extra: Vec<String>,
19}
20
21impl NvrtcOptions {
22    /// Build the cudarc equivalent without leaking: every flag funnels into
23    /// the raw `options` vector (cudarc passes those to the compiler as-is).
24    fn to_cudarc(&self) -> cudarc::nvrtc::CompileOptions {
25        let mut options = cudarc::nvrtc::CompileOptions::default();
26        if let Some(arch) = &self.arch {
27            options.options.push(format!("--gpu-architecture={arch}"));
28        }
29        if let Some(std_flag) = &self.std {
30            options.options.push(format!("--std={std_flag}"));
31        }
32        options.options.extend(self.extra.iter().cloned());
33        options
34    }
35}
36
37/// Compile CUDA source to PTX on the host using NVRTC.
38///
39/// # Errors
40///
41/// Returns the compiler's typed error via [`crate::Error::BackendSource`]
42/// (with the NVRTC log) when compilation fails, or a validation error when the
43/// source contains a NUL byte.
44pub fn compile_nvrtc(src: &str, opts: &NvrtcOptions) -> crate::Result<cudarc::nvrtc::Ptx> {
45    if src.as_bytes().contains(&0) {
46        return Err(crate::Error::invalid_argument(
47            "nvrtc.compile",
48            "source",
49            "CUDA source cannot contain NUL bytes",
50        ));
51    }
52    cudarc::nvrtc::compile_ptx_with_opts(src, opts.to_cudarc())
53        .map_err(|err| crate::Error::backend_source("nvrtc.compile", err))
54}