Digital Drosophila Archive
276 lines
16 KiB
Python
276 lines
16 KiB
Python
"""Small-case diagnostic CUDA backend; SI states, explicit dtype rounding.
|
|
|
|
run_gpu is serial, records CPU diagnostics, and is not a performance runtime.
|
|
Coefficients and optional weights must come from the same-precision Brian2 oracle.
|
|
"""
|
|
from functools import lru_cache
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
from pathlib import Path
|
|
import time
|
|
import traceback
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torch.utils.cpp_extension import load
|
|
|
|
|
|
SOURCE_ROOT = Path(__file__).resolve().parent
|
|
COEFFICIENT_NAMES = ("gd", "c0", "cg", "cv", "threshold", "reset", "seed_volt")
|
|
|
|
|
|
def require(condition, message):
|
|
if not condition:
|
|
raise ValueError(message)
|
|
|
|
|
|
def write_json(path, data):
|
|
with Path(path).open("x", encoding="utf-8") as output:
|
|
json.dump(data, output, indent=2, allow_nan=False)
|
|
output.write("\n")
|
|
|
|
|
|
def digest(path):
|
|
checksum = hashlib.sha256()
|
|
with Path(path).open("rb") as source:
|
|
for block in iter(lambda: source.read(1 << 20), b""):
|
|
checksum.update(block)
|
|
return checksum.hexdigest()
|
|
|
|
|
|
@lru_cache(maxsize=None)
|
|
def _load_extension(build_path, signature):
|
|
build = Path(build_path) / ("author_cuda_" + signature[:16])
|
|
build.mkdir(parents=True, exist_ok=True)
|
|
module = load(name="author_cuda_" + signature[:16],
|
|
sources=[str(SOURCE_ROOT / "author_ext.cpp"), str(SOURCE_ROOT / "author_ext.cu")],
|
|
build_directory=str(build), extra_cflags=["-O2"],
|
|
extra_cuda_cflags=["-O2", "--fmad=false", "--ftz=false", "--prec-div=true", "--prec-sqrt=true"],
|
|
verbose=False)
|
|
return module
|
|
|
|
|
|
def build_extension(build_root):
|
|
provenance = {"sources": {name: digest(SOURCE_ROOT / name)
|
|
for name in ("author_ext.cpp", "author_ext.cu")},
|
|
"torch": torch.__version__, "torch_cuda": torch.version.cuda,
|
|
"arch_list": os.environ.get("TORCH_CUDA_ARCH_LIST"),
|
|
"cuda_flags": ["-O2", "--fmad=false", "--ftz=false", "--prec-div=true", "--prec-sqrt=true"]}
|
|
signature = hashlib.sha256(json.dumps(provenance, sort_keys=True).encode()).hexdigest()
|
|
return _load_extension(str(Path(build_root).resolve()), signature), provenance
|
|
|
|
|
|
def read_case(case_dir, precision, coefficients, weights):
|
|
case_dir = Path(case_dir)
|
|
require(precision in ("float32", "float64"), "precision must be float32 or float64")
|
|
dtype = np.dtype(precision)
|
|
with (case_dir / "spec.json").open(encoding="utf-8") as source:
|
|
spec = json.load(source)
|
|
for key in ("dt_ms", "steps", "duration_ms", "model"):
|
|
require(key in spec, "Missing spec field: " + key)
|
|
model = spec["model"]
|
|
for key in ("threshold_mv", "reset_potential_mv", "resting_potential_mv",
|
|
"synaptic_delay_ms", "refractory_ms", "w_syn_mv_per_anatomical_synapse"):
|
|
require(key in model and math.isfinite(float(model[key])), "Missing/nonfinite model field: " + key)
|
|
steps = spec["steps"]
|
|
require(type(steps) is int and 0 < steps < np.iinfo(np.int32).max, "Invalid step count")
|
|
dt = float(spec["dt_ms"])
|
|
require(math.isfinite(dt) and dt > 0, "Invalid dt")
|
|
require(math.isfinite(float(spec["duration_ms"])) and
|
|
round(float(spec["duration_ms"]) / dt) == steps, "Duration/steps mismatch")
|
|
delay = round(float(model["synaptic_delay_ms"]) / dt)
|
|
refractory = round(float(model["refractory_ms"]) / dt)
|
|
require(delay == 18 and refractory == 22, "This declared backend requires delay=18 and refractory=22")
|
|
require(float(model["w_syn_mv_per_anatomical_synapse"]) == 0.275, "W_syn differs from frozen value")
|
|
coeff = {}
|
|
for key in COEFFICIENT_NAMES:
|
|
require(key in coefficients and math.isfinite(float(coefficients[key])), "Missing/nonfinite coefficient: " + key)
|
|
coeff[key] = float(coefficients[key])
|
|
require(bool(np.isfinite(dtype.type(coeff[key]))), "Coefficient overflows requested dtype: " + key)
|
|
require(dtype.type(coeff["threshold"]) == dtype.type(float(model["threshold_mv"]) * .001),
|
|
"Oracle threshold differs from case spec")
|
|
require(dtype.type(coeff["reset"]) == dtype.type(float(model["reset_potential_mv"]) * .001),
|
|
"Oracle reset differs from case spec")
|
|
with np.load(case_dir / "arrays.npz", allow_pickle=False) as inputs:
|
|
arrays = {key: inputs[key].copy() for key in inputs.files}
|
|
required_types = {"nodes": "int32", "edge_id": "int64", "src": "int32", "dst": "int32",
|
|
"signed_raw": "int64", "requests": "uint8", "external_mv": "float64",
|
|
"initial_v_mv": "float64", "initial_g_mv": "float64"}
|
|
for key, expected in required_types.items():
|
|
require(key in arrays and arrays[key].dtype == np.dtype(expected), "Missing/wrong array dtype: " + key)
|
|
nodes = arrays["nodes"]
|
|
require(nodes.ndim == 1 and len(np.unique(nodes)) == len(nodes), "Invalid neuron identities")
|
|
n = len(nodes)
|
|
require(n < 512, "This diagnostic backend is restricted to fewer than 512 cells")
|
|
edge_id = arrays["edge_id"]
|
|
require(edge_id.ndim == 1 and len(np.unique(edge_id)) == len(edge_id), "Invalid edge identities")
|
|
e = len(edge_id)
|
|
for key in ("src", "dst", "signed_raw"):
|
|
require(arrays[key].shape == (e,), "Edge vector shape mismatch: " + key)
|
|
for key in ("src", "dst"):
|
|
require(bool(np.all((arrays[key] >= 0) & (arrays[key] < n))), "Local edge endpoint out of bounds")
|
|
require(np.array_equal(np.lexsort((edge_id, arrays["src"])), np.arange(e)),
|
|
"Edges must already be canonical source/edge_id sorted; indices will not be silently reordered")
|
|
for key in ("requests", "external_mv"):
|
|
require(arrays[key].shape == (steps, n), "Trial array shape mismatch: " + key)
|
|
require(bool(np.all(arrays["requests"] <= 1)), "Request array is not Boolean uint8")
|
|
for key in ("initial_v_mv", "initial_g_mv"):
|
|
require(arrays[key].shape == (n,), "Initial state shape mismatch: " + key)
|
|
for key in ("external_mv", "initial_v_mv", "initial_g_mv"):
|
|
require(bool(np.isfinite(arrays[key]).all()), "Nonfinite case array: " + key)
|
|
if weights is None:
|
|
weight_array = np.asarray(arrays["signed_raw"] * (.275 * .001), dtype=dtype)
|
|
weight_origin = "signed_raw * (0.275 * 0.001), then cast to requested dtype"
|
|
else:
|
|
weight_array = np.asarray(weights, dtype=dtype)
|
|
weight_origin = "caller supplied oracle SI weight array"
|
|
require(weight_array.shape == (e,) and bool(np.isfinite(weight_array).all()), "Invalid SI weights")
|
|
for key in ("initial_v_mv", "initial_g_mv", "external_mv"):
|
|
converted = np.asarray(arrays[key] * .001, dtype=dtype)
|
|
require(bool(np.isfinite(converted).all()), "Conversion overflow: " + key)
|
|
return spec, arrays, coeff, np.ascontiguousarray(weight_array), weight_origin
|
|
|
|
|
|
def run_gpu(case_dir, out_dir, precision, coefficients, build_root, weights=None):
|
|
"""Run one case; return state dict and save SI state plus GPU-selected arrivals.
|
|
|
|
``weights`` is optional SI volts in the original arrays.npz edge-row order.
|
|
Repeated calls share a cached extension; each needs a new exclusive out_dir.
|
|
"""
|
|
case_dir, out_dir = Path(case_dir), Path(out_dir)
|
|
require(not out_dir.exists(), "Refusing existing output directory: " + str(out_dir))
|
|
spec, arrays, coeff, weight_array, weight_origin = read_case(case_dir, precision, coefficients, weights)
|
|
require(torch.cuda.is_available(), "CUDA unavailable")
|
|
require(torch.cuda.device_count() == 1 and "4090" in torch.cuda.get_device_name(), "Only the single declared RTX4090 is permitted")
|
|
out_dir.mkdir(parents=True, exist_ok=False)
|
|
start = time.perf_counter()
|
|
try:
|
|
ext, build_provenance = build_extension(build_root)
|
|
build_done = time.perf_counter()
|
|
device = torch.device("cuda", torch.cuda.current_device())
|
|
torch.cuda.reset_peak_memory_stats(device)
|
|
dtype = torch.float32 if precision == "float32" else torch.float64
|
|
numpy_dtype = np.dtype(precision)
|
|
n, steps = len(arrays["nodes"]), spec["steps"]
|
|
|
|
def upload(value, target_dtype=None):
|
|
return torch.as_tensor(np.ascontiguousarray(value), device=device, dtype=target_dtype).contiguous()
|
|
|
|
v = upload(np.asarray(arrays["initial_v_mv"] * .001, dtype=numpy_dtype), dtype)
|
|
g = upload(np.asarray(arrays["initial_g_mv"] * .001, dtype=numpy_dtype), dtype)
|
|
ref = torch.zeros(n, dtype=torch.int32, device=device)
|
|
counts = torch.zeros(n, dtype=torch.int32, device=device)
|
|
requests = upload(arrays["requests"], torch.uint8)
|
|
accepted_requests = torch.empty_like(requests)
|
|
rejected_requests = torch.empty_like(requests)
|
|
external = upload(np.asarray(arrays["external_mv"] * .001, dtype=numpy_dtype), dtype)
|
|
synaptic_weights = upload(weight_array, dtype)
|
|
dst = upload(arrays["dst"], torch.int32)
|
|
degrees_np = np.bincount(arrays["src"], minlength=n).astype(np.int64)
|
|
indptr_np = np.concatenate((np.zeros(1, np.int64), np.cumsum(degrees_np, dtype=np.int64)))
|
|
indptr = upload(indptr_np, torch.int64)
|
|
ring = torch.zeros((19, n), dtype=torch.uint8, device=device)
|
|
history = torch.empty((steps, 3, n), dtype=dtype, device=device)
|
|
event_steps, event_nodes = [], []
|
|
arrival_steps, arrival_indices, arrival_accepted = [], [], []
|
|
torch.cuda.synchronize(device)
|
|
simulation_start = time.perf_counter()
|
|
for step in range(steps):
|
|
spikes = ring[step % 19]
|
|
# The state kernel decrements refractory before considering a request.
|
|
# These diagnostics use that exact pre-kernel eligibility on the GPU.
|
|
available_after_decrement = ref <= 1
|
|
accepted_requests[step].copy_(requests[step] & available_after_decrement.to(torch.uint8))
|
|
rejected_requests[step].copy_(requests[step] & (~available_after_decrement).to(torch.uint8))
|
|
ext.state_step(v, g, ref, requests[step], spikes, counts,
|
|
*[coeff[key] for key in COEFFICIENT_NAMES], 22)
|
|
if step >= 18:
|
|
due = ring[(step - 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]
|
|
order = torch.argsort(destination, stable=True)
|
|
edge_indices = expanded[order].contiguous()
|
|
ordered_targets = destination[order]
|
|
boundary = torch.ones(edge_count, dtype=torch.bool, device=device)
|
|
boundary[1:] = ordered_targets[1:] != ordered_targets[:-1]
|
|
starts = torch.nonzero(boundary, as_tuple=False).flatten().contiguous()
|
|
targets = ordered_targets[starts].contiguous()
|
|
accepted = (ref[ordered_targets.long()] == 0).to(torch.uint8)
|
|
# Capture the actual CUDA-expanded and GPU-stable-sorted edge stream.
|
|
arrival_steps.append(np.full(edge_count, step, dtype=np.int32))
|
|
arrival_indices.append(edge_indices.cpu().numpy().copy())
|
|
arrival_accepted.append(accepted.cpu().numpy().copy())
|
|
ext.apply_arrivals(edge_indices, starts, targets, synaptic_weights, g, ref)
|
|
ext.apply_external(external[step], g, ref)
|
|
history[step, 0].copy_(v)
|
|
history[step, 1].copy_(g)
|
|
history[step, 2].copy_(ref)
|
|
fired = torch.nonzero(spikes, as_tuple=False).flatten().cpu().numpy()
|
|
if len(fired):
|
|
event_steps.append(np.full(len(fired), step, dtype=np.int32))
|
|
event_nodes.append(arrays["nodes"][fired])
|
|
torch.cuda.synchronize(device)
|
|
simulation_done = time.perf_counter()
|
|
|
|
def joined(chunks, dt):
|
|
return np.concatenate(chunks).astype(dt, copy=False) if chunks else np.empty(0, dtype=dt)
|
|
|
|
result = {"nodes": arrays["nodes"], "history": history.cpu().numpy(),
|
|
"counts": counts.cpu().numpy(), "v": v.cpu().numpy(), "g": g.cpu().numpy(),
|
|
"ref": ref.cpu().numpy(), "weights": synaptic_weights.cpu().numpy(),
|
|
"step": joined(event_steps, np.int32), "neuron_idx": joined(event_nodes, np.int32),
|
|
"request_attempted": arrays["requests"].copy(),
|
|
"request_accepted": accepted_requests.cpu().numpy(),
|
|
"request_rejected": rejected_requests.cpu().numpy()}
|
|
arrival_index = joined(arrival_indices, np.int64)
|
|
arrivals = {"step": joined(arrival_steps, np.int32), "edge_index": arrival_index,
|
|
"edge_id": arrays["edge_id"][arrival_index],
|
|
"target_local": arrays["dst"][arrival_index],
|
|
"accepted": joined(arrival_accepted, np.uint8)}
|
|
result.update(arrival_step=arrivals["step"], arrival_edge_index=arrivals["edge_index"],
|
|
arrival_accepted=arrivals["accepted"])
|
|
with (out_dir / "state.npz").open("xb") as target:
|
|
np.savez_compressed(target, **result)
|
|
with (out_dir / "arrival-order.npz").open("xb") as target:
|
|
np.savez_compressed(target, **arrivals)
|
|
finite = bool(all(np.isfinite(value).all() for value in result.values()))
|
|
metrics = {"precision": precision, "state_units": "SI volts; refractory is steps",
|
|
"coefficients_double_input": coeff,
|
|
"coefficients_effective_dtype": {key: float(numpy_dtype.type(value)) for key, value in coeff.items()},
|
|
"weights_origin": weight_origin, "nodes": n, "edges": len(weight_array), "steps": steps,
|
|
"total_spikes": len(result["step"]), "arriving_edges": len(arrival_index),
|
|
"accepted_arriving_edges": int(np.count_nonzero(arrivals["accepted"])),
|
|
"requested": int(np.count_nonzero(result["request_attempted"])),
|
|
"accepted_requests": int(np.count_nonzero(result["request_accepted"])),
|
|
"rejected_requests": int(np.count_nonzero(result["request_rejected"])),
|
|
"all_finite": finite, "delay_steps": 18, "refractory_steps": 22,
|
|
"build_or_cache_seconds": build_done - start,
|
|
"diagnostic_simulation_seconds": simulation_done - simulation_start,
|
|
"wall_seconds": time.perf_counter() - start,
|
|
"gpu": torch.cuda.get_device_name(device), "device": str(device),
|
|
"peak_allocated_vram_bytes": int(torch.cuda.max_memory_allocated(device)),
|
|
"peak_reserved_vram_bytes": int(torch.cuda.max_memory_reserved(device)),
|
|
"build": build_provenance,
|
|
"input_sha256": {name: digest(case_dir / name) for name in ("spec.json", "arrays.npz")},
|
|
"arrival_order_scope": "GPU expanded edges, stable target grouping; compare relative order per tick and target",
|
|
"performance_runtime_claim": False, "oracle_validation_claim": False,
|
|
"whole_brain_run": False, "learning": False,
|
|
"cpu_transfer_each_step_for_diagnostics": True}
|
|
write_json(out_dir / "metrics.json", metrics)
|
|
require(finite, "Nonfinite GPU state; raw trace preserved")
|
|
return result
|
|
except Exception as error:
|
|
write_json(out_dir / "error.json", {"type": type(error).__name__, "message": str(error),
|
|
"traceback": traceback.format_exc(), "precision": precision,
|
|
"validation_pass": False})
|
|
raise
|