Digital Drosophila Archive
347 lines
21 KiB
Python
347 lines
21 KiB
Python
"""Diagnostic author-float64 LIF plus frozen legacy-float32 APL overlay.
|
|
|
|
The sealed Phase 4AE CUDA operators are reused without modification. APL is
|
|
an explicitly mixed-precision engineering overlay, not an author-model claim.
|
|
All CUDA/CPU synchronization here is diagnostic; no performance claim is made.
|
|
"""
|
|
from functools import lru_cache
|
|
import csv
|
|
import hashlib
|
|
import json
|
|
import math
|
|
from pathlib import Path
|
|
import time
|
|
import traceback
|
|
import types
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from legacy_apl_worker import LegacyAPL, N_NEURONS, file_sha256, array_sha256
|
|
|
|
|
|
SOURCE_ROOT = Path(__file__).resolve().parent
|
|
AE_MODEL = SOURCE_ROOT.parent / "author-runtime-lab-v1"
|
|
AF_MODEL = SOURCE_ROOT.parent / "author-fullgraph-lab-v1"
|
|
COEFFICIENT_NAMES = ("gd", "c0", "cg", "cv", "threshold", "reset", "seed_volt")
|
|
REQUEST_FIELDS = ("tick", "neuron_idx", "attempted", "accepted", "rejected")
|
|
|
|
|
|
def require(condition, message):
|
|
if not condition:
|
|
raise ValueError(message)
|
|
|
|
|
|
def write_json(path, value):
|
|
with Path(path).open("x", encoding="utf-8") as output:
|
|
json.dump(value, output, indent=2, allow_nan=False)
|
|
output.write("\n")
|
|
|
|
|
|
def raw_hash(value, dtype):
|
|
return hashlib.sha256(np.ascontiguousarray(value, dtype=dtype).tobytes()).hexdigest()
|
|
|
|
|
|
def load_source(path, name):
|
|
module = types.ModuleType(name)
|
|
module.__file__ = str(path)
|
|
exec(compile(Path(path).read_bytes(), str(path), "exec"), module.__dict__)
|
|
return module
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def sealed_backend():
|
|
return load_source(AE_MODEL / "gpu_backend.py", "phase4ag_sealed_ae_backend")
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def queue_function():
|
|
return load_source(AF_MODEL / "oracle_full.py", "phase4ag_sealed_af_queue").queue_summary
|
|
|
|
|
|
@lru_cache(maxsize=4)
|
|
def load_graph(graph_directory):
|
|
root = Path(graph_directory)
|
|
require(root.is_absolute(), "graph_dir must be absolute")
|
|
graph = {}
|
|
for name, dtype in (("nodes", np.int32), ("edge_id", np.int64), ("src", np.int32),
|
|
("dst", np.int32), ("signed_raw", np.int64), ("ptr", np.int64),
|
|
("weights", np.float64)):
|
|
value = np.load(root / (name + ".npy"), mmap_mode="r", allow_pickle=False)
|
|
require(value.dtype == dtype and value.ndim == 1, "Graph dtype/rank mismatch: " + name)
|
|
graph[name] = value
|
|
nodes, edges = graph["nodes"], graph["edge_id"]
|
|
n, e = len(nodes), len(edges)
|
|
require(0 < n <= N_NEURONS and bool(np.all((nodes >= 0) & (nodes < N_NEURONS))) and
|
|
bool(np.all(np.diff(nodes) > 0)), "Global node IDs must be sorted unique")
|
|
require(n < 512 or (n == N_NEURONS and np.array_equal(nodes, np.arange(N_NEURONS, dtype=np.int32))),
|
|
"Only small fixtures or the intact frozen full graph are permitted")
|
|
require(n != N_NEURONS or e == 25582938, "Full graph edge count changed")
|
|
for name in ("src", "dst", "signed_raw", "weights"):
|
|
require(graph[name].shape == (e,), "Graph edge array length mismatch: " + name)
|
|
ptr = graph["ptr"]
|
|
require(ptr.shape == (n + 1,) and ptr[0] == 0 and ptr[-1] == e and
|
|
bool(np.all(np.diff(ptr) >= 0)), "Invalid graph CSR")
|
|
require(np.array_equal(graph["src"], np.repeat(np.arange(n, dtype=np.int32), np.diff(ptr))),
|
|
"Source order does not match CSR")
|
|
# Global edge IDs need not equal local row indices in a fixture.
|
|
for start in range(0, e, 1 << 20):
|
|
stop = min(start + (1 << 20), e)
|
|
require(bool(np.all((graph["dst"][start:stop] >= 0) & (graph["dst"][start:stop] < n))),
|
|
"Local destination out of range")
|
|
require(bool(np.isfinite(graph["weights"][start:stop]).all()), "Nonfinite base SI weights")
|
|
require(np.array_equal(np.sign(graph["weights"][start:stop]), np.sign(graph["signed_raw"][start:stop])),
|
|
"Base weight sign/zero pattern differs from frozen raw contacts")
|
|
lo = max(0, start - 1)
|
|
require(bool(np.all(edges[start:stop] >= 0)) and bool(np.all(np.diff(edges[lo:stop]) > 0)),
|
|
"Global edge IDs must retain canonical increasing source/edge order")
|
|
return graph
|
|
|
|
|
|
def load_inputs(case_dir):
|
|
spec = json.loads((case_dir / "spec.json").read_text(encoding="utf-8"))
|
|
for name in ("name", "steps", "dt_ms", "duration_ms", "model", "graph_dir", "apl",
|
|
"apl_bundle_dir", "coefficients_path", "record_history"):
|
|
require(name in spec, "Missing case field: " + name)
|
|
steps = spec["steps"]
|
|
require(type(steps) is int and 0 < steps <= 1000 and float(spec["dt_ms"]) == .1 and
|
|
float(spec["duration_ms"]) == steps * .1, "Unexpected steps/dt/duration")
|
|
require(type(spec["record_history"]) is bool, "record_history must be Boolean")
|
|
require(spec["apl"] == {"tau_ms": 20, "gain": .01, "normalization": 1}, "Frozen APL parameters changed")
|
|
for name in ("graph_dir", "apl_bundle_dir", "coefficients_path"):
|
|
require(Path(spec[name]).is_absolute(), "Input path must be absolute: " + name)
|
|
model = spec["model"]
|
|
for name in ("resting_potential_mv", "threshold_mv", "reset_potential_mv", "synaptic_delay_ms",
|
|
"refractory_ms", "w_syn_mv_per_anatomical_synapse"):
|
|
require(name in model and math.isfinite(float(model[name])), "Invalid model field: " + name)
|
|
require(float(model["synaptic_delay_ms"]) == 1.8 and float(model["refractory_ms"]) == 2.2 and
|
|
float(model["w_syn_mv_per_anatomical_synapse"]) == .275, "Frozen delay/ref/W_syn changed")
|
|
coeff = json.loads(Path(spec["coefficients_path"]).read_text(encoding="utf-8"))
|
|
for name in COEFFICIENT_NAMES:
|
|
require(name in coeff and math.isfinite(float(coeff[name])), "Invalid author coefficient: " + name)
|
|
require(coeff["threshold"] == float(model["threshold_mv"]) * .001 and
|
|
coeff["reset"] == float(model["reset_potential_mv"]) * .001, "Coefficient threshold/reset mismatch")
|
|
graph = load_graph(str(Path(spec["graph_dir"]).resolve()))
|
|
n = len(graph["nodes"])
|
|
require(spec["record_history"] == (n < 512), "Fixtures need dense history; full graph must not duplicate dense histories")
|
|
with np.load(case_dir / "requests.npz", allow_pickle=False) as source:
|
|
require(set(source.files) == {"tick", "neuron_idx"}, "Request sparse schema changed")
|
|
ticks, requested = source["tick"].copy(), source["neuron_idx"].copy()
|
|
require(ticks.dtype == np.int32 and requested.dtype == np.int32 and ticks.ndim == 1 and
|
|
ticks.shape == requested.shape, "Invalid sparse request dtype/shape")
|
|
require(bool(np.all((ticks >= 0) & (ticks < steps))), "Request tick outside trial")
|
|
if len(ticks) > 1:
|
|
require(bool(np.all((ticks[1:] > ticks[:-1]) |
|
|
((ticks[1:] == ticks[:-1]) & (requested[1:] > requested[:-1])))),
|
|
"Requests must be sorted unique by tick/global neuron")
|
|
local = np.searchsorted(graph["nodes"], requested)
|
|
require(bool(np.all(local < n)) and np.array_equal(graph["nodes"][local], requested),
|
|
"Requested neuron absent from graph")
|
|
initial = {"initial_v_mv": np.full(n, float(model["resting_potential_mv"]), dtype=np.float64),
|
|
"initial_g_mv": np.zeros(n, dtype=np.float64)}
|
|
initial_path = case_dir / "arrays.npz"
|
|
if initial_path.exists():
|
|
with np.load(initial_path, allow_pickle=False) as source:
|
|
require(set(source.files).issubset(initial), "Unexpected initial-state array")
|
|
for name in source.files:
|
|
value = source[name]
|
|
require(value.dtype == np.float64 and value.shape == (n,) and bool(np.isfinite(value).all()),
|
|
"Invalid initial state: " + name)
|
|
initial[name] = value.copy()
|
|
return spec, graph, coeff, ticks, requested, local, initial
|
|
|
|
|
|
def new_array(path, dtype, shape):
|
|
require(not Path(path).exists(), "Refusing existing output: " + str(path))
|
|
return np.lib.format.open_memmap(path, mode="w+", dtype=dtype, shape=shape)
|
|
|
|
|
|
def run_gpu(case_dir, out_dir, build_root):
|
|
"""Run fresh candidate state; never consume oracle/candidate APL traces.
|
|
|
|
pass_all denotes internal recording/finite integrity only. Independent
|
|
comparison and the original complete representation gate are separate.
|
|
"""
|
|
case_dir, out_dir = Path(case_dir), Path(out_dir)
|
|
require(not out_dir.exists(), "Refusing existing output directory: " + str(out_dir))
|
|
out_dir.mkdir(parents=True, exist_ok=False)
|
|
start, tick, completed = time.perf_counter(), -1, 0
|
|
maps = {}
|
|
try:
|
|
spec, graph, coeff, request_ticks, request_nodes, request_local, initial = load_inputs(case_dir)
|
|
require(torch.cuda.is_available() and torch.cuda.device_count() == 1 and
|
|
"4090" in torch.cuda.get_device_name(), "Only the declared single visible RTX4090 is permitted")
|
|
ext, build = sealed_backend().build_extension(build_root)
|
|
build_done = time.perf_counter()
|
|
apl = LegacyAPL(spec["apl_bundle_dir"], spec["dt_ms"])
|
|
write_json(out_dir / "apl-component.json", apl.metadata)
|
|
device = apl.device
|
|
torch.cuda.reset_peak_memory_stats(device)
|
|
n, e, steps = len(graph["nodes"]), len(graph["edge_id"]), spec["steps"]
|
|
|
|
def upload(value, dtype):
|
|
return torch.from_numpy(np.array(value, copy=True, order="C")).to(device=device, dtype=dtype).contiguous()
|
|
|
|
indptr, dst = upload(graph["ptr"], torch.int64), upload(graph["dst"], torch.int32)
|
|
global_nodes_gpu = upload(graph["nodes"], torch.int64)
|
|
base_weights = upload(graph["weights"], torch.float64)
|
|
weights = base_weights.clone()
|
|
del base_weights
|
|
local_apl = np.flatnonzero(np.isin(graph["nodes"], apl.apl_globals)).astype(np.int64)
|
|
apl_out_mask = np.isin(graph["src"], local_apl)
|
|
if bool(np.any(apl_out_mask)):
|
|
weights[upload(apl_out_mask, torch.bool)] = 0.0
|
|
effective_hash = hashlib.sha256()
|
|
for offset in range(0, e, 1 << 20):
|
|
effective_hash.update(weights[offset:offset + (1 << 20)].cpu().numpy().astype("<f8", copy=False).tobytes())
|
|
common_kc, local_kc, full_kc = np.intersect1d(graph["nodes"], apl.kc_globals, assume_unique=True, return_indices=True)
|
|
kc_local_gpu, kc_full_gpu = upload(local_kc, torch.int64), upload(full_kc, torch.int64)
|
|
v = upload(initial["initial_v_mv"] * .001, torch.float64)
|
|
g = upload(initial["initial_g_mv"] * .001, torch.float64)
|
|
ref = torch.zeros(n, device=device, dtype=torch.int32)
|
|
counts = torch.zeros(n, device=device, dtype=torch.int32)
|
|
requests = torch.zeros(n, device=device, dtype=torch.uint8)
|
|
ring = torch.zeros((19, n), device=device, dtype=torch.uint8)
|
|
for name, width in (("drive", 2), ("state", 2), ("inhibition", 4064)):
|
|
maps["apl_" + name] = new_array(out_dir / ("apl-" + name + ".float32.npy"), np.float32, (steps, width))
|
|
if spec["record_history"]:
|
|
for name, dtype in (("v", np.float64), ("g", np.float64), ("ref", np.int32)):
|
|
maps[name] = new_array(out_dir / (name + "-history." + np.dtype(dtype).name + ".npy"), dtype, (steps, n))
|
|
maps["pre_" + name] = new_array(out_dir / ("preAPL-" + name + "." + np.dtype(dtype).name + ".npy"), dtype, (steps, n))
|
|
event_steps, event_nodes = [], []
|
|
accepted_requests_total, due_total, accepted_due_total = 0, 0, 0
|
|
queue_summary = queue_function()
|
|
simulation_start = time.perf_counter()
|
|
with (out_dir / "queue.jsonl").open("x", encoding="utf-8", buffering=1) as queue_output, \
|
|
(out_dir / "tick-state-hashes.jsonl").open("x", encoding="utf-8", buffering=1) as hash_output, \
|
|
(out_dir / "requests.csv").open("x", encoding="utf-8", newline="", buffering=1) as request_output:
|
|
writer = csv.DictWriter(request_output, fieldnames=REQUEST_FIELDS)
|
|
writer.writeheader()
|
|
for tick in range(steps):
|
|
requests.zero_()
|
|
begin, end = np.searchsorted(request_ticks, (tick, tick + 1))
|
|
if end > begin:
|
|
requested_local_gpu = upload(request_local[begin:end], torch.int64)
|
|
requests[requested_local_gpu] = 1
|
|
accepted = (ref[requested_local_gpu] <= 1).to(torch.uint8).cpu().numpy()
|
|
accepted_requests_total += int(accepted.sum())
|
|
writer.writerows({"tick": tick, "neuron_idx": int(node), "attempted": 1,
|
|
"accepted": int(ok), "rejected": 1 - int(ok)}
|
|
for node, ok in zip(request_nodes[begin:end], accepted))
|
|
spikes = ring[tick % 19]
|
|
ext.state_step(v, g, ref, requests, spikes, counts,
|
|
*[float(coeff[name]) for name in COEFFICIENT_NAMES], 22)
|
|
expanded_global = np.empty(0, dtype=np.int64)
|
|
accepted_cpu = np.empty(0, dtype=np.uint8)
|
|
if tick >= 18:
|
|
due = ring[(tick - 18) % 19]
|
|
sources = torch.nonzero(due, as_tuple=False).flatten().contiguous()
|
|
if sources.numel():
|
|
degrees = indptr[sources + 1] - indptr[sources]
|
|
offsets = (torch.cumsum(degrees, dim=0) - degrees).contiguous()
|
|
edge_count = int(degrees.sum().item())
|
|
expanded = torch.empty(edge_count, device=device, dtype=torch.int64)
|
|
ext.fill_edges(indptr, sources, offsets, expanded)
|
|
if edge_count:
|
|
destination = dst[expanded]
|
|
# Actual active fanout order, BEFORE stable target grouping.
|
|
expanded_global = graph["edge_id"][expanded.cpu().numpy()].copy()
|
|
accepted_cpu = (ref[destination.long()] == 0).to(torch.uint8).cpu().numpy().copy()
|
|
order = torch.argsort(destination, stable=True)
|
|
ordered_edges, ordered_targets = expanded[order].contiguous(), destination[order]
|
|
boundary = torch.ones(edge_count, device=device, dtype=torch.bool)
|
|
boundary[1:] = ordered_targets[1:] != ordered_targets[:-1]
|
|
starts = torch.nonzero(boundary, as_tuple=False).flatten().contiguous()
|
|
targets = ordered_targets[starts].contiguous()
|
|
ext.apply_arrivals(ordered_edges, starts, targets, weights, g, ref)
|
|
queue = queue_summary(tick, expanded_global, accepted_cpu)
|
|
queue_output.write(json.dumps(queue, allow_nan=False) + "\n")
|
|
due_total += queue["count"]
|
|
accepted_due_total += queue["accepted_count"]
|
|
if spec["record_history"]:
|
|
for name, value in (("v", v), ("g", g), ("ref", ref)):
|
|
maps["pre_" + name][tick] = value.cpu().numpy()
|
|
if n == N_NEURONS:
|
|
full_spikes = spikes
|
|
else:
|
|
apl.spikes.zero_()
|
|
apl.spikes[global_nodes_gpu] = spikes
|
|
full_spikes = apl.spikes
|
|
drive, apl_state, h = apl.advance(full_spikes)
|
|
# Mixed precision contract: widen f32 h EXACTLY, THEN multiply
|
|
# in f64, THEN subtract in f64. Do not guard refractory KCs.
|
|
delta_si = h[kc_full_gpu].to(torch.float64) * 0.001
|
|
g[kc_local_gpu] = g[kc_local_gpu] - delta_si
|
|
apl_cpu = {"drive": drive.cpu().numpy().copy(), "state": apl_state.cpu().numpy().copy(),
|
|
"inhibition": h.cpu().numpy().copy()}
|
|
for name, value in apl_cpu.items():
|
|
maps["apl_" + name][tick] = value
|
|
actual = {"v": v.cpu().numpy().copy(), "g": g.cpu().numpy().copy(),
|
|
"ref": ref.cpu().numpy().copy()}
|
|
finite = all(bool(np.isfinite(value).all()) for value in list(actual.values()) + list(apl_cpu.values()))
|
|
row = {"tick": tick, **{name + "_sha256": raw_hash(value, "<i4" if name == "ref" else "<f8")
|
|
for name, value in actual.items()}}
|
|
hash_output.write(json.dumps(row, allow_nan=False) + "\n")
|
|
if spec["record_history"]:
|
|
for name, value in actual.items():
|
|
maps[name][tick] = value
|
|
fired_local = torch.nonzero(spikes, as_tuple=False).flatten().cpu().numpy()
|
|
if len(fired_local):
|
|
event_steps.append(np.full(len(fired_local), tick, dtype=np.int32))
|
|
event_nodes.append(graph["nodes"][fired_local].copy())
|
|
completed += 1
|
|
if not finite:
|
|
with (out_dir / "first-failure-state.npz").open("xb") as target:
|
|
np.savez_compressed(target, tick=np.array(tick, dtype=np.int32), **actual,
|
|
**{"apl_" + name: value for name, value in apl_cpu.items()})
|
|
require(False, "Nonfinite state; first failure saved")
|
|
torch.cuda.synchronize(device)
|
|
|
|
def joined(values):
|
|
return np.concatenate(values).astype(np.int32, copy=False) if values else np.empty(0, dtype=np.int32)
|
|
|
|
state = {"step": joined(event_steps), "neuron_idx": joined(event_nodes), "counts": counts.cpu().numpy(),
|
|
"v": v.cpu().numpy(), "g": g.cpu().numpy(), "ref": ref.cpu().numpy()}
|
|
with (out_dir / "state.npz").open("xb") as target:
|
|
np.savez_compressed(target, **state)
|
|
event_locals = np.searchsorted(graph["nodes"], state["neuron_idx"])
|
|
counts_exact = np.array_equal(np.bincount(event_locals, minlength=n), state["counts"])
|
|
require(counts_exact and completed == steps, "Sparse events/counts/completed ticks mismatch")
|
|
for value in maps.values():
|
|
value.flush()
|
|
summary = {"case": spec["name"], "complete": True, "pass_all": True,
|
|
"pass_all_scope": "internal finite/recording/event-count integrity only; independent equivalence not claimed",
|
|
"steps": steps, "completed_ticks": completed, "nodes": n, "edges": e, "precision": "float64",
|
|
"apl_precision": "float32", "state_units": "SI volts", "total_spikes": len(state["step"]),
|
|
"requested": len(request_nodes), "accepted_requests": accepted_requests_total,
|
|
"rejected_requests": len(request_nodes) - accepted_requests_total,
|
|
"arriving_edges": due_total, "accepted_arriving_edges": accepted_due_total,
|
|
"counts_match_sparse_events": bool(counts_exact), "all_finite": True,
|
|
"record_history": spec["record_history"], "apl_fast_edges_zeroed": int(apl_out_mask.sum()),
|
|
"kc_targets_present": len(common_kc), "effective_weights_raw_sha256": effective_hash.hexdigest(),
|
|
"apl_component": apl.metadata, "apl_schedule": "end, after threshold/reset and guarded arrivals, unconditional all KC",
|
|
"apl_si_conversion": "h32 -> float64 -> multiply float64 0.001 -> separate float64 subtraction",
|
|
"build": build, "sealed_ae_backend_sha256": file_sha256(AE_MODEL / "gpu_backend.py"),
|
|
"shared_apl_component_oracle_limitation": True, "oracle_validation_claim": False,
|
|
"learning": False, "performance_runtime_claim": False, "cpu_sync_each_tick": True,
|
|
"build_or_cache_seconds": build_done - start,
|
|
"diagnostic_simulation_seconds": time.perf_counter() - simulation_start,
|
|
"wall_seconds": time.perf_counter() - start,
|
|
"peak_allocated_vram_bytes": int(torch.cuda.max_memory_allocated(device)),
|
|
"peak_reserved_vram_bytes": int(torch.cuda.max_memory_reserved(device)),
|
|
"input_sha256": {"spec.json": file_sha256(case_dir / "spec.json"),
|
|
"requests.npz": file_sha256(case_dir / "requests.npz"),
|
|
"base_weights.npy": file_sha256(Path(spec["graph_dir"]) / "weights.npy"),
|
|
"coefficients.json": file_sha256(spec["coefficients_path"])}}
|
|
if (case_dir / "arrays.npz").exists():
|
|
summary["input_sha256"]["arrays.npz"] = file_sha256(case_dir / "arrays.npz")
|
|
write_json(out_dir / "summary.json", summary)
|
|
return summary
|
|
except Exception as error:
|
|
for value in maps.values():
|
|
value.flush()
|
|
write_json(out_dir / "error.json", {"complete": False, "pass_all": False,
|
|
"tick_in_progress": tick, "completed_ticks": completed,
|
|
"type": type(error).__name__, "message": str(error), "traceback": traceback.format_exc()})
|
|
raise
|