foxhunt/crates/ml/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_pipeline");
    let common_header = kernel_dir.join("common_device_functions.cuh");
    let trade_physics_header = kernel_dir.join("trade_physics.cuh");

    // Rebuild cubins when shared headers change
    println!("cargo:rerun-if-changed={}", trade_physics_header.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 -- gracefully skip if not (non-CUDA builds)
    let nvcc_path = find_nvcc();
    let nvcc = match nvcc_path {
        Some(p) => p,
        None => {
            eprintln!("  warning: nvcc not found, skipping 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)
        .unwrap_or_else(|e| panic!("Failed to read {}: {e}", common_header.display()));
    println!("cargo:rerun-if-changed={}", common_header.display());

    // All kernels to precompile.
    // Kernels marked "standalone" have their own helpers and don't need common header.
    // All others get common_device_functions.cuh prepended.
    let kernels_with_common = [
        // Original 24 kernels
        "epsilon_greedy_kernel.cu",
        "backtest_env_kernel.cu",
        "backtest_forward_ppo_kernel.cu",
        "backtest_forward_supervised_kernel.cu",
        "backtest_metrics_kernel.cu",
        "dt_kernels.cu",
        "ensemble_kernels.cu",
        "her_episode_kernel.cu",
        "her_relabel_kernel.cu",
        "signal_adapter_kernel.cu",
        "statistics_kernel.cu",
        "training_guard_kernel.cu",
        "c51_loss_kernel.cu",
        "mse_loss_kernel.cu",
        "curiosity_training_kernel.cu",
        "curiosity_inference_kernel.cu",
        "dqn_utility_kernels.cu",
        "attention_kernel.cu",
        "attention_backward_kernel.cu",
        "iql_value_kernel.cu",
        "iqn_dual_head_kernel.cu",
        "monitoring_kernel.cu",
        "nstep_kernel.cu",
        "reward_shaping_kernel.cu",
        "ppo_experience_kernel.cu",
        "bias_kernels.cu",
        // Formerly standalone — now need common header for BF16 types
        "experience_kernels.cu",
        "ema_kernel.cu",
        "relu_mask_kernel.cu",
        "c51_grad_kernel.cu",
        "mse_grad_kernel.cu",
        "q_stats_kernel.cu",
        "cql_grad_kernel.cu",
        "trade_stats_kernel.cu",
        "backward_kernels.cu",
        "iqn_cvar_kernel.cu",
        "mamba2_temporal_kernel.cu",
        "graph_utility_kernels.cu",
        "grad_decomp_kernel.cu",
        "branch_grad_balance_kernel.cu",
    ];

    // ALL kernels get common header (BF16 types + wrappers)
    let mut failed: Vec<&str> = Vec::new();
    for kernel_name in &kernels_with_common {
        if !try_compile_kernel(
            &nvcc, kernel_dir, kernel_name, &arch, &out_dir, Some(&common_src),
        ) {
            failed.push(kernel_name);
        }
    }
    let passed = kernels_with_common.len() - failed.len();
    eprintln!("  Precompiled {passed}/{} CUDA kernels ({arch}) — f32/TF32",
              kernels_with_common.len());
    if !failed.is_empty() {
        eprintln!("  FAILED: {}", failed.join(", "));
        panic!("nvcc failed to compile {} 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(header) => format!("{header}\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 (-I kernel_dir for shared .cuh includes)
    let include_flag = format!("-I{}", kernel_dir.display());
    let status = Command::new(nvcc)
        .args([
            "-cubin",
            &format!("-arch={arch}"),
            "-O3",
            "--ftz=true",
            "--fmad=true",
            "--prec-div=true",
            "--prec-sqrt=true",
            &include_flag,
            "-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> {
    // Try $CUDA_HOME/bin/nvcc first
    if let Ok(home) = std::env::var("CUDA_HOME") {
        let nvcc = PathBuf::from(home).join("bin/nvcc");
        if nvcc.exists() {
            return Some(nvcc);
        }
    }

    // Try common CUDA paths
    for path in &["/usr/local/cuda/bin/nvcc", "/usr/bin/nvcc"] {
        let p = PathBuf::from(path);
        if p.exists() {
            return Some(p);
        }
    }

    // Check if nvcc is in PATH
    match Command::new("nvcc").arg("--version").output() {
        Ok(output) if output.status.success() => Some(PathBuf::from("nvcc")),
        _ => None,
    }
}