Archive MaleCNS research sources, provenance, bibliography and immutable data inventory
This commit is contained in:
Digital Drosophila Archive
@@ -0,0 +1,37 @@
|
||||
{
|
||||
"artifact": "runtime-v1",
|
||||
"created_unix_time": 1789847798.995811,
|
||||
"files": [
|
||||
{
|
||||
"bytes": 8579,
|
||||
"path": "VERSION.json",
|
||||
"sha256": "1010e0f1258c07b94c6deaf71dd41882e6eb47cbe354323578acc0fad9eacd33"
|
||||
},
|
||||
{
|
||||
"bytes": 43906,
|
||||
"path": "phase4c.py",
|
||||
"sha256": "5ec83a132deb8ba20e62dab485537c2414016fe2fece546198d61c0668d75b6b"
|
||||
},
|
||||
{
|
||||
"bytes": 3143,
|
||||
"path": "phase4c.yaml",
|
||||
"sha256": "0734ffb3a99231b40db2052dea0f2d21efd5dad8d11d17d4731b952811ba5071"
|
||||
},
|
||||
{
|
||||
"bytes": 1307,
|
||||
"path": "phase4c_ext.cpp",
|
||||
"sha256": "187e268d38bb29ba271959c911d2112f9c84192b25d34ce2ef81ea89f1ae477d"
|
||||
},
|
||||
{
|
||||
"bytes": 6365,
|
||||
"path": "phase4c_ext.cu",
|
||||
"sha256": "ba8602151666c96a356fdc86e34deb66285a96efc0ec5e5280630bcaf43816c9"
|
||||
},
|
||||
{
|
||||
"bytes": 71,
|
||||
"path": "requirements.txt",
|
||||
"sha256": "1790ece7b1057db9efe2b24c74759b269b0e8911c4b78833b228200acfaa6277"
|
||||
}
|
||||
],
|
||||
"immutable_mode": "0444"
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
{
|
||||
"algorithm": "exact sparse event-driven CSR fan-out; no dense full-edge scan; exact 18-step GPU-resident ring buffer",
|
||||
"backend": "hybrid-exact",
|
||||
"eligibility": {
|
||||
"auditory-100hz": {
|
||||
"custom-cuda": {
|
||||
"candidate_event_sha256": "bf9472dfef89cd7c68f3d02d3098e988e094bcce1a75286461f38b61efdc948d",
|
||||
"eligible_for_runtime_v1": true,
|
||||
"final_g_max_abs_error": 0.001953125,
|
||||
"final_g_mean_abs_error": 2.3577847230171756e-07,
|
||||
"final_voltage_max_abs_error": 9.1552734375e-05,
|
||||
"final_voltage_mean_abs_error": 4.0572633253077584e-08,
|
||||
"per_neuron_spike_counts_exact": true,
|
||||
"propagation_statistics_exact": true,
|
||||
"reference_event_sha256": "bf9472dfef89cd7c68f3d02d3098e988e094bcce1a75286461f38b61efdc948d",
|
||||
"spike_events_exact": true,
|
||||
"state_within_tolerance": false
|
||||
},
|
||||
"hybrid-exact": {
|
||||
"candidate_event_sha256": "bf9472dfef89cd7c68f3d02d3098e988e094bcce1a75286461f38b61efdc948d",
|
||||
"eligible_for_runtime_v1": true,
|
||||
"final_g_max_abs_error": 0.0009765625,
|
||||
"final_g_mean_abs_error": 5.601728148008078e-08,
|
||||
"final_voltage_max_abs_error": 4.57763671875e-05,
|
||||
"final_voltage_mean_abs_error": 6.476060576687814e-09,
|
||||
"per_neuron_spike_counts_exact": true,
|
||||
"propagation_statistics_exact": true,
|
||||
"reference_event_sha256": "bf9472dfef89cd7c68f3d02d3098e988e094bcce1a75286461f38b61efdc948d",
|
||||
"spike_events_exact": true,
|
||||
"state_within_tolerance": false
|
||||
}
|
||||
},
|
||||
"auditory-10hz": {
|
||||
"custom-cuda": {
|
||||
"candidate_event_sha256": "5a2cfbe9168122bd153855efe645b65340437765a851a54ed5ee817da59630af",
|
||||
"eligible_for_runtime_v1": true,
|
||||
"final_g_max_abs_error": 0.00390625,
|
||||
"final_g_mean_abs_error": 2.360452100447219e-07,
|
||||
"final_voltage_max_abs_error": 6.103515625e-05,
|
||||
"final_voltage_mean_abs_error": 4.073282156014102e-08,
|
||||
"per_neuron_spike_counts_exact": true,
|
||||
"propagation_statistics_exact": true,
|
||||
"reference_event_sha256": "5a2cfbe9168122bd153855efe645b65340437765a851a54ed5ee817da59630af",
|
||||
"spike_events_exact": true,
|
||||
"state_within_tolerance": false
|
||||
},
|
||||
"hybrid-exact": {
|
||||
"candidate_event_sha256": "5a2cfbe9168122bd153855efe645b65340437765a851a54ed5ee817da59630af",
|
||||
"eligible_for_runtime_v1": true,
|
||||
"final_g_max_abs_error": 0.000244140625,
|
||||
"final_g_mean_abs_error": 3.382269753160472e-08,
|
||||
"final_voltage_max_abs_error": 3.0517578125e-05,
|
||||
"final_voltage_mean_abs_error": 5.308996797737109e-09,
|
||||
"per_neuron_spike_counts_exact": true,
|
||||
"propagation_statistics_exact": true,
|
||||
"reference_event_sha256": "5a2cfbe9168122bd153855efe645b65340437765a851a54ed5ee817da59630af",
|
||||
"spike_events_exact": true,
|
||||
"state_within_tolerance": false
|
||||
}
|
||||
},
|
||||
"compile-default": {
|
||||
"eligible_for_runtime_v1": false,
|
||||
"final_g_max_abs_error": 561.5676879882812,
|
||||
"final_g_mean_abs_error": 2.160905361175537,
|
||||
"final_voltage_max_abs_error": 111.6280746459961,
|
||||
"final_voltage_mean_abs_error": 0.47845765948295593,
|
||||
"per_neuron_spike_counts_exact": false,
|
||||
"spike_events_exact": false,
|
||||
"state_within_tolerance": false
|
||||
},
|
||||
"compile-reduce-overhead": {
|
||||
"eligible_for_runtime_v1": false,
|
||||
"final_g_max_abs_error": 561.5676879882812,
|
||||
"final_g_mean_abs_error": 2.160905361175537,
|
||||
"final_voltage_max_abs_error": 111.6280746459961,
|
||||
"final_voltage_mean_abs_error": 0.47845765948295593,
|
||||
"per_neuron_spike_counts_exact": false,
|
||||
"spike_events_exact": false,
|
||||
"state_within_tolerance": false
|
||||
},
|
||||
"cudagraph-custom-state": {
|
||||
"eligible_for_runtime_v1": false,
|
||||
"final_g_max_abs_error": 52968.98046875,
|
||||
"final_g_mean_abs_error": 17.617931365966797,
|
||||
"final_voltage_max_abs_error": 976.6533203125,
|
||||
"final_voltage_mean_abs_error": 6.446992874145508,
|
||||
"per_neuron_spike_counts_exact": false,
|
||||
"spike_events_exact": false,
|
||||
"state_within_tolerance": false
|
||||
},
|
||||
"gustatory-50hz": {
|
||||
"custom-cuda": {
|
||||
"candidate_event_sha256": "33f95e37e27a8ef6de8368948d388a1ed9e03f29ff331c266025646889f88c48",
|
||||
"eligible_for_runtime_v1": false,
|
||||
"final_g_max_abs_error": 0.00390625,
|
||||
"final_g_mean_abs_error": 2.401351650860306e-07,
|
||||
"final_voltage_max_abs_error": 6.103515625e-05,
|
||||
"final_voltage_mean_abs_error": 3.691125627369729e-08,
|
||||
"per_neuron_spike_counts_exact": true,
|
||||
"propagation_statistics_exact": true,
|
||||
"reference_event_sha256": "169bd42735d5357300546c37f826b6e5e216f43136c16262364bd6a991e66d53",
|
||||
"spike_events_exact": false,
|
||||
"state_within_tolerance": false
|
||||
},
|
||||
"hybrid-exact": {
|
||||
"candidate_event_sha256": "169bd42735d5357300546c37f826b6e5e216f43136c16262364bd6a991e66d53",
|
||||
"eligible_for_runtime_v1": true,
|
||||
"final_g_max_abs_error": 0.00390625,
|
||||
"final_g_mean_abs_error": 8.467309697834935e-08,
|
||||
"final_voltage_max_abs_error": 6.103515625e-05,
|
||||
"final_voltage_mean_abs_error": 6.498944493671388e-09,
|
||||
"per_neuron_spike_counts_exact": true,
|
||||
"propagation_statistics_exact": true,
|
||||
"reference_event_sha256": "169bd42735d5357300546c37f826b6e5e216f43136c16262364bd6a991e66d53",
|
||||
"spike_events_exact": true,
|
||||
"state_within_tolerance": false
|
||||
}
|
||||
},
|
||||
"mechanosensory-50hz": {
|
||||
"custom-cuda": {
|
||||
"candidate_event_sha256": "6e6e7bb80bedace02bc60790838f5d535854ddd8567b62e421a193f94818eae5",
|
||||
"eligible_for_runtime_v1": false,
|
||||
"final_g_max_abs_error": 604.771240234375,
|
||||
"final_g_mean_abs_error": 3.3535778522491455,
|
||||
"final_voltage_max_abs_error": 351.72637939453125,
|
||||
"final_voltage_mean_abs_error": 1.0685756206512451,
|
||||
"per_neuron_spike_counts_exact": false,
|
||||
"propagation_statistics_exact": false,
|
||||
"reference_event_sha256": "0cd6399bb437771a0901575674450fae32cd296df710f0a7ec0da26804176999",
|
||||
"spike_events_exact": false,
|
||||
"state_within_tolerance": false
|
||||
},
|
||||
"hybrid-exact": {
|
||||
"candidate_event_sha256": "0cd6399bb437771a0901575674450fae32cd296df710f0a7ec0da26804176999",
|
||||
"eligible_for_runtime_v1": true,
|
||||
"final_g_max_abs_error": 0.00048828125,
|
||||
"final_g_mean_abs_error": 4.223790739388278e-08,
|
||||
"final_voltage_max_abs_error": 6.103515625e-05,
|
||||
"final_voltage_mean_abs_error": 9.290744173995336e-09,
|
||||
"per_neuron_spike_counts_exact": true,
|
||||
"propagation_statistics_exact": true,
|
||||
"reference_event_sha256": "0cd6399bb437771a0901575674450fae32cd296df710f0a7ec0da26804176999",
|
||||
"spike_events_exact": true,
|
||||
"state_within_tolerance": false
|
||||
}
|
||||
},
|
||||
"olfactory-50hz": {
|
||||
"custom-cuda": {
|
||||
"candidate_event_sha256": "0a96a16ab0f4ea559238206e3feab22954bb090ad4e791f47a4ca530a23a460d",
|
||||
"eligible_for_runtime_v1": true,
|
||||
"final_g_max_abs_error": 0.00390625,
|
||||
"final_g_mean_abs_error": 2.6316607204535103e-07,
|
||||
"final_voltage_max_abs_error": 6.103515625e-05,
|
||||
"final_voltage_mean_abs_error": 3.892501609925603e-08,
|
||||
"per_neuron_spike_counts_exact": true,
|
||||
"propagation_statistics_exact": true,
|
||||
"reference_event_sha256": "0a96a16ab0f4ea559238206e3feab22954bb090ad4e791f47a4ca530a23a460d",
|
||||
"spike_events_exact": true,
|
||||
"state_within_tolerance": false
|
||||
},
|
||||
"hybrid-exact": {
|
||||
"candidate_event_sha256": "0a96a16ab0f4ea559238206e3feab22954bb090ad4e791f47a4ca530a23a460d",
|
||||
"eligible_for_runtime_v1": true,
|
||||
"final_g_max_abs_error": 0.00048828125,
|
||||
"final_g_mean_abs_error": 5.61670994159158e-08,
|
||||
"final_voltage_max_abs_error": 6.103515625e-05,
|
||||
"final_voltage_mean_abs_error": 7.231219623093921e-09,
|
||||
"per_neuron_spike_counts_exact": true,
|
||||
"propagation_statistics_exact": true,
|
||||
"reference_event_sha256": "0a96a16ab0f4ea559238206e3feab22954bb090ad4e791f47a4ca530a23a460d",
|
||||
"spike_events_exact": true,
|
||||
"state_within_tolerance": false
|
||||
}
|
||||
}
|
||||
},
|
||||
"fallback_backend": "eager",
|
||||
"name": "MaleCNS-LIF-runtime-v1",
|
||||
"reference": "MaleCNS-LIF-reference-v1",
|
||||
"spike_behavior_requirement": "exact on all five 1 s frozen reference trials for custom runtime; exact auditory trial for compiler variants",
|
||||
"ten_second_biological_realtime_speed": 0.12564218388193749
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
1010e0f1258c07b94c6deaf71dd41882e6eb47cbe354323578acc0fad9eacd33 VERSION.json
|
||||
5ec83a132deb8ba20e62dab485537c2414016fe2fece546198d61c0668d75b6b phase4c.py
|
||||
0734ffb3a99231b40db2052dea0f2d21efd5dad8d11d17d4731b952811ba5071 phase4c.yaml
|
||||
187e268d38bb29ba271959c911d2112f9c84192b25d34ce2ef81ea89f1ae477d phase4c_ext.cpp
|
||||
ba8602151666c96a356fdc86e34deb66285a96efc0ec5e5280630bcaf43816c9 phase4c_ext.cu
|
||||
1790ece7b1057db9efe2b24c74759b269b0e8911c4b78833b228200acfaa6277 requirements.txt
|
||||
4761723f4b65c3adca28c5b09e1cfcb77b18521ef31ac1e220244f1b35138e9a MANIFEST.json
|
||||
@@ -0,0 +1,756 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Phase 4C: freeze reference, audit visual semantics, and optimize exact runtime."""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import platform
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pyarrow as pa
|
||||
import pyarrow.parquet as pq
|
||||
import torch
|
||||
import yaml
|
||||
from torch.utils.cpp_extension import load as load_extension
|
||||
|
||||
import phase4b as p4b
|
||||
|
||||
LOG = logging.getLogger("phase4c")
|
||||
|
||||
|
||||
def sha256(path: Path, chunk: int = 8 << 20) -> str:
|
||||
h = hashlib.sha256()
|
||||
with path.open("rb") as f:
|
||||
while data := f.read(chunk):
|
||||
h.update(data)
|
||||
return h.hexdigest()
|
||||
|
||||
|
||||
def write_json(path: Path, value: Any) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
||||
|
||||
|
||||
def copy_immutable(source: Path, destination: Path) -> None:
|
||||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||||
if destination.exists():
|
||||
if sha256(source) != sha256(destination):
|
||||
raise RuntimeError(f"immutable destination differs: {destination}")
|
||||
return
|
||||
shutil.copy2(source, destination)
|
||||
destination.chmod(0o444)
|
||||
|
||||
|
||||
def exact_paths(cfg: dict[str, Any]) -> dict[str, Path]:
|
||||
root = Path(cfg["paths"]["connectome_root"]).resolve()
|
||||
allowed = Path("/data/digital-drosophila/datasets/male-cns-v1.0/processed/neuron-only").resolve()
|
||||
if root != allowed:
|
||||
raise RuntimeError(f"only allowed input is {allowed}")
|
||||
paths = {
|
||||
"indptr": root / "gpu/csr-indptr.int64.npy",
|
||||
"indices": root / "gpu/csr-indices.int32.npy",
|
||||
"raw_weights": root / "gpu/csr-weights.int64.npy",
|
||||
"metadata": root / "neuron-index/unified-metadata.parquet",
|
||||
"mapping": root / "neuron-index/body-id-neuron-idx.parquet",
|
||||
"csr_metadata": root / "gpu/csr-metadata.json",
|
||||
}
|
||||
for path in paths.values():
|
||||
if not path.is_file():
|
||||
raise FileNotFoundError(path)
|
||||
return paths
|
||||
|
||||
|
||||
def freeze_reference(cfg: dict[str, Any], config_path: Path, paths: dict[str, Path]) -> dict[str, Any]:
|
||||
model_root = Path(cfg["paths"]["model_root"])
|
||||
ref = model_root / "reference-v1"
|
||||
ref.mkdir(parents=True, exist_ok=True)
|
||||
source_dir = Path(__file__).resolve().parent
|
||||
phase4b_root = Path(cfg["paths"]["phase4b_root"])
|
||||
source_files = {
|
||||
"model.py": source_dir / "phase4b.py",
|
||||
"config.yaml": source_dir / "phase4b.yaml",
|
||||
}
|
||||
for name, source in source_files.items():
|
||||
copy_immutable(source, ref / name)
|
||||
phase4b_seal = phase4b_root / "manifests/checksums.sha256"
|
||||
phase4b_manifest = phase4b_root / "manifests/manifest.json"
|
||||
if not phase4b_seal.is_file() or not phase4b_manifest.is_file():
|
||||
raise RuntimeError("Phase 4B manifest/seal missing")
|
||||
seal_check = subprocess.run(
|
||||
["sha256sum", "-c", "--quiet", str(phase4b_seal)], cwd=phase4b_root,
|
||||
text=True, capture_output=True,
|
||||
)
|
||||
if seal_check.returncode:
|
||||
raise RuntimeError(f"Phase 4B seal failed: {seal_check.stderr}")
|
||||
frozen = {
|
||||
"name": cfg["names"]["reference"],
|
||||
"description": cfg["names"]["description"],
|
||||
"not_claimed": "strict Shiu 2024 original-model reproduction",
|
||||
"reason": "MaleCNS v1.0 transmitter inventory includes histamine and differs from the original dataset/inventory",
|
||||
"parameters": cfg["model"],
|
||||
"sign_policy": cfg["sign_policy"],
|
||||
"graph": cfg["expected_graph"],
|
||||
"source_phase": "Phase 4B PASS",
|
||||
"phase4b_manifest_sha256": sha256(phase4b_manifest),
|
||||
"phase4b_checksums_sha256": sha256(phase4b_seal),
|
||||
"phase4b_checksum_verification": "PASS",
|
||||
"local_repository_commit": cfg["reproducibility"]["local_repository_commit"],
|
||||
"remote_model_directory_is_git_repository": False,
|
||||
"code": {name: {"path": str(ref / name), "sha256": sha256(ref / name)} for name in source_files},
|
||||
"inputs": {name: {"path": str(path), "sha256": sha256(path)} for name, path in paths.items()},
|
||||
}
|
||||
version_path = ref / "VERSION.json"
|
||||
if version_path.exists():
|
||||
previous = json.loads(version_path.read_text(encoding="utf-8"))
|
||||
if previous != frozen:
|
||||
raise RuntimeError("existing immutable reference VERSION differs")
|
||||
else:
|
||||
write_json(version_path, frozen)
|
||||
version_path.chmod(0o444)
|
||||
return frozen
|
||||
|
||||
|
||||
def full_metadata(path: Path) -> pa.Table:
|
||||
table = pq.read_table(path)
|
||||
required = [
|
||||
"neuron_idx", "body_id", "annotation_superclass", "annotation_class",
|
||||
"annotation_subclass", "annotation_type", "annotation_instance", "side",
|
||||
"neurotransmitter",
|
||||
]
|
||||
missing = [x for x in required if x not in table.column_names]
|
||||
if missing:
|
||||
raise RuntimeError(f"metadata columns missing: {missing}")
|
||||
return table
|
||||
|
||||
|
||||
def list_mask(values: list[Any], accepted: set[str]) -> np.ndarray:
|
||||
return np.asarray([x in accepted for x in values], dtype=bool)
|
||||
|
||||
|
||||
def populations(table: pa.Table) -> dict[str, np.ndarray]:
|
||||
cls = table["annotation_class"].to_pylist()
|
||||
sub = table["annotation_subclass"].to_pylist()
|
||||
mech = {"mechanosensory", "mechanosensory_tactile", "mechanosensory_proprioceptive"}
|
||||
return {
|
||||
"visual": np.asarray([x == "visual" for x in cls]),
|
||||
"mechanosensory": np.asarray([(x in mech) and y != "auditory" for x, y in zip(cls, sub)]),
|
||||
"olfactory": np.asarray([x == "olfactory" for x in cls]),
|
||||
"gustatory": np.asarray([x == "gustatory" for x in cls]),
|
||||
"auditory": np.asarray([x == "auditory" for x in sub]),
|
||||
}
|
||||
|
||||
|
||||
def layer_edge_stats(indptr: np.ndarray, indices: np.ndarray, raw: np.ndarray,
|
||||
layers: dict[str, np.ndarray]) -> dict[str, Any]:
|
||||
n = len(indptr) - 1
|
||||
label = np.full(n, -1, dtype=np.int16)
|
||||
names = list(layers)
|
||||
for j, name in enumerate(names):
|
||||
label[layers[name]] = j
|
||||
counts = np.zeros((len(names), len(names)), dtype=np.int64)
|
||||
weights = np.zeros_like(counts)
|
||||
for src in range(n):
|
||||
a, b = int(indptr[src]), int(indptr[src + 1])
|
||||
x = label[src]
|
||||
if x < 0 or a == b:
|
||||
continue
|
||||
y = label[indices[a:b]]
|
||||
valid = y >= 0
|
||||
if np.any(valid):
|
||||
counts[x] += np.bincount(y[valid], minlength=len(names))
|
||||
weights[x] += np.bincount(y[valid], weights=raw[a:b][valid], minlength=len(names)).astype(np.int64)
|
||||
return {
|
||||
"layer_order": names,
|
||||
"edge_counts": counts.tolist(),
|
||||
"synapse_weight_sums": weights.tolist(),
|
||||
}
|
||||
|
||||
|
||||
def visual_audit(exp: Path, table: pa.Table, indptr: np.ndarray,
|
||||
indices: np.ndarray, raw: np.ndarray) -> tuple[dict[str, Any], dict[str, np.ndarray]]:
|
||||
out = exp / "visual-audit"
|
||||
out.mkdir(parents=True, exist_ok=True)
|
||||
pops = populations(table)
|
||||
visual = pops["visual"]
|
||||
nodes = np.flatnonzero(visual).astype(np.int32)
|
||||
fields = [
|
||||
"neuron_idx", "body_id", "annotation_superclass", "annotation_class",
|
||||
"annotation_subclass", "annotation_type", "annotation_instance", "side", "neurotransmitter",
|
||||
]
|
||||
visual_table = table.select(fields).filter(pa.array(visual))
|
||||
pq.write_table(visual_table, out / "phase4b-visual-6091.parquet", compression="zstd")
|
||||
arrays = {name: table[name].to_pylist() for name in fields}
|
||||
breakdown: dict[str, Any] = {}
|
||||
for field in fields[2:]:
|
||||
c = Counter("<null>" if arrays[field][i] is None else str(arrays[field][i]) for i in nodes)
|
||||
breakdown[field] = dict(sorted(c.items(), key=lambda x: (-x[1], x[0])))
|
||||
visual_types = sorted({str(arrays["annotation_type"][i]) for i in nodes if arrays["annotation_type"][i] is not None})
|
||||
exact_photoreceptor = visual & list_mask(arrays["annotation_superclass"], {"ol_sensory"})
|
||||
|
||||
# These anatomical memberships are derived only from explicit metadata fields.
|
||||
super_values = arrays["annotation_superclass"]
|
||||
type_values = ["" if x is None else str(x) for x in arrays["annotation_type"]]
|
||||
ol_intrinsic = list_mask(super_values, {"ol_intrinsic"})
|
||||
visual_projection = list_mask(super_values, {"visual_projection", "visual_projection_tbc"})
|
||||
lamina_types = {"L1", "L2", "L3", "L4", "L5", "C2", "C3", "T1", "Lawf1", "Lawf2"}
|
||||
lamina = ol_intrinsic & np.asarray([x in lamina_types for x in type_values])
|
||||
lobula_plate = ol_intrinsic & np.asarray([
|
||||
x.startswith(("T4", "T5", "LPi", "LPLC", "LLPC")) for x in type_values
|
||||
])
|
||||
lobula = ol_intrinsic & ~lobula_plate & np.asarray([
|
||||
x.startswith(("LC", "LT", "Li", "Lobula")) for x in type_values
|
||||
])
|
||||
medulla = ol_intrinsic & ~lamina & ~lobula & ~lobula_plate & np.asarray([
|
||||
x.startswith(("Mi", "Dm", "Tm", "TmY", "Pm", "Cm", "Me")) for x in type_values
|
||||
])
|
||||
optic_other = ol_intrinsic & ~(lamina | medulla | lobula | lobula_plate)
|
||||
anatomy = {
|
||||
"photoreceptor": exact_photoreceptor,
|
||||
"retinal_sensory": exact_photoreceptor,
|
||||
"lamina": lamina,
|
||||
"medulla": medulla,
|
||||
"lobula": lobula,
|
||||
"lobula_plate": lobula_plate,
|
||||
"optic_lobe_interneurons": ol_intrinsic,
|
||||
"visual_projection_neurons": visual_projection,
|
||||
"optic_lobe_interneurons_unassigned_to_named_region": optic_other,
|
||||
}
|
||||
membership_rows: list[dict[str, Any]] = []
|
||||
for role, mask in anatomy.items():
|
||||
for i in np.flatnonzero(mask):
|
||||
membership_rows.append({"anatomical_role": role, "neuron_idx": int(i), "body_id": int(arrays["body_id"][i])})
|
||||
pq.write_table(pa.Table.from_pylist(membership_rows), out / "anatomical-memberships.parquet", compression="zstd")
|
||||
|
||||
# Exact graph hierarchy uses shortest directed distance, not inferred body IDs.
|
||||
distance = p4b.directed_distances(indptr, indices, nodes.astype(np.int64), len(visual))
|
||||
optic_super = {"ol_intrinsic", "visual_projection", "visual_projection_tbc", "visual_centrifugal"}
|
||||
optic = list_mask(super_values, optic_super)
|
||||
layers = {
|
||||
"visual_sensory": distance == 0,
|
||||
"first_order": distance == 1,
|
||||
"second_order": distance == 2,
|
||||
"higher_visual": (distance >= 3) & optic,
|
||||
"central_brain": (distance >= 3) & ~optic,
|
||||
}
|
||||
for name, mask in layers.items():
|
||||
ids = np.flatnonzero(mask).astype(np.int32)
|
||||
pq.write_table(pa.table({
|
||||
"neuron_idx": pa.array(ids, type=pa.int32()),
|
||||
"body_id": pa.array(np.asarray(arrays["body_id"], dtype=np.uint64)[ids], type=pa.uint64()),
|
||||
"shortest_visual_path_depth": pa.array(distance[ids], type=pa.int16()),
|
||||
}), out / f"hierarchy-{name}.parquet", compression="zstd")
|
||||
edge_stats = layer_edge_stats(indptr, indices, raw, layers)
|
||||
write_json(out / "hierarchy-edge-statistics.json", edge_stats)
|
||||
summary = {
|
||||
"phase4b_visual_count": int(visual.sum()),
|
||||
"all_histaminergic": all(arrays["neurotransmitter"][i] == "histamine" for i in nodes),
|
||||
"all_ol_sensory": all(arrays["annotation_superclass"][i] == "ol_sensory" for i in nodes),
|
||||
"raw_visual_types": visual_types,
|
||||
"conclusion": "The 6091 annotation_class=visual bodies are photoreceptor/retinal sensory input neurons, not the whole visual processing network.",
|
||||
"field_breakdown": breakdown,
|
||||
"anatomical_membership_counts": {k: int(v.sum()) for k, v in anatomy.items()},
|
||||
"anatomical_rule_provenance": {
|
||||
"photoreceptor_retinal_sensory": "annotation_class=visual AND annotation_superclass=ol_sensory; observed types recorded verbatim",
|
||||
"optic_lobe_interneurons": "annotation_superclass=ol_intrinsic",
|
||||
"visual_projection_neurons": "annotation_superclass in {visual_projection, visual_projection_tbc}",
|
||||
"named_optic_regions": "explicit annotation_type prefix/set rules, preserved in this audit; no body_id was guessed",
|
||||
},
|
||||
"hierarchy_definition": "exclusive shortest directed connectome distance from the 6091 sensory neurons: depth 0 sensory, depth 1 first-order, depth 2 second-order, depth >=3 split into optic/visual-superclass higher_visual versus non-optic downstream central_brain; unreachable bodies are excluded",
|
||||
"hierarchy_counts": {k: int(v.sum()) for k, v in layers.items()},
|
||||
"edge_statistics_file": str(out / "hierarchy-edge-statistics.json"),
|
||||
}
|
||||
write_json(out / "summary.json", summary)
|
||||
return summary, layers
|
||||
|
||||
|
||||
def generate_forced(steps: int, target_count: int, rate_hz: float, dt: float,
|
||||
seed: int, device: torch.device) -> torch.Tensor:
|
||||
probability = rate_hz * dt / 1000.0
|
||||
generator = torch.Generator(device=device)
|
||||
generator.manual_seed(seed)
|
||||
forced = torch.empty((steps, target_count), dtype=torch.uint8, device=device)
|
||||
# Deliberately preserve Phase 4B's one rand() call per timestep.
|
||||
for step in range(steps):
|
||||
forced[step].copy_((torch.rand(target_count, generator=generator, device=device) < probability).to(torch.uint8))
|
||||
return forced
|
||||
|
||||
|
||||
def expand_rows(indptr: torch.Tensor, sources: torch.Tensor) -> torch.Tensor:
|
||||
if sources.numel() == 0:
|
||||
return torch.empty(0, dtype=torch.int64, device=indptr.device)
|
||||
starts = indptr[sources]
|
||||
counts = indptr[sources + 1] - starts
|
||||
cumulative = torch.cumsum(counts, 0)
|
||||
total = int(cumulative[-1].item()) if cumulative.numel() else 0
|
||||
if total == 0:
|
||||
return torch.empty(0, dtype=torch.int64, device=indptr.device)
|
||||
seq = torch.arange(total, dtype=torch.int64, device=indptr.device)
|
||||
row = torch.searchsorted(cumulative, seq, right=True)
|
||||
return starts[row] + seq - (cumulative - counts)[row]
|
||||
|
||||
|
||||
class Runtime:
|
||||
def __init__(self, cfg: dict[str, Any], table: pa.Table, indptr_np: np.ndarray,
|
||||
indices_np: np.ndarray, raw_np: np.ndarray, device: torch.device,
|
||||
extension: Any | None):
|
||||
self.cfg, self.table, self.device, self.extension = cfg, table, device, extension
|
||||
self.n = int(cfg["expected_graph"]["neurons"])
|
||||
self.indptr = torch.as_tensor(np.asarray(indptr_np), dtype=torch.int64, device=device)
|
||||
self.indices = torch.as_tensor(np.asarray(indices_np), dtype=torch.int32, device=device)
|
||||
nt = table["neurotransmitter"].to_pylist()
|
||||
policy = cfg["sign_policy"]
|
||||
signs = np.empty(self.n, dtype=np.int8)
|
||||
for i, value in enumerate(nt):
|
||||
key = "" if value is None else str(value).strip().lower()
|
||||
if not key:
|
||||
signs[i] = int(policy["unknown_sign"])
|
||||
elif key == "unclear":
|
||||
signs[i] = int(policy["unclear_sign"])
|
||||
elif key in policy["neurotransmitter_sign"]:
|
||||
signs[i] = int(policy["neurotransmitter_sign"][key])
|
||||
else:
|
||||
raise RuntimeError(f"unmapped neurotransmitter: {value!r}")
|
||||
effective = np.empty(len(raw_np), dtype=np.float32)
|
||||
for src in range(self.n):
|
||||
a, b = int(indptr_np[src]), int(indptr_np[src + 1])
|
||||
effective[a:b] = raw_np[a:b].astype(np.float32) * signs[src] * float(cfg["model"]["w_syn_mv_per_anatomical_synapse"])
|
||||
self.weights = torch.as_tensor(effective, dtype=torch.float32, device=device)
|
||||
self.signs = signs
|
||||
self.superclass = np.asarray(table["annotation_superclass"].to_pylist(), dtype=object)
|
||||
|
||||
def compiled_state(self, mode: str):
|
||||
m = self.cfg["model"]
|
||||
gd = 1.0 - float(m["dt_ms"]) / float(m["synaptic_tau_ms"])
|
||||
alpha = float(m["dt_ms"]) / float(m["membrane_tau_ms"])
|
||||
rest, reset, threshold = (float(m[x]) for x in ("resting_potential_mv", "reset_potential_mv", "threshold_mv"))
|
||||
|
||||
def fn(v, g, refractory, arrival):
|
||||
g2 = g * gd
|
||||
g2 = g2 + arrival
|
||||
r2 = torch.clamp_min(refractory, 1) - 1
|
||||
available = r2 == 0
|
||||
leak = v - rest
|
||||
drive = g2 - leak
|
||||
integrated = v + alpha * drive
|
||||
v2 = torch.where(available, integrated, torch.full_like(v, reset))
|
||||
natural = available & (v2 >= threshold)
|
||||
return v2, g2, r2, natural, available
|
||||
|
||||
kwargs: dict[str, Any] = {"backend": "inductor", "fullgraph": True}
|
||||
if mode == "compile-reduce-overhead":
|
||||
kwargs["mode"] = "reduce-overhead"
|
||||
return torch.compile(fn, **kwargs), fn
|
||||
|
||||
def run(self, backend: str, targets_np: np.ndarray, rate_hz: float, seed: int,
|
||||
duration_ms: float, capture_events: bool, monitor: bool = False,
|
||||
forced_override: torch.Tensor | None = None) -> dict[str, Any]:
|
||||
m = self.cfg["model"]
|
||||
dt = float(m["dt_ms"]); steps = int(round(duration_ms / dt))
|
||||
delay_steps = int(round(float(m["synaptic_delay_ms"]) / dt)); qlen = delay_steps + 1
|
||||
ref_steps = int(round(float(m["refractory_ms"]) / dt))
|
||||
alpha = dt / float(m["membrane_tau_ms"]); gd = 1.0 - dt / float(m["synaptic_tau_ms"])
|
||||
rest = float(m["resting_potential_mv"]); reset = float(m["reset_potential_mv"]); threshold = float(m["threshold_mv"])
|
||||
targets = torch.as_tensor(targets_np, dtype=torch.int64, device=self.device)
|
||||
forced = forced_override if forced_override is not None else generate_forced(steps, len(targets_np), rate_hz, dt, seed, self.device)
|
||||
if forced.shape != (steps, len(targets_np)):
|
||||
raise RuntimeError("forced input shape mismatch")
|
||||
v = torch.full((self.n,), rest, dtype=torch.float32, device=self.device)
|
||||
g = torch.zeros(self.n, dtype=torch.float32, device=self.device)
|
||||
refractory = torch.zeros(self.n, dtype=torch.int32, device=self.device)
|
||||
queue = torch.zeros((qlen, self.n), dtype=torch.float32, device=self.device)
|
||||
counts = torch.zeros(self.n, dtype=torch.int32, device=self.device)
|
||||
event_steps: list[np.ndarray] = []; event_nodes: list[np.ndarray] = []
|
||||
compiled = eager_fn = None
|
||||
if backend.startswith("compile-"):
|
||||
compiled, eager_fn = self.compiled_state(backend)
|
||||
# Exclude compile latency and record a graph-break diagnostic separately.
|
||||
tv, tg, tr, ta = v.clone(), g.clone(), refractory.clone(), queue[0].clone()
|
||||
compiled(tv, tg, tr, ta); torch.cuda.synchronize()
|
||||
del tv, tg, tr, ta
|
||||
if backend.startswith("custom") or backend in ("cudagraph-custom-state", "hybrid-exact"):
|
||||
if self.extension is None:
|
||||
raise RuntimeError("custom extension unavailable")
|
||||
spike_u8 = torch.zeros(self.n, dtype=torch.uint8, device=self.device)
|
||||
graph = staging = None
|
||||
if backend == "cudagraph-custom-state":
|
||||
staging = torch.zeros(self.n, dtype=torch.float32, device=self.device)
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
# Warm-up on disposable state, then capture only the graph-safe fixed-shape state kernel.
|
||||
sv, sg = v.clone(), g.clone(); sr = refractory.clone(); ss = spike_u8.clone(); sc = counts.clone(); sa = staging.clone()
|
||||
self.extension.state_step(sv, sg, sr, sa, ss, sc, gd, alpha, rest, reset, threshold, ref_steps)
|
||||
torch.cuda.synchronize()
|
||||
del sv, sg, sr, ss, sc, sa
|
||||
with torch.cuda.graph(graph):
|
||||
self.extension.state_step(v, g, refractory, staging, spike_u8, counts, gd, alpha, rest, reset, threshold, ref_steps)
|
||||
sampler = p4b.GPUSampler(self.cfg["device"]["launcher_visible_gpu_uuid"], float(self.cfg["benchmark"]["gpu_sample_interval_seconds"])) if monitor else None
|
||||
if sampler: sampler.start()
|
||||
torch.cuda.synchronize(); torch.cuda.reset_peak_memory_stats(); start = time.perf_counter()
|
||||
cpu_start = time.process_time()
|
||||
for step in range(steps):
|
||||
slot = step % qlen
|
||||
if backend in ("eager", "hybrid-exact"):
|
||||
g.mul_(gd).add_(queue[slot]); queue[slot].zero_()
|
||||
refractory.clamp_min_(1).sub_(1)
|
||||
available = refractory == 0
|
||||
v[available] += alpha * (g[available] - (v[available] - rest))
|
||||
v[~available] = reset
|
||||
spikes = available & (v >= threshold)
|
||||
accepted = forced[step].bool() & available[targets]
|
||||
forced_sources = targets[accepted]
|
||||
if forced_sources.numel(): spikes[forced_sources] = True
|
||||
sources = torch.nonzero(spikes, as_tuple=False).flatten()
|
||||
counts.add_(spikes.to(torch.int32))
|
||||
if sources.numel():
|
||||
v[sources] = reset; refractory[sources] = ref_steps
|
||||
if backend == "hybrid-exact":
|
||||
starts = self.indptr[sources]
|
||||
degrees = self.indptr[sources + 1] - starts
|
||||
cumulative = torch.cumsum(degrees, 0)
|
||||
total = int(cumulative[-1].item()) if cumulative.numel() else 0
|
||||
edges = torch.empty(total, dtype=torch.int64, device=self.device)
|
||||
if total:
|
||||
self.extension.fill_edges(self.indptr, sources, cumulative - degrees, edges)
|
||||
else:
|
||||
edges = expand_rows(self.indptr, sources)
|
||||
if edges.numel():
|
||||
queue[(step + delay_steps) % qlen].scatter_add_(0, self.indices[edges].long(), self.weights[edges])
|
||||
elif backend.startswith("compile-"):
|
||||
assert compiled is not None
|
||||
arrival = queue[slot].clone(); queue[slot].zero_()
|
||||
v, g, refractory, spikes, available = compiled(v, g, refractory, arrival)
|
||||
accepted = forced[step].bool() & available[targets]
|
||||
forced_sources = targets[accepted]
|
||||
if forced_sources.numel(): spikes[forced_sources] = True
|
||||
sources = torch.nonzero(spikes, as_tuple=False).flatten()
|
||||
counts.add_(spikes.to(torch.int32))
|
||||
if sources.numel():
|
||||
v[sources] = reset; refractory[sources] = ref_steps
|
||||
edges = expand_rows(self.indptr, sources)
|
||||
if edges.numel(): queue[(step + delay_steps) % qlen].scatter_add_(0, self.indices[edges].long(), self.weights[edges])
|
||||
else:
|
||||
assert self.extension is not None
|
||||
if backend == "cudagraph-custom-state":
|
||||
assert staging is not None and graph is not None
|
||||
staging.copy_(queue[slot]); queue[slot].zero_(); graph.replay()
|
||||
else:
|
||||
self.extension.state_step(v, g, refractory, queue[slot], spike_u8, counts, gd, alpha, rest, reset, threshold, ref_steps)
|
||||
self.extension.force_step(v, refractory, spike_u8, counts, targets, forced[step], reset, ref_steps)
|
||||
self.extension.fanout(spike_u8, self.indptr, self.indices, self.weights, queue[(step + delay_steps) % qlen])
|
||||
sources = torch.nonzero(spike_u8, as_tuple=False).flatten() if capture_events else None
|
||||
if capture_events:
|
||||
if not (backend.startswith("custom") or backend == "cudagraph-custom-state"):
|
||||
pass
|
||||
assert sources is not None
|
||||
src = sources.detach().cpu().numpy().astype(np.int32, copy=True)
|
||||
event_nodes.append(src); event_steps.append(np.full(len(src), step, dtype=np.int32))
|
||||
torch.cuda.synchronize(); wall = time.perf_counter() - start; cpu = time.process_time() - cpu_start
|
||||
resources = sampler.stop() if sampler else {}
|
||||
v_np = v.cpu().numpy().copy(); g_np = g.cpu().numpy().copy(); counts_np = counts.cpu().numpy().copy()
|
||||
es = np.concatenate(event_steps) if event_steps else np.empty(0, dtype=np.int32)
|
||||
en = np.concatenate(event_nodes) if event_nodes else np.empty(0, dtype=np.int32)
|
||||
active = counts_np > 0
|
||||
super_values = self.superclass
|
||||
desc = np.isin(super_values, ["descending_neuron", "descending_neuron_tbc"])
|
||||
motor = np.isin(super_values, ["vnc_motor", "cb_motor"])
|
||||
bio_s = duration_ms / 1000.0
|
||||
result = {
|
||||
"backend": backend, "duration_ms": duration_ms, "steps": steps,
|
||||
"wall_clock_seconds": wall, "steps_per_second": steps / wall,
|
||||
"biological_realtime_speed": bio_s / wall, "realtime_factor_wall_over_biological": wall / bio_s,
|
||||
"total_spikes": int(counts_np.sum(dtype=np.int64)), "active_neurons": int(active.sum()),
|
||||
"mean_firing_rate_hz": float(counts_np.sum(dtype=np.int64) / (self.n * bio_s)),
|
||||
"descending_active_neurons": int(np.count_nonzero(active & desc)),
|
||||
"descending_spikes": int(counts_np[desc].sum(dtype=np.int64)),
|
||||
"motor_active_neurons": int(np.count_nonzero(active & motor)),
|
||||
"motor_spikes": int(counts_np[motor].sum(dtype=np.int64)),
|
||||
"peak_vram_bytes": int(torch.cuda.max_memory_allocated()),
|
||||
"cpu_process_utilization_percent_of_one_core": 100.0 * cpu / wall,
|
||||
"resource_sampling": resources,
|
||||
"event_count": int(len(en)),
|
||||
"event_sha256": hashlib.sha256(es.tobytes() + en.tobytes()).hexdigest() if capture_events else None,
|
||||
"spike_counts_sha256": hashlib.sha256(counts_np.tobytes()).hexdigest(),
|
||||
"final_voltage_sha256": hashlib.sha256(v_np.tobytes()).hexdigest(),
|
||||
"final_g_sha256": hashlib.sha256(g_np.tobytes()).hexdigest(),
|
||||
}
|
||||
return {"metrics": result, "event_steps": es, "event_nodes": en, "counts": counts_np, "v": v_np, "g": g_np}
|
||||
|
||||
|
||||
def compare(reference: dict[str, Any], candidate: dict[str, Any], tolerance: float) -> dict[str, Any]:
|
||||
events = np.array_equal(reference["event_steps"], candidate["event_steps"]) and np.array_equal(reference["event_nodes"], candidate["event_nodes"])
|
||||
counts = np.array_equal(reference["counts"], candidate["counts"])
|
||||
vdiff = np.abs(reference["v"] - candidate["v"])
|
||||
gdiff = np.abs(reference["g"] - candidate["g"])
|
||||
return {
|
||||
"spike_events_exact": bool(events), "per_neuron_spike_counts_exact": bool(counts),
|
||||
"final_voltage_max_abs_error": float(vdiff.max(initial=0)),
|
||||
"final_voltage_mean_abs_error": float(vdiff.mean()),
|
||||
"final_g_max_abs_error": float(gdiff.max(initial=0)), "final_g_mean_abs_error": float(gdiff.mean()),
|
||||
"state_within_tolerance": bool(vdiff.max(initial=0) <= tolerance and gdiff.max(initial=0) <= tolerance),
|
||||
"eligible_for_runtime_v1": bool(events and counts),
|
||||
}
|
||||
|
||||
|
||||
def save_reference_trial(exp: Path, spec: dict[str, Any], run: dict[str, Any]) -> None:
|
||||
base = exp / "regression/reference" / spec["name"]
|
||||
base.mkdir(parents=True, exist_ok=True)
|
||||
pq.write_table(pa.table({
|
||||
"step": pa.array(run["event_steps"], type=pa.int32()),
|
||||
"timestamp_ms": pa.array(run["event_steps"].astype(np.float64) * 0.1, type=pa.float64()),
|
||||
"neuron_idx": pa.array(run["event_nodes"], type=pa.int32()),
|
||||
}), base / "spike-events.parquet", compression="zstd")
|
||||
np.save(base / "per-neuron-spike-counts.int32.npy", run["counts"], allow_pickle=False)
|
||||
np.save(base / "final-voltage.float32.npy", run["v"], allow_pickle=False)
|
||||
np.save(base / "final-g.float32.npy", run["g"], allow_pickle=False)
|
||||
write_json(base / "statistics.json", run["metrics"] | {"trial": spec})
|
||||
|
||||
|
||||
def propagation_stats(run: dict[str, Any], targets: np.ndarray, distance: np.ndarray,
|
||||
superclass: np.ndarray) -> dict[str, Any]:
|
||||
active = run["counts"] > 0
|
||||
nonnegative = distance[active & (distance >= 0)]
|
||||
return {
|
||||
"active_neurons": int(active.sum()),
|
||||
"downstream_active_neurons": int(np.count_nonzero(active & ~np.isin(np.arange(len(active)), targets))),
|
||||
"propagation_depth": int(nonnegative.max(initial=0)),
|
||||
"descending_active_neurons": int(np.count_nonzero(active & np.isin(superclass, ["descending_neuron", "descending_neuron_tbc"]))),
|
||||
"motor_active_neurons": int(np.count_nonzero(active & np.isin(superclass, ["vnc_motor", "cb_motor"]))),
|
||||
}
|
||||
|
||||
|
||||
def visual_dynamics(exp: Path, cfg: dict[str, Any], runtime: Runtime,
|
||||
layers: dict[str, np.ndarray], indptr: np.ndarray, indices: np.ndarray) -> dict[str, Any]:
|
||||
spec = cfg["visual_dynamics"]; duration = float(spec["duration_ms"]); rate = float(spec["rate_hz"]); seed = int(spec["seed"])
|
||||
results: dict[str, Any] = {}
|
||||
tests: list[tuple[str, np.ndarray, float, int, torch.Tensor | None, str]] = []
|
||||
sensory = np.flatnonzero(layers["visual_sensory"]).astype(np.int64)
|
||||
tests.append(("A-forced-histaminergic-sensory", sensory, rate, seed, None, "engineering forcing of sensory output spikes"))
|
||||
steps = int(round(duration / float(cfg["model"]["dt_ms"])))
|
||||
tonic = generate_forced(steps, len(sensory), float(spec["tonic_histamine_baseline_hz"]), float(cfg["model"]["dt_ms"]), seed + 1, runtime.device)
|
||||
mod_start = int(round(float(spec["modulation_start_ms"]) / float(cfg["model"]["dt_ms"])))
|
||||
mod = generate_forced(steps - mod_start, len(sensory), float(spec["modulation_rate_hz"]), float(cfg["model"]["dt_ms"]), seed + 2, runtime.device)
|
||||
tonic[mod_start:] = mod
|
||||
tests.append(("B-tonic-histamine-rate-modulation", sensory, 0.0, seed + 1, tonic, "exploratory tonic output plus rate modulation; not baseline"))
|
||||
first = np.flatnonzero(layers["first_order"]).astype(np.int64)
|
||||
second = np.flatnonzero(layers["second_order"]).astype(np.int64)
|
||||
tests.append(("C-first-order-external-input", first, rate, seed + 3, None, "engineering input interface moved to real first-order targets"))
|
||||
tests.append(("C-second-order-external-input", second, rate, seed + 4, None, "engineering input interface moved to real second-order targets"))
|
||||
for j, name in enumerate(("visual_sensory", "first_order", "second_order", "higher_visual", "central_brain")):
|
||||
targets = np.flatnonzero(layers[name]).astype(np.int64)
|
||||
# Keep the audit bounded without inventing body IDs: deterministic full layer up to 25k, otherwise fixed index stride.
|
||||
if len(targets) > 25000:
|
||||
targets = targets[np.linspace(0, len(targets) - 1, 25000, dtype=np.int64)]
|
||||
tests.append((f"D-layer-{name}", targets, rate, seed + 10 + j, None, f"engineering external input at hierarchy layer {name}"))
|
||||
for name, targets, hz, trial_seed, override, interpretation in tests:
|
||||
LOG.info("visual dynamics %s targets=%d", name, len(targets))
|
||||
run = runtime.run("eager", targets, hz, trial_seed, duration, False, forced_override=override)
|
||||
distance = p4b.directed_distances(indptr, indices, targets, runtime.n)
|
||||
record = run["metrics"] | propagation_stats(run, targets, distance, runtime.superclass) | {
|
||||
"test": name, "stimulated_population_size": int(len(targets)), "interpretation": interpretation,
|
||||
"biological_interpretation": "histamine remains inhibitory; forced sensory spikes are an engineering interface, not phototransduction/receptor-potential biology",
|
||||
"baseline_status": "exploratory only" if name.startswith("B-") else "semantic test; no automatic selection",
|
||||
}
|
||||
results[name] = record
|
||||
write_json(exp / "visual-audit/dynamics" / f"{name}.json", record)
|
||||
write_json(exp / "visual-audit/dynamics-summary.json", results)
|
||||
return results
|
||||
|
||||
|
||||
def compiler_report(exp: Path, runtime: Runtime) -> dict[str, Any]:
|
||||
report: dict[str, Any] = {"torch_version": torch.__version__, "modes": {}}
|
||||
n = runtime.n
|
||||
m = runtime.cfg["model"]
|
||||
for mode in ("compile-default", "compile-reduce-overhead"):
|
||||
try:
|
||||
compiled, fn = runtime.compiled_state(mode)
|
||||
v = torch.full((n,), float(m["resting_potential_mv"]), device=runtime.device)
|
||||
g = torch.zeros(n, device=runtime.device); r = torch.zeros(n, dtype=torch.int32, device=runtime.device); a = torch.zeros(n, device=runtime.device)
|
||||
compiled(v, g, r, a); torch.cuda.synchronize()
|
||||
report["modes"][mode] = {"status": "compiled", "fullgraph": True, "graph_breaks_in_fused_static_state": 0}
|
||||
except Exception as exc:
|
||||
report["modes"][mode] = {"status": "failed", "error": repr(exc)}
|
||||
report["fullgraph_dynamic_csr"] = {
|
||||
"status": "not eligible",
|
||||
"graph_break_reasons": [
|
||||
"torch.nonzero produces data-dependent dynamic shape",
|
||||
"Python branching on sources.numel() and edge_ids.numel()",
|
||||
"expand_rows requires cumulative[-1].item() host synchronization",
|
||||
"dynamic torch.arange allocation and event-dependent CSR fan-out",
|
||||
],
|
||||
"action": "kept eager reference; custom CUDA prototype keeps spike detection/fan-out/accumulation on GPU",
|
||||
}
|
||||
write_json(exp / "compiler-reports/torch-compile.json", report)
|
||||
return report
|
||||
|
||||
|
||||
def build_extension(source_dir: Path, exp: Path) -> tuple[Any | None, dict[str, Any]]:
|
||||
cpp, cu = source_dir / "phase4c_ext.cpp", source_dir / "phase4c_ext.cu"
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
ext = load_extension(
|
||||
name="male_cns_lif_runtime_v1_ext", sources=[str(cpp), str(cu)],
|
||||
extra_cflags=["-O3"], extra_cuda_cflags=["-O3"], verbose=True,
|
||||
build_directory=str(exp / "compiler-reports/extension-build"),
|
||||
)
|
||||
result = {"status": "PASS", "build_seconds": time.perf_counter() - started, "sources": {str(cpp): sha256(cpp), str(cu): sha256(cu)}}
|
||||
return ext, result
|
||||
except Exception as exc:
|
||||
result = {"status": "FAILED", "build_seconds": time.perf_counter() - started, "error": repr(exc), "sources": {str(cpp): sha256(cpp), str(cu): sha256(cu)}}
|
||||
return None, result
|
||||
|
||||
|
||||
def seal_outputs(exp: Path, cfg: dict[str, Any], paths: dict[str, Path], code_paths: list[Path], summary: dict[str, Any]) -> None:
|
||||
controls = {exp / "manifests/manifest.json", exp / "manifests/checksums.sha256"}
|
||||
outputs = sorted(x for x in exp.rglob("*") if x.is_file() and x not in controls)
|
||||
manifest = {
|
||||
"phase": "4C", "created_unix_time": time.time(), "status": summary["status"],
|
||||
"software": {"python": platform.python_version(), "torch": torch.__version__, "numpy": np.__version__, "pyarrow": pa.__version__},
|
||||
"code": [{"path": str(x), "bytes": x.stat().st_size, "sha256": sha256(x)} for x in code_paths],
|
||||
"inputs": {name: {"path": str(path), "bytes": path.stat().st_size, "sha256": sha256(path)} for name, path in paths.items()},
|
||||
"outputs": [{"path": str(x), "bytes": x.stat().st_size, "sha256": sha256(x)} for x in outputs],
|
||||
"prohibitions_respected": cfg["prohibitions"],
|
||||
}
|
||||
write_json(exp / "manifests/manifest.json", manifest)
|
||||
targets = outputs + [exp / "manifests/manifest.json"]
|
||||
(exp / "manifests/checksums.sha256").write_text("".join(f"{sha256(x)} {x.relative_to(exp)}\n" for x in targets), encoding="utf-8")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
ap = argparse.ArgumentParser(); ap.add_argument("--config", required=True, type=Path); args = ap.parse_args()
|
||||
cfg = yaml.safe_load(args.config.read_text(encoding="utf-8"))
|
||||
exp = Path(cfg["paths"]["experiment_root"])
|
||||
for sub in ("configs", "visual-audit", "benchmarks", "compiler-reports", "profiler", "regression", "manifests", "audit"):
|
||||
(exp / sub).mkdir(parents=True, exist_ok=True)
|
||||
p4b.setup_logging(exp / "audit/phase4c.log")
|
||||
shutil.copy2(args.config, exp / "configs/phase4c.yaml")
|
||||
paths = exact_paths(cfg); hashes_before = {k: sha256(v) for k, v in paths.items()}
|
||||
indptr, indices, raw = p4b.validate_graph(cfg, paths)
|
||||
table = full_metadata(paths["metadata"]); pops = populations(table)
|
||||
expected_pop = {"visual": 6091, "auditory": 114}
|
||||
for name, count in expected_pop.items():
|
||||
if int(pops[name].sum()) != count: raise RuntimeError(f"{name} count mismatch")
|
||||
frozen = freeze_reference(cfg, args.config, paths)
|
||||
visual_summary, layers = visual_audit(exp, table, indptr, indices, raw)
|
||||
if not torch.cuda.is_available() or torch.cuda.device_count() != 1: raise RuntimeError("exactly one visible CUDA GPU required")
|
||||
device = torch.device("cuda:0"); props = torch.cuda.get_device_properties(device)
|
||||
if cfg["device"]["required_name_substring"] not in props.name: raise RuntimeError(f"wrong GPU: {props.name}")
|
||||
source_dir = Path(__file__).resolve().parent
|
||||
extension, extension_report = build_extension(source_dir, exp)
|
||||
write_json(exp / "compiler-reports/custom-extension.json", extension_report)
|
||||
runtime = Runtime(cfg, table, indptr, indices, raw, device, extension)
|
||||
compiler = compiler_report(exp, runtime)
|
||||
|
||||
reference_runs: dict[str, dict[str, Any]] = {}
|
||||
regression: dict[str, Any] = {}
|
||||
tol = float(cfg["regression"]["float_absolute_tolerance"])
|
||||
for spec in cfg["reference_trials"]:
|
||||
targets = np.flatnonzero(pops[spec["population"]]).astype(np.int64)
|
||||
LOG.info("reference trial %s", spec["name"])
|
||||
ref = runtime.run("eager", targets, float(spec["rate_hz"]), int(spec["seed"]), float(spec["duration_ms"]), True)
|
||||
dist = p4b.directed_distances(indptr, indices, targets, runtime.n)
|
||||
ref["metrics"].update(propagation_stats(ref, targets, dist, runtime.superclass))
|
||||
save_reference_trial(exp, spec, ref); reference_runs[spec["name"]] = ref
|
||||
candidates: dict[str, Any] = {}
|
||||
for backend in (["hybrid-exact", "custom-cuda"] if extension is not None else []):
|
||||
cand = runtime.run(backend, targets, float(spec["rate_hz"]), int(spec["seed"]), float(spec["duration_ms"]), True)
|
||||
comp = compare(ref, cand, tol)
|
||||
comp.update({
|
||||
"propagation_statistics_exact": propagation_stats(ref, targets, dist, runtime.superclass) == propagation_stats(cand, targets, dist, runtime.superclass),
|
||||
"reference_event_sha256": ref["metrics"]["event_sha256"], "candidate_event_sha256": cand["metrics"]["event_sha256"],
|
||||
})
|
||||
candidates[backend] = comp
|
||||
regression[spec["name"]] = candidates
|
||||
write_json(exp / "regression" / f"{spec['name']}.json", candidates)
|
||||
visual_results = visual_dynamics(exp, cfg, runtime, layers, indptr, indices)
|
||||
|
||||
bench: dict[str, Any] = {}
|
||||
auditory = np.flatnonzero(pops["auditory"]).astype(np.int64)
|
||||
baseline_by_duration: dict[float, dict[str, Any]] = {}
|
||||
for duration in cfg["benchmark"]["durations_ms"]:
|
||||
duration = float(duration); bench[str(duration)] = {}
|
||||
for backend in cfg["benchmark"]["backends"]:
|
||||
if (backend.startswith("custom") or backend in ("cudagraph-custom-state", "hybrid-exact")) and extension is None:
|
||||
bench[str(duration)][backend] = {"status": "unavailable", "reason": extension_report.get("error")}; continue
|
||||
LOG.info("benchmark backend=%s duration_ms=%s", backend, duration)
|
||||
try:
|
||||
run = runtime.run(backend, auditory, float(cfg["benchmark"]["rate_hz"]), int(cfg["benchmark"]["seed"]), duration, False, monitor=True)
|
||||
record = run["metrics"] | {"status": "PASS"}
|
||||
if backend == "eager": baseline_by_duration[duration] = run
|
||||
else:
|
||||
ref = baseline_by_duration.get(duration)
|
||||
if ref is not None:
|
||||
state_cmp = compare({**ref, "event_steps": np.empty(0,np.int32), "event_nodes": np.empty(0,np.int32)}, {**run, "event_steps": np.empty(0,np.int32), "event_nodes": np.empty(0,np.int32)}, tol)
|
||||
# Performance runs deliberately disable event capture. Do not
|
||||
# mislabel two empty arrays as an event-sequence regression.
|
||||
state_cmp["spike_events_exact"] = None
|
||||
state_cmp["event_sequence_checked"] = False
|
||||
record["benchmark_state_regression"] = state_cmp
|
||||
bench[str(duration)][backend] = record
|
||||
except Exception as exc:
|
||||
bench[str(duration)][backend] = {"status": "FAILED", "error": repr(exc)}
|
||||
write_json(exp / "benchmarks/runtime-matrix.json", bench)
|
||||
|
||||
eligible = ["eager"]
|
||||
if extension is not None and all(regression[x].get("hybrid-exact", {}).get("eligible_for_runtime_v1") for x in regression):
|
||||
eligible.append("hybrid-exact")
|
||||
if extension is not None and all(regression[x].get("custom-cuda", {}).get("eligible_for_runtime_v1") for x in regression):
|
||||
eligible.append("custom-cuda")
|
||||
# Compiled/CUDA-graph paths need an exact 1 s event regression before eligibility.
|
||||
audit_spec = cfg["reference_trials"][0]; audit_ref = reference_runs[audit_spec["name"]]
|
||||
targets = np.flatnonzero(pops[audit_spec["population"]]).astype(np.int64)
|
||||
backend_regression: dict[str, Any] = {}
|
||||
for backend in ("compile-default", "compile-reduce-overhead", "cudagraph-custom-state"):
|
||||
if backend == "cudagraph-custom-state" and extension is None: continue
|
||||
try:
|
||||
candidate = runtime.run(backend, targets, float(audit_spec["rate_hz"]), int(audit_spec["seed"]), 1000.0, True)
|
||||
backend_regression[backend] = compare(audit_ref, candidate, tol)
|
||||
if backend_regression[backend]["eligible_for_runtime_v1"]: eligible.append(backend)
|
||||
except Exception as exc:
|
||||
backend_regression[backend] = {"eligible_for_runtime_v1": False, "error": repr(exc)}
|
||||
write_json(exp / "regression/backend-eligibility.json", backend_regression)
|
||||
|
||||
def speed(backend: str) -> float:
|
||||
x = bench.get("10000.0", {}).get(backend, {})
|
||||
return float(x.get("biological_realtime_speed", 0.0)) if x.get("status") == "PASS" else 0.0
|
||||
fastest = max(eligible, key=speed)
|
||||
runtime_root = Path(cfg["paths"]["model_root"]) / "runtime-v1"
|
||||
runtime_root.mkdir(parents=True, exist_ok=True)
|
||||
runtime_sources = (Path(__file__).resolve(), source_dir / "phase4c_ext.cpp", source_dir / "phase4c_ext.cu", source_dir / "requirements.txt", args.config.resolve())
|
||||
for source in runtime_sources:
|
||||
copy_immutable(source, runtime_root / source.name)
|
||||
runtime_version = {
|
||||
"name": cfg["names"]["runtime"], "backend": fastest,
|
||||
"reference": cfg["names"]["reference"], "eligibility": backend_regression | regression,
|
||||
"ten_second_biological_realtime_speed": speed(fastest),
|
||||
"spike_behavior_requirement": "exact on all five 1 s frozen reference trials for custom runtime; exact auditory trial for compiler variants",
|
||||
"algorithm": "exact sparse event-driven CSR fan-out; no dense full-edge scan; exact 18-step GPU-resident ring buffer",
|
||||
"fallback_backend": "eager",
|
||||
}
|
||||
write_json(runtime_root / "VERSION.json", runtime_version); (runtime_root / "VERSION.json").chmod(0o444)
|
||||
hashes_after = {k: sha256(v) for k, v in paths.items()}
|
||||
if hashes_before != hashes_after: raise RuntimeError("Phase 3B inputs changed")
|
||||
summary = {
|
||||
"status": "PASS" if fastest else "FAILED", "reference_frozen": True,
|
||||
"reference": frozen, "visual_audit": visual_summary, "visual_dynamics": visual_results,
|
||||
"compiler_report": compiler, "custom_extension": extension_report,
|
||||
"regression": regression, "backend_regression": backend_regression,
|
||||
"runtime_v1_backend": fastest, "benchmarks": bench,
|
||||
"device": {"name": props.name, "total_memory_bytes": int(props.total_memory)},
|
||||
"input_hashes_unchanged": True, "entered_phase4d": False, "game_rl_web_3d_training": False,
|
||||
}
|
||||
write_json(exp / "benchmarks/summary.json", summary)
|
||||
seal_outputs(exp, cfg, paths, list(runtime_sources), summary)
|
||||
LOG.info("PHASE4C_COMPLETE runtime=%s", fastest)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,103 @@
|
||||
phase: 4C
|
||||
names:
|
||||
reference: MaleCNS-LIF-reference-v1
|
||||
runtime: MaleCNS-LIF-runtime-v1
|
||||
description: Shiu-like LIF dynamics adapted to MaleCNS v1.0
|
||||
|
||||
paths:
|
||||
connectome_root: /data/digital-drosophila/datasets/male-cns-v1.0/processed/neuron-only
|
||||
phase4b_root: /data/digital-drosophila/experiments/phase4b-dynamics
|
||||
model_root: /data/digital-drosophila/models/male-cns-lif
|
||||
experiment_root: /data/digital-drosophila/experiments/phase4c-runtime
|
||||
|
||||
expected_graph:
|
||||
neurons: 166700
|
||||
directed_edges: 25582938
|
||||
synaptic_contact_weight_sum: 124177617
|
||||
isolated_neurons: 217
|
||||
self_loops: 101
|
||||
|
||||
device:
|
||||
required_name_substring: RTX 4090
|
||||
launcher_visible_gpu_uuid: GPU-79f3a85d-adce-7833-7003-6488140ca596
|
||||
mixed_multi_gpu: false
|
||||
|
||||
model:
|
||||
dt_ms: 0.1
|
||||
membrane_tau_ms: 20.0
|
||||
synaptic_tau_ms: 5.0
|
||||
refractory_ms: 2.2
|
||||
resting_potential_mv: -52.0
|
||||
reset_potential_mv: -52.0
|
||||
threshold_mv: -45.0
|
||||
synaptic_delay_ms: 1.8
|
||||
w_syn_mv_per_anatomical_synapse: 0.275
|
||||
integration: forward_euler
|
||||
refractory_g_policy: g continues to decay and receive exact delayed events while v is clamped at V_reset
|
||||
|
||||
sign_policy:
|
||||
neurotransmitter_sign:
|
||||
acetylcholine: 1
|
||||
gaba: -1
|
||||
glutamate: -1
|
||||
histamine: -1
|
||||
dopamine: 1
|
||||
octopamine: 1
|
||||
serotonin: 1
|
||||
unknown_sign: 0
|
||||
unclear_sign: 0
|
||||
unmapped_policy: error
|
||||
|
||||
sensory_populations:
|
||||
visual:
|
||||
expression: annotation_class == 'visual'
|
||||
mechanosensory:
|
||||
expression: annotation_class in {'mechanosensory','mechanosensory_tactile','mechanosensory_proprioceptive'} and annotation_subclass != 'auditory'
|
||||
olfactory:
|
||||
expression: annotation_class == 'olfactory'
|
||||
gustatory:
|
||||
expression: annotation_class == 'gustatory'
|
||||
auditory:
|
||||
expression: annotation_subclass == 'auditory'
|
||||
|
||||
reference_trials:
|
||||
- {name: auditory-10hz, population: auditory, rate_hz: 10, seed: 42001, duration_ms: 1000.0}
|
||||
- {name: auditory-100hz, population: auditory, rate_hz: 100, seed: 42001, duration_ms: 1000.0}
|
||||
- {name: olfactory-50hz, population: olfactory, rate_hz: 50, seed: 42003, duration_ms: 1000.0}
|
||||
- {name: mechanosensory-50hz, population: mechanosensory, rate_hz: 50, seed: 42003, duration_ms: 1000.0}
|
||||
- {name: gustatory-50hz, population: gustatory, rate_hz: 50, seed: 42003, duration_ms: 1000.0}
|
||||
|
||||
visual_dynamics:
|
||||
duration_ms: 100.0
|
||||
rate_hz: 50
|
||||
seed: 43100
|
||||
tonic_histamine_baseline_hz: 50
|
||||
modulation_rate_hz: 10
|
||||
modulation_start_ms: 50.0
|
||||
|
||||
benchmark:
|
||||
population: auditory
|
||||
rate_hz: 10
|
||||
seed: 43200
|
||||
durations_ms: [10.0, 100.0, 1000.0, 10000.0]
|
||||
gpu_sample_interval_seconds: 0.1
|
||||
backends: [eager, hybrid-exact, compile-default, compile-reduce-overhead, cudagraph-custom-state, custom-cuda]
|
||||
warmup_ms: 10.0
|
||||
|
||||
regression:
|
||||
require_exact_spike_events: true
|
||||
float_absolute_tolerance: 1.0e-5
|
||||
reject_on_spike_divergence: true
|
||||
|
||||
reproducibility:
|
||||
global_seed: 43000
|
||||
local_repository_commit: d826733f1c620a2ad0b096c846b9b8e9eb3ac0fa
|
||||
deterministic_forced_inputs: true
|
||||
source_phase: Phase 4B PASS
|
||||
|
||||
prohibitions:
|
||||
alter_connectome: false
|
||||
alter_raw_synapse_counts: false
|
||||
alter_transmitter_signs: false
|
||||
train_weights: false
|
||||
game_rl_web_3d: false
|
||||
@@ -0,0 +1,23 @@
|
||||
#include <torch/extension.h>
|
||||
|
||||
void state_step_cuda(torch::Tensor v, torch::Tensor g, torch::Tensor refractory,
|
||||
torch::Tensor arrival, torch::Tensor spikes,
|
||||
torch::Tensor spike_counts, double g_decay, double alpha,
|
||||
double rest, double reset, double threshold,
|
||||
int64_t refractory_steps);
|
||||
void force_step_cuda(torch::Tensor v, torch::Tensor refractory,
|
||||
torch::Tensor spikes, torch::Tensor spike_counts,
|
||||
torch::Tensor targets, torch::Tensor forced,
|
||||
double reset, int64_t refractory_steps);
|
||||
void fanout_cuda(torch::Tensor spikes, torch::Tensor indptr,
|
||||
torch::Tensor indices, torch::Tensor weights,
|
||||
torch::Tensor destination);
|
||||
void fill_edges_cuda(torch::Tensor indptr, torch::Tensor sources,
|
||||
torch::Tensor output_offsets, torch::Tensor output);
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("state_step", &state_step_cuda, "Fused exact LIF state update (CUDA)");
|
||||
m.def("force_step", &force_step_cuda, "Apply external forced spikes (CUDA)");
|
||||
m.def("fanout", &fanout_cuda, "Event-driven CSR fan-out (CUDA)");
|
||||
m.def("fill_edges", &fill_edges_cuda, "Exact source-ordered CSR edge expansion (CUDA)");
|
||||
}
|
||||
@@ -0,0 +1,152 @@
|
||||
#include <torch/extension.h>
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cstdint>
|
||||
|
||||
namespace {
|
||||
|
||||
__global__ void state_kernel(float* v, float* g, int* refractory,
|
||||
float* arrival, unsigned char* spikes,
|
||||
int* spike_counts, int n, float g_decay,
|
||||
float alpha, float rest, float reset,
|
||||
float threshold, int refractory_steps) {
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i >= n) return;
|
||||
// Phase 4B executes mul_ and add_ as separate kernels: prevent contraction.
|
||||
float gi = __fadd_rn(__fmul_rn(g[i], g_decay), arrival[i]);
|
||||
g[i] = gi;
|
||||
arrival[i] = 0.0f;
|
||||
int r = refractory[i];
|
||||
r = r > 0 ? r - 1 : 0;
|
||||
refractory[i] = r;
|
||||
float vi = v[i];
|
||||
bool available = r == 0;
|
||||
if (available) {
|
||||
// Match the separate float32 operations used by the eager reference.
|
||||
float leak = __fsub_rn(vi, rest);
|
||||
float drive = __fsub_rn(gi, leak);
|
||||
vi = __fadd_rn(vi, __fmul_rn(alpha, drive));
|
||||
} else {
|
||||
vi = reset;
|
||||
}
|
||||
bool fired = available && vi >= threshold;
|
||||
if (fired) {
|
||||
vi = reset;
|
||||
refractory[i] = refractory_steps;
|
||||
spike_counts[i] += 1;
|
||||
}
|
||||
v[i] = vi;
|
||||
spikes[i] = fired ? 1 : 0;
|
||||
}
|
||||
|
||||
__global__ void force_kernel(float* v, int* refractory,
|
||||
unsigned char* spikes, int* spike_counts,
|
||||
const int64_t* targets,
|
||||
const unsigned char* forced, int count,
|
||||
float reset, int refractory_steps) {
|
||||
int j = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (j >= count || !forced[j]) return;
|
||||
int64_t i = targets[j];
|
||||
// A natural spike already recorded at this target represents the same event.
|
||||
if (spikes[i]) return;
|
||||
if (refractory[i] == 0) {
|
||||
spikes[i] = 1;
|
||||
v[i] = reset;
|
||||
refractory[i] = refractory_steps;
|
||||
spike_counts[i] += 1;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void fanout_kernel(const unsigned char* spikes,
|
||||
const int64_t* indptr, const int* indices,
|
||||
const float* weights, float* destination,
|
||||
int n) {
|
||||
int src = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (src >= n || !spikes[src]) return;
|
||||
int64_t start = indptr[src];
|
||||
int64_t stop = indptr[src + 1];
|
||||
for (int64_t e = start; e < stop; ++e) {
|
||||
atomicAdd(destination + indices[e], weights[e]);
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void fill_edges_kernel(const int64_t* indptr,
|
||||
const int64_t* sources,
|
||||
const int64_t* output_offsets,
|
||||
int64_t* output, int source_count) {
|
||||
int j = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (j >= source_count) return;
|
||||
int64_t src = sources[j];
|
||||
int64_t begin = indptr[src];
|
||||
int64_t degree = indptr[src + 1] - begin;
|
||||
int64_t out = output_offsets[j];
|
||||
for (int64_t k = 0; k < degree; ++k) output[out + k] = begin + k;
|
||||
}
|
||||
|
||||
void check_cuda(const torch::Tensor& x) {
|
||||
TORCH_CHECK(x.is_cuda(), "all tensors must be CUDA tensors");
|
||||
TORCH_CHECK(x.is_contiguous(), "all tensors must be contiguous");
|
||||
}
|
||||
} // namespace
|
||||
|
||||
void state_step_cuda(torch::Tensor v, torch::Tensor g, torch::Tensor refractory,
|
||||
torch::Tensor arrival, torch::Tensor spikes,
|
||||
torch::Tensor spike_counts, double g_decay, double alpha,
|
||||
double rest, double reset, double threshold,
|
||||
int64_t refractory_steps) {
|
||||
check_cuda(v); check_cuda(g); check_cuda(refractory); check_cuda(arrival);
|
||||
check_cuda(spikes); check_cuda(spike_counts);
|
||||
TORCH_CHECK(v.scalar_type() == torch::kFloat32 && g.scalar_type() == torch::kFloat32,
|
||||
"v/g must be float32");
|
||||
TORCH_CHECK(refractory.scalar_type() == torch::kInt32 && spike_counts.scalar_type() == torch::kInt32,
|
||||
"refractory/counts must be int32");
|
||||
TORCH_CHECK(spikes.scalar_type() == torch::kUInt8, "spikes must be uint8");
|
||||
int n = v.numel();
|
||||
state_kernel<<<(n + 255) / 256, 256>>>(
|
||||
v.data_ptr<float>(), g.data_ptr<float>(), refractory.data_ptr<int>(),
|
||||
arrival.data_ptr<float>(), spikes.data_ptr<unsigned char>(),
|
||||
spike_counts.data_ptr<int>(), n, (float)g_decay, (float)alpha,
|
||||
(float)rest, (float)reset, (float)threshold, (int)refractory_steps);
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
}
|
||||
|
||||
void force_step_cuda(torch::Tensor v, torch::Tensor refractory,
|
||||
torch::Tensor spikes, torch::Tensor spike_counts,
|
||||
torch::Tensor targets, torch::Tensor forced,
|
||||
double reset, int64_t refractory_steps) {
|
||||
check_cuda(targets); check_cuda(forced);
|
||||
int count = targets.numel();
|
||||
force_kernel<<<(count + 255) / 256, 256>>>(
|
||||
v.data_ptr<float>(), refractory.data_ptr<int>(),
|
||||
spikes.data_ptr<unsigned char>(), spike_counts.data_ptr<int>(),
|
||||
targets.data_ptr<int64_t>(), forced.data_ptr<unsigned char>(), count,
|
||||
(float)reset, (int)refractory_steps);
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
}
|
||||
|
||||
void fanout_cuda(torch::Tensor spikes, torch::Tensor indptr,
|
||||
torch::Tensor indices, torch::Tensor weights,
|
||||
torch::Tensor destination) {
|
||||
check_cuda(spikes); check_cuda(indptr); check_cuda(indices);
|
||||
check_cuda(weights); check_cuda(destination);
|
||||
int n = spikes.numel();
|
||||
fanout_kernel<<<(n + 255) / 256, 256>>>(
|
||||
spikes.data_ptr<unsigned char>(), indptr.data_ptr<int64_t>(),
|
||||
indices.data_ptr<int>(), weights.data_ptr<float>(),
|
||||
destination.data_ptr<float>(), n);
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
}
|
||||
|
||||
void fill_edges_cuda(torch::Tensor indptr, torch::Tensor sources,
|
||||
torch::Tensor output_offsets, torch::Tensor output) {
|
||||
check_cuda(indptr); check_cuda(sources); check_cuda(output_offsets); check_cuda(output);
|
||||
TORCH_CHECK(indptr.scalar_type() == torch::kInt64 && sources.scalar_type() == torch::kInt64 &&
|
||||
output_offsets.scalar_type() == torch::kInt64 && output.scalar_type() == torch::kInt64,
|
||||
"fill_edges tensors must be int64");
|
||||
int count = sources.numel();
|
||||
fill_edges_kernel<<<(count + 255) / 256, 256>>>(
|
||||
indptr.data_ptr<int64_t>(), sources.data_ptr<int64_t>(),
|
||||
output_offsets.data_ptr<int64_t>(), output.data_ptr<int64_t>(), count);
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
numpy==2.2.6
|
||||
pyarrow==25.0.1
|
||||
PyYAML==6.0.3
|
||||
torch==2.14.0
|
||||
ninja==1.13.2
|
||||
Reference in New Issue
Block a user