CelerisLab/tests/validation/multigrid/test_refinement_blocked_base.py

"""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.

Dense refinement storage was removed: every level is blocked and level_blocks
is required (TileContractCPU). The retired dense/blocked route parity gates
(BlockedBaseGPU exact parity, DenseBlockedRouteParity incl. MRT and two
bodies, the route capability matrix) passed and are recorded at 1573942 and
2469af7.
"""
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
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
from CelerisLab.lbm.refinement_tables import blocked_coarse_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_get_info", return_value=(1 << 40, 1 << 40)))
        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 blocks in ([(12, 8), None], [(7, 8), FINE_BLOCK], [(12,), FINE_BLOCK]):
            with self.subTest(blocks=blocks), self.assertRaises(ValueError):
                FixedPatchStepper(config("periodic"), context=context, refinements=[BOUNDS],
                                  level_blocks=blocks)
        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_coarse_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_coarse_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_coarse_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"), context=m["context"],
                                        coarse_boundary="slot_reservoir", refinements=[BOUNDS],
                                        level_blocks=[(3, 4), FINE_BLOCK],
                                        level_objects=[(), (circle(actuation(.002)),)])
            try:
                backend.initialize(state(NY, NX, 1), state(*backend.hierarchy.pairs[0].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)
                self.assertEqual([call.kwargs["block_root_bc"] for call in
                                  compiler.generate_config.call_args_list], [1, 0])
                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"), context=m["context"], refinements=[BOUNDS],
                                      level_blocks=[(3, 4), FINE_BLOCK])
                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"), context=m["context"], refinements=[BOUNDS],
                                        level_blocks=[(3, 4), FINE_BLOCK])
            backend.initialize(state(NY, NX, 1), state(*backend.hierarchy.pairs[0].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 TileContractCPU(unittest.TestCase):
    """level_blocks is required: one positive (by, bx) tile per level, checked before CUDA."""

    def test_level_blocks_required_and_validated(self):
        context = Mock(activate=Mock(side_effect=AssertionError("CUDA reached")))
        cfg = config("periodic")
        for label, kw in (("omitted", {}),
                          ("None list", dict(level_blocks=None)),
                          ("too few levels", dict(level_blocks=[(12, 8)])),
                          ("too many levels", dict(level_blocks=[(12, 8), FINE_BLOCK, FINE_BLOCK])),
                          ("None root", dict(level_blocks=[None, FINE_BLOCK])),
                          ("None fine", dict(level_blocks=[(12, 8), None])),
                          ("wrong arity", dict(level_blocks=[(12,), FINE_BLOCK])),
                          ("non-integer", dict(level_blocks=[(12, 8.0), FINE_BLOCK])),
                          ("bool", dict(level_blocks=[(True, 8), FINE_BLOCK])),
                          ("non-positive", dict(level_blocks=[(0, 8), FINE_BLOCK])),
                          ("non-tuple", dict(level_blocks=[12, FINE_BLOCK])),
                          ("not dividing", dict(level_blocks=[(7, 8), FINE_BLOCK]))):
            with self.subTest(label), self.assertRaises(ValueError):
                FixedPatchStepper(cfg, context=context, refinements=[BOUNDS], **kw)
        with self.assertRaisesRegex(ValueError, "whole-level tile"):
            FixedPatchStepper(cfg, context=context, refinements=[BOUNDS])
        context.activate.assert_not_called()
        # The documented whole-level tiles are always valid (construction reaches CUDA).
        whole = [(NY, NX), FixedPatch2D.from_config(cfg, BOUNDS).fine_shape]
        with self.assertRaisesRegex(AssertionError, "CUDA reached"):
            FixedPatchStepper(cfg, context=context, refinements=[BOUNDS], level_blocks=whole)


if __name__ == "__main__":
    unittest.main()