openvm_cuda_builder/
lib.rs

1use std::{env, path::Path, process::Command};
2
3/// CUDA builder configuration
4#[derive(Debug, Clone)]
5pub struct CudaBuilder {
6    include_paths: Vec<String>,
7    source_files: Vec<String>,
8    watch_paths: Vec<String>,
9    watch_globs: Vec<String>,
10    library_name: String,
11    cuda_arch: Vec<String>,
12    cuda_opt_level: Option<String>,
13    lineinfo: bool,
14    custom_flags: Vec<String>,
15    link_libraries: Vec<String>,
16    link_search_paths: Vec<String>,
17}
18
19impl Default for CudaBuilder {
20    fn default() -> Self {
21        let mut link_search_paths = Vec::new();
22        if let Ok(cuda_lib_dir) = env::var("CUDA_LIB_DIR") {
23            link_search_paths.push(cuda_lib_dir);
24        } else {
25            link_search_paths.push("/usr/local/cuda/lib64".to_string());
26        }
27
28        Self {
29            include_paths: Vec::new(),
30            source_files: Vec::new(),
31            watch_paths: vec!["build.rs".to_string()],
32            watch_globs: Vec::new(),
33            library_name: String::new(),
34            cuda_arch: Vec::new(),
35            cuda_opt_level: None,
36            lineinfo: false,
37            custom_flags: vec![
38                "--std=c++17".to_string(),
39                "--expt-relaxed-constexpr".to_string(),
40                "-Xfatbin=-compress-all".to_string(),
41            ],
42            link_libraries: vec!["cudart".to_string(), "cuda".to_string()],
43            link_search_paths,
44        }
45    }
46}
47
48impl CudaBuilder {
49    /// Create a new CudaBuilder
50    pub fn new() -> Self {
51        Self::default()
52    }
53
54    /// Set the library name (useful when cloning from a template)
55    pub fn library_name(mut self, name: &str) -> Self {
56        self.library_name = name.to_string();
57        self
58    }
59
60    /// Add include path
61    pub fn include<P: AsRef<Path>>(mut self, path: P) -> Self {
62        let path_str = path.as_ref().to_string_lossy().to_string();
63        self.include_paths.push(path_str.clone());
64        self.watch_paths.push(path_str);
65        self
66    }
67
68    /// Add include path from another crate's exported include
69    pub fn include_from_dep(mut self, dep_env_var: &str) -> Self {
70        if let Ok(path) = env::var(dep_env_var) {
71            self.include_paths.push(path);
72        }
73        self
74    }
75
76    /// Add source file
77    pub fn file<P: AsRef<Path>>(mut self, path: P) -> Self {
78        let path_str = path.as_ref().to_string_lossy().to_string();
79        self.source_files.push(path_str.clone());
80        self.watch_paths.push(path_str);
81        self
82    }
83
84    /// Add multiple source files
85    pub fn files<P: AsRef<Path>, I: IntoIterator<Item = P>>(mut self, paths: I) -> Self {
86        for path in paths {
87            let path_str = path.as_ref().to_string_lossy().to_string();
88            self.source_files.push(path_str.clone());
89            self.watch_paths.push(path_str);
90        }
91        self
92    }
93
94    /// Add multiple source files matching a glob pattern
95    pub fn files_from_glob(mut self, pattern: &str) -> Self {
96        self.watch_globs.push(pattern.to_string());
97        for path in glob::glob(pattern).expect("Invalid glob pattern").flatten() {
98            if path.is_file() && path.extension().is_some_and(|ext| ext == "cu") {
99                self.source_files.push(path.to_string_lossy().to_string());
100            }
101        }
102        self
103    }
104
105    /// Watch a specific path for changes
106    pub fn watch<P: AsRef<Path>>(mut self, path: P) -> Self {
107        self.watch_paths
108            .push(path.as_ref().to_string_lossy().to_string());
109        self
110    }
111
112    /// Watch paths matching a glob pattern
113    pub fn watch_glob(mut self, pattern: &str) -> Self {
114        self.watch_globs.push(pattern.to_string());
115        self
116    }
117
118    /// Set CUDA architecture (e.g., "75", "80")
119    pub fn cuda_arch(mut self, arch: &str) -> Self {
120        self.cuda_arch = vec![arch.to_string()];
121        self
122    }
123
124    /// Set multiple CUDA architectures  
125    pub fn cuda_archs(mut self, archs: Vec<&str>) -> Self {
126        self.cuda_arch = archs.iter().map(|s| s.to_string()).collect();
127        self
128    }
129
130    /// Set CUDA optimization level (0-3)
131    pub fn cuda_opt_level(mut self, level: u8) -> Self {
132        self.cuda_opt_level = Some(level.to_string());
133        self
134    }
135
136    /// Enable line info for profiling (NCU). Keeps other optimizations intact.
137    pub fn lineinfo(mut self, enabled: bool) -> Self {
138        self.lineinfo = enabled;
139        self
140    }
141
142    /// Add custom compiler flag
143    pub fn flag(mut self, flag: &str) -> Self {
144        self.custom_flags.push(flag.to_string());
145        self
146    }
147
148    /// Add library to link
149    pub fn link_lib(mut self, lib: &str) -> Self {
150        self.link_libraries.push(lib.to_string());
151        self
152    }
153
154    /// Add library search path
155    pub fn link_search<P: AsRef<Path>>(mut self, path: P) -> Self {
156        self.link_search_paths
157            .push(path.as_ref().to_string_lossy().to_string());
158        self
159    }
160
161    /// Build the CUDA library
162    pub fn build(self) {
163        // Validation
164        self.validate();
165
166        // Set up rerun conditions
167        self.setup_rerun_conditions();
168
169        // Get or detect CUDA architecture
170        let cuda_archs = self.get_cuda_arch();
171
172        // Create cc::Build
173        let mut builder = cc::Build::new();
174        builder.cuda(true);
175
176        // Handle CUDA_DEBUG=1
177        self.handle_debug_shortcuts(&mut builder);
178
179        // Get optimization level
180        let cuda_opt_level = self.get_cuda_opt_level();
181
182        // Add include paths
183        for include in &self.include_paths {
184            builder.include(include);
185        }
186
187        // Add CUDA_PATH include if available
188        if let Ok(cuda_path) = env::var("CUDA_PATH") {
189            builder.include(format!("{}/include", cuda_path));
190        }
191
192        // Add custom flags
193        for flag in &self.custom_flags {
194            builder.flag(flag);
195        }
196
197        // Add SASS code for each architecture
198        for arch in &cuda_archs {
199            builder
200                .flag("-gencode")
201                .flag(format!("arch=compute_{},code=sm_{}", arch, arch));
202        }
203
204        // Add PTX for the highest architecture (forward compatibility)
205        // This allows the code to run on future GPUs
206        if let Some(max_arch) = cuda_archs.iter().max() {
207            builder.flag("-gencode").flag(format!(
208                "arch=compute_{},code=compute_{}",
209                max_arch, max_arch
210            ));
211        }
212
213        // Add parallel jobs flag
214        builder.flag(nvcc_parallel_jobs());
215
216        // Set optimization and debug flags
217        if cuda_opt_level == "0" {
218            builder.debug(true).flag("-O0");
219        } else {
220            builder
221                .debug(false)
222                .flag(format!("--ptxas-options=-O{}", cuda_opt_level));
223        }
224
225        // Add line info for profiling (NCU) if enabled
226        if self.get_lineinfo() {
227            builder.flag("-lineinfo");
228        }
229
230        // Add source files
231        for file in &self.source_files {
232            builder.file(file);
233        }
234
235        // Compile
236        builder.compile(&self.library_name);
237    }
238
239    /// Validate the builder configuration
240    fn validate(&self) {
241        if self.library_name.is_empty() {
242            panic!(
243                "Library name must be set using .library_name(\"name\") before calling .build()"
244            );
245        }
246
247        if self.source_files.is_empty() {
248            panic!("At least one source file must be added using .file() or .files() before calling .build()");
249        }
250
251        // Validate that source files exist (optional, but helpful)
252        for file in &self.source_files {
253            if !Path::new(file).exists() {
254                eprintln!("cargo:warning=CUDA source file does not exist: {}", file);
255            }
256        }
257
258        // Validate include paths exist (optional warning)
259        for include in &self.include_paths {
260            if !Path::new(include).exists() {
261                eprintln!("cargo:warning=Include path does not exist: {}", include);
262            }
263        }
264    }
265
266    pub fn emit_link_directives(&self) {
267        for path in &self.link_search_paths {
268            println!("cargo:rustc-link-search=native={}", path);
269        }
270        for lib in &self.link_libraries {
271            println!("cargo:rustc-link-lib={}", lib);
272        }
273    }
274
275    fn setup_rerun_conditions(&self) {
276        // Standard rerun conditions
277        println!("cargo:rerun-if-env-changed=CUDA_ARCH");
278        println!("cargo:rerun-if-env-changed=CUDA_OPT_LEVEL");
279        println!("cargo:rerun-if-env-changed=CUDA_DEBUG");
280        println!("cargo:rerun-if-env-changed=CUDA_LINEINFO");
281        println!("cargo:rerun-if-env-changed=NVCC_THREADS");
282
283        // Watch specific paths
284        for path in &self.watch_paths {
285            println!("cargo:rerun-if-changed={}", path);
286        }
287
288        // Watch glob patterns
289        for pattern in &self.watch_globs {
290            watch_glob(pattern);
291        }
292    }
293
294    fn get_cuda_arch(&self) -> Vec<String> {
295        if !self.cuda_arch.is_empty() {
296            return self.cuda_arch.clone();
297        }
298
299        // Check environment variable
300        if let Ok(env_archs) = env::var("CUDA_ARCH") {
301            return env_archs
302                .split(',')
303                .map(|s| s.trim().to_string())
304                .filter(|s| !s.is_empty())
305                .collect();
306        }
307
308        // Auto-detect current GPU
309        vec![detect_cuda_arch()]
310    }
311
312    fn get_cuda_opt_level(&self) -> String {
313        if let Some(level) = &self.cuda_opt_level {
314            return level.clone();
315        }
316
317        env::var("CUDA_OPT_LEVEL").unwrap_or_else(|_| "3".to_string())
318    }
319
320    fn get_lineinfo(&self) -> bool {
321        self.lineinfo || env::var("CUDA_LINEINFO").map(|v| v == "1").unwrap_or(false)
322    }
323
324    fn handle_debug_shortcuts(&self, builder: &mut cc::Build) {
325        if env::var("CUDA_DEBUG").map(|v| v == "1").unwrap_or(false) {
326            env::set_var("CUDA_OPT_LEVEL", "0");
327            env::set_var("CUDA_LAUNCH_BLOCKING", "1");
328            env::set_var("RUST_BACKTRACE", "full");
329            env::set_var("CUDA_ENABLE_COREDUMP_ON_EXCEPTION", "1");
330            env::set_var("CUDA_DEVICE_WAITS_ON_EXCEPTION", "1");
331
332            println!("cargo:warning=CUDA_DEBUG=1 → Enabling comprehensive debugging:");
333            println!("cargo:warning=  → CUDA_OPT_LEVEL=0 (no optimization)");
334            println!("cargo:warning=  → CUDA_LAUNCH_BLOCKING=1 (synchronous kernels)");
335            println!("cargo:warning=  → Line info and device debug symbols enabled");
336            println!("cargo:warning=  → CUDA_DEBUG macro defined for preprocessor");
337
338            builder.flag("-G"); // Device debug symbols
339            builder.flag("-Xcompiler=-fno-omit-frame-pointer"); // Better stack traces
340            builder.flag("-Xptxas=-v"); // Verbose PTX compilation
341            builder.define("CUDA_DEBUG", "1"); // Define CUDA_DEBUG macro
342        }
343    }
344}
345
346/// Check if CUDA is available on the system
347pub fn cuda_available() -> bool {
348    Command::new("nvcc").arg("--version").output().is_ok()
349}
350
351/// Detect CUDA architecture using nvidia-smi
352pub fn detect_cuda_arch() -> String {
353    let output = Command::new("nvidia-smi")
354        .args(["--query-gpu=compute_cap", "--format=csv,noheader"])
355        .output()
356        .expect("Failed to run nvidia-smi - make sure NVIDIA drivers are installed");
357
358    let full_output =
359        String::from_utf8(output.stdout).expect("nvidia-smi output is not valid UTF-8");
360
361    let arch = full_output
362        .lines()
363        .next()
364        .expect("nvidia-smi failed to return compute capability")
365        .trim()
366        .replace('.', ""); // Convert "7.5" to "75"
367
368    // Set both cargo env and process env
369    println!("cargo:rustc-env=CUDA_ARCH={}", arch);
370    env::set_var("CUDA_ARCH", &arch);
371    arch
372}
373
374/// Calculate optimal number of parallel NVCC jobs
375pub fn nvcc_parallel_jobs() -> String {
376    let threads = std::thread::available_parallelism()
377        .map(|n| n.get())
378        .unwrap_or(1);
379
380    let threads = env::var("NVCC_THREADS")
381        .ok()
382        .and_then(|v| v.parse::<usize>().ok())
383        .unwrap_or(threads);
384
385    format!("-t{}", threads)
386}
387
388/// Watch files matching a glob pattern
389fn watch_glob(pattern: &str) {
390    for path in glob::glob(pattern).expect("Invalid glob pattern").flatten() {
391        if path.is_file() {
392            println!("cargo:rerun-if-changed={}", path.display());
393        }
394    }
395}
openvm_cuda_builder/lib.rs

openvm_cuda_builder/
lib.rs
