rustytorch/build.rs

//! Build script for RustyTorch++
//! Configures GPU compilation and links with rustg

use std::env;
use std::path::PathBuf;
use std::process::Command;

fn main() {
    println!("cargo:rerun-if-changed=build.rs");

    // Check for CUDA installation
    check_cuda();

    // Set up rustg paths
    setup_rustg_paths();

    // Configure GPU architecture
    // Default to sm_86 (Ampere - RTX 30xx series) which is widely compatible
    // Use GPU_ARCH env var to override for specific hardware:
    //   sm_75 = Turing (RTX 20xx)
    //   sm_80 = Ampere (A100)
    //   sm_86 = Ampere (RTX 30xx, RTX A-series)
    //   sm_89 = Ada Lovelace (RTX 40xx)
    //   sm_90 = Hopper (H100)
    //   sm_100/sm_120 = Blackwell (RTX 50xx) - requires CUDA 13+
    let gpu_arch = env::var("GPU_ARCH").unwrap_or_else(|_| detect_gpu_arch());
    println!("cargo:rustc-env=GPU_ARCH={}", gpu_arch);

    // Set kernel cache directory
    let kernel_cache = PathBuf::from("target/kernel_cache");
    std::fs::create_dir_all(&kernel_cache).expect("Failed to create kernel cache directory");
    println!("cargo:rustc-env=KERNEL_CACHE_DIR={}", kernel_cache.display());
}

fn check_cuda() {
    // Source shell config and check nvcc
    let output = Command::new("sh")
    .arg("-c")
    .arg("source ~/.zshrc 2>/dev/null || source ~/.bashrc 2>/dev/null; which nvcc")
    .output();

    if let Ok(output) = output {
        if output.status.success() {
            let nvcc_path = String::from_utf8_lossy(&output.stdout);
            println!("cargo:warning=Found NVCC at: {}", nvcc_path.trim());

            // Check CUDA version
            let version_output = Command::new("nvcc")
            .arg("--version")
            .output();

            if let Ok(version) = version_output {
                let version_str = String::from_utf8_lossy(&version.stdout);
                if version_str.contains("release 13.0") {
                    println!("cargo:warning=CUDA 13.0 detected - RTX 5090 support enabled");
                }
            }
        }
        } else {
            println!("cargo:warning=NVCC not found - GPU compilation may be limited");
        }
    }

    fn setup_rustg_paths() {
        // Set up paths to rustg components
        let rustg_path = PathBuf::from("../rust/rustg");

        if rustg_path.exists() {
            println!("cargo:rustc-env=RUSTG_PATH={}", rustg_path.display());

            // Add cargo-g to PATH if available
            let cargo_g_path = rustg_path.join("cargo-g");
            if cargo_g_path.exists() {
                println!("cargo:rustc-env=CARGO_G_PATH={}", cargo_g_path.display());
            }

            // Add gpu-dev-tools to PATH if available
            let gpu_tools_path = rustg_path.join("gpu-dev-tools");
            if gpu_tools_path.exists() {
                println!("cargo:rustc-env=GPU_DEV_TOOLS_PATH={}", gpu_tools_path.display());
            }
            } else {
                println!("cargo:warning=rustg not found at expected path - using fallback compilation");
            }
        }

/// Auto-detect GPU architecture from nvidia-smi
fn detect_gpu_arch() -> String {
    // Try to detect GPU compute capability using nvidia-smi
    let output = Command::new("nvidia-smi")
        .args(["--query-gpu=compute_cap", "--format=csv,noheader"])
        .output();

    if let Ok(output) = output {
        if output.status.success() {
            let compute_cap = String::from_utf8_lossy(&output.stdout);
            let compute_cap = compute_cap.trim();

            // Convert compute capability (e.g., "8.6") to sm arch (e.g., "sm_86")
            if let Some(arch) = compute_cap_to_sm(compute_cap) {
                println!("cargo:warning=Auto-detected GPU architecture: {} (compute {})", arch, compute_cap);
                return arch;
            }
        }
    }

    // Fallback to sm_86 (RTX 30xx series) - widely compatible
    println!("cargo:warning=Could not detect GPU, defaulting to sm_86 (Ampere)");
    "sm_86".to_owned()
}

/// Convert compute capability string to sm_XX format
fn compute_cap_to_sm(compute_cap: &str) -> Option<String> {
    // Parse "X.Y" format
    let parts: Vec<&str> = compute_cap.split('.').collect();
    if parts.len() == 2 {
        if let (Ok(major), Ok(minor)) = (parts[0].parse::<u32>(), parts[1].parse::<u32>()) {
            return Some(format!("sm_{}{}", major, minor));
        }
    }
    None
}