CelerisLab/tests/validation/multigrid/test_block_stepper.py

"""Periodic body-free block backend: CPU contract doubles plus one short GPU case.

Frozen GPU gates (32x24, 8 steps, positive low-Mach non-equilibrium state):
finite cores; bit-identical 1x1, 2x2 and 4x2 splits; max |G - FP64 oracle|
<= 3e-6. Mass/momentum drift is reported, not gated.
"""
from contextlib import nullcontext
from pathlib import Path
import resource
from threading import RLock
import unittest
from unittest.mock import MagicMock, Mock, patch

import numpy as np
import pycuda.driver as cuda

from CelerisLab.config import LBMConfig
from CelerisLab.cuda import compiler_v2 as compiler
from CelerisLab.lbm import block_stepper as implementation
from CelerisLab.lbm.blocks import BlockLayout2D
from CelerisLab.lbm.block_stepper import PeriodicBlockStepper
from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY

ROOT = Path(__file__).resolve().parents[3]
NX, NY, NU, STEPS, ORACLE_ATOL = 32, 24, .1, 8, 3e-6
SPLITS = ((24, 32), (12, 16), (12, 8))  # 1x1, 2x2, 4x2 blocks (by, bx)
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(**kwargs):
    return LBMConfig(**dict(nx=NX, ny=NY, viscosity=NU, **kwargs))


def feq(rho, ux, uy):
    cu = CX[:, None, None] * ux + CY[:, None, None] * uy
    return W[:, None, None] * rho * (1 + 3 * cu + 4.5 * cu * cu - 1.5 * (ux * ux + uy * uy))


def initial_state():
    """Positive low-Mach equilibrium plus a seeded non-equilibrium part."""
    y, x = np.mgrid[0:NY, 0:NX]
    rho = 1 + 1e-3 * np.cos(2 * np.pi * x / NX)
    ux = .02 * np.sin(2 * np.pi * y / NY)
    uy = .01 * np.cos(2 * np.pi * x / NX)
    noise = np.random.default_rng(2026).uniform(-1, 1, (9, NY, NX)) * 1e-3 * W[:, None, None]
    return (feq(rho, ux, uy) + noise).astype(np.float32)


def oracle(g, steps):
    """FP64 periodic pull + SRT, from the same FP32-rounded G_0."""
    g = g.astype(np.float64)
    omega = 1 / (3 * NU + .5)
    for _ in range(steps):
        f = np.stack([np.roll(g[q], (D2Q9_EY[q], D2Q9_EX[q]), axis=(0, 1)) for q in range(9)])
        rho = f.sum(axis=0)
        ux, uy = (CX[:, None, None] * f).sum(axis=0) / rho, (CY[:, None, None] * f).sum(axis=0) / rho
        g = f + omega * (feq(rho, ux, uy) - f)
    return g


def totals(g):
    g = g.astype(np.float64)
    return g.sum(), (CX[:, None, None] * g).sum(), (CY[:, None, None] * g).sum()


