foxhunt/crates/ml/build.rs

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

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

    // ── Feature schema fingerprint ────────────────────────────────────────────
    // Stamped into every fxcache file so caches built against a different
    // feature-extractor / state-layout / fxcache-format version fail
    // header-validation at load time and trigger automatic regeneration via
    // the Argo `ensure-fxcache` step. Removes the manual "remember to bump
    // FXCACHE_VERSION on schema change" ritual that broke the L40S deploy.
    //
    // Hashing strategy: FNV-1a 64-bit over the raw bytes of these source
    // files. Stable across rust versions and machines (unlike
    // `std::hash::DefaultHasher`). Cost: harmless cosmetic edits to these
    // files (whitespace / comments) trigger one cache regen on next deploy
    // (~5min). That cost is acceptable; missed schema drift is not.
    emit_feature_schema_hash();

    // 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",
        "backtest_plan_kernel.cu",
        "tau_update_kernel.cu",
        "epsilon_update_kernel.cu",
        "per_branch_gamma_update_kernel.cu",
        "kelly_cap_update_kernel.cu",
        "atoms_update_kernel.cu",
        "q_quantile_kernel.cu",
        // D.8 Plan 2 Task 6C: TLOB attention kernels (SDP + state scatter/read)
        "tlob_kernel.cu",
        // C.2 Plan 3 Task 1: per-component |reward| EMA into ISV[63..69)
        "reward_component_ema_kernel.cu",
        // B.2 Plan 3 Task 3: Flat→Positioned transition rate EMA into ISV[71]
        "trade_rate_ema_kernel.cu",
        // B.4 Plan 3 Task 4: per-batch readiness EMA + derived plan_threshold
        // (producer for ISV[PLAN_THRESHOLD_INDEX=49] + ISV[READINESS_EMA_INDEX=75])
        "plan_threshold_update_kernel.cu",
        // C.3 Plan 3 Task 7: train-vs-val state-distribution KL EMA + Flat-trap
        // escape amplifier (producer for ISV[78]+ISV[79]). Adaptive amp
        // multiplies B.1 opp_cost and B.2 bonus consumers.
        "state_kl_divergence_kernel.cu",
        // B.3 Plan 3 Task 8: GPU-only seeded warm-start scripted policies.
        // Replaces network-Q action source during the seed phase. 4 policies
        // (uniform/momentum/mean-rev/vwap-deviation) deterministically mixed
        // 40/20/20/20 by `i % 5`. Driven by ISV[SEED_STEPS_DONE/TARGET]
        // dispatch at the launcher boundary.
        "scripted_policy_kernel.cu",
        // B.3 Plan 3 Task 8: per-collect_experiences seed-phase progress
        // counter. Increments ISV[SEED_STEPS_DONE], EMAs the derived
        // seed-fraction into ISV[SEED_FRAC_EMA] (consumed by Task 9).
        "seed_step_counter_update_kernel.cu",
        // C.5 Plan 3 Task 9: CQL α ramp coupled to ISV[SEED_FRAC_EMA].
        // EMAs ISV[CQL_ALPHA_INDEX=48] toward `final × (1 - seed_frac)`.
        // Producer-only upgrade for slot 48; CQL gradient kernel consumer
        // (already reading ISV[48] per Plan 1 Task 12) unchanged.
        "cql_alpha_seed_update_kernel.cu",
        // E.5 Plan 4 Task 5 Mode A: attention-focus interpretability EMAs.
        // Multi-block GPU reductions — 3 blocks reducing VSN mag/dir param
        // slices + Mamba2 enriched-hidden buffer. EMA-updates ISV[87..90).
        // Diagnostic only; no consumer reads these slots in Mode A.
        "attention_focus_ema_kernel.cu",
        // Plan 4 follow-up: target-drift EMA, replaces legacy CPU-DtoH
        // per_branch_target_drift(). 2-block kernel computing
        // RMS(target - online) for mag + dir branches → ISV[92,93].
        "target_drift_kernel.cu",
        // Plan 4 Task 2c.1: Gated Residual Network kernels (forward + backward).
        // Module is additive — no production callers in this commit; Task 2c.3+4
        // wires it into the trunk encoder.
        "grn_kernel.cu",
        // Plan 4 Task 2c.3c.5: per-batch RMS EMA of `save_h_s2` into ISV[96].
        // Single-block shmem-reduce kernel (256 threads, no atomicAdd) launched
        // alongside `reward_component_ema` from `training_loop.rs`. Producer-only
        // in this commit — 2c.3c.6 wires the consumer in `mag_concat_qdir`'s
        // adaptive-scale path.
        "h_s2_rms_ema_kernel.cu",
        // Plan 4 Task 3 (E.3): IQN multi-quantile diagnostic EMAs into
        // ISV[99..103) (Q_p05/Q_p25/Q_p75/Q_p95; median τ=0.50 is the existing
        // greedy-Q diagnostic and not duplicated). 4-block kernel, one block
        // per off-median quantile, shmem-reduce over (B × TBA) of
        // `save_q_online`. No atomicAdd. Producer-only — diagnostic only.
        "iqn_quantile_ema_kernel.cu",
        // Plan 4 Task 1B-i: VSN feature-selection kernels (forward + backward).
        // Per-sample softmax-over-6-groups + gate-multiply. Caller-side cuBLAS GEMMs
        // for the per-group MLPs (Linear_1 + ReLU + Linear_2). Module is additive
        // in this commit — ZERO production callers; consumers wired in 1B-iii/iv.
        "vsn_feature_selection_kernel.cu",
        // Plan 4 Task 1B-i: VSN per-group mask EMA producer (single-block shmem
        // reduction, no atomicAdd). Reads vsn_mask saved by the forward kernel,
        // EMA-updates 6 ISV slots. Producer-only; consumer-side ISV slot
        // allocation lands in 1B-ii.
        "vsn_mask_ema_kernel.cu",
        // Plan 4 Task 6 Commit A: multi-task auxiliary heads (E.6).
        // Two `Linear(SH2 → 32) → ELU → Linear(32 → K)` MLPs branching off
        // h_s2 — next-bar return regression (K=1, MSE) + 5-class regime
        // classification (K=5, cross-entropy). Forward + backward + loss-
        // reduce + label-builder + per-tensor batch-reduce kernels. Module
        // is additive — ZERO production callers in this commit; Commit B
        // wires the forward/backward + training-loop loss accumulation.
        "aux_heads_kernel.cu",
        // Plan 4 Task 6 Commit A: aux-head loss EMA producer (single-thread
        // single-block kernel mirroring h_s2_rms_ema_kernel's footprint).
        // Producer-only; ISV[AUX_NEXT_BAR_MSE_EMA_INDEX] +
        // ISV[AUX_REGIME_CE_EMA_INDEX] consumer wires in Commit B.
        "aux_heads_loss_ema_kernel.cu",
        // Plan A Phase 0 (Thompson sampling spec 2026-04-26): standalone test
        // kernel exercised only by `distributional_q_tests.rs`. Implements the
        // Thompson direction sampling math (inverse-CDF over C51 atoms,
        // uniform-τ over IQN quantiles, argmax of E[Q]) that Phase 2 will
        // integrate into `experience_action_select`. Test-only; no production
        // callers in this commit.
        "thompson_test_kernel.cu",
        // MoE Regime Redesign Phase 2 Task 2.1: mixture-forward kernel.
        // Single-thread-per-(b,c), no atomicAdd, capture-friendly.
        // h_s2[b,c] = Σ_k g[b,k] · expert_outputs[k,b,c].
        // Subsequent tasks extend this cubin with moe_mixture_backward,
        // moe_load_balance_loss, moe_expert_util_ema_update.
        "moe_kernels.cu",
        // MoE adaptive load-balance λ controller (2026-04-27): single-block
        // single-thread cold-path producer kernel matching kelly_cap_update /
        // cql_alpha_seed_update precedent. Reads ISV[MOE_GATE_ENTROPY_EMA_INDEX=126]
        // (entropy EMA from moe_expert_util_ema_update), writes
        // ISV[MOE_LAMBDA_EFF_INDEX=128] = λ_floor + λ_max_extra × deficit.
        // Consumer (moe_load_balance_loss in moe_kernels.cu) reads slot 128
        // at runtime — no DtoH per feedback_isv_for_adaptive_bounds.md.
        "moe_lambda_eff_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
        }
    }
}

