diff --git a/crates/ml/tests/sp21_per_trade_predicted_q_test.rs b/crates/ml/tests/sp21_per_trade_predicted_q_test.rs new file mode 100644 index 000000000..611c65f9a --- /dev/null +++ b/crates/ml/tests/sp21_per_trade_predicted_q_test.rs @@ -0,0 +1,521 @@ +//! SP21 T2.2 Phase 1.5 (2026-05-11) — `predicted_q` per-trade tape GPU oracle test. +//! +//! Validates the `entry_q` tracking added to `backtest_env_step_batch`: +//! +//! 1. **predicted_q populated on close** — open a Long position at step 0 +//! with a controlled `q_values_per_window[Long_Full] = 2.5`, hold for two +//! bars, close with Flat at step 3. Read the per-trade tape and assert +//! the single emitted trade carries `predicted_q ≈ 2.5` (the entry_q +//! captured at open and snapshotted before the close-detection block). +//! 2. **NULL q_values_per_window keeps predicted_q at 0** — same scenario +//! with `q_values_per_window = NULL` (single-step `evaluate()` semantics). +//! Kernel's open-time write is gated on the NULL guard, so +//! `entry_q_state[w]` stays 0.0 and the tape's `predicted_q` is 0.0. +//! 3. **entry_q persists across hold bars** — the persistent `entry_q_state` +//! buffer is per-window, written ONCE at open, read back at close. Hold +//! bars don't touch it. The test indirectly proves persistence by the +//! single-trade success of test (1). +//! +//! These cover the entry_q capture path end-to-end at the kernel level. The +//! eval pipeline (cuBLAS forward → action select → env step) integration test +//! lives in `gpu_backtest_validation.rs`; this file focuses on the per-trade +//! tape's new `predicted_q` field as a kernel-direct oracle. +//! +//! Run on a GPU host: +//! +//! SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 \ +//! cargo test -p ml --test sp21_per_trade_predicted_q_test --features cuda \ +//! -- --ignored --nocapture + +#![allow(clippy::tests_outside_test_module)] +#![allow(clippy::too_many_lines)] + +#[cfg(feature = "cuda")] +#[allow(unsafe_code)] // CUDA kernel launch. +mod gpu { + use std::sync::Arc; + + use cudarc::driver::{ + CudaContext, CudaFunction, CudaStream, DevicePtr, LaunchConfig, PushKernelArg, + }; + + /// Same cubin loaded by `gpu_backtest_evaluator.rs`. Test-side + /// `include_bytes!` reads from the per-crate `OUT_DIR` set by + /// `crates/ml/build.rs` — the same OUT_DIR used by the lib build. + const ENV_CUBIN: &[u8] = include_bytes!(concat!( + env!("OUT_DIR"), + "/backtest_env_kernel.cubin" + )); + + /// Branch sizes — match `DqnBacktestConfig::from_network_dims` defaults. + const B0: i32 = 4; // Short / Hold / Long / Flat + const B1: i32 = 3; // Quarter / Half / Full + const B2: i32 = 3; // Order type + const B3: i32 = 3; // Patient / Normal / Aggressive + /// Total factored action count. Matches the kernel's + /// `num_actions` arg. + const NUM_ACTIONS: i32 = B0 * B1 * B2 * B3; // 108 + + /// Direction constants (mirrors `state_layout.cuh:123-126`). + const DIR_LONG: i32 = 2; + const DIR_HOLD: i32 = 1; + const DIR_FLAT: i32 = 3; + + /// Magnitude constants (mirrors `state_layout.cuh:133-135`). + const MAG_FULL: i32 = 2; + const MAG_QUARTER: i32 = 0; + + /// Encode a factored action: `dir * b1*b2*b3 + mag * b2*b3 + order * b3 + urgency`. + /// Mirrors `trade_physics.cuh::decode_direction_4b` etc. exactly so the + /// kernel decoder maps our test actions back to the canonical DIR/MAG. + const fn encode_action(dir: i32, mag: i32, order: i32, urgency: i32) -> i32 { + dir * (B1 * B2 * B3) + mag * (B2 * B3) + order * B3 + urgency + } + + /// `MAX_TRADES_PER_WINDOW` matches the kernel guard at the per-trade + /// emission site. We size at 16 here — far smaller than production's + /// 200_000 — because the test only emits a handful of trades and the + /// per-window buffer alloc is `n_windows × max_trades`. + const MAX_TRADES_PER_WINDOW: i32 = 16; + + /// Tolerance for f32 comparisons against expected entry_q values. + /// The kernel writes the Q-value verbatim from `q_values_per_window` + /// (no math), so the round-trip is exact under IEEE-754 — but a small + /// epsilon defends against any future float-rewrite of the open path. + const Q_TOL: f32 = 1e-5; + + /// 4-bar single-window scenario: open Long Full at step 0, hold for two + /// bars, Flat-close at step 3. Returns (n_windows=1, max_len=4, chunk_len=4). + const N_WINDOWS: usize = 1; + const MAX_LEN: usize = 4; + const CHUNK_LEN: usize = 4; + const FEATURE_DIM: i32 = 4; // < 41 → kernel uses fixed-width trail (no features deref) + + fn try_cuda_stream() -> Option> { + match CudaContext::new(0) { + Ok(ctx) => Some(ctx.default_stream()), + Err(_) => { + eprintln!("CUDA not available, skipping per-trade-tape predicted_q test"); + None + } + } + } + + fn load_env_batch_kernel(stream: &Arc) -> CudaFunction { + let module = stream + .context() + .load_cubin(ENV_CUBIN.to_vec()) + .expect("load backtest_env_kernel.cubin"); + module + .load_function("backtest_env_step_batch") + .expect("load backtest_env_step_batch function") + } + + /// One window of OHLC. All bars at price 100.0 (flat market) → trade + /// closes at break-even, isolating the entry_q signal from any P&L + /// noise. The close emission predicate `pre_entry_price > 0.0f && + /// value > 1.0f` only requires entry_price was non-zero (set on open + /// at step 0) and capital is positive — both satisfied with flat + /// prices and `initial_capital = 100_000`. + fn flat_market_prices() -> Vec { + let mut prices = Vec::with_capacity(N_WINDOWS * MAX_LEN * 4); + for _ in 0..(N_WINDOWS * MAX_LEN) { + // OHLC = [100, 100, 100, 100] + prices.push(100.0_f32); + prices.push(100.0_f32); + prices.push(100.0_f32); + prices.push(100.0_f32); + } + prices + } + + /// Initial portfolio state matching `GpuBacktestEvaluator::init_portfolio_state`. + fn init_portfolio_state(initial_capital: f32) -> Vec { + let mut state = vec![0.0_f32; N_WINDOWS * 8]; + state[0] = initial_capital; // value + state[2] = initial_capital; // cash + state[4] = initial_capital; // max_equity + state + } + + /// Action sequence: open Long Full at step 0, hold, hold, close Flat. + /// Layout `[chunk_len, n_windows]` flattened — the kernel reads via + /// `chunked_actions[s * n_windows + w]`. + fn open_hold_close_actions() -> Vec { + // step 0: Long Full Market Normal (open) + let a_open = encode_action(DIR_LONG, MAG_FULL, 0, 1); + // step 1, 2: Hold Quarter Market Normal (carry the position; Hold = no-op, + // see `project_hold_action_design.md`) + let a_hold = encode_action(DIR_HOLD, MAG_QUARTER, 0, 1); + // step 3: Flat Quarter Market Normal (close all positions to zero) + let a_close = encode_action(DIR_FLAT, MAG_QUARTER, 0, 1); + vec![a_open, a_hold, a_hold, a_close] + } + + /// Q-values buffer layout `[chunk_len, n_windows, num_actions]` flattened. + /// The kernel indexes via `q[s * n_windows * num_actions + w * num_actions + /// + action_taken]` — at step 0 we set the slot for the open action only; + /// other slots stay at 0.0 (the kernel never reads them since it always + /// indexes by the chosen action). + fn q_values_with_open_at(open_action: i32, open_q: f32) -> Vec { + let n = N_WINDOWS; + let mut q = vec![0.0_f32; CHUNK_LEN * n * NUM_ACTIONS as usize]; + let s = 0_usize; + let w = 0_usize; + let a = open_action as usize; + q[s * n * NUM_ACTIONS as usize + w * NUM_ACTIONS as usize + a] = open_q; + q + } + + /// Configuration constants for the kernel — match `GpuBacktestConfig::default` + /// where it matters and zero out the optional shaping params. + struct KernelConfig { + max_position: f32, + tx_cost_bps: f32, + spread_cost: f32, + contract_multiplier: f32, + margin_pct: f32, + initial_capital: f32, + } + + impl KernelConfig { + fn for_test() -> Self { + Self { + max_position: 1.0, + // Zero costs so any P&L deviation comes from the kernel logic, + // not transaction friction. The test asserts on `predicted_q` + // (independent of P&L) but the close emission predicate's + // `value > 1.0f` guard is more comfortably satisfied with + // no cost drain. + tx_cost_bps: 0.0, + spread_cost: 0.0, + contract_multiplier: 1.0, + margin_pct: 0.0, + initial_capital: 100_000.0, + } + } + } + + /// Run the batch kernel one chunk over the 4-bar window with the given + /// q_values_per_window pointer (real or NULL=0u64) and `num_actions`. + /// Returns (per_trade_count[0], per_trade_predicted_q[0..count]). + #[allow(clippy::too_many_arguments)] + fn run_one_chunk( + stream: &Arc, + kernel: &CudaFunction, + cfg: &KernelConfig, + actions: &[i32], + q_values_dev_ptr: u64, // 0 = NULL + num_actions: i32, + ) -> (u32, Vec) { + // ── Allocate input buffers ── + let mut prices_dev = stream + .alloc_zeros::(N_WINDOWS * MAX_LEN * 4) + .expect("alloc prices"); + let prices_host = flat_market_prices(); + stream.memcpy_htod(&prices_host, &mut prices_dev).expect("HtoD prices"); + + let mut window_lens_dev = stream + .alloc_zeros::(N_WINDOWS) + .expect("alloc window_lens"); + let window_lens_host = vec![MAX_LEN as i32; N_WINDOWS]; + stream + .memcpy_htod(&window_lens_host, &mut window_lens_dev) + .expect("HtoD window_lens"); + + // Features buffer — sized `[n * max_len * feature_dim]`. With + // FEATURE_DIM=4 (< 41), the kernel's `compute_regime_trail_scales` + // takes the fallback (no feature deref), so zeros are safe. + let features_dev = stream + .alloc_zeros::(N_WINDOWS * MAX_LEN * FEATURE_DIM as usize) + .expect("alloc features"); + + let mut actions_dev = stream + .alloc_zeros::(CHUNK_LEN * N_WINDOWS) + .expect("alloc chunked_actions"); + stream + .memcpy_htod(actions, &mut actions_dev) + .expect("HtoD chunked_actions"); + + let mut portfolio_dev = stream + .alloc_zeros::(N_WINDOWS * 8) + .expect("alloc portfolio"); + let portfolio_host = init_portfolio_state(cfg.initial_capital); + stream + .memcpy_htod(&portfolio_host, &mut portfolio_dev) + .expect("HtoD portfolio"); + + let step_rewards_dev = stream + .alloc_zeros::(N_WINDOWS) + .expect("alloc step_rewards"); + let step_returns_dev = stream + .alloc_zeros::(N_WINDOWS * MAX_LEN) + .expect("alloc step_returns"); + let done_dev = stream + .alloc_zeros::(N_WINDOWS) + .expect("alloc done_flags"); + let actions_history_dev = stream + .alloc_zeros::(N_WINDOWS * MAX_LEN) + .expect("alloc actions_history"); + let kelly_stats_dev = stream + .alloc_zeros::(N_WINDOWS * 4) + .expect("alloc kelly_stats"); + + // Per-trade tape buffers (5 SoA + count + 1 new SoA for predicted_q). + let per_trade_pnl_dev = stream + .alloc_zeros::(N_WINDOWS * MAX_TRADES_PER_WINDOW as usize) + .expect("alloc per_trade_pnl"); + let per_trade_holding_bars_dev = stream + .alloc_zeros::(N_WINDOWS * MAX_TRADES_PER_WINDOW as usize) + .expect("alloc per_trade_holding_bars"); + let per_trade_bar_index_dev = stream + .alloc_zeros::(N_WINDOWS * MAX_TRADES_PER_WINDOW as usize) + .expect("alloc per_trade_bar_index"); + let per_trade_dir_mag_dev = stream + .alloc_zeros::(N_WINDOWS * MAX_TRADES_PER_WINDOW as usize) + .expect("alloc per_trade_dir_mag"); + let per_trade_count_dev = stream + .alloc_zeros::(N_WINDOWS) + .expect("alloc per_trade_count"); + let per_trade_predicted_q_dev = stream + .alloc_zeros::(N_WINDOWS * MAX_TRADES_PER_WINDOW as usize) + .expect("alloc per_trade_predicted_q"); + // Phase 1.5 entry_q persistent state — one float per window. + // Initialized to 0.0; the kernel writes it on open / reverse. + let entry_q_state_dev = stream + .alloc_zeros::(N_WINDOWS) + .expect("alloc entry_q_state"); + + // Scalar args. + let n_i32 = N_WINDOWS as i32; + let max_len_i32 = MAX_LEN as i32; + let start_step_i32 = 0_i32; + let chunk_len_i32 = CHUNK_LEN as i32; + let max_trades_i32 = MAX_TRADES_PER_WINDOW; + let feature_dim_i32 = FEATURE_DIM; + + // ── Launch ── + let env_blocks = ((N_WINDOWS as u32) + 255) / 256; + let launch_cfg = LaunchConfig { + grid_dim: (env_blocks.max(1), 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + + unsafe { + stream + .launch_builder(kernel) + .arg(&prices_dev) + .arg(&window_lens_dev) + .arg(&features_dev) + .arg(&feature_dim_i32) + .arg(&actions_dev) + .arg(&portfolio_dev) + .arg(&step_rewards_dev) + .arg(&step_returns_dev) + .arg(&done_dev) + .arg(&actions_history_dev) + .arg(&n_i32) + .arg(&max_len_i32) + .arg(&cfg.max_position) + .arg(&cfg.tx_cost_bps) + .arg(&cfg.spread_cost) + .arg(&start_step_i32) + .arg(&chunk_len_i32) + .arg(&B0) + .arg(&B1) + .arg(&B2) + .arg(&B3) + .arg(&0.0_f32) // holding_cost_rate (unused per Phase 3) + .arg(&0.0_f32) // churn_threshold (unused) + .arg(&0.0_f32) // churn_penalty_scale (unused) + .arg(&cfg.contract_multiplier) + .arg(&cfg.margin_pct) + .arg(&0.0_f32) // opp_cost_scale (unused) + .arg(&kelly_stats_dev) + .arg(&0u64) // isv_signals NULL → health=0.5 fallback + .arg(&0u64) // exploration_scale_ptr NULL + .arg(&0u64) // shaping_scale_ptr NULL + .arg(&0u64) // conviction_ptr NULL → 1.0 fallback + .arg(&0u64) // magnitude_conviction_ptr NULL → 0.0 fallback + .arg(&per_trade_pnl_dev) + .arg(&per_trade_holding_bars_dev) + .arg(&per_trade_bar_index_dev) + .arg(&per_trade_dir_mag_dev) + .arg(&per_trade_count_dev) + .arg(&max_trades_i32) + // SP21 T2.2 Phase 1.5 args + .arg(&entry_q_state_dev) + .arg(&q_values_dev_ptr) // 0 = NULL for the NULL-tolerance test + .arg(&num_actions) + .arg(&per_trade_predicted_q_dev) + .launch(launch_cfg) + .expect("backtest_env_step_batch launch"); + } + + stream.synchronize().expect("post-launch sync"); + + // ── Read tape ── + let mut count_host = vec![0u32; N_WINDOWS]; + stream + .memcpy_dtoh(&per_trade_count_dev, &mut count_host) + .expect("DtoH per_trade_count"); + let mut predicted_q_host = + vec![0.0_f32; N_WINDOWS * MAX_TRADES_PER_WINDOW as usize]; + stream + .memcpy_dtoh(&per_trade_predicted_q_dev, &mut predicted_q_host) + .expect("DtoH per_trade_predicted_q"); + + // Drop device buffers explicitly to free GPU memory before the + // next test runs (each #[test] gets its own isolated alloc set). + drop(prices_dev); + drop(window_lens_dev); + drop(features_dev); + drop(actions_dev); + drop(portfolio_dev); + drop(step_rewards_dev); + drop(step_returns_dev); + drop(done_dev); + drop(actions_history_dev); + drop(kelly_stats_dev); + drop(per_trade_pnl_dev); + drop(per_trade_holding_bars_dev); + drop(per_trade_bar_index_dev); + drop(per_trade_dir_mag_dev); + drop(per_trade_count_dev); + drop(per_trade_predicted_q_dev); + drop(entry_q_state_dev); + + let count_w0 = count_host[0]; + let predicted_q_w0: Vec = predicted_q_host[0..(count_w0 as usize)].to_vec(); + (count_w0, predicted_q_w0) + } + + /// Test 1: production path (real q_values) populates `predicted_q` from + /// the entry_q snapshot at trade open. + #[test] + #[ignore = "requires GPU"] + fn predicted_q_populated_on_close_with_real_q_values() { + let stream = match try_cuda_stream() { + Some(s) => s, + None => return, + }; + let kernel = load_env_batch_kernel(&stream); + + let cfg = KernelConfig::for_test(); + let actions = open_hold_close_actions(); + let open_action = actions[0]; + + // The Q-value we expect predicted_q to carry across the trade. + let expected_entry_q: f32 = 2.5; + let q_host = q_values_with_open_at(open_action, expected_entry_q); + let mut q_dev = stream + .alloc_zeros::(q_host.len()) + .expect("alloc q_values"); + stream.memcpy_htod(&q_host, &mut q_dev).expect("HtoD q_values"); + let q_dev_ptr = { + let (ptr, _g) = q_dev.device_ptr(&stream); + ptr + }; + + let (count, predicted_q) = run_one_chunk( + &stream, &kernel, &cfg, &actions, q_dev_ptr, NUM_ACTIONS, + ); + + // Drop the q_values buffer ourselves — it's owned by this test only. + drop(q_dev); + + assert_eq!( + count, 1, + "expected exactly 1 trade closed (open Long@step 0, Flat@step 3); got count={count}" + ); + assert!( + !predicted_q.is_empty(), + "predicted_q tape should have ≥1 entry when count={count}" + ); + let actual_q = predicted_q[0]; + assert!( + (actual_q - expected_entry_q).abs() < Q_TOL, + "predicted_q[0] = {actual_q}, expected {expected_entry_q} (entry_q captured from \ + q_values_per_window[w=0, a={open_action}] at step-0 open and snapshotted before the \ + step-3 close emission)" + ); + } + + /// Test 2: NULL `q_values_per_window` (single-step `evaluate()` semantics) + /// keeps `predicted_q` at the buffer-init 0.0 value. The kernel's + /// open-time write is gated on the NULL guard at the bottom of the + /// per-step loop, so `entry_q_state[w]` stays 0.0 and the close emit + /// snapshots 0.0 into the tape. + #[test] + #[ignore = "requires GPU"] + fn predicted_q_stays_zero_when_q_values_is_null() { + let stream = match try_cuda_stream() { + Some(s) => s, + None => return, + }; + let kernel = load_env_batch_kernel(&stream); + + let cfg = KernelConfig::for_test(); + let actions = open_hold_close_actions(); + + // q_values_per_window = NULL (0u64). num_actions = 0 so the kernel's + // open-time guard `num_actions > 0` short-circuits independently. + let (count, predicted_q) = run_one_chunk( + &stream, &kernel, &cfg, &actions, 0u64, 0, + ); + + assert_eq!( + count, 1, + "trade still closes regardless of q_values availability; got count={count}" + ); + let actual_q = predicted_q[0]; + assert!( + actual_q.abs() < Q_TOL, + "predicted_q[0] = {actual_q}, expected ~0.0 (NULL q_values_per_window means the \ + kernel skips the entry_q write at open; the persistent entry_q_state[w] stays at \ + buffer-init 0.0 and the close emission snapshots 0.0 into the tape)" + ); + } + + /// Test 3: explicit verification that the entry_q write is GATED on + /// `is_open || is_reverse_open` — a sequence with NO open (e.g., + /// flat-flat-flat-flat) should produce zero trades and never touch + /// `entry_q_state`. Provides an additional path-coverage check beyond + /// tests (1) and (2). + #[test] + #[ignore = "requires GPU"] + fn no_trades_no_predicted_q_emission() { + let stream = match try_cuda_stream() { + Some(s) => s, + None => return, + }; + let kernel = load_env_batch_kernel(&stream); + + let cfg = KernelConfig::for_test(); + // All Flat actions: no position ever opens, no close emission. + let a_flat = encode_action(DIR_FLAT, MAG_QUARTER, 0, 1); + let actions = vec![a_flat; CHUNK_LEN]; + + let q_host = q_values_with_open_at(0, 99.9); // won't be read + let mut q_dev = stream + .alloc_zeros::(q_host.len()) + .expect("alloc q_values"); + stream.memcpy_htod(&q_host, &mut q_dev).expect("HtoD q_values"); + let q_dev_ptr = { + let (ptr, _g) = q_dev.device_ptr(&stream); + ptr + }; + + let (count, _predicted_q) = run_one_chunk( + &stream, &kernel, &cfg, &actions, q_dev_ptr, NUM_ACTIONS, + ); + + drop(q_dev); + + assert_eq!( + count, 0, + "all-Flat action sequence should yield zero trades; got count={count}" + ); + } +}