Support blocked parent storage in fixed patch stepping

tests/validation/multigrid/test_refinement_blocked_base.py

@@ -0,0 +1,324 @@
+"""All-blocked fixed patch: transfer ordering, failure ownership and CUDA parity.
+
+Run CPU and GPU classes in separate processes. GPU callers must first check
+resources/device availability; this file never selects or reserves a GPU.
+"""
+from contextlib import ExitStack, contextmanager, nullcontext
+from dataclasses import replace
+from pathlib import Path
+import struct
+import unittest
+from unittest.mock import Mock, patch
+
+import numpy as np
+import pycuda.driver as cuda
+
+from CelerisLab.cuda import compiler_v2 as compiler
+from CelerisLab.lbm import refinement_stepper
+from CelerisLab.lbm.blocks import BlockLayout2D
+from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY
+from CelerisLab.lbm.refinement import ACTIVE, OUTER_GHOST, FixedPatch2D
+from CelerisLab.lbm.refinement_stepper import FixedPatchStepper, blocked_base_tables
+from test_refinement_blocked_fine import NX, NY, BOUNDS, config, state, circle, actuation
+
+BASE_BLOCKS = ((24, 32), (12, 8), (3, 4), (1, 1))
+FINE_BLOCK = (16, 20)
+
+
+def raw(backend):
+    arrays = []
+    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)
+        arrays.append(array if level.layout is None else level.layout.to_dense(array))
+    return tuple(arrays)
+
+
+def exact(test, actual, expected, label):
+    """Compare the complete invalid mask and every finite FP32 bit (incl. -0)."""
+    np.testing.assert_array_equal(np.isnan(actual), np.isnan(expected), err_msg=label)
+    test.assertFalse(np.isinf(actual).any(), label)
+    test.assertFalse(np.isinf(expected).any(), label)
+    finite = np.isfinite(expected)
+    np.testing.assert_array_equal(actual[finite].view(np.uint32),
+                                  expected[finite].view(np.uint32), err_msg=label)
+
+
+@contextmanager
+def mocked_backend(mode="periodic", fail_alloc=None):
+    """Use the real constructor, uploads and launch logic, recording driver calls."""
+    calls, allocations, uploads, functions = [], [], [], {}
+    stream = Mock()
+    context = Mock(activate=nullcontext, sm_arch="sm_70")
+
+    def function(name):
+        if name not in functions:
+            functions[name] = Mock(side_effect=lambda *a, **kw: calls.append((name, a, kw)))
+        return functions[name]
+
+    def allocate(size):
+        if len(allocations) == fail_alloc:
+            raise MemoryError("injected allocation failure")
+        if size <= 0:
+            raise AssertionError("empty allocation")
+        pointer = Mock()
+        allocations.append(pointer)
+        return pointer
+
+    module = Mock(get_function=Mock(side_effect=function))
+    module.get_global.side_effect = lambda name: (name, 16 if name == "d_block_wall" else 20)
+    with ExitStack() as stack:
+        for name in ("generate_config", "compile_kernel", "cleanup_build_dir"):
+            stack.enter_context(patch.object(compiler, name))
+        stack.enter_context(patch.object(compiler, "load_module", return_value=module))
+        stack.enter_context(patch.object(cuda, "Stream", return_value=stream))
+        stack.enter_context(patch.object(cuda, "mem_alloc", side_effect=allocate))
+        stack.enter_context(patch.object(cuda, "memcpy_htod",
+                                        side_effect=lambda p, a: uploads.append((p, a))))
+        yield dict(context=context, calls=calls, allocations=allocations, uploads=uploads,
+                   functions=functions, stream=stream)
+
+
+class BlockedBaseCPU(unittest.TestCase):
+    def test_reject_before_cuda(self):
+        context = Mock(activate=Mock(side_effect=AssertionError("CUDA reached")))
+        for kwargs in ({"base_block": (12, 8)},
+                       {"base_block": (7, 8), "fine_block": FINE_BLOCK},
+                       {"base_block": (12,), "fine_block": FINE_BLOCK}):
+            with self.subTest(kwargs=kwargs), self.assertRaises(ValueError):
+                FixedPatchStepper(config("periodic"), BOUNDS, context=context, **kwargs)
+        context.activate.assert_not_called()
+
+    def test_all_arrival_rows_donors_and_kind_masks(self):
+        p = FixedPatch2D.from_config(config("periodic"), BOUNDS)
+        order = np.argsort(p.coalesce_coarse)
+        self.assertEqual(len(order), 172)
+        for mode in ("periodic", "slot_reservoir"):
+            leaves = np.ones((NY, NX), bool)
+            x0, y0, x1, y1 = BOUNDS
+            leaves[y0:y1, x0:x1] = False
+            if mode != "periodic":
+                leaves[[0, -1]] = False
+            for by, bx in BASE_BLOCKS:
+                layout = BlockLayout2D.build(NY, NX, by, bx, periodic=(True, mode == "periodic"))
+                donors, targets, kind = blocked_base_tables(p, layout, order, mode != "periodic")
+                ly, lx = layout.local_shape
+                L = ly * lx
+                expected, source_cells = [], []
+                for slot in p.coalesce_coarse[order]:
+                    q, receiver = divmod(int(slot), NX * NY)
+                    y, x = divmod(receiver, NX)
+                    block = (y // by) * (NX // bx) + x // bx
+                    # Scalar coordinate construction independent of layout.owner / slot.
+                    sy, sx = y % by + 1 - D2Q9_EY[q], x % bx + 1 - D2Q9_EX[q]
+                    expected.append((block * 9 + q) * L + sy * lx + sx)
+                    source_cells.append((y - D2Q9_EY[q], x - D2Q9_EX[q]))
+                np.testing.assert_array_equal(targets, expected)
+                self.assertEqual(len(set(targets.tolist())), 172)
+                self.assertEqual(sum(not (y0 <= y < y1 and x0 <= x < x1)
+                                     for y, x in source_cells), 8)
+                expected_donors = []
+                for cell in p.ghost_coarse:
+                    y, x = divmod(int(cell), NX)
+                    block = (y // by) * (NX // bx) + x // bx
+                    expected_donors.append(block * L + (y % by + 1) * lx + x % bx + 1)
+                np.testing.assert_array_equal(donors, expected_donors)
+                expected_kind = np.full(layout.nblocks * L, OUTER_GHOST, np.uint8)
+                expected_kind[layout.owner[leaves].astype(np.intp)] = ACTIVE
+                np.testing.assert_array_equal(kind, expected_kind)
+
+    def test_staging_negative_controls_detect_seam_and_donor_corruption(self):
+        p = FixedPatch2D.from_config(config("periodic"), BOUNDS)
+        order = np.argsort(p.coalesce_coarse)
+        layout = BlockLayout2D.build(NY, NX, 3, 4)
+        donors, targets, _ = blocked_base_tables(p, layout, order)
+        dense = np.arange(9 * NY * NX, dtype=np.float32).reshape(9, NY, NX)
+        initial = layout.to_blocked(dense)
+        arrivals = np.arange(len(targets), dtype=np.float32) + 10000
+        staged = initial.copy()
+        staged.ravel()[targets.astype(np.intp)] = arrivals
+        def pull(storage):
+            out = np.empty_like(dense)
+            by, bx = layout.block
+            ly, lx = layout.local_shape
+            for y in range(NY):
+                for x in range(NX):
+                    b = (y // by) * (NX // bx) + x // bx
+                    for q in range(9):
+                        out[q, y, x] = storage[b, q, y % by + 1 - D2Q9_EY[q],
+                                               x % bx + 1 - D2Q9_EX[q]]
+            return out
+
+        expected = np.array([np.roll(dense[q], (D2Q9_EY[q], D2Q9_EX[q]), axis=(0, 1))
+                             for q in range(9)])
+        expected.ravel()[p.coalesce_coarse[order].astype(np.intp)] = arrivals
+        leaves = np.ones((NY, NX), bool)
+        leaves[6:18, 8:24] = False
+        # Compare ALL parent arrivals, not only the slots deliberately written.
+        np.testing.assert_array_equal(pull(staged)[:, leaves], expected[:, leaves])
+        # A late halo gather erases receiver-local arrivals at seams.
+        late_gather = staged.copy()
+        layout.exchange(late_gather)
+        self.assertGreater(np.count_nonzero(pull(late_gather)[:, leaves] != expected[:, leaves]), 0)
+        # Canonical source owners are not the slots read by each receiver's pull.
+        q, receiver = np.divmod(p.coalesce_coarse[order], NX * NY)
+        y, x = np.divmod(receiver.astype(np.intp), NX)
+        owners = layout.owner[y - np.asarray(D2Q9_EY)[q.astype(np.intp)],
+                              x - np.asarray(D2Q9_EX)[q.astype(np.intp)]]
+        wrong = initial.copy()
+        wrong.ravel()[layout.slot(owners, q).astype(np.intp)] = arrivals
+        self.assertGreater(np.count_nonzero(pull(wrong)[:, leaves] != expected[:, leaves]), 0)
+        # With one parent block, all eight grazing sources are owner cores.
+        # Each is read for four fine ghost children: late explosion corrupts32.
+        single = BlockLayout2D.build(NY, NX, NY, NX)
+        donors, targets, _ = blocked_base_tables(p, single, order)
+        initial = single.to_blocked(dense)
+        staged = initial.copy()
+        staged.ravel()[targets.astype(np.intp)] = arrivals
+        donor_slots = single.slot(donors[None, :], np.arange(9)[:, None]).astype(np.intp)
+        self.assertEqual(np.count_nonzero(staged.ravel()[donor_slots]
+                                          != initial.ravel()[donor_slots]), 32)
+
+    def test_raw_mask_gate_detects_leak_hidden_by_download_mask(self):
+        valid = state(NY, NX, 1)
+        valid[:, 6:18, 8:24] = np.nan
+        leaked = valid.copy()
+        leaked[5, 6, 8] = 42  # skipped invalid-cell store leaves stale staging data
+        public = leaked.copy()
+        public[:, 6:18, 8:24] = np.nan
+        exact(self, public, valid, "masked download cannot detect stale staging")
+        with self.assertRaises(AssertionError):
+            exact(self, leaked, valid, "raw owner-core must detect stale staging")
+
+    def test_launch_order_stream_wall_constants_and_buffer_swaps(self):
+        with mocked_backend("slot_reservoir") as m:
+            backend = FixedPatchStepper(config("slot_reservoir"), BOUNDS, context=m["context"],
+                                        coarse_boundary="slot_reservoir", fine_block=FINE_BLOCK,
+                                        base_block=(3, 4), fine_objects=[circle(actuation(.002))])
+            try:
+                backend.initialize(state(NY, NX, 1), state(*backend.patch.fine_shape, 2))
+                coarse, fine = backend._levels
+                current, spare, fine_current = coarse.current, coarse.spare, fine.current
+                backend.step(2)
+                names = [name for name, _, _ in m["calls"]]
+                interval = ["BlockExplode", "BlockHaloCopy", "BlockCurvedBoundary", "BlockStep",
+                            "BlockHaloCopy", "BlockCurvedBoundary", "BlockStep",
+                            "BlockHaloCopy", "BlockCoalesce", "BlockStep"]
+                self.assertEqual(names, interval * 2)
+                self.assertEqual(backend.step_count, 2)
+                self.assertIs(coarse.current, current)
+                self.assertIs(coarse.spare, spare)
+                self.assertIs(fine.current, fine_current)
+                for _, _, kw in m["calls"]:
+                    self.assertIs(kw["stream"], m["stream"])
+                for offset, src, dst in ((0, current, spare), (10, spare, current)):
+                    explode, gather, stage, collide = [m["calls"][offset + i][1] for i in (0, 7, 8, 9)]
+                    self.assertIs(explode[0], src)
+                    self.assertIs(gather[0], src)
+                    self.assertIs(stage[0], fine_current)
+                    self.assertIs(stage[1], src)
+                    self.assertIs(collide[0], src)
+                    self.assertIs(collide[1], dst)
+                    self.assertIs(collide[4], coarse.kind)
+                walls = [a for ptr, a in m["uploads"] if isinstance(ptr, str) and ptr == "d_block_wall"]
+                self.assertEqual(len(walls), 1)
+                width, height, slot, speed = struct.unpack("=IIIf", walls[0])
+                self.assertEqual((width, height, slot), (NX, NY, 8))
+                self.assertAlmostEqual(speed, .001)
+            finally:
+                backend.close()
+            for allocation in m["allocations"]:
+                allocation.free.assert_called_once_with()
+
+    def test_allocation_failure_releases_every_owned_buffer(self):
+        # Fail during the first level, second level and transfer-table uploads.
+        for fail_at in (1, 6, 10):
+            with self.subTest(fail_at=fail_at), mocked_backend(fail_alloc=fail_at) as m:
+                with self.assertRaisesRegex(MemoryError, "injected allocation"):
+                    FixedPatchStepper(config("periodic"), BOUNDS, context=m["context"],
+                                      fine_block=FINE_BLOCK, base_block=(3, 4))
+                self.assertEqual(len(m["allocations"]), fail_at)
+                for allocation in m["allocations"]:
+                    allocation.free.assert_called_once_with()
+
+    def test_failed_parent_stage_cannot_publish_or_resume(self):
+        with mocked_backend() as m:
+            backend = FixedPatchStepper(config("periodic"), BOUNDS, context=m["context"],
+                                        fine_block=FINE_BLOCK, base_block=(3, 4))
+            backend.initialize(state(NY, NX, 1), state(*backend.patch.fine_shape, 2))
+            m["functions"]["BlockCoalesce"].side_effect = RuntimeError("injected parent stage failure")
+            current = backend._levels[0].current
+            with self.assertRaisesRegex(RuntimeError, "injected parent stage"):
+                backend.step()
+            self.assertEqual(backend.step_count, 0)
+            self.assertEqual(backend._status, "failed")
+            self.assertIs(backend._levels[0].current, current)
+            with self.assertRaisesRegex(RuntimeError, "failed"):
+                backend.download()
+            with self.assertRaisesRegex(RuntimeError, "failed"):
+                backend.step()
+            backend.close()
+            for allocation in m["allocations"]:
+                allocation.free.assert_called_once_with()
+
+
+class BlockedBaseGPU(unittest.TestCase):
+    def test_bodyfree_exact_parity_and_raw_invalid_slots(self):
+        self.compare([(None, ((24, 32), (3, 4), (1, 1)), 8)])
+
+    def test_actuated_circle_exact_parity_and_raw_invalid_slots(self):
+        self.compare([(.002, ((12, 8),), 8), (-.002, ((12, 8),), 8)])
+
+    def test_non_power_two_slot_width_exact_parity(self):
+        self.compare([(None, ((3, 4),), 7), (None, ((3, 4),), 10)], ("slot_reservoir",))
+
+    def compare(self, cases, modes=("periodic", "slot_reservoir")):
+        from CelerisLab.cuda.context import CudaContext
+        expected_source = Path(__file__).resolve().parents[3] / "src/CelerisLab/lbm/refinement_stepper.py"
+        self.assertEqual(Path(refinement_stepper.__file__).resolve(), expected_source)
+        with CudaContext(0) as context:
+            for mode in modes:
+                for speed, layouts, width in cases:
+                    cfg = config(mode)
+                    if mode == "slot_reservoir":
+                        cfg = replace(cfg, wall_slot_width=width)
+                    p = FixedPatch2D.from_config(cfg, BOUNDS)
+                    arrays = state(NY, NX, 1), state(*p.fine_shape, 2)
+                    reference = None
+                    for base_block in (None, *layouts):
+                        label = f"{mode} width={width} speed={speed} base={base_block}"
+                        snapshots = []
+                        with FixedPatchStepper(cfg, BOUNDS, context=context, coarse_boundary=mode,
+                                               fine_block=FINE_BLOCK, base_block=base_block,
+                                               fine_objects=() if speed is None else [circle(actuation(speed))]) as backend:
+                            backend.initialize(*arrays)
+                            for step in (0, 1, 2, 8):
+                                if step == 1 and base_block is not None:
+                                    # Blocked staging dirties invalid input slots; its output
+                                    # must poison them. Dense parent does not stage there.
+                                    coarse = backend._levels[0]
+                                    spare = coarse.layout.to_blocked(np.full(coarse.shape, 42, np.float32))
+                                    cuda.memcpy_htod(coarse.spare, spare)
+                                backend.step(step - backend.step_count)
+                                public, owner = backend.download(), raw(backend)
+                                self.assertEqual(backend.step_count, step)
+                                self.assertTrue(np.isnan(owner[0][:, 6:18, 8:24]).all(), label)
+                                if mode != "periodic":
+                                    self.assertTrue(np.isnan(owner[0][:, (0, -1), :]).all(), label)
+                                snapshots.append(public + owner)
+                            backend.initialize(*arrays)
+                            self.assertEqual(backend.step_count, 0)
+                            for actual, expected in zip(backend.download() + raw(backend), snapshots[0]):
+                                exact(self, actual, expected, label + " reset")
+                        if reference is None:
+                            reference = snapshots
+                        else:
+                            for step, actual, expected in zip((0, 1, 2, 8), snapshots, reference):
+                                for index, (a, b) in enumerate(zip(actual, expected)):
+                                    exact(self, a, b, f"{label} step={step} array={index}")
+                            print(f"PASS {label}: full finite bits/NaN masks at 0,1,2,8 + reset", flush=True)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/validation/multigrid/test_refinement_blocked_fine.py

@@ -128,8 +128,8 @@ class BlockedFineCPU(unittest.TestCase):
             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
+            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)
@@ -187,7 +187,7 @@ class BlockedFineCPU(unittest.TestCase):
             backend, calls = self.mocked_backend([circle(spec)])
             c, f = backend._levels
             coarse, cur, spare = c.current, f.current, f.spare
-            flags, kind = backend._fine_flags, backend._kind
+            flags, kind = backend._fine_flags, backend._levels[1].kind
             velocity = backend._link_tables.get("surface_velocity")
             self.assertEqual(velocity is None, spec is None)
             backend.step(1)

tests/validation/multigrid/test_refinement_followup.py

@@ -111,7 +111,7 @@ class FaultGuardsCPU(unittest.TestCase):
         backend._lock = RLock()
         backend._status = "ready"
         backend._stream = Mock()
-        backend._levels = [Mock(shape=(9, 1, 1)), Mock(shape=(9, 1, 1))]
+        backend._levels = [Mock(shape=(9, 1, 1), layout=None), Mock(shape=(9, 1, 1), layout=None)]
         backend._allocations = [Mock(), Mock()]
         return backend