Digital Drosophila Archive
161 lines
11 KiB
Python
161 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""Canonical deterministic MaleCNS LIF reference backend and validation tasks."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import time
|
|
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 phase4c
|
|
|
|
|
|
def sha_bytes(*arrays: np.ndarray) -> str:
|
|
h=hashlib.sha256()
|
|
for x in arrays:h.update(np.ascontiguousarray(x).tobytes())
|
|
return h.hexdigest()
|
|
|
|
|
|
def save(path: Path,x: object):
|
|
path.parent.mkdir(parents=True,exist_ok=True)
|
|
path.write_text(json.dumps(x,indent=2,sort_keys=True,default=str)+'\n')
|
|
|
|
|
|
def build_extension(model:Path,exp:Path):
|
|
build=exp/'gpu-reference/extension-build';build.mkdir(parents=True,exist_ok=True)
|
|
start=time.perf_counter()
|
|
ext=load_extension(name='male_cns_lif_deterministic_v2_ext',
|
|
sources=[str(model/'phase4p_ext.cpp'),str(model/'phase4p_ext.cu')],
|
|
build_directory=str(build),extra_cflags=['-O3','-ffp-contract=off'],
|
|
extra_cuda_cflags=['-O3','--fmad=false','-std=c++20'],verbose=True)
|
|
report={'status':'PASS','build_seconds':time.perf_counter()-start,'binary':str(Path(ext.__file__)),
|
|
'binary_sha256':phase4c.sha256(Path(ext.__file__)),'flags':{'cxx':['-O3','-ffp-contract=off'],'cuda':['-O3','--fmad=false','-std=c++20']}}
|
|
save(exp/'gpu-reference/extension.json',report);return ext
|
|
|
|
|
|
class CanonicalRuntime:
|
|
def __init__(self,cfg:dict[str,Any],model:Path,exp:Path):
|
|
self.cfg=cfg;self.exp=exp;self.device=torch.device('cuda:0')
|
|
if torch.cuda.device_count()!=1 or cfg['device']['required_name_substring'] not in torch.cuda.get_device_name(0):
|
|
raise RuntimeError('exactly one visible RTX 4090 required')
|
|
p4cfg=yaml.safe_load((Path(cfg['paths']['historical_runtime_root'])/'phase4c.yaml').read_text())
|
|
paths=phase4c.exact_paths(p4cfg)
|
|
self.indptr_np,self.indices_np,self.raw_np=phase4c.p4b.validate_graph(p4cfg,paths)
|
|
self.table=phase4c.full_metadata(paths['metadata']);self.pops=phase4c.populations(self.table)
|
|
base=phase4c.Runtime(p4cfg,self.table,self.indptr_np,self.indices_np,self.raw_np,self.device,None)
|
|
self.indptr,self.indices,self.weights=base.indptr,base.indices,base.weights
|
|
self.ext=build_extension(model,exp)
|
|
m=cfg['model'];self.n=int(cfg['expected_graph']['neurons']);self.dt=float(m['dt_ms'])
|
|
self.delay_steps=int(round(m['synaptic_delay_ms']/self.dt));self.qlen=self.delay_steps+1
|
|
self.ref_steps=int(round(m['refractory_ms']/self.dt));self.gd=np.float32(1-self.dt/m['synaptic_tau_ms'])
|
|
self.alpha=np.float32(self.dt/m['membrane_tau_ms']);self.rest=np.float32(m['resting_potential_mv'])
|
|
self.resetv=np.float32(m['reset_potential_mv']);self.threshold=np.float32(m['threshold_mv'])
|
|
self.superclass=np.asarray(self.table['annotation_superclass'].to_pylist(),object)
|
|
|
|
def reset(self):
|
|
d=self.device;self.v=torch.full((self.n,),float(self.rest),dtype=torch.float32,device=d)
|
|
self.g=torch.zeros(self.n,dtype=torch.float32,device=d);self.ref=torch.zeros(self.n,dtype=torch.int32,device=d)
|
|
self.queue=torch.zeros((self.qlen,self.n),dtype=torch.float32,device=d)
|
|
self.spike_u8=torch.zeros(self.n,dtype=torch.uint8,device=d);self.counts=torch.zeros(self.n,dtype=torch.int32,device=d)
|
|
self.step_index=0;self.event_steps=[];self.event_nodes=[]
|
|
|
|
def canonical_fanout(self,sources:torch.Tensor,destination:torch.Tensor)->dict[str,float|int]:
|
|
if not sources.numel():return {'contributions':0,'segments':0,'sort_seconds':0.,'reduce_seconds':0.}
|
|
starts=self.indptr[sources];degree=self.indptr[sources+1]-starts;cum=torch.cumsum(degree,0);total=int(cum[-1].item())
|
|
edges=torch.empty(total,dtype=torch.int64,device=self.device)
|
|
self.ext.fill_edges(self.indptr,sources,cum-degree,edges)
|
|
t0=time.perf_counter();targets=self.indices[edges]
|
|
order=torch.argsort(targets,stable=True);sorted_edges=edges[order];sorted_targets=targets[order]
|
|
boundary=torch.ones(total,dtype=torch.bool,device=self.device)
|
|
if total>1:boundary[1:]=sorted_targets[1:]!=sorted_targets[:-1]
|
|
seg_starts=torch.nonzero(boundary,as_tuple=False).flatten().to(torch.int64)
|
|
unique_targets=sorted_targets[seg_starts].to(torch.int32);sort_sec=time.perf_counter()-t0
|
|
t1=time.perf_counter();self.ext.segmented_sum(sorted_edges,seg_starts,unique_targets,self.weights,destination)
|
|
return {'contributions':total,'segments':int(len(unique_targets)),'sort_seconds':sort_sec,'reduce_seconds':time.perf_counter()-t1}
|
|
|
|
def run(self,targets_np:np.ndarray,rate:float,seed:int,duration_ms:float,capture=True,forced_override=None):
|
|
self.reset();steps=int(round(duration_ms/self.dt));targets=torch.as_tensor(targets_np,dtype=torch.int64,device=self.device)
|
|
forced=forced_override if forced_override is not None else phase4c.generate_forced(steps,len(targets_np),rate,self.dt,seed,self.device)
|
|
torch.cuda.synchronize();t0=time.perf_counter();sort_t=reduce_t=0.;contrib=0
|
|
for step in range(steps):
|
|
slot=step%self.qlen
|
|
self.ext.state_step(self.v,self.g,self.ref,self.queue[slot],self.spike_u8,self.counts,
|
|
float(self.gd),float(self.alpha),float(self.rest),float(self.resetv),float(self.threshold),self.ref_steps)
|
|
self.ext.force_step(self.v,self.ref,self.spike_u8,self.counts,targets,forced[step],float(self.resetv),self.ref_steps)
|
|
sources=torch.nonzero(self.spike_u8,as_tuple=False).flatten()
|
|
stats=self.canonical_fanout(sources,self.queue[(step+self.delay_steps)%self.qlen])
|
|
sort_t+=stats['sort_seconds'];reduce_t+=stats['reduce_seconds'];contrib+=stats['contributions']
|
|
if capture and sources.numel():
|
|
x=sources.cpu().numpy().astype(np.int32,copy=True);self.event_nodes.append(x);self.event_steps.append(np.full(len(x),step,np.int32))
|
|
self.step_index+=1
|
|
torch.cuda.synchronize();wall=time.perf_counter()-t0
|
|
es=np.concatenate(self.event_steps) if self.event_steps else np.empty(0,np.int32)
|
|
en=np.concatenate(self.event_nodes) if self.event_nodes else np.empty(0,np.int32)
|
|
return {'step':es,'neuron_idx':en,'counts':self.counts.cpu().numpy().copy(),'v':self.v.cpu().numpy().copy(),'g':self.g.cpu().numpy().copy(),
|
|
'metrics':{'wall_seconds':wall,'biological_ms':duration_ms,'realtime_factor':wall/(duration_ms/1000),
|
|
'peak_vram_bytes':int(torch.cuda.max_memory_allocated()),'contributions':int(contrib),
|
|
'sort_submit_seconds':sort_t,'reduce_submit_seconds':reduce_t,'events':len(en),'active':int((self.counts>0).sum().item())}}
|
|
|
|
|
|
def cpu_ordered_sum(edge_ids:np.ndarray,indices:np.ndarray,weights:np.ndarray):
|
|
edge_ids=np.asarray(edge_ids,np.int64);targets=np.asarray(indices[edge_ids],np.int32)
|
|
order=np.argsort(targets,kind='stable');ordered=edge_ids[order];ordered_target=targets[order]
|
|
unique=[];sums=[];begin=0
|
|
while begin<len(ordered):
|
|
target=int(ordered_target[begin]);end=begin+1;acc=np.float32(0)
|
|
while end<len(ordered) and ordered_target[end]==target:end+=1
|
|
for e in ordered[begin:end]:acc=np.float32(acc+np.float32(weights[e]))
|
|
unique.append(target);sums.append(acc);begin=end
|
|
return np.asarray(unique,np.int32),np.asarray(sums,np.float32),ordered
|
|
|
|
|
|
def save_trial(root:Path,run:dict):
|
|
root.mkdir(parents=True,exist_ok=True)
|
|
pq.write_table(pa.table({'step':run['step'],'timestamp_ms':run['step'].astype(np.float64)*.1,'neuron_idx':run['neuron_idx']}),root/'spike-events.parquet',compression='zstd')
|
|
for name,key,dtype in [('per-neuron-spike-counts.int32.npy','counts',np.int32),('final-voltage.float32.npy','v',np.float32),('final-g.float32.npy','g',np.float32)]:np.save(root/name,run[key].astype(dtype),allow_pickle=False)
|
|
save(root/'statistics.json',run['metrics']|{'event_sha256':sha_bytes(run['step'],run['neuron_idx']),
|
|
'counts_sha256':sha_bytes(run['counts']),'v_sha256':sha_bytes(run['v']),'g_sha256':sha_bytes(run['g'])})
|
|
|
|
|
|
def phase4c_trials(cfg,rt,repeats):
|
|
exp=Path(cfg['paths']['experiment_root']);rows=[]
|
|
for spec in cfg['reference_trials']:
|
|
target=np.flatnonzero(rt.pops[spec['population']]).astype(np.int64)
|
|
baseline=None
|
|
for repeat in range(repeats):
|
|
run=rt.run(target,float(spec['rate_hz']),int(spec['seed']),float(spec['duration_ms']),True)
|
|
if baseline is None:
|
|
baseline={k:run[k].copy() for k in ('step','neuron_idx','counts','v','g')};save_trial(exp/'phase4c-v2/reference'/spec['name'],run)
|
|
exact={k:bool(np.array_equal(run[k],baseline[k])) for k in ('step','neuron_idx','counts','v','g')}
|
|
row={'trial':spec['name'],'repeat':repeat,**exact,**run['metrics']};rows.append(row)
|
|
save(exp/'phase4c-v2/reproducibility-matrix.json',rows)
|
|
print(json.dumps({'trial':spec['name'],'repeat':repeat,'exact':all(exact.values()),'wall':run['metrics']['wall_seconds']}),flush=True)
|
|
summary={'repeats':repeats,'matrix':{s['name']:{k:sum(r[k] for r in rows if r['trial']==s['name']) for k in ('step','neuron_idx','counts','v','g')} for s in cfg['reference_trials']},
|
|
'all_exact':all(all(r[k] for k in ('step','neuron_idx','counts','v','g')) for r in rows),
|
|
'performance':{s['name']:{'mean_wall_seconds':float(np.mean([r['wall_seconds'] for r in rows if r['trial']==s['name']])),
|
|
'mean_realtime_factor':float(np.mean([r['realtime_factor'] for r in rows if r['trial']==s['name']])),
|
|
'peak_vram_bytes':max(r['peak_vram_bytes'] for r in rows if r['trial']==s['name']),
|
|
'mean_contributions':float(np.mean([r['contributions'] for r in rows if r['trial']==s['name']]))} for s in cfg['reference_trials']}}
|
|
save(exp/'phase4c-v2/summary.json',summary)
|
|
|
|
|
|
def main():
|
|
ap=argparse.ArgumentParser();ap.add_argument('--config',type=Path,required=True);ap.add_argument('task',choices=['phase4c']);ap.add_argument('--repeats',type=int)
|
|
a=ap.parse_args();cfg=yaml.safe_load(a.config.read_text());exp=Path(cfg['paths']['experiment_root']);model=Path(cfg['paths']['model_root'])
|
|
for sub in ('numerical-contract','cpu-reference','gpu-reference','threshold-3942','fanout-tests','phase4c-v2','historical-comparison','brainio','activity-stream','eligibility-replay','performance','revalidation','regression','audit','manifests'):(exp/sub).mkdir(parents=True,exist_ok=True)
|
|
rt=CanonicalRuntime(cfg,model,exp)
|
|
phase4c_trials(cfg,rt,a.repeats or int(cfg['repeats_per_trial']))
|
|
|
|
|
|
if __name__=='__main__':main()
|