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

    // Common header lives in the ml crate — use relative path from ml-core
    let common_header_path = Path::new("../../ml/src/cuda_pipeline/common_device_functions.cuh")
        .canonicalize()
        .unwrap_or_else(|_| {
            // Fallback: try from workspace root
            let workspace_root = std::env::var("CARGO_MANIFEST_DIR")
                .map(PathBuf::from)
                .unwrap_or_default();
            workspace_root
                .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-core 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 to precompile — all get common header prepended
    let kernels = [
        "activation_kernels.cu",
        "elementwise_kernels.cu",
        "linear_kernels.cu",
        "loss_kernels.cu",
        "reduction_kernels.cu",
        "dropout_kernels.cu",
        "layer_norm_kernels.cu",
        "optimizer_kernels.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, &common_src) {
            failed.push(kernel_name);
        }
    }

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

/// Compile a single .cu kernel file to a .cubin via nvcc.
/// Common header 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: &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: common header + kernel
    let full_source = format!("{common_header}\n{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",
            "--use_fast_math",
            "--ftz=true",
            "--fmad=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,
    }
}