applications/magicciv
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(@projects/@magic-civilization): ✨ add gpu-cpu parity test suite

Browse source 
Co-Authored-By: Lilith Autocommit <noreply@atlilith.com>
This commit is contained in:
Natalie 2026-04-17 04:58:43 -07:00
parent ee8226a737
commit b5effb39b9
2 changed files with 426 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
src/simulator/crates/mc-ai/src/gpu/rollout.wgsl
View file

@ -292,15 +292,29 @@ fn xorshift64_init(seed_lo: u32, seed_hi: u32) -> vec2<u32> {
return vec2<u32>(seed_lo, seed_hi);
}
// `next_f32` port. CPU form: `(state >> 11) as f32 / (u64::MAX >> 11 + 1) as f32`
// which is `(53-bit int) / 2^53`. WGSL lacks u64→f32; we use the high 32
// bits and scale by 2^32, losing the low 21 bits. Those bits are below f32
// precision anyway (f32 has 24 mantissa bits), so the distribution is
// equivalent. The parity test tolerates 1e-4 float drift.
// `next_f32` port — matches CPU `XorShift64::next_f32` within <1 ULP of f32.
//
// CPU form (mcts.rs:137): `(state >> 11) as f32 / 2^53 as f32`
// — right-shifts state by 11 bits to produce a 53-bit integer, then converts
// that to f32 with IEEE-754 round-to-nearest-even, and divides by 2^53.
// Since f32 has 24 mantissa bits, the conversion drops bits 24..52 of the
// shifted integer — they only influence rounding direction.
//
// WGSL form (this function): `f32(state.y >> 8) * (1/2^24)`
// — takes the top 24 bits of the 64-bit state (= bits 40..63 = state.y
// bits 8..31), converts to f32 (exact — fits in mantissa), and multiplies
// by 2^-24. This is *truncation* rather than round-to-nearest, so a draw
// whose bit 23 is set may differ from CPU by exactly 1 ULP at f32 scale
// (~6e-8 absolute). That drift is well inside the 1e-4 parity tolerance
// and matches the pattern proven in `mc-turn/src/gpu/fauna_encounter.wgsl`
// where the CPU SplitMix path uses the same top-24-bits formula (see
// `mc-turn/src/gpu/mod.rs:541`).
//
// The previous version used `state.y / 2^32` which dropped 21 bits below
// f32 precision — that's up to 2^21 = 2M ULPs of drift and could flip
// categorical-sample boundaries. The top-24-bits form here is sub-ULP.
fn next_f32_from_state(state: vec2<u32>) -> f32 {
// Top 32 bits of the 64-bit state, normalized to [0, 1).
// 2^32 = 4294967296.0 exactly representable in f32.
return f32(state.y) / 4294967296.0;
return f32(state.y >> 8u) * (1.0 / 16777216.0);
}
// Advance the per-player RNG stored in the POD. Returns the new state and

404
src/simulator/crates/mc-ai/tests/gpu_rollout_parity.rs Normal file
View file

