Add isolated exact eligibility GPU replay prototype and fixed benchmark protocol
This commit is contained in:
Digital Drosophila Archive
@@ -0,0 +1,233 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user