class CpuContracts(unittest.TestCase):
    """No device: driver doubles prove rejection before allocation and lifecycle."""

    def forbidden_context(self):
        return Mock(activate=Mock(side_effect=AssertionError("CUDA reached")), sm_arch="sm_70")

    def test_unsupported_inputs_rejected_before_cuda(self):
        layout = BlockLayout2D.build(NY, NX, 12, 16)
        cases = ((config(), BlockLayout2D.build(NY, NX, 12, 16, periodic=(True, False))),
                 (config(), BlockLayout2D.build(NY, NX * 2, 12, 16)),
                 (config(collision="MRT"), layout),
                 (config(), "layout"))
        for cfg, lay in cases:
            with self.subTest(cfg=cfg.collision, layout=getattr(lay, "periodic", lay)):
                with self.assertRaises(ValueError):
                    PeriodicBlockStepper(cfg, lay, context=self.forbidden_context())
        with patch.object(implementation, "MAX_GRID_DIM", 3):
            with self.assertRaisesRegex(ValueError, "launch limits"):
                PeriodicBlockStepper(config(), layout, context=self.forbidden_context())

    def test_memory_and_param_layout_rejected_before_allocation(self):
        layout = BlockLayout2D.build(NY, NX, 12, 16)
        context = Mock(activate=nullcontext, sm_arch="sm_70")
        with patch.object(cuda, "mem_get_info", return_value=(1024, 1 << 30)), \
                patch.object(cuda, "Stream") as stream, patch.object(cuda, "mem_alloc") as alloc:
            with self.assertRaises(MemoryError):
                PeriodicBlockStepper(config(), layout, context=context)
            stream.assert_not_called()
            alloc.assert_not_called()
        module = Mock()
        module.get_global.return_value = (Mock(), 16)
        with patch.object(cuda, "mem_get_info", return_value=(1 << 30, 1 << 30)), \
                patch.object(cuda, "Stream"), patch.object(compiler, "generate_config"), \
                patch.object(compiler, "compile_kernel") as compile_kernel, \
                patch.object(compiler, "cleanup_build_dir"), \
                patch.object(compiler, "load_module", return_value=module), \
                patch.object(cuda, "mem_alloc") as alloc, patch.object(cuda, "memcpy_htod") as upload:
            with self.assertRaisesRegex(RuntimeError, "layout mismatch"):
                PeriodicBlockStepper(config(), layout, context=context)
            self.assertEqual(compile_kernel.call_args.kwargs["source"], "blocks.cu")
            alloc.assert_not_called()
            upload.assert_not_called()

    def fake_backend(self):
        backend = PeriodicBlockStepper.__new__(PeriodicBlockStepper)
        backend.layout = BlockLayout2D.build(NY, NX, 12, 16)
        backend.ctx = Mock(activate=nullcontext)
        backend._lock = RLock()
        backend._status = "ready"
        backend._stream = Mock()
        backend._copy = backend._step = Mock()
        backend._dst, backend._src = Mock(), Mock()
        backend._block, backend._grid, backend._copy_grid = (32, 1, 1), (1, 1, 1), (1, 1)
        backend._allocations = [Mock(), Mock()]
        backend._current, backend._spare = backend._allocations
        return backend

    def test_initialize_rejects_bad_state_before_upload(self):
        backend = self.fake_backend()
        good = initial_state()
        bad = (good[:, :-1], np.where(np.arange(NX) == 3, np.nan, good), -good)
        with patch.object(cuda, "memcpy_htod") as upload:
            for g in bad:
                with self.assertRaises(ValueError):
                    backend.initialize(g)
            upload.assert_not_called()
            backend.initialize(good)
            self.assertEqual(upload.call_count, 2)
            self.assertTrue(np.isnan(upload.call_args_list[1].args[1]).all())
        backend.close()

    def test_failed_launch_and_download_make_backend_unusable(self):
        for failure in ("step", "download"):
            backend = self.fake_backend()
            if failure == "step":
                backend._step = Mock(side_effect=RuntimeError("launch failed"))
                with self.assertRaisesRegex(RuntimeError, "launch failed"):
                    backend.step()
            else:
                with patch.object(cuda, "memcpy_dtoh", side_effect=RuntimeError("read failed")):
                    with self.assertRaisesRegex(RuntimeError, "read failed"):
                        backend.download()
            for operation in (backend.download, backend.step):
                with self.assertRaisesRegex(RuntimeError, "failed"):
                    operation()
            backend.close()
            for allocation in backend_allocations(backend):
                allocation.free.assert_called_once()

    def test_activation_failure_retains_resources_for_retry(self):
        backend = self.fake_backend()
        allocations, stream = backend._allocations[:], backend._stream
        unavailable = MagicMock()
        unavailable.__enter__.side_effect = RuntimeError("activation failed")
        backend.ctx.activate = Mock(side_effect=[unavailable, nullcontext()])
        with self.assertRaisesRegex(RuntimeError, "activation failed"):
            backend.close()
        self.assertEqual(backend._status, "failed")
        self.assertEqual(backend._allocations, allocations)
        self.assertIs(backend._stream, stream)
        backend.close()
        self.assertEqual(backend._status, "closed")
        for allocation in allocations:
            allocation.free.assert_called_once()


def backend_allocations(backend):
    return (backend._current, backend._spare)


class PeriodicBlockGPU(unittest.TestCase):
    def setUp(self):
        self.assertEqual(Path(implementation.__file__).resolve(),
                         ROOT / "src/CelerisLab/lbm/block_stepper.py")

    def tearDown(self):
        rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024
        print(f"host_peak_rss={rss:.3f} MiB", flush=True)

    def test_splits_bit_identical_and_fp64_oracle(self):
        from CelerisLab.cuda.context import CudaContext
        g0 = initial_state()
        results = []
        with CudaContext(0) as context:
            for by, bx in SPLITS:
                layout = BlockLayout2D.build(NY, NX, by, bx)
                with PeriodicBlockStepper(config(), layout, context=context) as backend:
                    backend.initialize(g0)
                    backend.step(STEPS)
                    self.assertEqual(backend.step_count, STEPS)
                    results.append(backend.download())
        for (by, bx), g in zip(SPLITS, results):
            self.assertTrue(np.isfinite(g).all(), (by, bx))
            np.testing.assert_array_equal(g, results[0], err_msg=f"split {(by, bx)}")
        error = np.abs(results[0].astype(np.float64) - oracle(g0, STEPS)).max()
        before, after = totals(g0), totals(results[0])
        print(f"max_abs_oracle_error={error:.3e} gate={ORACLE_ATOL:.0e}", flush=True)
        print("drift mass={:.3e} jx={:.3e} jy={:.3e}".format(
            *(a - b for a, b in zip(after, before))), flush=True)
        self.assertLessEqual(error, ORACLE_ATOL)


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