ptx_details/
ptx_details.rs

1use std::{
2    env,
3    path::{Path, PathBuf},
4    process::Command,
5};
6
7use clap::Parser;
8use openvm_scripts::{find_cuda_include_dirs, get_cuda_dep_common_include_dirs};
9use tempfile::Builder as TempfileBuilder;
10
11#[derive(Parser)]
12#[command(author, version, about = "Compile CUDA code with nvcc")]
13struct Args {
14    /// Write output to this path. If omitted, a temporary file is used and deleted.
15    #[arg(long)]
16    out: Option<PathBuf>,
17
18    /// Path to nvcc (default: nvcc from PATH)
19    #[arg(long, default_value = "nvcc")]
20    nvcc: String,
21
22    /// Path to CUDA installation (used for include paths)
23    #[arg(long, env = "CUDA_PATH", default_value = "/usr/local/cuda")]
24    cuda_path: String,
25
26    /// CUDA architecture (e.g. 89 -> sm_89 / compute_89)
27    #[arg(long, default_value = "89")]
28    cuda_arch: String,
29
30    /// nvcc thread count (-t)
31    #[arg(long, default_value_t = 16)]
32    threads: u32,
33
34    /// Input .cu file to compile
35    input: PathBuf,
36}
37
38fn add_if_exists(cmd: &mut Command, include_dir: &Path) {
39    if include_dir.exists() {
40        cmd.arg(format!("-I{}", include_dir.display()));
41    }
42}
43
44fn main() -> eyre::Result<()> {
45    let args = Args::parse();
46    let workspace_root = env::current_dir()?;
47
48    // Reuse the same include discovery approach as other scripts.
49    let mut include_dirs = find_cuda_include_dirs(&workspace_root);
50    include_dirs.extend(get_cuda_dep_common_include_dirs());
51
52    let (out_path, _tmp_file_guard) = match args.out {
53        Some(p) => (p, None),
54        None => {
55            let tmp = TempfileBuilder::new().suffix(".ptx").tempfile()?;
56            let path = tmp.path().to_path_buf();
57            (path, Some(tmp))
58        }
59    };
60
61    let cuda_include = PathBuf::from(&args.cuda_path).join("include");
62    let cccl_include = cuda_include.join("cccl");
63
64    let mut cmd = Command::new(&args.nvcc);
65    cmd.env("LC_ALL", "C");
66
67    // Keep flags matched to the original ptx.sh as closely as possible.
68    cmd.args([
69        "-ccbin=c++",
70        "-Xcompiler",
71        "-O3",
72        "-Xcompiler",
73        "-ffunction-sections",
74        "-Xcompiler",
75        "-fdata-sections",
76        "-Xcompiler",
77        "-fPIC",
78        "-m64",
79    ]);
80
81    for dir in &include_dirs {
82        cmd.arg(format!("-I{}", dir.display()));
83    }
84
85    add_if_exists(&mut cmd, &cuda_include);
86    add_if_exists(&mut cmd, &cccl_include);
87
88    cmd.args([
89        "-Xcompiler",
90        "-Wall",
91        "-Xcompiler",
92        "-Wextra",
93        "--std=c++17",
94        "--expt-relaxed-constexpr",
95        "-Xfatbin=-compress-all",
96    ]);
97
98    cmd.arg("-gencode")
99        .arg(format!("arch=compute_{a},code=sm_{a}", a = args.cuda_arch));
100    cmd.arg("-gencode").arg(format!(
101        "arch=compute_{a},code=compute_{a}",
102        a = args.cuda_arch
103    ));
104    cmd.arg(format!("-t{}", args.threads));
105
106    cmd.args(["--ptxas-options=-v", "-O3"]);
107    cmd.arg("-o").arg(&out_path);
108    cmd.args(["-c", "--device-c"]);
109
110    // Make input absolute (so callers can run from workspace root like other scripts).
111    let input = if args.input.is_absolute() {
112        args.input
113    } else {
114        workspace_root.join(args.input)
115    };
116    cmd.arg(input);
117
118    let status = cmd.status()?;
119    if !status.success() {
120        return Err(eyre::eyre!("nvcc failed with status: {status}"));
121    }
122
123    Ok(())
124}
ptx_details/ptx_details.rs

ptx_details/
ptx_details.rs
