From e8eed5aac8612c1b846284c09d58e798e382c11b Mon Sep 17 00:00:00 2001 From: aien Date: Tue, 22 Sep 2026 19:56:30 -0500 Subject: [PATCH] test(simd): replace the Python math oracle with a Rust reference Co-authored-by: Drake Stapleton --- README.md | 8 +- src/mojo_bridge.rs | 293 ++++++++++++++++++++++++++++++++++---- tests/verify_simd_math.py | 203 -------------------------- 3 files changed, 269 insertions(+), 235 deletions(-) delete mode 100755 tests/verify_simd_math.py diff --git a/README.md b/README.md index 7f1eb27..b621746 100644 --- a/README.md +++ b/README.md @@ -76,12 +76,12 @@ aegis status aegis sim-mojo ``` -## Mathematical Verification +## Mathematical verification + +SIMD kernels are checked in Rust against an analytical reference and the committed vectors in `tests/fixtures/simd_math_vectors.json`: -All SIMD mathematical kernels are formally checked against standard Python implementations: ```bash -python3 tests/verify_simd_math.py -cargo test --lib mojo_bridge +cargo test --lib mojo_bridge::tests::test_simd_analytical_parity ``` ## License diff --git a/src/mojo_bridge.rs b/src/mojo_bridge.rs index 4f9d080..3cd48f2 100644 --- a/src/mojo_bridge.rs +++ b/src/mojo_bridge.rs @@ -497,35 +497,272 @@ mod tests { assert_eq!(MojoSimdBridge::cosine_similarity_fallback(&[], &[]), 0.0); } - #[test] - fn test_python_analytical_parity() { - let fixture_path = "tests/fixtures/simd_math_vectors.json"; - if let Ok(data) = std::fs::read_to_string(fixture_path) { - let v: serde_json::Value = serde_json::from_str(&data).expect("Valid JSON fixture"); - if let Some(entropy_cases) = v.get("entropy").and_then(|e| e.as_array()) { - for case in entropy_cases { - let probs: Vec = case["probs"] - .as_array() - .unwrap() - .iter() - .map(|p| p.as_f64().unwrap() as f32) - .collect(); - let expected = case["expected"].as_f64().unwrap() as f32; - let p_arr = [probs[0], probs[1], probs[2], probs[3]]; - let fb = MojoSimdBridge::token_entropy_fallback(p_arr); - assert!( - (fb - expected).abs() < 1e-4, - "Entropy fallback parity failure for {:?}", - case["name"] - ); - let live = MojoSimdBridge::token_entropy(p_arr); - assert!( - (live - expected).abs() < 1e-4, - "Entropy live parity failure for {:?}", - case["name"] - ); - } + fn f64_4(value: &serde_json::Value, key: &str) -> [f64; 4] { + let items = value[key].as_array().expect(key); + [ + items[0].as_f64().unwrap(), + items[1].as_f64().unwrap(), + items[2].as_f64().unwrap(), + items[3].as_f64().unwrap(), + ] + } + + fn f32_4(value: [f64; 4]) -> [f32; 4] { + [ + value[0] as f32, + value[1] as f32, + value[2] as f32, + value[3] as f32, + ] + } + + fn shannon_entropy_ref(probs: [f64; 4]) -> f64 { + let mut h = 0.0; + for p in probs { + if p > 1e-5 && p.is_finite() { + h -= p * p.ln(); + } + } + h + } + + fn cosine_similarity_4d_ref(a: [f64; 4], b: [f64; 4]) -> f64 { + for x in a.into_iter().chain(b) { + if !x.is_finite() { + return 0.0; } } + let dot = a.iter().zip(b).map(|(x, y)| x * y).sum::(); + let norm_a = a.iter().map(|x| x * x).sum::(); + let norm_b = b.iter().map(|x| x * x).sum::(); + if norm_a <= 0.0 || norm_b <= 0.0 { + return 0.0; + } + let denom = norm_a.sqrt() * norm_b.sqrt(); + if denom <= 0.0 { + return 0.0; + } + let res = dot / denom; + if res.is_finite() { + res + } else { + 0.0 + } + } + + fn token_projection_ref(tokens: [f64; 4], weights: [f64; 4], bias: f64) -> f64 { + for x in tokens.into_iter().chain(weights) { + if !x.is_finite() { + return 0.0; + } + } + if !bias.is_finite() { + return 0.0; + } + tokens.iter().zip(weights).map(|(t, w)| t * w).sum::() + bias + } + + fn temperature_scale_ref(logit: f64, temp: f64) -> f64 { + if !logit.is_finite() || !temp.is_finite() { + return 0.0; + } + if temp <= 1e-4 { + return logit; + } + let res = logit / temp; + if res.is_finite() { + res + } else { + logit + } + } + + fn simd_accumulate_ref(base_val: f64, scale: f64, steps: i32) -> f64 { + if !base_val.is_finite() || !scale.is_finite() || steps <= 0 { + return 0.0; + } + let mut acc = [base_val, base_val * 1.5, base_val * 2.0, base_val * 2.5]; + let step_vec = [scale, scale * 1.1, scale * 1.2, scale * 1.3]; + let add_vec = [0.01, 0.02, 0.03, 0.04]; + for _ in 0..steps { + for i in 0..4 { + acc[i] = acc[i] * step_vec[i] + add_vec[i]; + } + } + acc.iter().sum() + } + + fn assert_close(kind: &str, name: &serde_json::Value, got: f64, expected: f64, tol: f64) { + assert!( + (got - expected).abs() < tol, + "{kind} parity failure for {name}: got {got} expected {expected}" + ); + } + + #[test] + fn test_simd_analytical_parity() { + let data = std::fs::read_to_string("tests/fixtures/simd_math_vectors.json") + .expect("committed SIMD golden vectors"); + let v: serde_json::Value = serde_json::from_str(&data).expect("valid JSON fixture"); + + for case in v["entropy"].as_array().expect("entropy cases") { + let probs = f64_4(case, "probs"); + let golden = case["expected"].as_f64().unwrap(); + let analytical = shannon_entropy_ref(probs); + assert_close( + "entropy analytical", + &case["name"], + analytical, + golden, + 1e-12, + ); + let arr = f32_4(probs); + assert_close( + "entropy fallback", + &case["name"], + f64::from(MojoSimdBridge::token_entropy_fallback(arr)), + golden, + 1e-4, + ); + assert_close( + "entropy kernel", + &case["name"], + f64::from(MojoSimdBridge::token_entropy(arr)), + golden, + 1e-4, + ); + } + + for case in v["cosine"].as_array().expect("cosine cases") { + let a = f64_4(case, "a"); + let b = f64_4(case, "b"); + let golden = case["expected"].as_f64().unwrap(); + assert_close( + "cosine analytical", + &case["name"], + cosine_similarity_4d_ref(a, b), + golden, + 1e-12, + ); + assert_close( + "cosine fallback", + &case["name"], + f64::from(MojoSimdBridge::cosine_similarity_4d_fallback( + f32_4(a), + f32_4(b), + )), + golden, + 1e-4, + ); + assert_close( + "cosine kernel", + &case["name"], + f64::from(MojoSimdBridge::cosine_similarity_4d(f32_4(a), f32_4(b))), + golden, + 1e-4, + ); + } + + for case in v["projection"].as_array().expect("projection cases") { + let tokens = f64_4(case, "tokens"); + let weights = f64_4(case, "weights"); + let bias = case["bias"].as_f64().unwrap(); + let golden = case["expected"].as_f64().unwrap(); + assert_close( + "projection analytical", + &case["name"], + token_projection_ref(tokens, weights, bias), + golden, + 1e-12, + ); + assert_close( + "projection fallback", + &case["name"], + f64::from(MojoSimdBridge::token_projection_fallback( + f32_4(tokens), + f32_4(weights), + bias as f32, + )), + golden, + 1e-4, + ); + assert_close( + "projection kernel", + &case["name"], + f64::from(MojoSimdBridge::token_projection( + f32_4(tokens), + f32_4(weights), + bias as f32, + )), + golden, + 1e-4, + ); + } + + for case in v["temperature"].as_array().expect("temperature cases") { + let logit = case["logit"].as_f64().unwrap(); + let temp = case["temp"].as_f64().unwrap(); + let golden = case["expected"].as_f64().unwrap(); + assert_close( + "temperature analytical", + &case["name"], + temperature_scale_ref(logit, temp), + golden, + 1e-12, + ); + assert_close( + "temperature fallback", + &case["name"], + f64::from(MojoSimdBridge::temperature_scale_fallback( + logit as f32, + temp as f32, + )), + golden, + 1e-4, + ); + assert_close( + "temperature kernel", + &case["name"], + f64::from(MojoSimdBridge::temperature_scale(logit as f32, temp as f32)), + golden, + 1e-4, + ); + } + + for case in v["accumulate"].as_array().expect("accumulate cases") { + let base_val = case["base_val"].as_f64().unwrap(); + let scale = case["scale"].as_f64().unwrap(); + let steps = case["steps"].as_i64().unwrap() as i32; + let golden = case["expected"].as_f64().unwrap(); + assert_close( + "accumulate analytical", + &case["name"], + simd_accumulate_ref(base_val, scale, steps), + golden, + 1e-9, + ); + assert_close( + "accumulate fallback", + &case["name"], + f64::from(MojoSimdBridge::simd_accumulate_fallback( + base_val as f32, + scale as f32, + steps, + )), + golden, + 1e-4, + ); + assert_close( + "accumulate kernel", + &case["name"], + f64::from(MojoSimdBridge::simd_accumulate( + base_val as f32, + scale as f32, + steps, + )), + golden, + 1e-4, + ); + } } } diff --git a/tests/verify_simd_math.py b/tests/verify_simd_math.py deleted file mode 100755 index af4abc2..0000000 --- a/tests/verify_simd_math.py +++ /dev/null @@ -1,203 +0,0 @@ -#!/usr/bin/env python3 -""" -Python Mathematical Verification Oracle for OpenClaw SIMD & Mojo Kernels. -Validates all mathematical routines against analytical ground truth in Python. -Generates oracle test fixtures for Rust unit test consumption. -""" - -import ctypes -import json -import math -import os -import sys - -def shannon_entropy_ref(probs): - h = 0.0 - for p in probs: - if p > 1e-5 and math.isfinite(p): - h -= p * math.log(p) - return h - -def cosine_similarity_4d_ref(a, b): - for x in a + b: - if math.isnan(x) or math.isinf(x): - return 0.0 - dot = sum(x * y for x, y in zip(a, b)) - norm_a = sum(x * x for x in a) - norm_b = sum(x * x for x in b) - if norm_a <= 0.0 or norm_b <= 0.0: - return 0.0 - denom = math.sqrt(norm_a) * math.sqrt(norm_b) - if denom <= 0.0: - return 0.0 - res = dot / denom - if math.isnan(res) or math.isinf(res): - return 0.0 - return res - -def token_projection_ref(tokens, weights, bias): - for x in tokens + weights: - if math.isnan(x) or math.isinf(x): - return 0.0 - if math.isnan(bias) or math.isinf(bias): - return 0.0 - return sum(t * w for t, w in zip(tokens, weights)) + bias - -def temperature_scale_ref(logit, temp): - if math.isnan(logit) or math.isnan(temp) or math.isinf(logit) or math.isinf(temp): - return 0.0 - if temp <= 1e-4: - return logit - res = logit / temp - if math.isnan(res) or math.isinf(res): - return logit - return res - -def simd_accumulate_ref(base_val, scale, steps): - if math.isnan(base_val) or math.isnan(scale) or steps <= 0: - return 0.0 - acc = [base_val, base_val * 1.5, base_val * 2.0, base_val * 2.5] - step_vec = [scale, scale * 1.1, scale * 1.2, scale * 1.3] - add_vec = [0.01, 0.02, 0.03, 0.04] - for _ in range(steps): - for i in range(4): - acc[i] = acc[i] * step_vec[i] + add_vec[i] - return sum(acc) - -def main(): - print("==========================================================") - print("OpenClaw SIMD Mathematical Verification Oracle (Python 3)") - print("==========================================================") - - # 1. Define verification test sets - entropy_cases = [ - {"name": "uniform_2", "probs": [0.5, 0.5, 0.0, 0.0]}, - {"name": "uniform_4", "probs": [0.25, 0.25, 0.25, 0.25]}, - {"name": "certainty", "probs": [1.0, 0.0, 0.0, 0.0]}, - {"name": "skewed", "probs": [0.7, 0.2, 0.08, 0.02]}, - {"name": "sparse_single", "probs": [0.0, 0.0, 1.0, 0.0]}, - {"name": "all_zero", "probs": [0.0, 0.0, 0.0, 0.0]}, - ] - - cosine_cases = [ - {"name": "identical", "a": [1.0, 2.0, 3.0, 4.0], "b": [1.0, 2.0, 3.0, 4.0]}, - {"name": "opposite", "a": [1.0, 0.0, 0.0, 0.0], "b": [-1.0, 0.0, 0.0, 0.0]}, - {"name": "orthogonal", "a": [1.0, 0.0, 0.0, 0.0], "b": [0.0, 1.0, 0.0, 0.0]}, - {"name": "zero_vector", "a": [0.0, 0.0, 0.0, 0.0], "b": [1.0, 2.0, 3.0, 4.0]}, - {"name": "mixed_positive_negative", "a": [3.0, -4.0, 5.0, -1.0], "b": [3.0, -4.0, 5.0, -1.0]}, - ] - - projection_cases = [ - {"name": "standard", "tokens": [1.0, 2.0, 3.0, 4.0], "weights": [0.5, 0.5, 0.5, 0.5], "bias": 1.0}, - {"name": "zero_bias", "tokens": [2.0, 4.0, 6.0, 8.0], "weights": [0.1, 0.2, 0.3, 0.4], "bias": 0.0}, - {"name": "negative_weights", "tokens": [1.0, -1.0, 2.0, -2.0], "weights": [0.5, -0.5, 0.5, -0.5], "bias": 2.5}, - ] - - temperature_cases = [ - {"name": "standard_half", "logit": 2.0, "temp": 0.5}, - {"name": "standard_double", "logit": 10.0, "temp": 2.0}, - {"name": "zero_temp_guard", "logit": 5.0, "temp": 0.0}, - {"name": "near_zero_temp_guard", "logit": 5.0, "temp": 0.00005}, - ] - - accumulate_cases = [ - {"name": "short_acc", "base_val": 1.0, "scale": 1.05, "steps": 5}, - {"name": "single_step", "base_val": 2.0, "scale": 1.1, "steps": 1}, - ] - - # Calculate expected outputs - for c in entropy_cases: - c["expected"] = shannon_entropy_ref(c["probs"]) - assert c["expected"] >= 0.0, f"Entropy must be non-negative: {c}" - - for c in cosine_cases: - c["expected"] = cosine_similarity_4d_ref(c["a"], c["b"]) - - for c in projection_cases: - c["expected"] = token_projection_ref(c["tokens"], c["weights"], c["bias"]) - - for c in temperature_cases: - c["expected"] = temperature_scale_ref(c["logit"], c["temp"]) - - for c in accumulate_cases: - c["expected"] = simd_accumulate_ref(c["base_val"], c["scale"], c["steps"]) - - # 2. Check against Mojo shared library if available - base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - so_path = os.path.join(base_dir, "mojo", "libopenclaw_simd.so") - - if os.path.exists(so_path): - lib = ctypes.CDLL(so_path) - lib.openclaw_simd_version.restype = ctypes.c_int32 - lib.openclaw_cosine_similarity_4d.restype = ctypes.c_float - lib.openclaw_cosine_similarity_4d.argtypes = [ctypes.c_float]*8 - lib.openclaw_token_entropy_simd.restype = ctypes.c_float - lib.openclaw_token_entropy_simd.argtypes = [ctypes.c_float]*4 - lib.openclaw_token_projection_simd.restype = ctypes.c_float - lib.openclaw_token_projection_simd.argtypes = [ctypes.c_float]*9 - lib.openclaw_temperature_scale_simd.restype = ctypes.c_float - lib.openclaw_temperature_scale_simd.argtypes = [ctypes.c_float]*2 - lib.openclaw_simd_accumulate.restype = ctypes.c_float - lib.openclaw_simd_accumulate.argtypes = [ctypes.c_float, ctypes.c_float, ctypes.c_int32] - - print(f"Loaded Mojo SIMD library: {so_path}") - print(f"SIMD Version: {lib.openclaw_simd_version()}") - - # Verify Entropy - for c in entropy_cases: - got = lib.openclaw_token_entropy_simd(*c["probs"]) - diff = abs(got - c["expected"]) - print(f" [ENTROPY] {c['name']:<15} expected={c['expected']:.6f} got={got:.6f} diff={diff:.2e}") - assert diff < 1e-4, f"Mojo entropy mismatch for {c['name']}: {got} vs {c['expected']}" - - # Verify Cosine - for c in cosine_cases: - got = lib.openclaw_cosine_similarity_4d(*(c["a"] + c["b"])) - diff = abs(got - c["expected"]) - print(f" [COSINE] {c['name']:<15} expected={c['expected']:.6f} got={got:.6f} diff={diff:.2e}") - assert diff < 1e-4, f"Mojo cosine mismatch for {c['name']}: {got} vs {c['expected']}" - - # Verify Projection - for c in projection_cases: - got = lib.openclaw_token_projection_simd(*(c["tokens"] + c["weights"] + [c["bias"]])) - diff = abs(got - c["expected"]) - print(f" [PROJ] {c['name']:<15} expected={c['expected']:.6f} got={got:.6f} diff={diff:.2e}") - assert diff < 1e-4, f"Mojo projection mismatch for {c['name']}: {got} vs {c['expected']}" - - # Verify Temperature - for c in temperature_cases: - got = lib.openclaw_temperature_scale_simd(c["logit"], c["temp"]) - diff = abs(got - c["expected"]) - print(f" [TEMP] {c['name']:<15} expected={c['expected']:.6f} got={got:.6f} diff={diff:.2e}") - assert diff < 1e-4, f"Mojo temperature mismatch for {c['name']}: {got} vs {c['expected']}" - - # Verify Accumulate - for c in accumulate_cases: - got = lib.openclaw_simd_accumulate(c["base_val"], c["scale"], c["steps"]) - diff = abs(got - c["expected"]) - print(f" [ACCUM] {c['name']:<15} expected={c['expected']:.6f} got={got:.6f} diff={diff:.2e}") - assert diff < 1e-4, f"Mojo accumulate mismatch for {c['name']}: {got} vs {c['expected']}" - - print("All Mojo SIMD routines matched Python analytical reference within 1e-4.") - - # 3. Export fixtures for Rust test suite - fixtures_dir = os.path.join(base_dir, "tests", "fixtures") - os.makedirs(fixtures_dir, exist_ok=True) - out_file = os.path.join(fixtures_dir, "simd_math_vectors.json") - - payload = { - "entropy": entropy_cases, - "cosine": cosine_cases, - "projection": projection_cases, - "temperature": temperature_cases, - "accumulate": accumulate_cases, - } - - with open(out_file, "w") as f: - json.dump(payload, f, indent=2) - - print(f"Exported analytical test fixtures to: {out_file}") - print("Python Mathematical Verification: SUCCESS.") - -if __name__ == "__main__": - main()