"""Host-only tests; CUDA equality requires separate real-device evidence.""" from __future__ import annotations from dataclasses import replace import importlib.util import math import os from pathlib import Path import sys import unittest import numpy as np from gpu_backend import (CUDA_OPTIONS, CUDA_SOURCE, COMPARTMENTS, CudaEligibility, cpu_reference, prepare_events) ROOT = Path(__file__).resolve().parents[1] PHASE4J = Path(os.environ.get("PHASE4J_SOURCE", str(ROOT / "models" / "male-cns-lif" / "phase4j.py"))) BRIDGE_PATH = Path(os.environ.get("ELIGIBILITY_BRIDGE_SOURCE", str(ROOT / "wholebrain-plastic-interface-v1" / "eligibility_bridge.py"))) spec = importlib.util.spec_from_file_location("_gpu_math_frozen_bridge", BRIDGE_PATH) bridge = importlib.util.module_from_spec(spec) sys.modules[spec.name] = bridge spec.loader.exec_module(bridge) def assert_bits(test, left, right): test.assertEqual(left.dtype, right.dtype) test.assertEqual(left.shape, right.shape) test.assertEqual(left.tobytes(), right.tobytes()) class HostMathTests(unittest.TestCase): def setUp(self): self.kc = np.array([2, 5, 9], np.int32) self.step = np.array([0, 0, 1, 2, 2, 5, 7], np.int32) self.neurons = np.array([2, 4, 2, 5, 9, 9, 2], np.int32) self.sources = np.array([2, 2, 5, 9], np.int32) self.compartments = np.array([COMPARTMENTS[0], COMPARTMENTS[1], COMPARTMENTS[0], COMPARTMENTS[0]]) self.base = np.array([0.001, 0.01, 0.004, 0.08], np.float32) self.delta = np.zeros(4, np.float32) self.kw = dict(d=1.0, frozen=False, eta_mv=bridge.ETA_MV, floor_fraction=bridge.FLOOR_FRACTION, decay=math.exp(-bridge.DT_MS / bridge.TAU_MS)) def prepared(self, commands=(0, 2, 6), steps=8): return prepare_events(self.step, self.neurons, self.kc, steps, commands, n_neurons=12) def run_host(self, prepared=None, **overrides): args = self.kw | overrides return cpu_reference(self.prepared() if prepared is None else prepared, self.sources, self.compartments, COMPARTMENTS[0], self.delta, self.base, **args) def test_matches_frozen_trace_and_update_and_no_input_mutation(self): arrays = (self.step, self.neurons, self.kc, self.sources, self.compartments, self.base, self.delta) before = [(a.tobytes(), a.flags.writeable) for a in arrays] prepared = self.prepared() result = self.run_host(prepared) trace = bridge.trace_from_events(self.step, self.neurons, self.kc, 8, [0, 2, 6], phase4j_source=PHASE4J, n_neurons=12) expected, _ = bridge.apply_eligibility(trace.last_command_trace, self.kc, self.sources, self.compartments, COMPARTMENTS[0], self.delta, self.base, d=1.0, frozen=False, phase4j_source=PHASE4J, n_neurons=12) for left, right in ((result.snapshots, trace.snapshots), (result.final_trace, trace.final), (result.sampled_trace, trace.last_command_trace), (result.delta_after, expected)): assert_bits(self, left, right) self.assertFalse(left.flags.writeable) self.assertEqual(before, [(a.tobytes(), a.flags.writeable) for a in arrays]) self.assertEqual(prepared.kc_event_count, 6) self.assertEqual(prepared.indptr.tolist(), [0, 3, 4, 6]) self.assertEqual(prepared.ticks.tolist(), [0, 1, 7, 2, 2, 5]) def test_no_command_returns_positive_zero_but_nonzero_final(self): result = self.run_host(self.prepared(()), d=0) self.assertEqual(result.snapshots.shape, (0, 3)) assert_bits(self, result.sampled_trace, np.zeros(3, np.float32)) self.assertTrue(np.any(result.final_trace > 0)) assert_bits(self, result.delta_after, self.delta) def test_sampling_after_spike_and_last_command_only(self): result = self.run_host(self.prepared((0, 2, 7))) self.assertEqual(result.snapshots[0, 0], np.float32(1)) self.assertEqual(result.snapshots[1, 1], np.float32(1)) assert_bits(self, result.sampled_trace, result.snapshots[-1]) assert_bits(self, result.final_trace, result.snapshots[-1]) def test_frozen_d_zero_off_compartment_and_floor(self): self.delta = np.array([-0.0001, -0.0005, -0.001, 0], np.float32) for kwargs in (dict(d=0), dict(frozen=True)): assert_bits(self, self.run_host(**kwargs).delta_after, self.delta) result = self.run_host() assert_bits(self, result.delta_after[1:2], self.delta[1:2]) self.assertEqual(result.delta_after[0], np.float32(-0.25) * self.base[0]) self.assertEqual(result.delta_after[2], self.delta[2]) def test_empty_events_and_one_tick(self): prepared = prepare_events(np.empty(0, np.int32), np.empty(0, np.int32), self.kc, 1, [0], n_neurons=12) result = self.run_host(prepared) assert_bits(self, result.final_trace, np.zeros(3, np.float32)) assert_bits(self, result.delta_after, self.delta) def test_long_empty_gap_does_not_use_power(self): self.step = np.array([0], np.int32) self.neurons = np.array([2], np.int32) prepared = self.prepared((9999,), steps=10000) result = self.run_host(prepared) expected = np.float32(1) for _ in range(9999): expected = np.float32(expected * np.float32(self.kw["decay"])) assert_bits(self, result.sampled_trace[:1], np.array([expected], np.float32)) jumped = np.float32(float(np.float32(self.kw["decay"])) ** 9999) self.assertNotEqual(expected.view(np.uint32), jumped.view(np.uint32)) def test_random_sparse_case_matches_frozen(self): rng = np.random.default_rng(717) self.kc = np.arange(0, 80, 2, dtype=np.int32) encoded = np.unique(rng.integers(0, 80 * 257, size=1400)) self.step = (encoded // 80).astype(np.int32) self.neurons = (encoded % 80).astype(np.int32) self.sources = rng.choice(self.kc, 159).astype(np.int32) self.compartments = rng.choice(np.asarray(COMPARTMENTS), 159) self.base = rng.uniform(0.0001, 2, 159).astype(np.float32) self.delta = -(rng.uniform(0, 0.24, 159).astype(np.float32) * self.base) prepared = prepare_events(self.step, self.neurons, self.kc, 257, [0, 1, 20, 99, 256], n_neurons=80) result = self.run_host(prepared) trace = bridge.trace_from_events(self.step, self.neurons, self.kc, 257, [0, 1, 20, 99, 256], phase4j_source=PHASE4J, n_neurons=80) expected, _ = bridge.apply_eligibility(trace.last_command_trace, self.kc, self.sources, self.compartments, COMPARTMENTS[0], self.delta, self.base, d=1, phase4j_source=PHASE4J, n_neurons=80) assert_bits(self, result.snapshots, trace.snapshots) assert_bits(self, result.delta_after, expected) def test_signed_zero_source_floor_tie_is_preserved(self): # The bridge's identity guard rejects this pathological subnormal base. # Compare the original apply_v2 directly: np.maximum(+0,-0) gives -0. self.base[:] = np.nextafter(np.float32(0), np.float32(1)) result = self.run_host(d=0) original, _, _ = bridge._load_original(PHASE4J) from types import SimpleNamespace validator = SimpleNamespace(n=12, kc=self.kc, base_candidate=self.base, cand={"source_neuron_idx": bridge._Column(self.sources), "target_compartment_verbatim": bridge._Column(self.compartments)}) expected, _ = original(validator, result.sampled_trace, COMPARTMENTS[0], 0, self.delta, {"eligibility_v2": {"eta_effective_mv": bridge.ETA_MV, "max_fractional_depression": bridge.FLOOR_FRACTION}}) assert_bits(self, result.delta_after, expected) self.assertTrue(np.all(np.signbit(result.delta_after))) assert_bits(self, self.run_host(frozen=True).delta_after, self.delta) def test_input_validation(self): bad_events = [ (self.step.astype(np.int64), self.neurons, self.kc, 8, [0]), (self.step, self.neurons.astype(np.float32), self.kc, 8, [0]), (self.step, self.neurons, self.kc[::-1], 8, [0]), (self.step, self.neurons, self.kc, 8, [0, 0]), (self.step, self.neurons, self.kc, 8, [0.0]), (self.step, self.neurons, self.kc, 8, [8]), (self.step, self.neurons, self.kc, True, [0]), (self.step, self.neurons, self.kc, 0, []), (np.array([0, 0], np.int32), np.array([2, 2], np.int32), self.kc, 8, []), ] for args in bad_events: with self.subTest(args=args), self.assertRaises(ValueError): prepare_events(*args, n_neurons=12) for kwargs in (dict(d=True), dict(d=0.5), dict(d=float("nan")), dict(frozen=1), dict(eta_mv=float("inf")), dict(floor_fraction=0), dict(decay=float("nan"))): with self.subTest(kwargs=kwargs), self.assertRaises(ValueError): self.run_host(**kwargs) self.delta[0] = -0.0 with self.assertRaises(ValueError): self.run_host() self.delta[0] = 0 self.base[0] = np.nan with self.assertRaises(ValueError): self.run_host() def test_forged_prepared_rejected(self): prepared = self.prepared() bad = replace(prepared, ticks=np.array([0, 0, 7, 2, 2, 5], np.int32)) with self.assertRaises(ValueError): self.run_host(bad) def test_cuda_entry_rejects_bad_inputs_before_accessing_device(self): # An uninitialized shell has no cp/device fields. Rejection therefore # proves the CUDA entry's validation happens before any CUDA operation. backend = object.__new__(CudaEligibility) prepared = self.prepared() base_args = dict(prepared=prepared, candidate_sources=self.sources, candidate_compartments=self.compartments, compartment=COMPARTMENTS[0], delta_before=self.delta, base_candidate=self.base, **self.kw) bad = [dict(candidate_sources=self.sources.astype(np.int64)), dict(candidate_compartments=self.compartments.astype("S")), dict(delta_before=self.delta.astype(np.float64)), dict(base_candidate=np.array([np.inf, 1, 1, 1], np.float32)), dict(delta_before=np.array([np.nan, 0, 0, 0], np.float32)), dict(candidate_sources=np.array([1, 2, 5, 9], np.int32)), dict(candidate_sources=np.array([2, 2, 5, 11], np.int32)), dict(delta_before=np.array([-1, 0, 0, 0], np.float32)), dict(compartment="unknown"), dict(kernel_warmup=-1), dict(kernel_repeats=0), dict(kernel_repeats=True)] for overrides in bad: with self.subTest(overrides=overrides), self.assertRaises(ValueError): backend.run(**(base_args | overrides)) def test_no_dense_tick_raster_and_input_storage_owned(self): prepared = self.prepared() for arr in (prepared.kc, prepared.indptr, prepared.ticks, prepared.command_steps): self.assertEqual(arr.ndim, 1) self.assertFalse(arr.flags.writeable) self.kc[:] = 0 self.step[:] = 0 self.assertEqual(prepared.kc.tolist(), [2, 5, 9]) self.assertEqual(prepared.ticks.tolist(), [0, 1, 7, 2, 2, 5]) def test_kernel_explicit_rounding_controls(self): self.assertIn("--fmad=false", CUDA_OPTIONS) self.assertIn("--ftz=false", CUDA_OPTIONS) self.assertIn("__fmul_rn", CUDA_SOURCE) self.assertIn("__fadd_rn", CUDA_SOURCE) self.assertNotIn("pow(", CUDA_SOURCE) self.assertNotIn("__expf", CUDA_SOURCE) if __name__ == "__main__": unittest.main()