473 lines
24 KiB
Python
473 lines
24 KiB
Python
"""Test-only arrival-f GPU owner; CUDA compilation/execution requires admission."""
|
|
from dataclasses import replace
|
|
import hashlib
|
|
from numbers import Integral
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import spacetime_reference as ref
|
|
from CelerisLab.lbm.population_storage import _encode_physical, _decode_stored, _WEIGHTS
|
|
|
|
PINS = {"spacetime_reference.py": "1fe13242b9a41062a1227db4a6f9b695165aefab7f4c0b3b2adff0c9023552bf",
|
|
"cpu_reference.py": "bba2a06a9f6517b32079ffff15e4060d8c3b7ed35f16ca134e1877226c9e4e1e",
|
|
"../taylor_green/oracle.py": "83e4d7425d45538278e9fb181af450a36358bd9a5ed6be57ff611dd9ae5eb35b"}
|
|
|
|
|
|
def check_sources():
|
|
for name, expected in PINS.items():
|
|
if hashlib.sha256((Path(__file__).parent/name).read_bytes()).hexdigest() != expected:
|
|
raise RuntimeError(f"reference source changed: {name}")
|
|
|
|
|
|
def validate_config(config, *, uniform=False, allow_n128=False):
|
|
if type(uniform) is not bool or type(allow_n128) is not bool:
|
|
raise ValueError("admission flags must be bool")
|
|
cfg = replace(config)
|
|
cfg.validate()
|
|
if (cfg.nx != cfg.ny or not 8 <= cfg.nx <= (256 if uniform else (128 if allow_n128 else 64)) or cfg.nz != 1
|
|
or cfg.collision != "SRT" or not cfg.is_d2q9
|
|
or cfg.streaming != "double_buffer" or cfg.store_precision != "FP32"
|
|
or cfg.data_type != "FP32" or type(cfg.ddf_shifting) is not bool or cfg.les_enabled
|
|
or cfg.velocity != 0 or cfg.wall_slot_width or cfg.wall_normal_velocity):
|
|
raise ValueError("experiment requires periodic unforced square FP32 D2Q9/SRT within admitted size")
|
|
if uniform:
|
|
return (cfg,)
|
|
fine = replace(cfg, nx=2*cfg.nx, ny=2*cfg.ny, viscosity=2*cfg.viscosity)
|
|
fine.validate()
|
|
return cfg, fine
|
|
|
|
|
|
def pack_plan(plan):
|
|
"""Pack reviewed coordinate events; no production transfer builder imports."""
|
|
n, nf = plan.n, plan.nf
|
|
groups, receivers = plan.groups, plan.receivers
|
|
data = dict(samples=np.zeros((len(groups), 10), np.uint32),
|
|
origins=np.empty(len(groups), np.uint32), sizes=np.empty(len(groups), np.int32),
|
|
weights=np.zeros((len(groups), 4, 10), np.float64),
|
|
targets=(receivers[:, 0]*n*n+receivers[:, 1]).astype(np.uint32))
|
|
lookup = {tuple(row): i for i, row in enumerate(receivers)}
|
|
packets = [[] for _ in receivers]; incoming = [[], []]
|
|
for i, g in enumerate(groups):
|
|
size = len(g.samples)
|
|
if size not in (6, 10):
|
|
raise ValueError("unexpected interpolation stencil")
|
|
data['samples'][i, :size] = g.q*n*n+g.samples
|
|
data['origins'][i] = g.q*n*n+g.source
|
|
data['sizes'][i] = size
|
|
data['weights'][i, :, :size] = g.weights
|
|
for k, event in enumerate(g.events):
|
|
if event is None:
|
|
packets[lookup[(g.q, g.receiver)]].append((-1, 4*i+k))
|
|
else:
|
|
m, slot = event
|
|
incoming[m].append((g.q*nf*nf+slot, 4*i+k))
|
|
for m, outgoing in enumerate(plan.outgoing):
|
|
for q, source, target in outgoing:
|
|
packets[lookup[(int(q), int(target))]].append((m, int(q*nf*nf+source)))
|
|
if any(len(row) != 4 or len(set(row)) != 4 for row in packets):
|
|
raise ValueError("each receiver needs four distinct quarter-packets")
|
|
if len({packet for row in packets for packet in row}) != 4*len(receivers):
|
|
raise ValueError("packet has multiple receiver owners")
|
|
data['times'] = np.array([[t for t, _ in row] for row in packets], np.int32)
|
|
data['sources'] = np.array([[s for _, s in row] for row in packets], np.uint32)
|
|
for m, rows in enumerate(incoming):
|
|
data[f'incoming{m}'] = np.array(rows, np.uint32)
|
|
if len({slot for slot, _ in rows}) != len(rows):
|
|
raise ValueError("duplicate incoming writer")
|
|
return {key: np.ascontiguousarray(value) for key, value in data.items()}
|
|
|
|
|
|
def table_hash(data):
|
|
digest = hashlib.sha256()
|
|
for key, value in sorted(data.items()):
|
|
digest.update(f"{key}:{value.dtype}:{value.shape}".encode()+value.tobytes())
|
|
return digest.hexdigest()
|
|
|
|
|
|
def rounding_ledger(plan, tables, before, after, g0, predicted, fine_phases, pre, h, registers, *, shifted=False):
|
|
"""Audit actual stored FP32 phases; never call a missing packet 'rounding'."""
|
|
maps, n = plan.maps, plan.n
|
|
f64 = lambda a: np.asarray(a, dtype=np.float64)
|
|
difference = f64(g0[:, maps.coarse_active])-f64(before.coarse_f[:, maps.coarse_active])
|
|
coarse_collision = difference.sum(axis=1)
|
|
cell_mass = [difference.sum(axis=0)]
|
|
collision = coarse_collision.copy(); fine_collision = []
|
|
for f, g in fine_phases:
|
|
difference = f64(g[:, maps.fine_active])-f64(f[:, maps.fine_active])
|
|
fine_collision.append(.25*difference.sum(axis=1))
|
|
cell_mass.append(.25*difference.sum(axis=0))
|
|
collision += fine_collision[-1]
|
|
np.testing.assert_array_equal(h, pre.astype(np.float32))
|
|
origins = f64(g0).ravel()[tables['origins']]
|
|
np.testing.assert_allclose(pre.mean(axis=1), origins, rtol=0, atol=2e-14)
|
|
group = np.bincount(tables['origins']//(n*n), weights=f64(h).mean(axis=1)-origins, minlength=9)
|
|
interpolation = np.bincount(tables['origins']//(n*n), weights=pre.mean(axis=1)-origins, minlength=9)
|
|
child_cast = np.bincount(tables['origins']//(n*n), weights=f64(h).mean(axis=1)-pre.mean(axis=1), minlength=9)
|
|
packet = np.empty(tables['times'].shape, np.float64)
|
|
for when, source in ((-1, h), (0, fine_phases[0][1]), (1, fine_phases[1][1])):
|
|
mask = tables['times'] == when
|
|
packet[mask] = f64(source).ravel()[tables['sources'][mask]]
|
|
p = f64(predicted).ravel()[tables['targets']]
|
|
ideal = -p + .25*packet.sum(axis=1)
|
|
reg_error = registers-ideal
|
|
bound = 16*np.finfo(float).eps*(abs(p)+.25*abs(packet).sum(axis=1))
|
|
if np.any(abs(reg_error) > bound):
|
|
raise AssertionError("receiver packet register mismatch (not roundoff)")
|
|
actual = after.coarse_f.ravel()[tables['targets']]
|
|
np.testing.assert_array_equal(actual, (p+registers).astype(np.float32))
|
|
q = tables['targets']//(n*n)
|
|
store_error = np.bincount(q, weights=f64(actual)-(p+registers), minlength=9)
|
|
register_error = np.bincount(q, weights=reg_error, minlength=9)
|
|
# directional() must sum converted arrays, not accumulate in FP32.
|
|
total = lambda s: ref.directional(ref.cpu.State(f64(s.coarse_f), f64(s.fine_f)), maps)
|
|
net = total(after)-total(before)
|
|
raw = net-collision
|
|
unexplained = raw-group-store_error-register_error
|
|
limit = float(64*np.finfo(float).eps*(total(before).sum()+int(shifted)*n*n))
|
|
np.testing.assert_allclose(unexplained, 0, rtol=0, atol=limit)
|
|
offset = _WEIGHTS.ravel()[tables['origins']//(n*n),None] if shifted else 0.
|
|
return dict(raw=raw.tolist(), net_change=net.tolist(), physical_collision=collision.tolist(),
|
|
coarse_collision=coarse_collision.tolist(), fine0_collision=fine_collision[0].tolist(),
|
|
fine1_collision=fine_collision[1].tolist(), child_interpolation=interpolation.tolist(),
|
|
child_cast=child_cast.tolist(),
|
|
collision_cell_mass=[[float(a.sum()),float(a[a>0].sum()),float(a[a<0].sum()),float(abs(a).sum())]
|
|
for a in cell_mass],
|
|
child_rounding=group.tolist(), receiver_rounding=store_error.tolist(),
|
|
register_rounding=register_error.tolist(), unexplained=unexplained.tolist(),
|
|
unexplained_bound=limit, minimum_pre=float((pre+offset).min()),
|
|
minimum_stored=float((f64(h)+offset).min()))
|
|
|
|
|
|
class SpaceTimeGPU:
|
|
"""Caller-owned context, one stream, explicit synchronized arrival state."""
|
|
def __init__(self, config, bounds, *, context, mode='spacetime', identity=False, allow_n128=False):
|
|
check_sources()
|
|
self.configs = validate_config(config,allow_n128=allow_n128)
|
|
if mode not in ('spacetime', 'frozen') or type(identity) is not bool:
|
|
raise ValueError("unsupported interpolation or debug collision mode")
|
|
self.bounds = tuple(bounds)
|
|
if len(self.bounds) != 4 or any(isinstance(v, bool) or not isinstance(v, Integral) for v in self.bounds):
|
|
raise ValueError("rectangle bounds must be four integers")
|
|
self.plan = ref.Plan(config.nx, self.bounds)
|
|
self.tables = pack_plan(self.plan)
|
|
self.ctx, self.mode, self.identity = context, mode, identity
|
|
self.allocations, self.modules, self.build_metadata = [], [], []
|
|
self.status, self.step_count, self.device_bytes = 'new', 0, 0
|
|
self.health_counts = np.zeros(6, np.uint64)
|
|
import pycuda.driver as cuda
|
|
from CelerisLab.cuda import compiler_v2 as compiler
|
|
root = Path(__file__).resolve().parents[3]
|
|
if Path(compiler._KERNEL_DIR).resolve() != root/'src/CelerisLab/lbm/kernels':
|
|
raise RuntimeError("wrong production kernel checkout")
|
|
self.cuda = cuda
|
|
with context.activate():
|
|
self.stream = cuda.Stream()
|
|
try:
|
|
source = Path(__file__).with_suffix('.cu')
|
|
for cfg in self.configs:
|
|
build = compiler.generate_config(cfg)
|
|
try:
|
|
Path(build, source.name).write_bytes(source.read_bytes())
|
|
ptx = compiler.compile_kernel(context.sm_arch, source=source.name, build_dir=build)
|
|
self.modules.append(compiler.load_module(ptx))
|
|
import json
|
|
metadata = json.loads(Path(ptx+'.json').read_text())
|
|
metadata['test_entry_sha256'] = hashlib.sha256(source.read_bytes()).hexdigest()
|
|
self.build_metadata.append(metadata)
|
|
finally:
|
|
compiler.cleanup_build_dir(build)
|
|
self.device_tables = {key: self.upload(value) for key, value in self.tables.items()}
|
|
self.active = self.upload(self.plan.maps.fine_active.astype(np.uint8))
|
|
self.coarse = [self.allocate(9*config.nx**2*4) for _ in range(4)]
|
|
self.fine = [self.allocate(9*(2*config.nx)**2*4) for _ in range(3)]
|
|
self.pre = self.allocate(len(self.plan.groups)*4*8)
|
|
self.h = self.allocate(len(self.plan.groups)*4*4)
|
|
self.registers = self.allocate(len(self.plan.receivers)*8)
|
|
self.counts = self.allocate(6*4)
|
|
except BaseException:
|
|
try:
|
|
self.close()
|
|
except Exception:
|
|
pass
|
|
raise
|
|
|
|
def allocate(self, size):
|
|
pointer = self.cuda.mem_alloc(size)
|
|
self.allocations.append(pointer); self.device_bytes += size
|
|
return pointer
|
|
|
|
def upload(self, array):
|
|
pointer = self.allocate(array.nbytes)
|
|
self.cuda.memcpy_htod(pointer, array)
|
|
return pointer
|
|
|
|
def read(self, pointer, shape, dtype=np.float32):
|
|
self.stream.synchronize()
|
|
out = np.empty(shape, dtype)
|
|
self.cuda.memcpy_dtoh(out, pointer)
|
|
return out
|
|
|
|
def launch(self, level, name, *args, count=None):
|
|
cfg = self.configs[level]
|
|
nt = cfg.threads_per_block
|
|
grid = ((cfg.nx+nt-1)//nt, cfg.ny) if count is None else ((count+nt-1)//nt, 1)
|
|
self.modules[level].get_function(name)(*args, block=(nt,1,1), grid=grid, stream=self.stream)
|
|
|
|
def initialize(self, state):
|
|
if (self.status not in ('new', 'ready') or state.phase != 'arrival_f'
|
|
or isinstance(state.tick, bool) or not isinstance(state.tick, Integral) or state.tick != 0):
|
|
raise ValueError("initial synchronized arrival state with tick=0 required")
|
|
shifted = self.configs[0].ddf_shifting
|
|
arrays = []
|
|
for values, mask in zip((state.coarse_f,state.fine_f),
|
|
(self.plan.maps.coarse_active,self.plan.maps.fine_active)):
|
|
if np.shape(values) != (9,*mask.shape):
|
|
raise ValueError("positive finite leaf populations and exact shapes required")
|
|
a = _encode_physical(values,shifted)
|
|
physical = _decode_stored(a,shifted)
|
|
if not np.isfinite(physical[:,mask]).all() or (physical[:,mask] <= 0).any():
|
|
raise ValueError("positive finite leaf populations and exact shapes required")
|
|
a[:,~mask] = np.nan
|
|
arrays.append(a)
|
|
with self.ctx.activate():
|
|
try:
|
|
self.cuda.memcpy_htod(self.coarse[0], arrays[0])
|
|
self.cuda.memcpy_htod(self.fine[0], arrays[1])
|
|
except BaseException:
|
|
self.status = 'failed'
|
|
raise
|
|
self.status, self.step_count = 'ready', 0
|
|
self.health_counts.fill(0)
|
|
|
|
def _download_stored(self):
|
|
"""Private audit/reset snapshot; values remain in their storage representation."""
|
|
if self.status != 'ready':
|
|
raise RuntimeError(f"GPU owner is {self.status}")
|
|
try:
|
|
with self.ctx.activate():
|
|
arrays = [self.read(self.coarse[0], (9,self.plan.n,self.plan.n)),
|
|
self.read(self.fine[0], (9,self.plan.nf,self.plan.nf))]
|
|
except BaseException:
|
|
self.status = 'failed'
|
|
raise
|
|
for a, mask in zip(arrays, (self.plan.maps.coarse_active,self.plan.maps.fine_active)):
|
|
a[:,~mask] = np.nan
|
|
return ref.cpu.State(*arrays, tick=2*self.step_count)
|
|
|
|
def download(self):
|
|
"""Physical arrival populations; shifted values are reconstructed in FP64."""
|
|
stored = self._download_stored()
|
|
return ref.cpu.State(*(_decode_stored(a,self.configs[0].ddf_shifting)
|
|
for a in (stored.coarse_f,stored.fine_f)),tick=stored.tick)
|
|
|
|
def step(self, *, mutation=0, audit=True):
|
|
"""One interval; audit=False skips only host ledger snapshots, never health checks."""
|
|
if (self.status != 'ready' or type(mutation) is not int or mutation not in range(5)
|
|
or type(audit) is not bool or (mutation and not audit)):
|
|
raise RuntimeError("ready owner and known test mutation required")
|
|
before = self._download_stored() if audit else None
|
|
n, nf, t = self.plan.n, self.plan.nf, self.device_tables
|
|
ng, nr = len(self.plan.groups), len(self.plan.receivers)
|
|
u32, null = np.uint32, np.uintp(0)
|
|
omega = [np.float32(0 if self.identity else cfg.omega) for cfg in self.configs]
|
|
with self.ctx.activate():
|
|
try:
|
|
self.status = 'advancing'
|
|
self.cuda.memset_d32_async(self.counts, 0, 6, self.stream)
|
|
fc, g0, predicted, gp = self.coarse
|
|
self.launch(0, 'CoarseStart', fc,self.fine[0],g0,omega[0],
|
|
*(u32(v) for v in self.bounds),self.counts)
|
|
self.launch(0, 'Pull', g0,predicted,null,self.counts)
|
|
self.launch(0, 'Collide', predicted,gp,null,omega[0],self.counts)
|
|
self.launch(0, 'Interpolate', g0,gp,t['samples'],t['origins'],t['sizes'],t['weights'],
|
|
self.pre,self.h,u32(ng),u32(self.mode=='frozen'),self.counts,count=ng)
|
|
if audit:
|
|
coarse_g = self.read(g0,(9,n,n)); prediction = self.read(predicted,(9,n,n))
|
|
fine_phases = []
|
|
for m in (0,1):
|
|
f, out, g = self.fine
|
|
if audit:
|
|
old_f = self.read(f,(9,nf,nf))
|
|
self.launch(1, 'Collide', f,g,self.active,omega[1],self.counts)
|
|
if audit:
|
|
fine_phases.append((old_f,self.read(g,(9,nf,nf))))
|
|
self.launch(0, 'Register', predicted,g,self.h,t['targets'],t['sources'],t['times'],
|
|
self.registers,u32(nr),np.int32(m),np.int32(mutation),count=nr)
|
|
rows = self.tables[f'incoming{m}']
|
|
self.launch(1, 'Inject', g,self.h,t[f'incoming{m}'],u32(len(rows)),u32(mutation==4),count=len(rows))
|
|
self.launch(1, 'Pull', g,out,self.active,self.counts)
|
|
self.fine[0], self.fine[1] = out, f
|
|
self.launch(0, 'ApplyRegisters', predicted,self.registers,t['targets'],u32(nr),self.counts,count=nr)
|
|
counts = self.read(self.counts,(6,),np.uint32)
|
|
self.health_counts += counts.astype(np.uint64)
|
|
if counts.any():
|
|
raise AssertionError(f"boundary/leaf health failed: {counts.tolist()}")
|
|
ledger = dict(boundary_budget_status='not_measured',unexplained=None)
|
|
if audit:
|
|
coarse_after = self.read(predicted,(9,n,n))
|
|
coarse_after[:,~self.plan.maps.coarse_active] = np.nan
|
|
after = ref.cpu.State(coarse_after,self.read(self.fine[0],(9,nf,nf)),tick=2*(self.step_count+1))
|
|
ledger = rounding_ledger(self.plan,self.tables,before,after,coarse_g,prediction,fine_phases,
|
|
self.read(self.pre,(ng,4),np.float64),self.read(self.h,(ng,4)),
|
|
self.read(self.registers,(nr,),np.float64),
|
|
shifted=self.configs[0].ddf_shifting)
|
|
ledger['boundary_budget_status'] = 'measured'
|
|
ledger['health_counts'] = counts.tolist()
|
|
self.coarse[0], self.coarse[2] = predicted, fc
|
|
self.step_count += 1
|
|
self.status = 'ready'
|
|
return ledger
|
|
except BaseException:
|
|
self.status = 'failed'
|
|
raise
|
|
|
|
def close(self):
|
|
if self.status == 'closed':
|
|
return
|
|
with self.ctx.activate():
|
|
error = None
|
|
try:
|
|
self.stream.synchronize()
|
|
except Exception as exc:
|
|
error = exc
|
|
failed = []
|
|
for pointer in self.allocations:
|
|
try:
|
|
pointer.free()
|
|
except Exception as exc:
|
|
error = error or exc
|
|
failed.append(pointer)
|
|
self.allocations = failed
|
|
self.status = 'cleanup_failed' if failed else 'closed'
|
|
if not failed:
|
|
self.modules.clear()
|
|
if error:
|
|
raise error
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, kind, value, traceback):
|
|
try:
|
|
self.close()
|
|
except Exception:
|
|
if kind is None:
|
|
raise
|
|
|
|
|
|
def uniform_ledger(before, after, collided, *, shifted):
|
|
"""Periodic streaming permutes each direction; only collision changes its sum."""
|
|
a,b,g = [np.asarray(x,np.float64) for x in (before,after,collided)]
|
|
collision=(g-a).sum(axis=(1,2)); net=b.sum(axis=(1,2))-a.sum(axis=(1,2))
|
|
residual=net-collision
|
|
bound=float(64*np.finfo(float).eps*(a.sum()+int(shifted)*a.shape[1]*a.shape[2]))
|
|
np.testing.assert_allclose(residual,0,rtol=0,atol=bound)
|
|
cells=(g-a).sum(axis=0)
|
|
out={key:[0.]*9 for key in ('fine0_collision','fine1_collision','child_rounding',
|
|
'child_interpolation','child_cast','register_rounding','receiver_rounding')}
|
|
out.update(net_change=net.tolist(),physical_collision=collision.tolist(),coarse_collision=collision.tolist(),
|
|
unexplained=residual.tolist(),unexplained_bound=bound,
|
|
collision_cell_mass=[[float(cells.sum()),float(cells[cells>0].sum()),
|
|
float(cells[cells<0].sum()),float(abs(cells).sum())],[0.]*4,[0.]*4],
|
|
boundary_budget_status='not_applicable_uniform',inactive_components='fine, child, register, receiver')
|
|
return out
|
|
|
|
|
|
class UniformGPU(SpaceTimeGPU):
|
|
"""One explicit arrival-F level; reuse the reviewed kernels and owner helpers."""
|
|
def __init__(self, config, *, context, identity=False):
|
|
check_sources()
|
|
self.configs=validate_config(config,uniform=True)
|
|
if type(identity) is not bool:
|
|
raise ValueError('identity must be bool')
|
|
self.ctx,self.identity=context,identity
|
|
self.allocations,self.modules,self.build_metadata=[],[],[]
|
|
self.status,self.step_count,self.device_bytes='new',0,0
|
|
self.health_counts=np.zeros(6,np.uint64)
|
|
import pycuda.driver as cuda
|
|
from CelerisLab.cuda import compiler_v2 as compiler
|
|
if Path(compiler._KERNEL_DIR).resolve()!=Path(__file__).resolve().parents[3]/'src/CelerisLab/lbm/kernels':
|
|
raise RuntimeError('wrong production kernel checkout')
|
|
self.cuda=cuda
|
|
with context.activate():
|
|
self.stream=cuda.Stream()
|
|
try:
|
|
source=Path(__file__).with_suffix('.cu')
|
|
build=compiler.generate_config(self.configs[0])
|
|
try:
|
|
Path(build,source.name).write_bytes(source.read_bytes())
|
|
ptx=compiler.compile_kernel(context.sm_arch,source=source.name,build_dir=build)
|
|
self.modules.append(compiler.load_module(ptx))
|
|
import json
|
|
metadata=json.loads(Path(ptx+'.json').read_text())
|
|
metadata['test_entry_sha256']=hashlib.sha256(source.read_bytes()).hexdigest()
|
|
self.build_metadata.append(metadata)
|
|
finally:
|
|
compiler.cleanup_build_dir(build)
|
|
n=config.nx
|
|
self.buffers=[self.allocate(9*n*n*4) for _ in range(3)]
|
|
self.active=self.upload(np.ones((n,n),np.uint8))
|
|
self.counts=self.allocate(6*4)
|
|
except BaseException:
|
|
try:
|
|
self.close()
|
|
except BaseException:
|
|
pass
|
|
raise
|
|
|
|
def initialize(self, physical):
|
|
cfg=self.configs[0]
|
|
if self.status not in ('new','ready') or np.shape(physical)!=(9,cfg.ny,cfg.nx):
|
|
raise ValueError('ready owner and exact uniform population shape required')
|
|
stored=_encode_physical(physical,cfg.ddf_shifting)
|
|
decoded=_decode_stored(stored,cfg.ddf_shifting)
|
|
if not np.isfinite(decoded).all() or (decoded<=0).any():
|
|
raise ValueError('positive finite physical populations required')
|
|
try:
|
|
with self.ctx.activate():
|
|
self.cuda.memcpy_htod(self.buffers[0],stored)
|
|
self.health_counts[:]=0; self.step_count=0; self.status='ready'
|
|
except BaseException:
|
|
self.status='failed'
|
|
raise
|
|
|
|
def _download_stored(self):
|
|
if self.status!='ready':
|
|
raise RuntimeError('uniform owner is not ready')
|
|
try:
|
|
with self.ctx.activate():
|
|
n=self.configs[0].nx
|
|
return self.read(self.buffers[0],(9,n,n))
|
|
except BaseException:
|
|
self.status='failed'
|
|
raise
|
|
|
|
def download(self):
|
|
return _decode_stored(self._download_stored(),self.configs[0].ddf_shifting)
|
|
|
|
def step(self, *, audit=True):
|
|
if self.status!='ready' or type(audit) is not bool:
|
|
raise RuntimeError('ready owner and bool audit required')
|
|
cfg=self.configs[0]; shape=(9,cfg.ny,cfg.nx)
|
|
before=self._download_stored() if audit else None
|
|
with self.ctx.activate():
|
|
try:
|
|
self.status='advancing'
|
|
self.cuda.memset_d32_async(self.counts,0,6,self.stream)
|
|
f,g,out=self.buffers
|
|
self.launch(0,'Collide',f,g,self.active,np.float32(0 if self.identity else cfg.omega),self.counts)
|
|
self.launch(0,'Pull',g,out,self.active,self.counts)
|
|
counts=self.read(self.counts,(6,),np.uint32)
|
|
self.health_counts+=counts.astype(np.uint64)
|
|
if counts.any():
|
|
raise AssertionError(f'uniform population health failed: {counts.tolist()}')
|
|
ledger=uniform_ledger(before,self.read(out,shape),self.read(g,shape),shifted=cfg.ddf_shifting) if audit else dict(
|
|
boundary_budget_status='not_measured',unexplained=None)
|
|
ledger['health_counts']=counts.tolist()
|
|
self.buffers[0],self.buffers[2]=out,f
|
|
self.step_count+=1; self.status='ready'
|
|
return ledger
|
|
except BaseException:
|
|
self.status='failed'
|
|
raise
|