"""Blocked fine storage for FixedPatchStepper: CPU tables, staging and launch contracts. The former GPU gate (dense fine storage vs 1x1, 2x2 and 4x2 fine blocks, bit-identical) was retired with the dense route; its result and the dense/blocked route parity are recorded at 1573942 and 2469af7. """ 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 from CelerisLab.lbm import curved_links 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(spec=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=spec) def actuation(speed): """Astra's seam fixture: 21 aperture links; in each decomposed layout 3 of them write a halo image and 2 read a q<.5 donor from a halo. Coarse lattice units: prepare_fine_boundary does the (acoustic) fine scaling, never scale twice.""" return SurfaceActuation(0.7, 0.85, speed) def fine_circle(mode, spec=None): cfg = config(mode) patch = FixedPatch2D.from_config(cfg, BOUNDS) return patch, prepare_fine_boundary(cfg, patch, (circle(spec),), wall_aware=mode == "slot_reservoir") def fine_cfg(mode, patch): from dataclasses import replace cfg = replace(config(mode), nx=patch.fine_shape[1], ny=patch.fine_shape[0], viscosity=.2) if mode == "slot_reservoir": cfg = replace(cfg, y_wall_bc="bounce_back", top_wall_bc="velocity", wall_slot_width=0, wall_normal_velocity=0.0, velocity=0.0) cfg.validate() # as the stepper does before surface_link_velocity return cfg 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 of every level; blocked storage unpacked from owner cores.""" out = [] for level in backend._levels: array = np.empty(level.shape if level.layout is None else level.layout.blocked_shape, np.float32) cuda.memcpy_dtoh(array, level.current) out.append(array if level.layout is None else level.layout.to_dense(array)) 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): for kw in ({"level_blocks": [None, (7, 7)]}, {"level_blocks": [None, (16,)]}): with self.subTest(kw=kw), self.assertRaises(ValueError): FixedPatchStepper(config("periodic"), context=self.forbidden_context(), refinements=[BOUNDS], **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(cuda, "mem_get_info", return_value=(1 << 40, 1 << 40)), 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"), context=context, refinements=[BOUNDS], level_blocks=[(12, 8), FINE_BLOCKS[0]]) self.assertEqual(len(backend.layouts[1].halo_dst), 0) self.assertIsNone(backend._levels[1].halo_dst) backend._gather_halo(backend._levels[1], 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) self.assertTrue((fallback[q < .5] == 0).all(), mode) # every q<.5 link uses its donor for by, bx in FINE_BLOCKS: layout = BlockLayout2D.build(*patch.fine_shape, by, bx, periodic=(False, False)) _, mapped = blocked_link_tables(layout, boundary.flags, fluid_idx, dirs, q, fallback) single = layout.nblocks == 1 self.assertEqual(seam_cells(layout, mapped), (0, 0) if single else (8, 8), (mode, by, bx)) self.assertEqual(len(set(mapped.block.tolist())), 1 if single else 4, (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 mocked_backend(self, objects): """Real constructor/initialize with driver doubles; kernels record (name, args).""" 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(cuda, "mem_get_info", return_value=(1 << 40, 1 << 40)), 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"), \ patch.object(curved_links, "CurvedLinkSoA", side_effect=AssertionError("dense SoA built")): backend = FixedPatchStepper(config("periodic"), context=context, refinements=[BOUNDS], level_blocks=[(12, 8), FINE_BLOCKS[1]], level_objects=[(), tuple(objects)]) backend.initialize(state(NY, NX, 1), state(*backend.hierarchy.pairs[0].fine_shape, 2)) return backend, calls def test_same_buffer_halo_bfl_step_launch_order(self): for spec in (None, actuation(.002)): backend, calls = self.mocked_backend([circle(spec)]) c, f = backend._levels coarse, cur, spare = c.current, f.current, f.spare flags, kind = backend._level_boundaries[1].flags, backend._levels[1].kind velocity = backend._level_boundaries[1].device.get("surface_velocity") self.assertEqual(velocity is None, spec is None) 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), ("BlockHaloCopy", coarse), ("BlockCoalesce", cur), ("BlockStep", coarse)]) steps = [a for name, a in calls if name == "BlockStep"][:2] # the two fine hops 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) for a in (a for name, a in calls if name == "BlockCurvedBoundary"): if spec is None: # explicit null at the raw launch ABI self.assertIsInstance(a[7], np.uintp) self.assertEqual(int(a[7]), 0) else: self.assertIs(a[7], velocity) def test_actuation_without_velocity_buffer_fails_before_bfl(self): for speed in (0.0, .002): # any spec, even zero amplitude backend, calls = self.mocked_backend([circle(actuation(speed))]) self.assertIsNotNone(backend._level_boundaries[1].tables["surface_velocity"]) del backend._level_boundaries[1].device["surface_velocity"] with self.assertRaisesRegex(RuntimeError, "Surface velocity device buffer is missing"): backend.step(1) names = [name for name, _ in calls] self.assertNotIn("BlockCurvedBoundary", names) self.assertNotIn("BlockStep", names) self.assertEqual(backend._status, "failed") backend.close() def test_surface_velocity_helper_matches_dense_soa(self): for mode in ("periodic", "slot_reservoir"): patch, boundary = fine_circle(mode, actuation(-.002)) cfg = fine_cfg(mode, patch) soa = curved_links.CurvedLinkSoA() soa.assign_host(*boundary.curved_host) soa.assign_surface(boundary.objects, cfg, boundary.flags) fluid_idx, dirs, _, body_id, rx, ry, _, _ = boundary.curved_host velocity, aperture = curved_links.surface_link_velocity( boundary.objects, cfg, boundary.flags, fluid_idx, dirs, body_id, rx, ry) np.testing.assert_array_equal(velocity, soa.surface_velocity) np.testing.assert_array_equal(aperture, soa.aperture) self.assertEqual(curved_links.surface_link_velocity( (circle(),), cfg, boundary.flags, fluid_idx, dirs, body_id, rx, ry), (None, None)) # Actuated (not just geometric) seam coverage in both decomposed layouts. self.assertEqual(int(aperture.sum()), 21, mode) self.assertTrue(np.all(np.any(velocity[aperture] != 0, axis=1)), 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, boundary.curved_host[2], boundary.curved_host[7]) L = int(np.prod(layout.local_shape)) core = set(layout.owner.ravel().tolist()) halo_write = np.array([int(w) // (9 * L) * L + int(w) % L not in core for w in mapped.write]) halo_donor = np.array([d != NO_DONOR and int(d) not in core for d in mapped.donor]) self.assertEqual((int((halo_write & aperture).sum()), int((halo_donor & aperture).sum())), (3, 2), (mode, by, bx)) 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()) if __name__ == "__main__": unittest.main()