foxhunt/crates/ml-alpha/build.rs

//! Pre-compile all ml-alpha CUDA kernels into arch-specific cubins.
//!
//! Per `feedback_no_nvrtc.md`: no runtime kernel compilation.
//! Per `pearl_build_rs_rerun_if_env_changed.md`: every `std::env::var`
//! is paired with `cargo:rerun-if-env-changed`.

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

const KERNELS: &[&str] = &[
    "mamba2_alpha_kernel",          // Mamba2 SSM scan kernel (used by PerceptionTrainer's encoder prefix)
    "snap_feature_assemble",
    "cfc_step",
    "multi_horizon_heads",
    "projection",
    "bce_loss_multi_horizon",        // Kendall σ-weighted multi-horizon BCE (axis A)
    "adamw_step",
    "grad_norm",
    "horizon_lambda",                // ISV-driven per-horizon gradient scaler (EMA + lambda)
    "layer_norm",                    // Phase 1: trunk pre-CfC normalisation
    "variable_selection",            // Phase 2D: TFT-style per-feature gating
    "attention_pool",                // Phase 3: single-Q learned content summary at CfC k=0
    "reduce_axis0",                  // Phase B: cross-batch param-grad reducer
    "output_smoothness",             // CRT.train: per-horizon adjacent-position prob-jitter penalty
    "smoothness_lambda_controller",  // CRT.train: ISV-driven λ controller anchored on h30 jitter
    "gpu_log_ring",                  // GPU diagnostic log ring — tick kernel + log_record helper
    "bucket_transition_kernels",     // Per-horizon CfC Phase 1→2 transition: tau_sort, bucket_assign, bucket_iqr, channels_in_bucket, heads_compact, zero_off_bucket (ALPHA fix 2026-05-21)
    "cfc_step_per_branch",           // Per-horizon CfC Phase 2: fused per-(batch, branch) fwd + bwd over [25,25,25,25,28] buckets
    "heads_block_diagonal_fwd",      // Per-horizon CfC Phase 2: heads w_skip projection with compact ragged storage (640→128 floats)
    "aux_trunk",                     // SDD-3 Layer B3: smaller single-bucket CfC trunk (AUX_HIDDEN=64) for outcome-supervision (D-labels)
    "aux_heads",                     // SDD-3 Layer B4: per-direction linear regression heads on AuxTrunk output (long + short, N_AUX_HORIZONS each)
    "aux_loss",                      // SDD-3 Layer B4: Huber loss + grad for aux trade-outcome regression targets (NaN-masked)
    "aux_vec_add",                   // SDD-3 Layer B5: element-wise dst += src for aux→encoder gradient accumulation (lifted stop-grad)
    "dqn_distributional_q",          // RL Phase C: C51 distributional Q-head fwd + Bellman TD bwd for integrated RL trainer
    "rl_gamma_controller",           // RL Phase C: ISV controller emitting γ to ISV[RL_GAMMA_INDEX=400]
    "rl_target_tau_controller",      // RL Phase C: ISV controller emitting τ to ISV[RL_TARGET_TAU_INDEX=401]
    "ppo_clipped_surrogate",         // RL Phase D: PPO clipped-surrogate + entropy bonus + value MSE fwd/bwd
    "rl_ppo_clip_controller",        // RL Phase D: ISV controller emitting ε to ISV[RL_PPO_CLIP_INDEX=402]
    "rl_entropy_coef_controller",    // RL Phase D: ISV controller emitting entropy bonus weight to ISV[RL_ENTROPY_COEF_INDEX=403]
    "v_head_fwd_bwd",                // RL Phase E.2: scalar V(s) head fwd + MSE bwd (linear layer; per-batch scratch + reduce_axis0)
    "grad_h_accumulate",             // RL Phase E.2: element-wise grad_h_encoder += λ × grad_h_head accumulator (one head at a time, serialised by stream)
    "bellman_target_projection",     // RL Phase E.2-DEFER: C51 categorical projection of Bellman target Z(s_{t+1}, a*) onto the discrete support, reads γ from ISV[400]; replaces host-side build_synthetic_bellman_target stand-in
    "rl_lr_controller",              // RL Phase E.2-DEFER: per-head learning-rate ISV emitter — bootstraps ISV[412..417] with 1e-3 (BCE/Q/π/V/aux); replaces hardcoded PHASE_E2_DEFAULT_LR
    "rl_rollout_steps_controller",   // RL Phase E.3b: rollout-buffer-length ISV emitter — emits ISV[RL_N_ROLLOUT_STEPS_INDEX=404] from var(advantage)/|mean A| EMA; bootstraps 2048
    "rl_per_alpha_controller",       // RL Phase E.3b: PER priority-exponent ISV emitter — emits ISV[RL_PER_ALPHA_INDEX=405] from TD-error kurtosis EMA; bootstraps 0.6
    "rl_reward_scale_controller",    // RL Phase R1 (rebuild): reward-standardisation scale ISV emitter — emits ISV[RL_REWARD_SCALE_INDEX=406] from mean |realized_pnl_usd| EMA; bootstraps 1.0
    "ema_update_on_done",            // RL Phase R3: generic done-gated EMA producer (slot-parameterised) for closed-trade-magnitude EMAs (mean_abs_pnl, q_divergence, td_kurtosis)
    "ema_update_per_step",           // RL Phase R3: generic per-step EMA producer (slot-parameterised) for continuous EMAs (kl_pi, entropy_observed, advantage_var_ratio, trade_duration)
    "compute_advantage_return",      // RL Phase R3: element-wise A_t = r + γ(1-done)·V(s_{t+1}) − V(s_t), R_t = r + γ(1-done)·V(s_{t+1}); reads γ from ISV[400]
];

