foxhunt/crates/ml-alpha/build.rs

//! ml-alpha build.rs — precompile the Mamba2 temporal SSM scan kernel.
//!
//! Compiles `crates/ml/src/cuda_pipeline/mamba2_temporal_kernel.cu` (the only
//! kernel ml-alpha needs from the parent ml crate's library, for Phase 1d.1's
//! GPU-resident stateful encoder) into an arch-specific cubin in OUT_DIR.
//!
//! Discipline:
//! - **rerun-if-env-changed=CUDA_COMPUTE_CAP** is mandatory: cargo target PVCs
//!   are shared across L40S (sm_89) and H100 (sm_90) builds; without this,
//!   stale cubins surface as `CUDA_ERROR_NO_BINARY_FOR_GPU` at runtime.
//! - The kernel source has zero `#include`s, so no header prepending is
//!   needed (unlike ml-core's autograd kernels).
//! - On non-CUDA builds (`CARGO_FEATURE_CUDA` unset) we silently skip; the
//!   library still compiles, only the Mamba2Block::new() runtime path fails
//!   with a missing-cubin error.

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

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

    if std::env::var("CARGO_FEATURE_CUDA").is_err() {
        return;
    }

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

    // Kernel source lives in this crate's `cuda/` directory — purpose-built
    // simplified scan for the alpha supervised path (no ISV, no
    // temporal_weight, no NULL-pointer branches that the DQN trainer's
    // shared kernel carries).
    let kernel_src_path = Path::new("cuda/mamba2_alpha_kernel.cu")
        .canonicalize()
        .unwrap_or_else(|_| {
            let workspace = std::env::var("CARGO_MANIFEST_DIR")
                .map(PathBuf::from)
                .unwrap_or_default();
            workspace
                .join("cuda/mamba2_alpha_kernel.cu")
                .canonicalize()
                .expect("Cannot locate cuda/mamba2_alpha_kernel.cu")
        });
    println!("cargo:rerun-if-changed={}", kernel_src_path.display());

    println!("cargo:rerun-if-env-changed=CUDA_COMPUTE_CAP");
    let cuda_compute_cap = std::env::var("CUDA_COMPUTE_CAP").unwrap_or_else(|_| "80".to_string());
    let arch = format!("sm_{cuda_compute_cap}");

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

    let cubin_path = out_dir.join("mamba2_alpha_kernel.cubin");

    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_src_path.to_str().unwrap(),
        ])
        .status();

    match status {
        Ok(s) if s.success() => {
            eprintln!(
                "  ml-alpha: Compiled mamba2_alpha_kernel.cu -> mamba2_alpha_kernel.cubin ({arch})"
            );
        }
        Ok(s) => panic!(
            "ml-alpha: nvcc failed to compile mamba2_alpha_kernel.cu (exit={})",
            s.code().unwrap_or(-1)
        ),
        Err(e) => panic!("ml-alpha: nvcc invocation error: {e}"),
    }
}

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,
    }
}