foxhunt/crates/ml-dqn/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");

    // Common header lives in the ml crate — use relative path from ml-dqn
    let manifest_dir = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap());
    let common_header_path = manifest_dir
        .join("../ml/src/cuda_pipeline/common_device_functions.cuh")
        .canonicalize()
        .expect("Cannot find common_device_functions.cuh — is the ml crate present?");

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

    // 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-dqn CUDA kernel precompilation");
            eprintln!("           Install CUDA toolkit or set CUDA_HOME for GPU builds");
            return;
        }
    };

    // Read common header once
    let common_src = std::fs::read_to_string(&common_header_path)
        .unwrap_or_else(|e| panic!("Failed to read {}: {e}", common_header_path.display()));

    // All kernels get the common header (BF16 types + math wrappers)
    let bf16_kernels = [
        "rmsnorm_kernels.cu",
        "noisy_kernels.cu",
        "residual_kernels.cu",
        "cast_kernels.cu",
        "replay_buffer_kernels.cu",
        "seg_tree_kernel.cu",
    ];

    let standalone_kernels: [&str; 0] = [];

    let mut failed: Vec<&str> = Vec::new();

    for kernel_name in &bf16_kernels {
        if !try_compile_kernel(&nvcc, kernel_dir, kernel_name, &arch, &out_dir, Some(&common_src)) {
            failed.push(kernel_name);
        }
    }

    for kernel_name in &standalone_kernels {
        if !try_compile_kernel(&nvcc, kernel_dir, kernel_name, &arch, &out_dir, None) {
            failed.push(kernel_name);
        }
    }

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

/// Compile a single .cu kernel file to a .cubin via nvcc.
/// If common_header is Some, it is prepended to the kernel source.
fn try_compile_kernel(
    nvcc: &Path,
    kernel_dir: &Path,
    kernel_name: &str,
    arch: &str,
    out_dir: &Path,
    common_header: Option<&str>,
) -> 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 kernel_src = std::fs::read_to_string(&kernel_path)
        .unwrap_or_else(|e| panic!("Failed to read {}: {e}", kernel_path.display()));

    // Compose source: optional common header + kernel
    let full_source = match common_header {
        Some(hdr) => format!("{hdr}\n{kernel_src}"),
        None => kernel_src,
    };

    // Write composed source to temp file
    let tmp_src = out_dir.join(format!("_{kernel_name}"));
    std::fs::write(&tmp_src, &full_source).unwrap();

    // Compile with nvcc
    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(),
            tmp_src.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,
    }
}