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]
    "rl_action_kernel",              // RL Phase R4: Thompson sampler over C51 atoms; one block per batch, N_ACTIONS threads; per-batch xorshift32 PRNG state; replaces host Thompson loop per feedback_cpu_is_read_only
    "argmax_expected_q",             // RL Phase R4: argmax over expected Q per action; Bellman-target argmax (Double-DQN); pairs with rl_action_kernel per pearl_thompson_for_distributional_action_selection
    "log_pi_at_action",              // RL Phase R4: per-batch log π(action_b) via log-softmax + lookup; PPO importance-ratio path
    "dqn_target_soft_update",        // RL Phase R5: element-wise target[i] = (1-τ)·target + τ·current, reads τ from ISV[401]; closes defect #4 (no target-net soft update in flawed branch)
    "extract_realized_pnl_delta",    // RL Phase R6: GPU-pure reward + done extraction from device Pos array; replaces host read_pos loop per feedback_cpu_is_read_only
    "apply_reward_scale",            // RL Phase R6: element-wise rewards *= ISV[RL_REWARD_SCALE_INDEX=406]; closes the F.3b host roundtrip
    "actions_to_market_targets",     // RL Phase R6: 9-action grid → LobSim market_targets[B*2] on device; replaces host submit_market loop per feedback_cpu_is_read_only
    "abs_copy",                      // RL Phase R7a: element-wise dst[b] = fabsf(src[b]); feeds |reward| into ema_update_on_done for the MEAN_ABS_PNL_EMA slot
    "rl_var_over_abs_mean_b",        // R9 EMA wiring: var(x)/|mean(x)| reduction; feeds advantage_var_ratio_ema (ISV[421]) from advantages_d
    "rl_kurtosis_b",                 // R9 EMA wiring: E[(x-μ)⁴]/σ⁴ reduction; feeds td_kurtosis_ema (ISV[422]) from td_per_sample_d. (entropy_observed_ema uses existing ema_update_per_step's internal mean reduce directly on entropy_d — no separate mean kernel needed.)
];

// Cache bust v31 (2026-05-23): R9 EMA wiring follow-up — three new
// reduce kernels populate the input EMAs for 3 of the 6 previously
// frozen controllers (entropy_coef, rollout_steps, per_alpha). Each
// kernel is a single-block tree reduction over b_size floats:
//   * rl_reduce_mean_b — mean(x) — for entropy_observed
//   * rl_var_over_abs_mean_b — var(x)/|mean(x)| — for advantage_var_ratio
//   * rl_kurtosis_b — E[(x-μ)⁴]/σ⁴ — for td_kurtosis
// Block dim must be next pow2 ≥ b_size (trainer enforces); shared mem
// is reused across multi-pass reductions. The trainer launches each
// reducer after the source signal is populated (PPO surrogate fwd for
// entropy, compute_advantage_return for advantages, dqn backward_logits
// for td_per_sample) then feeds the scalar output to ema_update_per_step
// or ema_update_on_done targeting the right ISV slot.
//
// The remaining 3 EMAs (kl_pi, q_divergence, trade_duration) need new
// state/derivation kernels — separate follow-up commit.

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"))
}