"""Opt-in blocked fine storage for FixedPatchStepper: CPU tables plus GPU gate. Frozen GPU gate (same stream, 8 coarse steps, nonequilibrium start), in both periodic and slot_reservoir modes: dense fine storage vs 1x1, 2x2 and 4x2 fine blocks give bit-identical masked downloads AND bit-identical raw coarse and unpacked owner-core fine arrays, including patch-ghost values/NaN masks. """ from contextlib import nullcontext import unittest from unittest.mock import Mock, patch import numpy as np import pycuda.driver as cuda from CelerisLab.body.geometry.circle import CircleGeometry from CelerisLab.body.objects import SimObject, SurfaceActuation from CelerisLab.config import LBMConfig from CelerisLab.cuda import compiler_v2 as compiler from CelerisLab.lbm.blocks import NO_DONOR, BlockLayout2D, blocked_link_tables from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY from CelerisLab.lbm.descriptors import SOLID from CelerisLab.lbm.refinement import (ACTIVE, INNER_GHOST, OUTER_GHOST, FixedPatch2D, prepare_fine_boundary) from CelerisLab.lbm.refinement_stepper import FixedPatchStepper, blocked_fine_tables NX, NY, BOUNDS, STEPS = 32, 24, (8, 6, 24, 18), 8 FINE_BLOCKS = ((32, 40), (16, 20), (16, 10)) # fine window 32x40: 1x1, 2x2, 4x2 W = np.array([4 / 9] + [1 / 9] * 4 + [1 / 36] * 4) CX, CY = np.array(D2Q9_EX, float), np.array(D2Q9_EY, float) def config(mode): if mode == "periodic": return LBMConfig(nx=NX, ny=NY, viscosity=.1) return LBMConfig(nx=NX, ny=NY, viscosity=.1, velocity=0, wall_normal_velocity=.001, wall_slot_width=8, y_wall_bc="on_site_velocity", top_wall_bc="reservoir") def state(ny, nx, seed): """Positive low-Mach equilibrium plus a seeded nonequilibrium part.""" y, x = np.mgrid[0:ny, 0:nx] rho = 1 + 1e-3 * np.cos(2 * np.pi * x / nx) ux, uy = .02 * np.sin(2 * np.pi * y / ny), .01 * np.cos(2 * np.pi * x / nx) cu = CX[:, None, None] * ux + CY[:, None, None] * uy feq = W[:, None, None] * rho * (1 + 3 * cu + 4.5 * cu * cu - 1.5 * (ux * ux + uy * uy)) noise = np.random.default_rng(seed).uniform(-1, 1, feq.shape) * 1e-3 * W[:, None, None] return (feq + noise).astype(np.float32) def circle(actuation=None): """Astra fixture: coarse cell-edge C=(16,12), R=2.15 -> fine C=(19.5,15.5), R=4.3.""" return SimObject(0, geometry=CircleGeometry(16.0, 12.0, 2.15), center=(16.0, 12.0), radius=2.15, surface_actuation=actuation) def fine_circle(mode): cfg = config(mode) patch = FixedPatch2D.from_config(cfg, BOUNDS) return patch, prepare_fine_boundary(cfg, patch, (circle(),), wall_aware=mode == "slot_reservoir") def seam_cells(layout, mapped): """Distinct non-core cells written / read as q<.5 donors (receiver-local halos).""" L = int(np.prod(layout.local_shape)) core = set(layout.owner.ravel().tolist()) writes = {int(w) // (9 * L) * L + int(w) % L for w in mapped.write} donors = {int(d) for d in mapped.donor if d != NO_DONOR} return len(writes - core), len(donors - core) def raw(backend): """Unmasked device arrays; blocked fine storage unpacked from owner cores.""" coarse, fine = backend._levels out = [np.empty(coarse.shape, np.float32)] out.append(np.empty(fine.shape if backend.fine_layout is None else backend.fine_layout.blocked_shape, np.float32)) cuda.memcpy_dtoh(out[0], coarse.current) cuda.memcpy_dtoh(out[1], fine.current) if backend.fine_layout is not None: out[1] = backend.fine_layout.to_dense(out[1]) return out class BlockedFineCPU(unittest.TestCase): def forbidden_context(self): return Mock(activate=Mock(side_effect=AssertionError("CUDA reached")), sm_arch="sm_70") def test_rejections_before_cuda(self): actuated = circle(SurfaceActuation(0.0, 0.5, 0.002)) for kw in ({"fine_block": (7, 7)}, {"fine_block": (16,)}, {"fine_block": (16, 20), "fine_objects": [actuated]}): with self.subTest(kw=kw), self.assertRaises(ValueError): FixedPatchStepper(config("periodic"), BOUNDS, context=self.forbidden_context(), **kw) def test_single_block_layout_allocates_no_empty_halo_tables(self): # A 1x1 non-periodic fine layout has zero halo pairs; CUDA rejects mem_alloc(0). context = Mock(activate=nullcontext, sm_arch="sm_70") module = Mock() module.get_global.return_value = (Mock(), 20) sizes = [] def alloc(size): if size <= 0: raise AssertionError("zero-byte device allocation") sizes.append(size) return Mock() with patch.object(cuda, "Stream"), patch.object(compiler, "generate_config"), \ patch.object(compiler, "compile_kernel"), patch.object(compiler, "cleanup_build_dir"), \ patch.object(compiler, "load_module", return_value=module), \ patch.object(cuda, "mem_alloc", side_effect=alloc), patch.object(cuda, "memcpy_htod"): backend = FixedPatchStepper(config("periodic"), BOUNDS, context=context, fine_block=FINE_BLOCKS[0]) self.assertEqual(len(backend.fine_layout.halo_dst), 0) self.assertIsNone(backend._halo_dst) backend._fine_halo(Mock()) # no launch without pairs module.get_function.return_value.assert_not_called() backend.close() self.assertTrue(sizes) def test_circle_fixture_counts_and_seam_staging(self): for mode in ("periodic", "slot_reservoir"): patch, boundary = fine_circle(mode) fluid_idx, dirs, q, _, _, _, _, fallback = boundary.curved_host self.assertEqual(len(fluid_idx), 76, mode) self.assertEqual(int(np.count_nonzero(q < .5)), 68, mode) for by, bx in FINE_BLOCKS[1:]: layout = BlockLayout2D.build(*patch.fine_shape, by, bx, periodic=(False, False)) _, mapped = blocked_link_tables(layout, boundary.flags, fluid_idx, dirs, q, fallback) self.assertEqual(seam_cells(layout, mapped), (8, 8), (mode, by, bx)) def test_storage_adapter_flags_sizing(self): patch, boundary = fine_circle("periodic") columns = [boundary.curved_host[i] for i in (0, 1, 2, 7)] layout = BlockLayout2D.build(*patch.fine_shape, 16, 20, periodic=(False, False)) tables, mapped = blocked_link_tables(layout, boundary.flags, *columns) L = int(np.prod(layout.local_shape)) self.assertEqual(tables["flags"].shape, (layout.nblocks * L,)) np.testing.assert_array_equal(tables["flags"][layout.owner.ravel().astype(np.intp)], boundary.flags) np.testing.assert_array_equal(tables["cell"], mapped.receiver % np.uint64(L)) for bad in (boundary.flags[:-1], boundary.flags.astype(np.uint32)): with self.assertRaises(ValueError): blocked_link_tables(layout, bad, *columns) def test_same_buffer_halo_bfl_step_launch_order(self): context = Mock(activate=nullcontext, sm_arch="sm_70") calls, functions = [], {} def get_function(name): return functions.setdefault(name, Mock(side_effect=lambda *a, **k: calls.append((name, a)))) module = Mock(get_function=Mock(side_effect=get_function)) module.get_global.return_value = (Mock(), 20) with patch.object(cuda, "Stream"), patch.object(compiler, "generate_config"), \ patch.object(compiler, "compile_kernel"), patch.object(compiler, "cleanup_build_dir"), \ patch.object(compiler, "load_module", return_value=module), \ patch.object(cuda, "mem_alloc", side_effect=lambda size: Mock()), \ patch.object(cuda, "memcpy_htod"): backend = FixedPatchStepper(config("periodic"), BOUNDS, context=context, fine_objects=[circle()], fine_block=FINE_BLOCKS[1]) g_fine = state(*backend.patch.fine_shape, 2) backend.initialize(state(NY, NX, 1), g_fine) c, f = backend._levels coarse, cur, spare = c.current, f.current, f.spare flags, kind = backend._fine_flags, backend._kind backend.step(1) backend.close() order = [(name, a[0]) for name, a in calls] self.assertEqual(order, [("BlockExplode", coarse), ("BlockHaloCopy", cur), ("BlockCurvedBoundary", cur), ("BlockStep", cur), ("BlockHaloCopy", spare), ("BlockCurvedBoundary", spare), ("BlockStep", spare), ("PatchCoarseStep", coarse)]) steps = [a for name, a in calls if name == "BlockStep"] self.assertEqual([a[1] for a in steps], [spare, cur]) for a in steps: self.assertIs(a[3], flags) # blocked flags, not null self.assertIs(a[4], kind) def test_tables_map_back_to_dense_slots_in_order(self): patch = FixedPatch2D.from_config(config("periodic"), BOUNDS) fy, fx = patch.fine_shape cells = fy * fx order = np.argsort(patch.coalesce_coarse) for by, bx in FINE_BLOCKS: layout = BlockLayout2D.build(fy, fx, by, bx, periodic=(False, False)) core = set(layout.owner.ravel().tolist()) targets, children, kind = blocked_fine_tables(patch, layout, order) # Explosion targets: owner cores of exactly the non-ACTIVE fine cells. np.testing.assert_array_equal(targets, layout.owner.ravel()[patch.ghost_fine]) self.assertEqual(sorted(targets.tolist()), sorted(layout.owner[patch.fine_kind != ACTIVE].tolist())) # Children: same rows/child order, same (q, dense cell), owner cores. L = int(np.prod(layout.local_shape)) b, rest = np.divmod(children, np.uint64(9 * L)) q, c = np.divmod(rest, np.uint64(L)) dense = patch.coalesce_fine[order] np.testing.assert_array_equal(q, dense // cells) np.testing.assert_array_equal(b * np.uint64(L) + c, layout.owner.ravel()[(dense % cells).astype(np.intp)]) self.assertTrue(set((b * np.uint64(L) + c).ravel().tolist()) <= core) # Kinds follow owner cores; unused halo entries are OUTER. np.testing.assert_array_equal(kind[layout.owner.ravel().astype(np.intp)], patch.fine_kind.ravel()) self.assertEqual(set(np.unique(kind).tolist()) - {ACTIVE, INNER_GHOST, OUTER_GHOST}, set()) class BlockedFineGPU(unittest.TestCase): def test_dense_and_blocked_fine_storage_are_bit_identical(self): from CelerisLab.cuda.context import CudaContext for mode in ("periodic", "slot_reservoir"): cfg = config(mode) patch = FixedPatch2D.from_config(cfg, BOUNDS) g_coarse, g_fine = state(NY, NX, 1), state(*patch.fine_shape, 2) results = [] with CudaContext(0) as context: for fine_block in (None, *FINE_BLOCKS): with FixedPatchStepper(cfg, BOUNDS, context=context, coarse_boundary=mode, fine_block=fine_block) as backend: backend.initialize(g_coarse, g_fine) backend.step(STEPS) self.assertEqual(backend.step_count, STEPS) results.append((backend.download(), raw(backend))) (dense_public, dense_raw), blocked = results[0], results[1:] kind = patch.fine_kind self.assertTrue(np.isnan(dense_raw[1][:, kind == OUTER_GHOST]).all(), mode) # INNER populations may be NaN after hop 2 (pulled from poisoned OUTER); # their exact pattern is pinned by the dense-vs-blocked comparison below. self.assertTrue(np.isfinite(dense_raw[1][:, kind == ACTIVE]).all(), mode) for fine_block, (public, raw_arrays) in zip(FINE_BLOCKS, blocked): msg = f"{mode} fine_block={fine_block}" for a, b in zip(public + tuple(raw_arrays), dense_public + tuple(dense_raw)): np.testing.assert_array_equal(a, b, err_msg=msg) # values and NaN masks print(f"{mode}: dense vs {len(blocked)} fine-block layouts bit-identical " f"after {STEPS} coarse steps", flush=True) def test_passive_circle_dense_and_blocked_fine_storage_are_bit_identical(self): from CelerisLab.cuda.context import CudaContext for mode in ("periodic", "slot_reservoir"): cfg = config(mode) patch, boundary = fine_circle(mode) solid = (boundary.flags.reshape(patch.fine_shape) & SOLID) != 0 g_coarse, g_fine = state(NY, NX, 1), state(*patch.fine_shape, 2) results = [] with CudaContext(0) as context: for fine_block in (None, *FINE_BLOCKS): with FixedPatchStepper(cfg, BOUNDS, context=context, coarse_boundary=mode, fine_objects=[circle()], fine_block=fine_block) as backend: backend.initialize(g_coarse, g_fine) backend.step(STEPS) results.append((backend.download(), raw(backend))) (dense_public, dense_raw), blocked = results[0], results[1:] fine = dense_raw[1] self.assertTrue(np.isfinite(fine[:, (patch.fine_kind == ACTIVE) & ~solid]).all(), mode) self.assertTrue(np.isnan(fine[:, solid]).all(), mode) self.assertTrue(np.isnan(fine[:, patch.fine_kind == OUTER_GHOST]).all(), mode) for fine_block, (public, raw_arrays) in zip(FINE_BLOCKS, blocked): msg = f"circle {mode} fine_block={fine_block}" for a, b in zip(public + tuple(raw_arrays), dense_public + tuple(dense_raw)): np.testing.assert_array_equal(a, b, err_msg=msg) # values and NaN masks print(f"circle {mode}: dense vs {len(blocked)} fine-block layouts bit-identical " f"after {STEPS} coarse steps", flush=True) if __name__ == "__main__": unittest.main()