foxhunt/crates/ml-ensemble/build.rs

use std::path::{Path, PathBuf};
use std::process::Command;

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

    // Only compile CUDA kernels when the cuda feature is enabled
    if std::env::var("CARGO_FEATURE_CUDA").is_err() {
        return;
    }

    let out_dir = PathBuf::from(std::env::var("OUT_DIR").unwrap());
    let kernel_dir = Path::new("src");

    // Detect GPU architecture from env or default to sm_80
    let cuda_compute_cap = std::env::var("CUDA_COMPUTE_CAP").unwrap_or_else(|_| "80".to_string());
    let arch = format!("sm_{cuda_compute_cap}");

    // Check if nvcc is available
    let nvcc = match find_nvcc() {
        Some(p) => p,
        None => {
            eprintln!("  warning: nvcc not found, skipping ml-ensemble CUDA kernel precompilation");
            eprintln!("           Install CUDA toolkit or set CUDA_HOME for GPU builds");
            return;
        }
    };

    // Standalone kernels (plain f32, no common header needed)
    let kernels = ["ensemble_reduce_kernel.cu"];

    let mut failed: Vec<&str> = Vec::new();
    for kernel_name in &kernels {
        if !try_compile_kernel(&nvcc, kernel_dir, kernel_name, &arch, &out_dir) {
            failed.push(kernel_name);
        }
    }

    let passed = kernels.len() - failed.len();
    eprintln!(
        "  ml-ensemble: Precompiled {passed}/{} CUDA kernels ({arch})",
        kernels.len()
    );
    if !failed.is_empty() {
        eprintln!("  FAILED: {}", failed.join(", "));
        panic!(
            "nvcc failed to compile {} ml-ensemble kernel(s): {}",
            failed.len(),
            failed.join(", ")
        );
    }
}

/// Compile a single .cu kernel file to a .cubin via nvcc.
fn try_compile_kernel(
    nvcc: &Path,
    kernel_dir: &Path,
    kernel_name: &str,
    arch: &str,
    out_dir: &Path,
) -> bool {
    let kernel_path = kernel_dir.join(kernel_name);
    let cubin_name = kernel_name.replace(".cu", ".cubin");
    let cubin_path = out_dir.join(&cubin_name);

    println!("cargo:rerun-if-changed={}", kernel_path.display());

    let status = Command::new(nvcc)
        .args([
            "-cubin",
            &format!("-arch={arch}"),
            "-O3",
            "--ftz=true",
            "--fmad=true",
            "--prec-div=true",
            "--prec-sqrt=true",
            "-o",
            cubin_path.to_str().unwrap(),
            kernel_path.to_str().unwrap(),
        ])
        .status();

    match status {
        Ok(s) if s.success() => {
            eprintln!("  Compiled {kernel_name} -> {cubin_name} ({arch})");
            true
        }
        Ok(s) => {
            eprintln!(
                "  FAILED: {kernel_name} (exit={})",
                s.code().unwrap_or(-1)
            );
            false
        }
        Err(e) => {
            eprintln!("  FAILED: {kernel_name} (nvcc error: {e})");
            false
        }
    }
}

/// Find nvcc: prefer $CUDA_HOME/bin/nvcc, then check PATH
fn find_nvcc() -> Option<PathBuf> {
    if let Ok(home) = std::env::var("CUDA_HOME") {
        let nvcc = PathBuf::from(home).join("bin/nvcc");
        if nvcc.exists() {
            return Some(nvcc);
        }
    }

    for path in &["/usr/local/cuda/bin/nvcc", "/usr/bin/nvcc"] {
        let p = PathBuf::from(path);
        if p.exists() {
            return Some(p);
        }
    }

    match Command::new("nvcc").arg("--version").output() {
        Ok(output) if output.status.success() => Some(PathBuf::from("nvcc")),
        _ => None,
    }
}