test: prepare matched shifted GPU spacing ladder

tests/validation/multigrid/spacetime_gpu.py

@@ -19,15 +19,19 @@ def check_sources():
             raise RuntimeError(f"reference source changed: {name}")
 
 
-def validate_config(config):
+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 <= 64 or cfg.nz != 1
+    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 N8..64, FP32 D2Q9/SRT")
+        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
@@ -136,9 +140,9 @@ def rounding_ledger(plan, tables, before, after, g0, predicted, fine_phases, pre
 
 class SpaceTimeGPU:
     """Caller-owned context, one stream, explicit synchronized arrival state."""
-    def __init__(self, config, bounds, *, context, mode='spacetime', identity=False):
+    def __init__(self, config, bounds, *, context, mode='spacetime', identity=False, allow_n128=False):
         check_sources()
-        self.configs = validate_config(config)
+        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)
@@ -349,3 +353,120 @@ class SpaceTimeGPU:
         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