foxhunt/crates/ml/tests/sp20_aggregate_inputs_test.rs

//! SP20 Phase 1.4 (2026-05-09) — `sp20_aggregate_inputs_kernel` GPU oracle tests.
//!
//! Validates the per-env → SP20EmaInputs aggregation rules atomically
//! with the Path C buffer-arg refactor of `sp20_emas_compute_kernel`.
//! The aggregation kernel is the GPU-side replacement for what would
//! otherwise need a host-side reduction over `trade_close_per_sample`,
//! `step_ret_per_sample`, etc. — forbidden by `feedback_cpu_is_read_only`
//! + `feedback_no_htod_htoh_only_mapped_pinned`.
//!
//! Each test constructs synthetic per-env arrays in mapped-pinned f32/i32
//! buffers, primes the upstream stats / dir_acc inputs, launches the
//! aggregation kernel, then reads back the resulting `SP20EmaInputs`
//! struct via the same f32-aliased buffer trick the EMA tests use.
//!
//! Run on a GPU host:
//!
//!     SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 \
//!       cargo test -p ml --test sp20_aggregate_inputs_test --features cuda \
//!         -- --ignored --nocapture

#![allow(clippy::tests_outside_test_module)]

#[cfg(feature = "cuda")]
#[allow(unsafe_code)] // CUDA kernel launch + mapped-pinned memory.
mod gpu {
    use std::sync::Arc;

    use cudarc::driver::{CudaContext, CudaFunction, CudaStream};
    use ml::cuda_pipeline::mapped_pinned::{MappedF32Buffer, MappedI32Buffer};
    use ml::cuda_pipeline::sp20_aggregate_inputs::{
        launch_sp20_aggregate_inputs, SP20_EMA_INPUTS_F32_LEN,
    };

    const SP20_AGGREGATE_INPUTS_CUBIN: &[u8] = include_bytes!(concat!(
        env!("OUT_DIR"),
        "/sp20_aggregate_inputs_kernel.cubin"
    ));

    /// Production direction divisor (`NUM_MAGNITUDES * NUM_ORD * NUM_URG = 27`)
    /// matching `state_layout.cuh` and `experience_action_select`.
    const DIR_DIVISOR: i32 = 27;

    /// `DIR_HOLD = 1` per `state_layout.cuh`. The aggregation kernel
    /// flags `action_is_hold` when the majority of envs decode to this
    /// direction.
    const HOLD_ACTION_ID: i32 = 1;

    /// Encode a packed factored action with the given direction.
    /// `dir * 27 + 0` lands at the `mag=0, ord=0, urg=0` slot of the
    /// chosen direction — the aggregation kernel only cares about the
    /// direction component.
    fn encode_dir(dir: i32) -> i32 {
        dir * DIR_DIVISOR
    }

    fn make_test_stream() -> Arc<CudaStream> {
        let ctx = CudaContext::new(0).expect("CUDA context — is a GPU available?");
        ctx.default_stream()
    }