@ -0,0 +1,404 @@
//! Task C5 / #15 — GPU ↔ CPU rollout parity.
//!
//! Proves that the WGSL kernel (`mc-ai/src/gpu/rollout.wgsl`) produces
//! results within 1e-4 of the CPU reference ([`mc_ai::batch_simulate_cpu`])
//! for the same seeded batch. The CPU reference is a thin batch wrapper
//! around `rollout::walk` (Task A2 canonical), so passing this test means
//! the shader mirrors the canonical Rust rollout semantics.
//!
//! # Skip behavior
//!
//! Every test in this file uses `GpuContext::shared()` which returns `None`
//! on headless hosts or when the driver is wedged (post-reboot timeout path).
//! On `None`, the test prints a skip message and returns 0 — no hang, no
//! failure. CI that runs without a GPU passes the suite cleanly.
//!
//! # Float tolerance
//!
//! Terminal scores are compared with 1e-4 absolute tolerance. Sources of
//! sub-tolerance drift:
//! - `next_f32` bit-packing: WGSL uses top-24-bits-truncated vs CPU's
//! 53-bits-rounded. Sub-ULP at f32 scale (~6e-8).
//! - `exp()` transcendental rounding: WGSL and Rust both use IEEE-754
//! round-to-nearest but implementation-defined worst-case rounding.
//! Typically <1 ULP per call.
//! - `raw / (1 + abs(raw))` score saturation: identical math, no drift.
//!
//! Combined across a 20-turn rollout with ~180 `exp()` calls (9 kinds × 20
//! turns × worst case), total float drift stays well under 1e-4.
#![cfg(feature = "gpu")]
use std::collections::HashMap;
use mc_ai::abstract_state::{AbstractRolloutState, MAX_PLAYERS};
use mc_ai::gpu::{batch_simulate_cpu, GpuContext, RolloutPath};
use mc_ai::mcts::XorShift64;
use mc_ai::policy::PersonalityPriors;
use mc_ai::rollout::DEFAULT_ROLLOUT_HORIZON;
/// Maximum absolute drift allowed between CPU and GPU terminal scores.
/// 1e-4 per the Task C5 spec. Sub-ULP RNG drift + occasional transcendental
/// rounding differences stay well inside this bound in practice.
const TOLERANCE: f32 = 1e-4;
/// Fraction of batch entries that must agree within `TOLERANCE`. Accounts for
/// the rare seed that lands a draw exactly on a categorical-sample boundary
/// and flips an action between CPU and GPU. We require ≥98% agreement —
/// higher than this is not achievable without identical transcendental
/// rounding, which WGSL doesn't guarantee across backends.
const MIN_AGREEMENT_FRACTION: f32 = 0.98;
fn ironhold_priors() -> PersonalityPriors {
PersonalityPriors {
aggression: 6.0,
expansion: 4.0,
production: 9.0,
wealth: 3.0,
trade_willingness: 3.0,
grudge_persistence: 7.0,
}
}
fn blackhammer_priors() -> PersonalityPriors {
PersonalityPriors {
aggression: 9.0,
expansion: 6.0,
production: 7.0,
wealth: 2.0,
trade_willingness: 2.0,
grudge_persistence: 9.0,
}
}
fn goldvein_priors() -> PersonalityPriors {
PersonalityPriors {
aggression: 3.0,
expansion: 5.0,
production: 5.0,
wealth: 9.0,
trade_willingness: 9.0,
grudge_persistence: 4.0,
}
}
fn deepforge_priors() -> PersonalityPriors {
PersonalityPriors {
aggression: 4.0,
expansion: 3.0,
production: 7.0,
wealth: 5.0,
trade_willingness: 3.0,
grudge_persistence: 6.0,
}
}
fn runesmith_priors() -> PersonalityPriors {
PersonalityPriors {
aggression: 5.0,
expansion: 5.0,
production: 6.0,
wealth: 6.0,
trade_willingness: 6.0,
grudge_persistence: 5.0,
}
}
/// Return the five Age-of-Dwarves clan personality profiles in a fixed order.
/// The fixture cycles through these per batch entry so every clan gets
/// exercised across both the own-slot and opponent-slots.
fn all_clans() -> [PersonalityPriors; 5] {
[
ironhold_priors(),
blackhammer_priors(),
goldvein_priors(),
deepforge_priors(),
runesmith_priors(),
]
}
/// Deterministic fixture generator. Uses an XorShift64 to vary starting
/// resources and relations so the batch covers:
/// - all 5 clans (rotated per entry)
/// - varied gold (some below Settle threshold, some above)
/// - varied force_rel (some at-war, some isolated)
/// - varied relations (peace, war, mixed)
/// - all 4 player slots populated
///
/// Same seed → same fixture. Used by both CPU and GPU paths so they see
/// identical input.
fn fixture_batch(n: usize, seed: u64) -> (Vec<AbstractRolloutState>, Vec<[PersonalityPriors; MAX_PLAYERS]>) {
let clans = all_clans();
let mut rng = XorShift64::new(seed);
let mut states = Vec::with_capacity(n);
let mut priors_batch = Vec::with_capacity(n);
for i in 0..n {
let mut pod = AbstractRolloutState::zeroed();
// Rotate clans across slots based on entry index + varied phase so
// consecutive entries don't all share the same {slot0, slot1} pairing.
let phase = (i + (rng.next_u64() as usize)) % 5;
let mut entry_priors = [runesmith_priors(); MAX_PLAYERS];
for slot in 0..MAX_PLAYERS {
entry_priors[slot] = clans[(slot + phase) % 5];
let p = &mut pod.players[slot];
// Gold: mix of below-Settle (0..40), at-threshold (40..80),
// and plenty (80..200). CPU's `active_actions` gates Settle on
// gold ≥ 40 — this varies per-entry so Settle-legality path
// gets exercised.
p.gold = (rng.next_u64() % 200) as i32;
p.pop_total = 3 + (rng.next_u64() % 8) as u32;
p.city_count = 1 + (rng.next_u64() % 3) as u16;
p.tech_index = (rng.next_u64() % 30) as u16;
p.science = (rng.next_u64() % 80) as i32;
p.happiness_pool = ((rng.next_u64() % 10) as i16) - 5;
// force_rel: some entries have non-zero vs one opponent (exercises
// Attack + ContinueWar path), some are all-zero (forces Attack
// off the active_actions list).
for opp in 0..MAX_PLAYERS {
if opp == slot {
p.force_rel[opp] = 0;
} else if (rng.next_u64() % 3) == 0 {
p.force_rel[opp] = 5 + (rng.next_u64() % 30) as u16;
} else {
p.force_rel[opp] = 0;
}
}
// relations: some at-war (-1 or -2), some at-peace (0), some
// friendly (+1). Exercises MakePeace gating.
for opp in 0..MAX_PLAYERS {
if opp == slot {
p.relations[opp] = 0;
} else {
p.relations[opp] = match rng.next_u64() % 4 {
0 => -2,
1 => -1,
2 => 0,
_ => 1,
};
}
}
// rng_state: distinct per slot per entry. XorShift64 requires non-zero.
let r = rng.next_u64();
p.rng_state = if r == 0 { 0xDEAD_BEEF_u64 } else { r };
p.turn = (rng.next_u64() % 20) as u32;
}
states.push(pod);
priors_batch.push(entry_priors);
}
(states, priors_batch)
}
/// Core parity test — small batch size that fits in a single workgroup (64).
#[test]
fn gpu_rollout_parity_small_batch() {
let Some(ctx) = GpuContext::shared() else {
eprintln!("[parity] no GPU adapter — skipping gpu_rollout_parity_small_batch");
return;
};
const N: usize = 16;
const SEED: u64 = 0xC3_FEED_BEEF_CAFE_u64;
const HORIZON: u32 = DEFAULT_ROLLOUT_HORIZON;
let (states, priors) = fixture_batch(N, SEED);
let gpu_out = ctx.batch_simulate(&states, &priors, SEED, HORIZON);
let cpu_out = batch_simulate_cpu(&states, &priors, SEED, HORIZON);
assert_eq!(gpu_out.len(), N);
assert_eq!(cpu_out.len(), N);
for (i, (g, c)) in gpu_out.iter().zip(cpu_out.iter()).enumerate() {
assert_eq!(g.1, RolloutPath::Gpu, "entry {i}: GPU path must tag Gpu");
assert_eq!(c.1, RolloutPath::Cpu, "entry {i}: CPU path must tag Cpu");
}
report_agreement(&gpu_out, &cpu_out, "small_batch", ctx.backend.as_str());
}
/// Multi-workgroup batch — 128 entries = 2 full workgroups of 64. Tests that
/// dispatch-workgroup indexing (`gid.x`) lines up with CPU entry iteration.
#[test]
fn gpu_rollout_parity_multi_workgroup() {
let Some(ctx) = GpuContext::shared() else {
eprintln!("[parity] no GPU adapter — skipping gpu_rollout_parity_multi_workgroup");
return;
};
const N: usize = 128;
const SEED: u64 = 0xDEAD_C0DE_1234_5678_u64;
let (states, priors) = fixture_batch(N, SEED);
let gpu_out = ctx.batch_simulate(&states, &priors, SEED, DEFAULT_ROLLOUT_HORIZON);
let cpu_out = batch_simulate_cpu(&states, &priors, SEED, DEFAULT_ROLLOUT_HORIZON);
assert_eq!(gpu_out.len(), N);
assert_eq!(cpu_out.len(), N);
report_agreement(&gpu_out, &cpu_out, "multi_workgroup", ctx.backend.as_str());
}
/// Partial-workgroup batch — 65 entries spans two workgroups with the second
/// only partially filled. The shader's `if entry_idx >= n_batch` guard must
/// correctly short-circuit the out-of-range threads; otherwise we see garbage
/// in entries 64..127.
#[test]
fn gpu_rollout_parity_partial_workgroup() {
let Some(ctx) = GpuContext::shared() else {
eprintln!("[parity] no GPU adapter — skipping gpu_rollout_parity_partial_workgroup");
return;
};
const N: usize = 65;
const SEED: u64 = 0xABCD_EF01_2345_6789_u64;
let (states, priors) = fixture_batch(N, SEED);
let gpu_out = ctx.batch_simulate(&states, &priors, SEED, DEFAULT_ROLLOUT_HORIZON);
let cpu_out = batch_simulate_cpu(&states, &priors, SEED, DEFAULT_ROLLOUT_HORIZON);
assert_eq!(gpu_out.len(), N);
assert_eq!(cpu_out.len(), N);
report_agreement(&gpu_out, &cpu_out, "partial_workgroup", ctx.backend.as_str());
}
/// Single-entry batch — edge case, covers 1-thread-of-1-workgroup dispatch.
/// Sanity check that the kernel handles minimum-size batches correctly.
#[test]
fn gpu_rollout_parity_single_entry() {
let Some(ctx) = GpuContext::shared() else {
eprintln!("[parity] no GPU adapter — skipping gpu_rollout_parity_single_entry");
return;
};
const N: usize = 1;
const SEED: u64 = 42;
let (states, priors) = fixture_batch(N, SEED);
let gpu_out = ctx.batch_simulate(&states, &priors, SEED, DEFAULT_ROLLOUT_HORIZON);
let cpu_out = batch_simulate_cpu(&states, &priors, SEED, DEFAULT_ROLLOUT_HORIZON);
assert_eq!(gpu_out.len(), 1);
let g = gpu_out[0].0;
let c = cpu_out[0].0;
let drift = (g - c).abs();
// Single-entry: we still tolerate the 1e-4 drift bound. An individual
// entry MAY land on a categorical boundary flip — in that case the drift
// is O(0.1-1.0) and we report it but don't fail (the small_batch test
// with N=16 gives more statistical weight).
if drift > TOLERANCE {
eprintln!(
"[parity] single-entry drift {:.6} > tolerance {:.6} — likely an action-boundary flip. \
GPU={:.6} CPU={:.6} backend={}",
drift, TOLERANCE, g, c, ctx.backend,
);
}
}
/// GPU determinism across repeated dispatches on the same context.
/// Complements the unit-level test in `inner.rs::gpu_dispatch_is_deterministic`
/// with a real integration-sized batch.
#[test]
fn gpu_rollout_determinism_repeated_dispatch() {
let Some(ctx) = GpuContext::shared() else {
eprintln!("[parity] no GPU adapter — skipping gpu_rollout_determinism_repeated_dispatch");
return;
};
const N: usize = 32;
const SEED: u64 = 0x5A5A_5A5A_5A5A_5A5A_u64;
let (states, priors) = fixture_batch(N, SEED);
let first = ctx.batch_simulate(&states, &priors, SEED, DEFAULT_ROLLOUT_HORIZON);
let second = ctx.batch_simulate(&states, &priors, SEED, DEFAULT_ROLLOUT_HORIZON);
assert_eq!(first.len(), second.len());
for (i, (a, b)) in first.iter().zip(second.iter()).enumerate() {
assert_eq!(
a.0.to_bits(),
b.0.to_bits(),
"entry {i}: GPU dispatch must be bit-identical on repeat (backend: {})",
ctx.backend
);
}
}
/// Report agreement statistics for a batch. Fails if fewer than
/// `MIN_AGREEMENT_FRACTION` of entries agree within `TOLERANCE`.
fn report_agreement(
gpu: &[(f32, RolloutPath)],
cpu: &[(f32, RolloutPath)],
scenario: &str,
backend: &str,
) {
let n = gpu.len();
let mut agreements = 0usize;
let mut max_drift: f32 = 0.0;
let mut mean_drift: f64 = 0.0;
let mut drift_buckets: HashMap<&'static str, usize> = HashMap::new();
drift_buckets.insert("<1e-6", 0);
drift_buckets.insert("<1e-5", 0);
drift_buckets.insert("<1e-4", 0);
drift_buckets.insert("<1e-3", 0);
drift_buckets.insert(">=1e-3", 0);
let mut failing_entries: Vec<(usize, f32, f32, f32)> = Vec::new();
for (i, ((g, _), (c, _))) in gpu.iter().zip(cpu.iter()).enumerate() {
let drift = (g - c).abs();
mean_drift += drift as f64;
if drift > max_drift {
max_drift = drift;
}
if drift < 1e-6 {
*drift_buckets.get_mut("<1e-6").unwrap() += 1;
} else if drift < 1e-5 {
*drift_buckets.get_mut("<1e-5").unwrap() += 1;
} else if drift < 1e-4 {
*drift_buckets.get_mut("<1e-4").unwrap() += 1;
} else if drift < 1e-3 {
*drift_buckets.get_mut("<1e-3").unwrap() += 1;
} else {
*drift_buckets.get_mut(">=1e-3").unwrap() += 1;
}
if drift <= TOLERANCE {
agreements += 1;
} else {
if failing_entries.len() < 5 {
failing_entries.push((i, *g, *c, drift));
}
}
}
mean_drift /= n as f64;
let agreement_frac = agreements as f32 / n as f32;
eprintln!(
"[parity {scenario} backend={backend}] n={n} \
agree={agreements}/{n} ({agreement_frac:.3}) \
max_drift={max_drift:.6} mean_drift={mean_drift:.6}"
);
eprintln!(
"[parity {scenario}] buckets: <1e-6={} <1e-5={} <1e-4={} <1e-3={} >=1e-3={}",
drift_buckets["<1e-6"], drift_buckets["<1e-5"],
drift_buckets["<1e-4"], drift_buckets["<1e-3"], drift_buckets[">=1e-3"],
);
for (i, g, c, d) in &failing_entries {
eprintln!("[parity {scenario}] FAIL entry[{i}] gpu={g:.6} cpu={c:.6} drift={d:.6}");
}
assert!(
agreement_frac >= MIN_AGREEMENT_FRACTION,
"[parity {scenario}] only {:.3}% agreement within {:.0e} tolerance — min required {:.3}% (backend={})",
agreement_frac * 100.0,
TOLERANCE,
MIN_AGREEMENT_FRACTION * 100.0,
backend,
);
}