/// Compute and emit the feature-schema fingerprint as `FEATURE_SCHEMA_HASH`
/// (decimal u64 string) so `fxcache.rs` can pick it up via `env!`.
///
/// Hashes the raw bytes of every source file that defines the on-disk fxcache
/// schema. Any edit (including whitespace / comments) bumps the hash. The
/// fxcache header records this value at write time; `validate()` rejects
/// caches whose recorded hash mismatches the compiled-in const, and
/// `precompute_features` regenerates them.
fn emit_feature_schema_hash() {
    // Source files whose bytes determine the on-disk schema.
    // Touching ANY of these files implies "fxcache regen on next deploy".
    //
    // 2026-04-27: added `examples/precompute_features.rs` after diagnosing the
    // L40S `label_scale=5420` (raw-price magnitude) vs smoke `label_scale=0.05`
    // (z-score) divergence. Root cause: the z-score normalisation step landed
    // in `precompute_features.rs:625-630` on 2026-04-03 but that file was not
    // tracked here, so the hash didn't change. PVC fxcaches written by an
    // earlier build (no normalisation) still validated against the current
    // binary, silently feeding raw features through column 0 of next_states
    // and leaking future-bar prices into the policy via the aux head's
    // shared trunk → impossible Sharpe (141 at epoch 0, 0.32% max DD over
    // 214k bars). Adding the precompute pipeline to the hash forces every
    // fxcache to be regenerated when its writer changes — including any
    // future tweaks to feature normalisation, target ordering, or per-feature
    // post-processing.
    let schema_sources: &[&str] = &[
        "src/features/extraction.rs",
        "src/fxcache.rs",
        "examples/precompute_features.rs",
        "../ml-core/src/state_layout.rs",
    ];

    // FNV-1a 64-bit init constants (RFC-style).
    const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
    const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;

    let mut hash: u64 = FNV_OFFSET;
    for rel in schema_sources {
        let path = PathBuf::from(rel);
        let bytes = std::fs::read(&path)
            .unwrap_or_else(|e| panic!("Failed to read schema source {}: {e}", path.display()));
        // Mix file path bytes too so reordering / renaming bumps the hash even
        // if the bytes happen to match.
        for b in rel.as_bytes() {
            hash ^= *b as u64;
            hash = hash.wrapping_mul(FNV_PRIME);
        }
        // Length separator so concatenation boundaries can't collide.
        for b in (bytes.len() as u64).to_le_bytes() {
            hash ^= b as u64;
            hash = hash.wrapping_mul(FNV_PRIME);
        }
        for b in &bytes {
            hash ^= *b as u64;
            hash = hash.wrapping_mul(FNV_PRIME);
        }
        println!("cargo:rerun-if-changed={}", path.display());
    }

    // Decimal so the consumer can use the const-stable `u64::from_str_radix`
    // with radix 10 without any hex-prefix stripping.
    println!("cargo:rustc-env=FEATURE_SCHEMA_HASH={hash}");
}

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