309 lines
16 KiB
Python
309 lines
16 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
|
|
|
|
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):
|
|
cfg = replace(config)
|
|
cfg.validate()
|
|
if (cfg.nx != cfg.ny or not 8 <= cfg.nx <= 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 cfg.ddf_shifting 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 N8..64, unshifted FP32 D2Q9/SRT")
|
|
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):
|
|
"""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)
|
|
collision = (f64(g0[:, maps.coarse_active])-f64(before.coarse_f[:, maps.coarse_active])).sum(axis=1)
|
|
for f, g in fine_phases:
|
|
collision += .25*(f64(g[:, maps.fine_active])-f64(f[:, maps.fine_active])).sum(axis=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)
|
|
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)
|
|
raw = total(after)-total(before)-collision
|
|
unexplained = raw-group-store_error-register_error
|
|
limit = float(64*np.finfo(float).eps*total(before).sum())
|
|
np.testing.assert_allclose(unexplained, 0, rtol=0, atol=limit)
|
|
return dict(raw=raw.tolist(), physical_collision=collision.tolist(),
|
|
child_rounding=group.tolist(), receiver_rounding=store_error.tolist(),
|
|
register_rounding=register_error.tolist(), unexplained=unexplained.tolist(),
|
|
unexplained_bound=limit, minimum_pre=float(pre.min()), minimum_stored=float(h.min()))
|
|
|
|
|
|
class SpaceTimeGPU:
|
|
"""Caller-owned context, one stream, explicit synchronized arrival state."""
|
|
def __init__(self, config, bounds, *, context, mode='spacetime', identity=False):
|
|
check_sources()
|
|
self.configs = validate_config(config)
|
|
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
|
|
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")
|
|
arrays = [np.array(a, dtype=np.float32, copy=True, order='C') for a in (state.coarse_f,state.fine_f)]
|
|
for a, mask in zip(arrays, (self.plan.maps.coarse_active,self.plan.maps.fine_active)):
|
|
if a.shape != (9,*mask.shape) or not np.isfinite(a[:,mask]).all() or (a[:,mask] <= 0).any():
|
|
raise ValueError("positive finite leaf populations and exact shapes required")
|
|
a[:,~mask] = np.nan
|
|
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
|
|
|
|
def download(self):
|
|
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 step(self, *, mutation=0):
|
|
"""One audited interval; mutation 1/2/3 drops packets, 4 omits a ghost."""
|
|
if self.status != 'ready' or type(mutation) is not int or mutation not in range(5):
|
|
raise RuntimeError("ready owner and known test mutation required")
|
|
before = self.download()
|
|
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)
|
|
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
|
|
old_f = self.read(f,(9,nf,nf))
|
|
self.launch(1, 'Collide', f,g,self.active,omega[1],self.counts)
|
|
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)
|
|
if counts.any():
|
|
raise AssertionError(f"boundary/leaf health failed: {counts.tolist()}")
|
|
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))
|
|
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
|