test: add opt-in shifted storage to space-time GPU prototype

This commit is contained in:
Codex committed 2026-10-18 11:13:50 -07:00
1 parent c2496c2103
commit e8b8c58995
3 files changed
+164 -52

No files matched your search

+29 -11
View File
@@ -6,6 +6,7 @@ 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",
@@ -24,9 +25,9 @@ def validate_config(config):
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.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 N8..64, unshifted FP32 D2Q9/SRT")
raise ValueError("experiment requires periodic unforced square N8..64, FP32 D2Q9/SRT")
fine = replace(cfg, nx=2*cfg.nx, ny=2*cfg.ny, viscosity=2*cfg.viscosity)
fine.validate()
return cfg, fine
@@ -79,7 +80,7 @@ def table_hash(data):
return digest.hexdigest()
def rounding_ledger(plan, tables, before, after, g0, predicted, fine_phases, pre, h, registers):
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)
@@ -118,8 +119,9 @@ def rounding_ledger(plan, tables, before, after, g0, predicted, fine_phases, pre
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())
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(),
@@ -128,7 +130,8 @@ def rounding_ledger(plan, tables, before, after, g0, predicted, fine_phases, pre
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.min()), minimum_stored=float(h.min()))
unexplained_bound=limit, minimum_pre=float((pre+offset).min()),
minimum_stored=float((f64(h)+offset).min()))
class SpaceTimeGPU:
@@ -210,11 +213,18 @@ class SpaceTimeGPU:
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():
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])
@@ -225,7 +235,8 @@ class SpaceTimeGPU:
self.status, self.step_count = 'ready', 0
self.health_counts.fill(0)
def download(self):
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:
@@ -239,12 +250,18 @@ class SpaceTimeGPU:
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() if audit else None
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)
@@ -288,7 +305,8 @@ class SpaceTimeGPU:
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.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