Add opt-in shifted raw remap GPU lifecycle gate

2026-10-16 23:28:09 -07:00
parent b811b643e7
commit aaf35c6505
1 changed files with 113 additions and 0 deletions
@@ -0,0 +1,113 @@
+"""Opt-in GPU raw shifted remap lifecycle gates; no precision/performance claim.
+
+Run only on an authorized GPU:
+  PYTHONPATH=src python tests/validation/multigrid/test_shifted_remap_gpu.py -v
+Frozen gates: bitwise raw overlap/copy and identity continuation; moved-remap
+matches the existing host fill loop exactly; finite owned leaves after stepping.
+"""
+from contextlib import ExitStack
+import unittest
+import numpy as np
+import pycuda.driver as cuda
+from CelerisLab.config import LBMConfig
+from CelerisLab.cuda.context import CudaContext
+from CelerisLab.lbm.population_storage import _StoredPopulations
+from CelerisLab.lbm.refinement_stepper import FixedPatchStepper
+from CelerisLab.lbm.synchronized_remap import RemapTopology, _remap_stored
+
+CFG=LBMConfig(nx=48,ny=32,viscosity=.1,ddf_shifting=True)
+PARENTS=((4,4,20,28),(24,4,44,28))
+A=((16,16,32,40),)
+B=((20,16,36,40),)
+
+
+def make(ctx):
+    return FixedPatchStepper(CFG,PARENTS,context=ctx,base_block=(4,4),
+                             fine_block=(4,4),child_block=(4,4),
+                             child_rectangles_in_parent=A)
+
+
+def initial(run):
+    # Non-unit density, nonzero velocity and nonequilibrium perturbations.
+    w=np.array([4/9,*([1/9]*4),*([1/36]*4)],np.float64)
+    cx=np.array([0,1,-1,0,0,1,-1,1,-1]);cy=np.array([0,0,0,1,-1,1,-1,-1,1])
+    rho,u,v=.99,.013,-.007
+    cu=3*(cx*u+cy*v)
+    eq=w*((rho-1)+rho*(cu+.5*cu*cu-1.5*(u*u+v*v)))
+    rng=np.random.default_rng(8473)
+    arrays=tuple((np.broadcast_to(eq[:,None,None],level.shape)+
+                  rng.normal(0,1e-7,level.shape)).astype(np.float32) for level in run._levels)
+    return _StoredPopulations(arrays,True)
+
+
+class ShiftedRemapGPU(unittest.TestCase):
+    def exact(self,a,b,label):
+        self.assertEqual(a.shape,b.shape,label)
+        self.assertEqual(a.dtype,b.dtype,label)
+        self.assertEqual(a.tobytes(),b.tobytes(),label)
+
+    def snapshot_equal(self,a,b,label):
+        self.assertTrue(a.shifted and b.shifted)
+        self.assertEqual(len(a.values),len(b.values))
+        for x,y in zip(a.values,b.values):self.exact(x,y,label)
+
+    def check_reset(self,run,wanted):
+        for level,want in zip(run._levels,wanted):
+            blocked=level.layout.to_blocked(want)
+            for pointer in (level.current,level.spare):
+                actual=np.empty(level.layout.blocked_shape,np.float32)
+                cuda.memcpy_dtoh(actual,pointer)
+                self.exact(actual,blocked,'fresh current and spare')
+
+    def check_finite(self,run):
+        raw=run._download_raw()
+        self.assertTrue(raw.shifted)
+        for a,leaf in zip(raw.values,run._leaf):
+            self.assertTrue(np.isfinite(a[:,leaf]).all())
+            self.assertTrue(np.isnan(a[:,~leaf]).all())
+            self.assertTrue((1+a[:,leaf].sum(axis=0,dtype=np.float64)>0).all())
+
+    def test_identity_next_step_bitwise(self):
+        with CudaContext(0) as ctx,ExitStack() as owned:
+            old=owned.enter_context(make(ctx))
+            old._initialize_raw(initial(old));old.step(2)
+            before=old._download_raw()
+            new=owned.enter_context(old.remap_child(A))
+            self.check_reset(new,before.values)
+            self.snapshot_equal(old._download_raw(),before,'old unchanged')
+            self.snapshot_equal(new._download_raw(),before,'identity raw copy')
+            old.step();new.step()
+            self.assertEqual((old.step_count,new.step_count),(3,3))
+            self.snapshot_equal(old._download_raw(),new._download_raw(),'identity next step')
+            self.check_finite(old);self.check_finite(new)
+
+    def test_moved_remap_matches_raw_host_and_old_remains_usable(self):
+        with CudaContext(0) as ctx,ExitStack() as owned:
+            old=owned.enter_context(make(ctx))
+            old._initialize_raw(initial(old));old.step(2)
+            before=old._download_raw();count=old.step_count
+            old_mesh=RemapTopology.from_config(CFG,old.patch,A)
+            new_mesh=RemapTopology.from_config(CFG,old.patch,B)
+            expected=_remap_stored(old_mesh,new_mesh,before,shifted=True,
+                base_count=count,level_steps=(count,2*count,4*count),phase='postcollision_g')
+            new=owned.enter_context(old.remap_child(B))
+            self.check_reset(new,expected.g)
+            self.snapshot_equal(new._download_raw(),_StoredPopulations(expected.g,True),'host raw fill')
+            for a,b,om,nm in zip(before.values,new._download_raw().values,old_mesh.leaf_masks,new_mesh.leaf_masks):
+                self.exact(a[:,om&nm],b[:,om&nm],'physical same-level overlap')
+            self.snapshot_equal(old._download_raw(),before,'source unchanged')
+            old.step();new.step(2)
+            self.assertEqual((old.step_count,new.step_count),(3,4))
+            self.check_finite(old);self.check_finite(new)
+
+    def test_rejected_remap_does_not_mutate_source(self):
+        with CudaContext(0) as ctx,make(ctx) as run:
+            run._initialize_raw(initial(run));run.step(2)
+            before=run._download_raw()
+            with self.assertRaises(ValueError):run.remap_child([])
+            self.snapshot_equal(run._download_raw(),before,'rejection preserves source')
+            self.assertEqual(run.step_count,2)
+            run.step();self.check_finite(run)
+
+
+if __name__=='__main__':unittest.main()