"""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