// Cache bust v26 (2026-05-23): RL Phase R3 rebuild — three new GPU-resident
// kernels close defect #5 from the flawed Phase F+G arc
// (feedback_cpu_is_read_only violation in step_with_lobsim's host advantage
// + EMA loops). ema_update_on_done + ema_update_per_step are generic
// slot-parameterised EMA producers feeding the 7 RL controllers' ISV[417..424]
// inputs; compute_advantage_return is element-wise A_t + R_t on device,
// replacing the host loop. All three honor pearl_first_observation_bootstrap
// (bootstrap path defers if mean_obs == 0) + pearl_wiener_alpha_floor_for_nonstationary
// (host pre-floors α at 0.4) + feedback_no_atomicadd (tree reduction).

fn main() {
    println!("cargo:rerun-if-changed=build.rs");
    // Track shared headers so .cuh / .h edits trigger rebuilds of every
    // .cu that #includes them. Without these, an edit to a helper header
    // leaves a stale cubin.
    println!("cargo:rerun-if-changed=cuda/gpu_log_ids.h");
    println!("cargo:rerun-if-changed=cuda/gpu_log_helpers.cuh");

    println!("cargo:rerun-if-env-changed=CARGO_FEATURE_CUDA");
    if std::env::var("CARGO_FEATURE_CUDA").is_err() {
        eprintln!("  ml-alpha: cuda feature disabled, skipping kernel build");
        return;
    }

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

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

    let out = PathBuf::from(std::env::var("OUT_DIR").expect("OUT_DIR not set by cargo"));

    for k in KERNELS {
        let src = PathBuf::from(format!("cuda/{k}.cu"));
        if !src.exists() {
            eprintln!("  ml-alpha: skipping {k} — source not yet present");
            continue;
        }
        println!("cargo:rerun-if-changed={}", src.display());
        let cubin = out.join(format!("{k}.cubin"));
        compile(&nvcc, &src, &cubin, &arch);
    }
}

fn compile(nvcc: &Path, src: &Path, cubin: &Path, arch: &str) {
    let status = Command::new(nvcc)
        .args([
            "-cubin",
            &format!("-arch={arch}"),
            "-O3",
            "--use_fast_math",
            "--ftz=true",
            "--fmad=true",
            "-o",
            cubin.to_str().unwrap(),
            src.to_str().unwrap(),
        ])
        .status()
        .unwrap_or_else(|e| panic!("nvcc spawn failed for {}: {e}", src.display()));
    if !status.success() {
        panic!(
            "nvcc failed for {} (exit {})",
            src.display(),
            status.code().unwrap_or(-1)
        );
    }
    eprintln!(
        "  ml-alpha: compiled {} -> {} ({arch})",
        src.display(),
        cubin.display()
    );
}

fn find_nvcc() -> Option<PathBuf> {
    if let Ok(home) = std::env::var("CUDA_HOME") {
        let p = PathBuf::from(home).join("bin/nvcc");
        if p.exists() {
            return Some(p);
        }
    }
    for cand in ["/usr/local/cuda/bin/nvcc", "/usr/bin/nvcc"] {
        let p = PathBuf::from(cand);
        if p.exists() {
            return Some(p);
        }
    }
    Command::new("nvcc")
        .arg("--version")
        .output()
        .ok()
        .filter(|o| o.status.success())
        .map(|_| PathBuf::from("nvcc"))
}