foxhunt/crates/ml-backtesting/build.rs

//! Cubin compilation for the LOB simulator CUDA kernels.
//! Mirrors crates/ml-alpha/build.rs.
//! Per feedback_no_nvrtc.md: cubins only, no runtime nvrtc.

use std::path::PathBuf;
use std::process::Command;

fn main() -> Result<(), String> {
    let cuda_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("cuda");
    let out_dir = PathBuf::from(std::env::var("OUT_DIR").map_err(|e| e.to_string())?);

    // Per pearl_build_rs_rerun_if_env_changed.md: pair every env::var
    // with rerun-if-env-changed.
    //
    // Arch detection (mirrors crates/ml-alpha/build.rs:detect_arch — the
    // original "mirror" docstring above was aspirational, B-9 cluster
    // alpha-rl-mbg2n on H100 caught the gap: ml-alpha picked sm_90 via
    // CUDA_COMPUTE_CAP but ml-backtesting fell through to default sm_86
    // → no kernel image error at LobSimCuda::new). Priority:
    //   1. CUDA_COMPUTE_CAP (numeric, e.g. "90") — set by alpha-rl-template
    //      from nvidia-smi inside the compile pod
    //   2. FOXHUNT_CUDA_ARCH (sm_-prefixed, e.g. "sm_90") — set by
    //      lob-backtest-sweep-template
    //   3. nvidia-smi --query-gpu=compute_cap at build time
    //   4. Default sm_86 (RTX 3050 Ti local dev)
    let arch = detect_arch();
    eprintln!("  ml-backtesting: compiling kernels for {arch}");
    println!("cargo:rerun-if-env-changed=CUDA_COMPUTE_CAP");
    println!("cargo:rerun-if-env-changed=FOXHUNT_CUDA_ARCH");
    println!("cargo:rerun-if-env-changed=CUDA_PATH");
    println!("cargo:rerun-if-env-changed=NVCC");

    let nvcc = std::env::var("NVCC")
        .map(PathBuf::from)
        .unwrap_or_else(|_| PathBuf::from("nvcc"));

    if !cuda_dir.exists() {
        // No kernels yet (e.g. before C4 lands them) — quietly noop.
        return Ok(());
    }

    println!("cargo:rerun-if-changed={}", cuda_dir.display());
    let entries = std::fs::read_dir(&cuda_dir).map_err(|e| format!("read cuda dir: {e}"))?;
    for entry in entries {
        let entry = entry.map_err(|e| e.to_string())?;
        let path = entry.path();
        if path.extension().and_then(|s| s.to_str()) != Some("cu") {
            continue;
        }
        let stem = path
            .file_stem()
            .ok_or("cu file has no stem")?
            .to_string_lossy()
            .to_string();
        let cubin = out_dir.join(format!("{stem}.cubin"));
        println!("cargo:rerun-if-changed={}", path.display());

        let mut cmd = Command::new(&nvcc);
        cmd.args(["--cubin", "-O3", "-std=c++17"])
            .arg(format!("-arch={arch}"))
            .arg("-I")
            .arg(&cuda_dir)
            .arg("-o")
            .arg(&cubin)
            .arg(&path);
        let status = cmd
            .status()
            .map_err(|e| format!("spawn nvcc for {}: {e}", path.display()))?;
        if !status.success() {
            return Err(format!("nvcc failed for {}", path.display()));
        }
    }
    Ok(())
}

/// Returns the `sm_<NN>` arch string, honoring (in order):
/// `CUDA_COMPUTE_CAP` numeric env → `FOXHUNT_CUDA_ARCH` sm_-prefixed env
/// → `nvidia-smi --query-gpu=compute_cap` → default `sm_86`.
fn detect_arch() -> String {
    // 1. Numeric CUDA_COMPUTE_CAP (e.g. "90") from alpha-rl-template.
    if let Ok(cap) = std::env::var("CUDA_COMPUTE_CAP") {
        let trimmed = cap.trim();
        if !trimmed.is_empty() {
            return format!("sm_{trimmed}");
        }
    }
    // 2. sm_-prefixed FOXHUNT_CUDA_ARCH (e.g. "sm_90") from lob-backtest-sweep.
    if let Ok(arch) = std::env::var("FOXHUNT_CUDA_ARCH") {
        let trimmed = arch.trim();
        if !trimmed.is_empty() {
            return trimmed.to_string();
        }
    }
    // 3. Query the GPU on this machine (e.g. "9.0" → "90" → "sm_90").
    if let Ok(output) = Command::new("nvidia-smi")
        .args(["--query-gpu=compute_cap", "--format=csv,noheader"])
        .output()
    {
        if output.status.success() {
            let s = String::from_utf8_lossy(&output.stdout);
            let cap = s.trim().replace('.', "");
            if !cap.is_empty() {
                return format!("sm_{cap}");
            }
        }
    }
    // 4. Default — RTX 3050 Ti local dev.
    "sm_86".to_string()
}