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}