    fn load_aggregation_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
        let module = stream
            .context()
            .load_cubin(SP20_AGGREGATE_INPUTS_CUBIN.to_vec())
            .expect("load sp20_aggregate_inputs_kernel cubin");
        module
            .load_function("sp20_aggregate_inputs_kernel")
            .expect("load sp20_aggregate_inputs_kernel function")
    }

    /// Read the post-launch `SP20EmaInputs` struct out of the f32-aliased
    /// buffer. Returns `(is_close, is_win, trade_duration, action_is_hold,
    /// alpha, per_bar_hold_reward, p50, std, dir_acc)`.
    fn read_ema_inputs(buf: &MappedF32Buffer) -> (i32, i32, i32, i32, f32, f32, f32, f32, f32) {
        let v = buf.read_all();
        assert_eq!(v.len(), SP20_EMA_INPUTS_F32_LEN);
        // The kernel writes int fields in the first 4 slots and float
        // fields in the last 5. We cast back from f32::to_bits.
        let is_close       = v[0].to_bits() as i32;
        let is_win         = v[1].to_bits() as i32;
        let trade_duration = v[2].to_bits() as i32;
        let action_is_hold = v[3].to_bits() as i32;
        let alpha          = v[4];
        let per_bar_hr     = v[5];
        let p50            = v[6];
        let std            = v[7];
        let dir_acc        = v[8];
        (is_close, is_win, trade_duration, action_is_hold,
         alpha, per_bar_hr, p50, std, dir_acc)
    }

    /// Construct a 1-stride arrangement (one slice per env at offset
    /// `env_id * 1`). This mirrors a per-step slice that has already
    /// been offset to the current timestep — the production wire-up
    /// passes `base + current_t` so the kernel reads `[env*L + 0]`.
    /// In tests we use stride=1 and pack envs contiguously.
    fn run_kernel(
        trade_close: &[i32],
        step_ret: &[f32],
        hold_at_exit: &[f32],
        actions: &[i32],
        aux_logits_p50: f32,
        aux_logits_std: f32,
        aux_dir_acc: f32,
    ) -> (i32, i32, i32, i32, f32, f32, f32, f32, f32) {
        let stream = make_test_stream();
        let kernel = load_aggregation_kernel(&stream);
        let n = trade_close.len();
        assert!(n > 0, "test n_envs must be > 0");
        assert_eq!(step_ret.len(), n);
        assert_eq!(hold_at_exit.len(), n);
        assert_eq!(actions.len(), n);

        let tc_buf = unsafe { MappedI32Buffer::new(n) }.expect("alloc tc");
        tc_buf.write_from_slice(trade_close);
        let sr_buf = unsafe { MappedF32Buffer::new(n) }.expect("alloc sr");
        sr_buf.write_from_slice(step_ret);
        let he_buf = unsafe { MappedF32Buffer::new(n) }.expect("alloc he");
        he_buf.write_from_slice(hold_at_exit);
        let act_buf = unsafe { MappedI32Buffer::new(n) }.expect("alloc act");
        act_buf.write_from_slice(actions);

        let p50_buf = unsafe { MappedF32Buffer::new(1) }.expect("alloc p50");
        p50_buf.write_from_slice(&[aux_logits_p50]);
        let std_buf = unsafe { MappedF32Buffer::new(1) }.expect("alloc std");
        std_buf.write_from_slice(&[aux_logits_std]);
        let dir_buf = unsafe { MappedF32Buffer::new(1) }.expect("alloc dir");
        dir_buf.write_from_slice(&[aux_dir_acc]);

        let out_buf = unsafe { MappedF32Buffer::new(SP20_EMA_INPUTS_F32_LEN) }
            .expect("alloc out");
        // Pre-fill with sentinel so we can detect a kernel that fails to write.
        out_buf.write_from_slice(&[f32::NAN; SP20_EMA_INPUTS_F32_LEN]);

        unsafe {
            launch_sp20_aggregate_inputs(
                &stream,
                &kernel,
                tc_buf.dev_ptr,
                sr_buf.dev_ptr,
                he_buf.dev_ptr,
                act_buf.dev_ptr,
                /* env_stride = */ 1,
                p50_buf.dev_ptr,
                std_buf.dev_ptr,
                dir_buf.dev_ptr,
                /* n_envs = */ n as i32,
                DIR_DIVISOR,
                HOLD_ACTION_ID,
                out_buf.dev_ptr,
            ).expect("launch aggregate kernel");
        }
        stream.synchronize().expect("sync after aggregate");
        read_ema_inputs(&out_buf)
    }

    /// All envs flat (no trade close), all picking Long. Expected:
    /// is_close = 0, action_is_hold = 0, trade_duration = 0,
    /// alpha = per_bar_hold_reward = 0 (Phase 2 / 3.2 placeholders),
    /// upstream stats forwarded.
    #[test]
    #[ignore = "requires GPU"]
    fn no_trades_no_holds_emits_sentinel_state() {
        let trade_close = vec![0_i32; 8];
        let step_ret    = vec![0.0_f32; 8];
        let hold_at_exit = vec![0.0_f32; 8];
        let actions     = vec![encode_dir(2 /* DIR_LONG */); 8];
        let (ic, iw, td, ah, alpha, hr, p50, std, dir) =
            run_kernel(&trade_close, &step_ret, &hold_at_exit, &actions,
                       0.3, 0.1, 0.55);
        assert_eq!(ic, 0, "is_close must be 0 (no trades closed)");
        assert_eq!(iw, 0, "is_win must be 0 when is_close=0");
        assert_eq!(td, 0, "trade_duration must be 0 when is_close=0");
        assert_eq!(ah, 0, "action_is_hold must be 0 (all envs picked Long)");
        assert_eq!(alpha, 0.0, "alpha is Phase 2 forward ref placeholder = 0.0");
        assert_eq!(hr, 0.0, "per_bar_hold_reward is Phase 3.2 forward ref = 0.0");
        assert!((p50 - 0.3).abs() < 1e-6, "aux_logits_p50 forwarded");
        assert!((std - 0.1).abs() < 1e-6, "aux_logits_std forwarded");
        assert!((dir - 0.55).abs() < 1e-6, "aux_dir_acc forwarded");
    }

    /// Half envs win, half lose; all close. Expected:
    /// is_close = 1, is_win = 1 (4 wins / 8 closed = 0.5, threshold).
    /// trade_duration = round(mean(hold_at_exit)) = round(7.5) = 8
    /// (round-to-nearest-even at 7.5 ⇒ 8).
    #[test]
    #[ignore = "requires GPU"]
    fn half_wins_meets_threshold_and_duration_rounds() {
        // 4 wins, 4 losses. Hold-at-exit mean = (5+6+7+8+7+8+9+10)/8 = 7.5.
        let trade_close = vec![1; 8];
        let step_ret    = vec![ 0.5,  0.3,  0.2,  0.1, -0.5, -0.3, -0.2, -0.1];
        let hold_at_exit = vec![5.0, 6.0, 7.0, 8.0, 7.0, 8.0, 9.0, 10.0];
        let actions     = vec![encode_dir(0); 8]; // all DIR_SHORT — not Hold
        let (ic, iw, td, ah, _, _, _, _, _) =
            run_kernel(&trade_close, &step_ret, &hold_at_exit, &actions,
                       0.0, 0.0, 0.0);
        assert_eq!(ic, 1);
        assert_eq!(iw, 1, "fraction wins (4/8 = 0.5) meets >= 0.5 threshold");
        // round-to-nearest-even at 7.5 = 8 (banker's rounding); CUDA
        // __float2int_rn uses round-to-nearest-even. Either 7 or 8 are
        // structurally acceptable; the kernel uses CUDA's RN so we
        // accept both as a bound test.
        assert!(
            td == 7 || td == 8,
            "trade_duration must round 7.5 to 7 (down) or 8 (RN-even); got {td}",
        );
        assert_eq!(ah, 0);
    }

    /// Strict majority Hold over n_envs. With 5 of 8 picking Hold,
    /// 5 > 8/2 = 4 ⇒ action_is_hold = 1.
    #[test]
    #[ignore = "requires GPU"]
    fn majority_hold_sets_action_is_hold() {
        let trade_close = vec![0; 8];
        let step_ret    = vec![0.0; 8];
        let hold_at_exit = vec![0.0; 8];
        let mut actions = Vec::with_capacity(8);
        for i in 0..8 {
            actions.push(if i < 5 { encode_dir(HOLD_ACTION_ID) } else { encode_dir(0) });
        }
        let (_, _, _, ah, _, _, _, _, _) =
            run_kernel(&trade_close, &step_ret, &hold_at_exit, &actions,
                       0.0, 0.0, 0.0);
        assert_eq!(ah, 1, "5 of 8 = strict majority Hold => action_is_hold=1");
    }

    /// Tie at exactly half Hold. 4 of 8 = 4 = n_envs/2; the kernel uses
    /// `count * 2 > n_envs` (strict majority), so a tie ⇒ 0.
    #[test]
    #[ignore = "requires GPU"]
    fn tie_hold_count_does_not_set_action_is_hold() {
        let trade_close = vec![0; 8];
        let step_ret    = vec![0.0; 8];
        let hold_at_exit = vec![0.0; 8];
        let mut actions = Vec::with_capacity(8);
        for i in 0..8 {
            actions.push(if i < 4 { encode_dir(HOLD_ACTION_ID) } else { encode_dir(2) });
        }
        let (_, _, _, ah, _, _, _, _, _) =
            run_kernel(&trade_close, &step_ret, &hold_at_exit, &actions,
                       0.0, 0.0, 0.0);
        assert_eq!(ah, 0, "tie (4 of 8 Hold) => action_is_hold=0 (strict majority)");
    }

    /// Single env closed with a loss: is_close=1, is_win=0
    /// (0/1 < 0.5), trade_duration = round(3.0) = 3.
    #[test]
    #[ignore = "requires GPU"]
    fn single_env_loss_close() {
        let trade_close = vec![1, 0, 0, 0];
        let step_ret    = vec![-0.2, 0.0, 0.0, 0.0];
        let hold_at_exit = vec![3.0, 0.0, 0.0, 0.0];
        let actions     = vec![encode_dir(0); 4];
        let (ic, iw, td, _, _, _, _, _, _) =
            run_kernel(&trade_close, &step_ret, &hold_at_exit, &actions,
                       0.0, 0.0, 0.0);
        assert_eq!(ic, 1);
        assert_eq!(iw, 0, "loss (0 wins out of 1 closed) => is_win=0");
        assert_eq!(td, 3, "trade_duration = round(3.0) = 3");
    }

    /// Phase 2 / Phase 3.2 forward-reference placeholders are emitted
    /// as 0.0 regardless of any input combinations. This test
    /// **verifies the documented contract** — the kernel does NOT
    /// surface any per-env signal into `alpha` / `per_bar_hold_reward`
    /// in Phase 1.4; that wiring lands in Phase 2 / 3.2.
    #[test]
    #[ignore = "requires GPU"]
    fn alpha_and_per_bar_hold_reward_are_phase_2_and_3_2_placeholders() {
        let trade_close = vec![1, 1, 1];
        let step_ret    = vec![0.5, 0.4, 0.3];
        let hold_at_exit = vec![10.0, 10.0, 10.0];
        let actions     = vec![encode_dir(HOLD_ACTION_ID); 3];
        let (_, _, _, _, alpha, hr, _, _, _) =
            run_kernel(&trade_close, &step_ret, &hold_at_exit, &actions,
                       0.5, 0.5, 0.99);
        assert_eq!(
            alpha, 0.0,
            "alpha is Phase 2 placeholder; aggregation kernel must NOT \
             surface a per-env signal here regardless of inputs",
        );
        assert_eq!(
            hr, 0.0,
            "per_bar_hold_reward is Phase 3.2 placeholder; aggregation \
             kernel must NOT surface a per-env signal here",
        );
    }
}