|
| 1 | +use std::env; |
| 2 | +use std::fs; |
| 3 | +use std::path::PathBuf; |
| 4 | +use std::process::Command; |
| 5 | + |
| 6 | +fn cuda_home() -> PathBuf { |
| 7 | + env::var_os("CUDA_HOME") |
| 8 | + .or_else(|| env::var_os("CUDA_PATH")) |
| 9 | + .map(PathBuf::from) |
| 10 | + .unwrap_or_else(|| PathBuf::from("/usr/local/cuda")) |
| 11 | +} |
| 12 | + |
| 13 | +fn nvcc_path() -> PathBuf { |
| 14 | + cuda_home().join("bin").join("nvcc") |
| 15 | +} |
| 16 | + |
| 17 | +/// Query `nvidia-smi` for the local GPU's compute capability (e.g. "12.0" |
| 18 | +/// for Blackwell). Returns a `compute_XX` target on success, falling back |
| 19 | +/// to `compute_89` (Ada) when no GPU is visible or the query fails. |
| 20 | +fn detect_arch() -> String { |
| 21 | + const FALLBACK: &str = "compute_89"; |
| 22 | + let output = match Command::new("nvidia-smi") |
| 23 | + .args(["--query-gpu=compute_cap", "--format=csv,noheader"]) |
| 24 | + .output() |
| 25 | + { |
| 26 | + Ok(o) if o.status.success() => o, |
| 27 | + _ => return FALLBACK.to_string(), |
| 28 | + }; |
| 29 | + let line = match std::str::from_utf8(&output.stdout) { |
| 30 | + Ok(s) => s, |
| 31 | + Err(_) => return FALLBACK.to_string(), |
| 32 | + }; |
| 33 | + // First line, first comma-separated value (covers multi-GPU hosts). |
| 34 | + let cap = match line.lines().next() { |
| 35 | + Some(l) => l.split(',').next().unwrap_or("").trim(), |
| 36 | + None => return FALLBACK.to_string(), |
| 37 | + }; |
| 38 | + let (major, minor) = match cap.split_once('.') { |
| 39 | + Some((m, n)) => (m.trim(), n.trim()), |
| 40 | + None => return FALLBACK.to_string(), |
| 41 | + }; |
| 42 | + if major.chars().all(|c| c.is_ascii_digit()) && minor.chars().all(|c| c.is_ascii_digit()) { |
| 43 | + format!("compute_{major}{minor}") |
| 44 | + } else { |
| 45 | + FALLBACK.to_string() |
| 46 | + } |
| 47 | +} |
| 48 | + |
| 49 | +fn compile_ptx(src: &str, out_name: &str, have_nvcc: bool) { |
| 50 | + let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); |
| 51 | + let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap()); |
| 52 | + let src_path = manifest_dir.join("kernels").join(src); |
| 53 | + let out_path = out_dir.join(out_name); |
| 54 | + |
| 55 | + println!("cargo:rerun-if-changed=kernels/{src}"); |
| 56 | + println!("cargo:rerun-if-env-changed=CUDA_HOME"); |
| 57 | + println!("cargo:rerun-if-env-changed=CUDA_PATH"); |
| 58 | + println!("cargo:rerun-if-env-changed=CUDARC_NVCC_ARCH"); |
| 59 | + |
| 60 | + // When nvcc is missing from PATH, emit an empty PTX stub so the crate |
| 61 | + // still compiles. include_str! in src/device.rs needs the file to exist |
| 62 | + // at build time. Any runtime kernel call panics in cudarc when loading |
| 63 | + // the empty module. We can't run GPU code without nvcc on the build |
| 64 | + // host anyway. |
| 65 | + if !have_nvcc { |
| 66 | + fs::write(&out_path, "").expect("failed to write empty PTX stub"); |
| 67 | + return; |
| 68 | + } |
| 69 | + |
| 70 | + // Emit PTX for a virtual architecture; the CUDA driver JIT-compiles it for the |
| 71 | + // actual GPU at load time. Override with CUDARC_NVCC_ARCH to pin a specific |
| 72 | + // compute capability. If unset, try `nvidia-smi` to match the host GPU |
| 73 | + // (avoids JIT failures like nvcc-13.0 PTX rejected on Blackwell drivers); |
| 74 | + // fall back to compute_89 (Ada) when detection fails. |
| 75 | + let arch = env::var("CUDARC_NVCC_ARCH").unwrap_or_else(|_| detect_arch()); |
| 76 | + |
| 77 | + let status = Command::new(nvcc_path()) |
| 78 | + .args(["--ptx", "-O3", "-std=c++17", "-arch", &arch, "-o"]) |
| 79 | + .arg(&out_path) |
| 80 | + .arg(&src_path) |
| 81 | + .status() |
| 82 | + .expect("failed to invoke nvcc — is CUDA installed and CUDA_HOME set?"); |
| 83 | + |
| 84 | + if !status.success() { |
| 85 | + panic!("nvcc failed compiling {}", src_path.display()); |
| 86 | + } |
| 87 | +} |
| 88 | + |
| 89 | +fn main() { |
| 90 | + // Headers aren't compiled, so emit rerun-if-changed to rebuild on |
| 91 | + // header edits. |
| 92 | + println!("cargo:rerun-if-changed=kernels/goldilocks.cuh"); |
| 93 | + println!("cargo:rerun-if-changed=kernels/ext3.cuh"); |
| 94 | + |
| 95 | + // Probe for nvcc once. Workspace consumers (clippy, fmt, CPU-only test |
| 96 | + // runners) build math-cuda incidentally without using its kernels. Stub |
| 97 | + // out PTX when nvcc is unavailable so those builds succeed. |
| 98 | + let have_nvcc = Command::new(nvcc_path()) |
| 99 | + .arg("--version") |
| 100 | + .output() |
| 101 | + .map(|o| o.status.success()) |
| 102 | + .unwrap_or(false); |
| 103 | + if !have_nvcc { |
| 104 | + println!( |
| 105 | + "cargo:warning=math-cuda: nvcc not found at {} — emitting empty PTX stubs. \ |
| 106 | + Runtime GPU calls will panic. Install CUDA and rebuild for a working backend.", |
| 107 | + nvcc_path().display() |
| 108 | + ); |
| 109 | + } |
| 110 | + |
| 111 | + compile_ptx("arith.cu", "arith.ptx", have_nvcc); |
| 112 | + compile_ptx("ntt.cu", "ntt.ptx", have_nvcc); |
| 113 | +} |
0 commit comments