Digital Drosophila Archive
316 lines
16 KiB
Python
316 lines
16 KiB
Python
"""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 subprocess
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
import numpy as np
|
|
|
|
from gpu_backend import (CUDA_OPTIONS, CUDA_SOURCE, COMPARTMENTS, CudaEligibility, _nvcc_cubin, 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)
|
|
|
|
|
|
class CompilerCacheTests(unittest.TestCase):
|
|
"""Host-side compiler invocation/evidence checks; these do not compile CUDA."""
|
|
|
|
def test_explicit_flags_authenticated_cache_and_no_global_environment_mutation(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
compiler = root / "nvcc"
|
|
compiler.write_bytes(b"host-test-compiler-placeholder")
|
|
calls = []
|
|
|
|
def run(argv, **kwargs):
|
|
if argv[-1] == "--version":
|
|
return subprocess.CompletedProcess(argv, 0, stdout="nvcc test version", stderr="")
|
|
calls.append((argv, kwargs))
|
|
self.assertIn("--ftz=false", argv)
|
|
self.assertIn("--fmad=false", argv)
|
|
self.assertIn("-arch=sm_70", argv)
|
|
self.assertNotIn("--use_fast_math", argv)
|
|
self.assertNotIn("-ftz=true", argv)
|
|
Path(argv[argv.index("-o") + 1]).write_bytes(b"test cubin, not executable CUDA")
|
|
kwargs["stdout"].write(b"mock nvcc stdout\n")
|
|
kwargs["stderr"].write(b"mock nvcc stderr\n")
|
|
return subprocess.CompletedProcess(argv, 0)
|
|
|
|
with mock.patch.dict(os.environ, {"NVCC_PREPEND_FLAGS": "", "NVCC_APPEND_FLAGS": ""}), \
|
|
mock.patch("gpu_backend.subprocess.run", side_effect=run):
|
|
before = dict(os.environ)
|
|
cubin, evidence, timing = _nvcc_cubin(root / "cache", "70", compiler)
|
|
self.assertFalse(evidence["cache_hit"])
|
|
self.assertTrue(evidence["complete"])
|
|
self.assertIn("eligibility.cubin", evidence["files_sha256"])
|
|
self.assertIn("compiler.stdout.log", evidence["files_sha256"])
|
|
self.assertIn("compiler.stderr.log", evidence["files_sha256"])
|
|
self.assertEqual(before, dict(os.environ))
|
|
again, cached, second_timing = _nvcc_cubin(root / "cache", "70", compiler)
|
|
self.assertEqual(cubin, again)
|
|
self.assertTrue(cached["cache_hit"])
|
|
self.assertEqual(second_timing["nvcc_compile_seconds"], 0)
|
|
self.assertEqual(len(calls), 1)
|
|
self.assertEqual(evidence["receipt_sha256"], cached["receipt_sha256"])
|
|
cubin.write_bytes(b"corrupted test cache")
|
|
with self.assertRaisesRegex(ValueError, "artifact changed"):
|
|
_nvcc_cubin(root / "cache", "70", compiler)
|
|
|
|
def test_failed_compile_evidence_is_retained_and_not_overwritten(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
compiler = root / "nvcc"
|
|
compiler.write_bytes(b"host-test-compiler-placeholder")
|
|
|
|
def run(argv, **kwargs):
|
|
if argv[-1] == "--version":
|
|
return subprocess.CompletedProcess(argv, 0, stdout="nvcc test version", stderr="")
|
|
kwargs["stderr"].write(b"mock compiler failure\n")
|
|
return subprocess.CompletedProcess(argv, 1)
|
|
|
|
with mock.patch.dict(os.environ, {"NVCC_PREPEND_FLAGS": "", "NVCC_APPEND_FLAGS": ""}), \
|
|
mock.patch("gpu_backend.subprocess.run", side_effect=run):
|
|
with self.assertRaisesRegex(RuntimeError, "evidence retained"):
|
|
_nvcc_cubin(root / "cache", "70", compiler)
|
|
receipts = list((root / "cache").rglob("build-receipt.json"))
|
|
self.assertEqual(len(receipts), 1)
|
|
before = receipts[0].read_bytes()
|
|
with self.assertRaisesRegex(ValueError, "failed or mismatched"):
|
|
_nvcc_cubin(root / "cache", "70", compiler)
|
|
self.assertEqual(receipts[0].read_bytes(), before)
|
|
|
|
def test_injected_compiler_flags_rejected_without_clearing_global_environment(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
compiler = root / "nvcc"
|
|
compiler.write_bytes(b"host-test-compiler-placeholder")
|
|
with mock.patch.dict(os.environ, {"NVCC_APPEND_FLAGS": "--ftz=true"}):
|
|
with self.assertRaisesRegex(ValueError, "injected flags"):
|
|
_nvcc_cubin(root / "cache", "70", compiler)
|
|
self.assertEqual(os.environ["NVCC_APPEND_FLAGS"], "--ftz=true")
|
|
self.assertFalse((root / "cache").exists())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|