"""Short root-channel GPU gates. Caller owns GPU admission (CUDA_VISIBLE_DEVICES). Run: python -m unittest discover -s tests/validation/multigrid -p test_channel_root_gpu.py Uniform parity covers every ghost/corner and bx=1 outlet halo donors. AMR gates are split invariance and finite evolution, not external DFG benchmark validation. """ from dataclasses import replace import struct import unittest import numpy as np import pycuda.driver as cuda from CelerisLab.config import LBMConfig from CelerisLab.cuda import compiler from CelerisLab.cuda.context import CudaContext from CelerisLab.lbm.blocks import BlockLayout2D from CelerisLab.lbm.descriptors import FLUID, SOLID, BC_WALL, BC_INLET, BC_OUTLET from CelerisLab.lbm.population_storage import _encode_physical from CelerisLab.lbm.refinement_stepper import FixedPatchStepper def state(ny, nx): q, y, x = np.mgrid[:9, :ny, :nx] w = np.array([4/9]+[1/9]*4+[1/36]*4)[:, None, None] return w*(1+.001*np.sin(x+3*y+q)) def compile_module(cfg, context, source): cfg.validate() build = compiler.generate_config(cfg, block_root_bc=2 if source == "blocked.cu" else 0) try: module = compiler.load_module(compiler.compile_kernel(context.sm_arch, source=source, build_dir=build)) finally: compiler.cleanup_build_dir(build) cuda.memcpy_htod(module.get_global("d_params")[0], struct.pack("=ffffI", cfg.omega, 0, 0, 0, 0)) return module def run_root(cfg, context, initial, block=None, steps=8): """Launch the production uniform or blocked kernel, without a test-specific backend.""" allocations = [] def upload(a): a = np.ascontiguousarray(a) p = cuda.mem_alloc(a.nbytes); allocations.append(p); cuda.memcpy_htod(p, a) return p try: layout = None if block is None else BlockLayout2D.build(cfg.ny, cfg.nx, *block) local = cfg if layout is None else replace(cfg, nx=layout.local_shape[1], ny=layout.local_shape[0]) module = compile_module(local, context, "uniform.cu" if layout is None else "blocked.cu") data = initial if layout is None else layout.to_blocked(initial) current, spare = upload(data), upload(np.full_like(data, np.nan)) tpb = cfg.threads_per_block kw = dict(block=(tpb, 1, 1)) if layout is None: flags = np.full((cfg.ny,cfg.nx), FLUID, np.uint16) flags[:, 0] = SOLID|BC_INLET; flags[:, -1] = SOLID|BC_OUTLET flags[[0, -1], :] = SOLID|BC_WALL flag = upload(flags) step = module.get_function("OneStep") else: dst, src = upload(layout.halo_dst), upload(layout.halo_src) gather = module.get_function("BlockHaloCopy") step = module.get_function("BlockChannelStep") cuda.memcpy_htod(module.get_global("d_block_channel")[0], struct.pack("=IIf", cfg.nx, cfg.ny, cfg.velocity)) grid = ((local.nx+tpb-1)//tpb, local.ny, 1 if layout is None else layout.nblocks) for _ in range(steps): if layout is None: step(flag, current, spare, grid=grid, **kw) else: gather(current, dst, src, np.uint32(len(layout.halo_dst)), grid=((len(layout.halo_dst)+tpb-1)//tpb, 1), **kw) step(current, spare, np.uint32(1), np.uintp(0), np.uintp(0), grid=grid, **kw) current, spare = spare, current cuda.Context.synchronize() result = np.empty_like(data); cuda.memcpy_dtoh(result, current) return result if layout is None else layout.to_dense(result) finally: for pointer in allocations: pointer.free() class ChannelRootGPU(unittest.TestCase): def test_uniform_matched_root(self): with CudaContext(0) as context: for shifted in (False, True): for profile, clamp, damping in (("uniform", True, .5), ("parabolic", True, .5), ("parabolic", False, .8)): cfg = LBMConfig(nx=24, ny=20, viscosity=.1, velocity=.025, ddf_shifting=shifted, inlet_profile=profile, outlet_backflow_clamp=clamp, outlet_srt_neq_damp=damping) initial = _encode_physical(state(cfg.ny, cfg.nx), shifted) expected = run_root(cfg, context, initial) self.assertTrue(np.isfinite(expected).all()) for block in ((20,24), (5,6), (4,1)): actual = run_root(cfg, context, initial, block) np.testing.assert_array_equal(actual, expected) def test_amr_channel_body_jet_smoke(self): from CelerisLab.body.geometry.circle import CircleGeometry from CelerisLab.body.objects import SimObject, SurfaceActuation center = (24., 16.) body = SimObject(0, geometry=CircleGeometry(*center, 2.15), center=center, radius=2.15) jet = SurfaceActuation(.7, .85, .001).as_jet(center, 2.15) with CudaContext(0) as context: for shifted in (False, True): for bodies in (False, True): cfg = LBMConfig(nx=48, ny=32, viscosity=.1, velocity=.015, ddf_shifting=shifted) reference = None for block in ((32,48), (8,8), (4,1)): with FixedPatchStepper(cfg, context=context, coarse_boundary="channel", refinements=[(8,6,40,26)], level_blocks=[block, (8,8)], level_objects=[(), (body,) if bodies else ()], level_jets=[(), (jet,) if bodies else ()]) as backend: arrays = [state(*level.shape[1:]) for level in backend._levels] backend.initialize(*arrays) backend.step(4) if bodies: backend.set_jet_amplitude(0, -.001) backend.step(4) raw = backend._download_raw() physical = backend.download() for a, mask in zip(raw.values, backend._state): self.assertTrue(np.isfinite(a[:,mask]).all()) self.assertTrue(np.isfinite(raw.values[0][:,0,:]).all()) self.assertTrue(np.isnan(physical[0][:,0,:]).all()) backend._initialize_raw(raw) for a,b in zip(raw.values, backend._download_raw().values): np.testing.assert_array_equal(a,b) if reference is None: reference = raw.values else: for a,b in zip(reference, raw.values): np.testing.assert_array_equal(a,b) with self.assertRaisesRegex(ValueError, "periodic"): backend.remap_child([]) if __name__ == "__main__": unittest.main()