foxhunt/crates/ml-supervised/build.rs

//! Precompile CUDA kernels owned by `ml-supervised`.
//!
//! Modelled on `crates/ml-dqn/build.rs` (the canonical reference). Emits
//! `.cubin` files into `$OUT_DIR`; Rust callers `include_bytes!` them via
//! `concat!(env!("OUT_DIR"), "/<name>.cubin")` and load through
//! `stream.context().load_cubin(...)`.
//!
//! Kernels compiled here:
//! - `mamba/dropout_kernel.cu` → `dropout_kernel.cubin` (forward-only
//!   Bernoulli inverted dropout used by `Mamba2SSM::forward_with_gradients`).

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 (matches
    // ml-dqn's gating exactly). Non-CUDA builds skip kernel compilation.
    if std::env::var("CARGO_FEATURE_CUDA").is_err() {
        return;
    }

    let out_dir = PathBuf::from(std::env::var("OUT_DIR").unwrap());
    let manifest_dir = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap());

    // Detect GPU architecture from env or default to sm_80 (L40S / A100).
    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-supervised CUDA kernel precompilation");
            eprintln!("           Install CUDA toolkit or set CUDA_HOME for GPU builds");
            return;
        }
    };

    // Kernels owned by ml-supervised. Each entry is (source path relative to
    // crate root, output cubin base name). The kernel is self-contained and
    // does NOT consume ml/src/cuda_pipeline/common_device_functions.cuh —
    // that header carries DQN-specific constants (DQN_NUM_ACTIONS,
    // STATE_DIM, ...) that are irrelevant here.
    let kernels: &[(&str, &str)] = &[
        ("src/mamba/dropout_kernel.cu", "dropout_kernel.cu"),
    ];

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

    for (src_path, kernel_name) in kernels {
        let full_src = manifest_dir.join(src_path);
        if !try_compile_kernel(&nvcc, &full_src, kernel_name, &arch, &out_dir) {
            failed.push(kernel_name);
        }
    }

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

/// Compile a single standalone `.cu` kernel to a `.cubin` via nvcc.
///
/// No header prepending — the kernel is expected to be self-contained.
fn try_compile_kernel(
    nvcc: &Path,
    kernel_path: &Path,
    kernel_name: &str,
    arch: &str,
    out_dir: &Path,
) -> bool {
    println!("cargo:rerun-if-changed={}", kernel_path.display());

    if !kernel_path.exists() {
        eprintln!("  FAILED: {kernel_name} (source not found at {})", kernel_path.display());
        return false;
    }

    let cubin_name = kernel_name.replace(".cu", ".cubin");
    let cubin_path = out_dir.join(&cubin_name);

    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 common install paths,
/// then 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,
    }
}