Digital Drosophila Archive
176 lines
8.0 KiB
Python
176 lines
8.0 KiB
Python
"""Frozen eager Torch float32 APL component and binary oracle sidecar.
|
|
|
|
This is the Phase 4P operator sequence, not an author-neuron implementation.
|
|
No precision/determinism flags are changed here. The full matrix dimensions
|
|
are retained even when the caller simulates a small neuronal fixture.
|
|
"""
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import math
|
|
import os
|
|
from pathlib import Path
|
|
import struct
|
|
import sys
|
|
import traceback
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
|
|
PROTOCOL = "phase4ag-apl-v1"
|
|
N_NEURONS, N_APL, N_KC = 166700, 2, 4064
|
|
BUNDLE_FILES = ("input_matrix.npy", "output_matrix.npy", "kc_globals.npy", "apl_globals.npy")
|
|
FRAME = struct.Struct("<qQ")
|
|
PAYLOAD_BYTES = (N_APL + N_APL + N_KC) * 4
|
|
|
|
|
|
def require(condition, message):
|
|
if not condition:
|
|
raise ValueError(message)
|
|
|
|
|
|
def file_sha256(path):
|
|
digest = hashlib.sha256()
|
|
with Path(path).open("rb") as stream:
|
|
for block in iter(lambda: stream.read(8 << 20), b""):
|
|
digest.update(block)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def array_sha256(value):
|
|
return hashlib.sha256(np.ascontiguousarray(value).tobytes()).hexdigest()
|
|
|
|
|
|
def environment_metadata():
|
|
device = torch.cuda.current_device()
|
|
return {"torch": str(torch.__version__), "numpy": str(np.__version__),
|
|
"cuda": torch.version.cuda, "cudnn": torch.backends.cudnn.version(),
|
|
"device_name": torch.cuda.get_device_name(device), "device_index": device,
|
|
"device_capability": list(torch.cuda.get_device_capability(device)),
|
|
"deterministic_algorithms": torch.are_deterministic_algorithms_enabled(),
|
|
"cudnn_deterministic": torch.backends.cudnn.deterministic,
|
|
"cudnn_benchmark": torch.backends.cudnn.benchmark,
|
|
"cuda_matmul_allow_tf32": torch.backends.cuda.matmul.allow_tf32,
|
|
"cudnn_allow_tf32": torch.backends.cudnn.allow_tf32,
|
|
"float32_matmul_precision": torch.get_float32_matmul_precision(),
|
|
"CUBLAS_WORKSPACE_CONFIG": os.environ.get("CUBLAS_WORKSPACE_CONFIG"),
|
|
"precision_configuration_changed_by_component": False}
|
|
|
|
|
|
class LegacyAPL:
|
|
"""Independent zero-initialized APL state with original full-width matrices."""
|
|
|
|
def __init__(self, bundle_dir, dt_ms):
|
|
require(float(dt_ms) == 0.1, "Only frozen dt_ms=0.1 is allowed")
|
|
require(torch.cuda.is_available(), "CUDA is required for the legacy Torch component")
|
|
self.bundle_dir = Path(bundle_dir)
|
|
expected = {"input_matrix.npy": (np.float32, (N_APL, N_NEURONS)),
|
|
"output_matrix.npy": (np.float32, (N_APL, N_KC)),
|
|
"kc_globals.npy": (np.int32, (N_KC,)),
|
|
"apl_globals.npy": (np.int32, (N_APL,))}
|
|
arrays = {}
|
|
for name, (dtype, shape) in expected.items():
|
|
value = np.load(self.bundle_dir / name, allow_pickle=False)
|
|
require(value.dtype == dtype and value.shape == shape and value.flags.c_contiguous,
|
|
"APL bundle dtype/shape/layout mismatch: " + name)
|
|
require(bool(np.isfinite(value).all()), "Nonfinite APL bundle: " + name)
|
|
arrays[name] = value
|
|
for name in ("kc_globals.npy", "apl_globals.npy"):
|
|
ids = arrays[name]
|
|
require(bool(np.all((ids >= 0) & (ids < N_NEURONS))) and
|
|
bool(np.all(np.diff(ids) > 0)), "APL/KC globals must be sorted unique: " + name)
|
|
require(not np.intersect1d(arrays["kc_globals.npy"], arrays["apl_globals.npy"]).size,
|
|
"APL and KC identities overlap")
|
|
require(bool(np.all(arrays["input_matrix.npy"] >= 0)) and
|
|
bool(np.all(arrays["output_matrix.npy"] >= 0)), "APL raw-contact matrices must be nonnegative")
|
|
self.device = torch.device("cuda", torch.cuda.current_device())
|
|
self.input_gpu = torch.from_numpy(arrays["input_matrix.npy"].copy()).to(self.device)
|
|
self.output_gpu = torch.from_numpy(arrays["output_matrix.npy"].copy()).to(self.device)
|
|
self.kc_globals = arrays["kc_globals.npy"].copy()
|
|
self.apl_globals = arrays["apl_globals.npy"].copy()
|
|
self.state = torch.zeros(N_APL, dtype=torch.float32, device=self.device)
|
|
self.spikes = torch.zeros(N_NEURONS, dtype=torch.uint8, device=self.device)
|
|
self.decay = np.float32(math.exp(-float(dt_ms) / 20.0))
|
|
self.gain = np.float32(0.01)
|
|
self.metadata = {"ok": True, "protocol": PROTOCOL, "n_neurons": N_NEURONS,
|
|
"n_apl": N_APL, "n_kc": N_KC, "dt_ms": float(dt_ms),
|
|
"files_sha256": {name: file_sha256(self.bundle_dir / name) for name in BUNDLE_FILES},
|
|
"array_sha256": {name: array_sha256(value) for name, value in arrays.items()},
|
|
"decay_f32_bits": f"{int(self.decay.view(np.uint32)):08x}",
|
|
"gain_f32_bits": f"{int(self.gain.view(np.uint32)):08x}",
|
|
"environment": environment_metadata(),
|
|
"operator_sequence": "state.mul_(decay32); drive=mv(I32,spikes32); state.add_(drive); h=matmul(state,O32)*gain32",
|
|
"h_units": "legacy mV", "initial_state": "float32 zeros"}
|
|
|
|
@torch.no_grad()
|
|
def advance(self, spikes_u8):
|
|
require(spikes_u8.device == self.device and spikes_u8.dtype == torch.uint8 and
|
|
tuple(spikes_u8.shape) == (N_NEURONS,) and spikes_u8.is_contiguous(),
|
|
"APL expects full-width contiguous CUDA uint8 spikes")
|
|
# Preserve separate eager operations, including Python evaluation order
|
|
# in the frozen state.mul_(...).add_(torch.mv(...)) expression.
|
|
self.state.mul_(float(self.decay))
|
|
drive = torch.mv(self.input_gpu, spikes_u8.to(torch.float32))
|
|
self.state.add_(drive)
|
|
h = torch.matmul(self.state, self.output_gpu) * self.gain
|
|
return drive, self.state, h
|
|
|
|
def advance_global_fired(self, fired):
|
|
require(fired.dtype == np.int32 and fired.ndim == 1, "Global spikes must be int32 vector")
|
|
require(bool(np.all((fired >= 0) & (fired < N_NEURONS))) and
|
|
bool(np.all(np.diff(fired) > 0)), "Global spikes must be sorted unique and in range")
|
|
self.spikes.zero_()
|
|
if fired.size:
|
|
indices = torch.from_numpy(fired.copy()).to(device=self.device, dtype=torch.int64)
|
|
self.spikes[indices] = 1
|
|
return self.advance(self.spikes)
|
|
|
|
|
|
def read_exact(stream, count, boundary_eof=False):
|
|
result = bytearray()
|
|
while len(result) < count:
|
|
block = stream.read(count - len(result))
|
|
if not block:
|
|
if boundary_eof and not result:
|
|
return None
|
|
raise EOFError("Truncated APL request frame")
|
|
result.extend(block)
|
|
return bytes(result)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--bundle-dir", required=True)
|
|
parser.add_argument("--dt-ms", required=True, type=float)
|
|
args = parser.parse_args()
|
|
# stdout is exclusively the handshake line and binary response frames.
|
|
output, source = sys.stdout.buffer, sys.stdin.buffer
|
|
component = LegacyAPL(args.bundle_dir, args.dt_ms)
|
|
output.write((json.dumps(component.metadata, allow_nan=False) + "\n").encode("utf-8"))
|
|
output.flush()
|
|
expected_tick = 0
|
|
while True:
|
|
header = read_exact(source, FRAME.size, boundary_eof=True)
|
|
if header is None:
|
|
return
|
|
tick, count = FRAME.unpack(header)
|
|
require(tick == expected_tick and count <= N_NEURONS, "Invalid APL request tick/count")
|
|
fired = np.frombuffer(read_exact(source, count * 4), dtype="<i4").copy()
|
|
drive, state, h = component.advance_global_fired(fired)
|
|
arrays = [value.detach().cpu().numpy().astype("<f4", copy=False) for value in (drive, state, h)]
|
|
require(all(bool(np.isfinite(value).all()) for value in arrays), "Nonfinite APL output")
|
|
output.write(FRAME.pack(tick, PAYLOAD_BYTES))
|
|
for value in arrays:
|
|
output.write(value.tobytes())
|
|
output.flush()
|
|
expected_tick += 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
main()
|
|
except Exception:
|
|
traceback.print_exc(file=sys.stderr)
|
|
sys.exit(1)
|