Merge claude/block-gpu-periodic for bounded integration validation

src/CelerisLab/lbm/block_stepper.py

@@ -0,0 +1,201 @@
+# CelerisLab/lbm/block_stepper.py
+"""Opt-in body-free periodic backend for one BlockLayout2D level.
+
+D2Q9/SRT/FP32 unforced double buffers. Each step refreshes every halo from
+its owner core, then pulls/collides core cells into the spare buffer, then
+swaps. Only halo pair tables are uploaded; dense ownership stays on the host.
+The caller owns the CudaContext and must close this backend before it.
+No bodies, walls, refinement, checkpoint or runtime parameter mutation;
+wall/slot/inlet fields in cfg are ignored (fully periodic level).
+"""
+from dataclasses import replace
+from numbers import Integral
+import struct
+from threading import RLock
+
+import numpy as np
+import pycuda.driver as cuda
+
+from ..cuda import compiler_v2 as compiler
+from .blocks import NQ, BlockLayout2D
+
+MAX_GRID_DIM = 65535  # CUDA gridDim.y/z limit (rows, blocks)
+
+
+class PeriodicBlockStepper:
+    """G->G stepping of Q-major (9, ny, nx) postcollision populations.
+
+    On a device/launch failure the instance becomes unusable: close and
+    recreate it rather than assuming a partially advanced buffer rolled back.
+    """
+    def __init__(self, cfg, layout, *, context):
+        if not isinstance(layout, BlockLayout2D) or layout.periodic != (True, True):
+            raise ValueError("PeriodicBlockStepper requires a doubly periodic BlockLayout2D")
+        checked = replace(cfg)
+        checked.validate()
+        if (not checked.is_d2q9 or checked.collision != "SRT"
+                or checked.streaming != "double_buffer" or checked.store_precision != "FP32"
+                or checked.ddf_shifting or checked.les_enabled
+                or (checked.ny, checked.nx) != layout.shape):
+            raise ValueError("PeriodicBlockStepper requires matching D2Q9/SRT/FP32/"
+                             "unshifted/double_buffer config without LES")
+        ly, lx = layout.local_shape
+        pairs = len(layout.halo_dst)
+        tpb = checked.threads_per_block
+        if layout.nblocks > MAX_GRID_DIM or ly > MAX_GRID_DIM or not 0 < pairs < 2**31:
+            raise ValueError("block count, block rows or halo pairs exceed launch limits")
+        self.layout = layout
+        self.ctx = context
+        # Wall/slot fields are not applied; neutralize them like the fine patch level.
+        self._local_cfg = replace(checked, nx=lx, ny=ly, y_wall_bc="bounce_back", top_wall_bc="velocity",
+                                  wall_slot_width=0, wall_normal_velocity=0.0, velocity=0.0)
+        self._block = (tpb, 1, 1)
+        self._grid = ((lx + tpb - 1) // tpb, ly, layout.nblocks)
+        self._copy_grid = ((pairs + tpb - 1) // tpb, 1)
+        self._bytes = int(np.prod(layout.blocked_shape)) * 4
+        self._lock = RLock()
+        self._status = "new"
+        self._step_count = 0
+        self._allocations = []
+        self._stream = self._module = self._copy = self._step = None
+        with self.ctx.activate():
+            try:
+                free, _ = cuda.mem_get_info()
+                if 2 * self._bytes + layout.halo_dst.nbytes + layout.halo_src.nbytes > free:
+                    raise MemoryError("blocked buffers exceed free device memory")
+                self._stream = cuda.Stream()
+                config = self._local_cfg
+                config.validate()
+                if not np.isclose(config.omega, 1 / (3 * checked.viscosity + .5), rtol=1e-14, atol=0):
+                    raise RuntimeError("block omega/tau mismatch")
+                build = compiler.generate_config(config)
+                try:
+                    ptx = compiler.compile_kernel(self.ctx.sm_arch, source="blocks.cu", build_dir=build)
+                    module = compiler.load_module(ptx)
+                finally:
+                    compiler.cleanup_build_dir(build)
+                ptr, size = module.get_global("d_params")
+                if size != struct.calcsize("=ffffI"):
+                    raise RuntimeError("LBMParams host/device layout mismatch")
+                cuda.memcpy_htod(ptr, struct.pack("=ffffI", config.omega, 0, 0, 0, 0))
+                self._module = module
+                self._current = self._allocate(self._bytes)
+                self._spare = self._allocate(self._bytes)
+                self._dst = self._upload(layout.halo_dst)
+                self._src = self._upload(layout.halo_src)
+                self._copy = module.get_function("BlockHaloCopy")
+                self._step = module.get_function("BlockStep")
+            except BaseException:
+                try:
+                    self.close()
+                except Exception:
+                    pass  # Preserve the original construction failure.
+                raise
+
+    def _allocate(self, size):
+        pointer = cuda.mem_alloc(size)
+        self._allocations.append(pointer)
+        return pointer
+
+    def _upload(self, array):
+        array = np.ascontiguousarray(array)
+        pointer = self._allocate(array.nbytes)
+        cuda.memcpy_htod(pointer, array)
+        return pointer
+
+    def _require(self, *states):
+        if self._status not in states:
+            raise RuntimeError(f"Block stepper is {self._status}; expected {' or '.join(states)}")
+
+    @property
+    def step_count(self):
+        return self._step_count
+
+    def initialize(self, g):
+        """Upload dense G_0; the spare buffer starts NaN so stale halos cannot hide."""
+        with self._lock, self.ctx.activate():
+            self._require("new", "ready")
+            dense = np.array(g, dtype=np.float32, order="C", copy=True)
+            if dense.shape != (NQ, *self.layout.shape):
+                raise ValueError("Expected a Q-major (9, ny, nx) array")
+            if not np.isfinite(dense).all() or (dense.sum(axis=0) <= 0).any():
+                raise ValueError("Populations must be finite with positive density")
+            blocked = self.layout.to_blocked(dense)
+            spare = np.full_like(blocked, np.nan)
+            try:
+                cuda.memcpy_htod(self._current, blocked)
+                cuda.memcpy_htod(self._spare, spare)
+                self._step_count = 0
+                self._status = "ready"
+            except BaseException:
+                self._status = "failed"
+                raise
+
+    def step(self, n=1):
+        """Advance n steps: halo refresh, then core pull/collide, per step."""
+        if isinstance(n, bool) or not isinstance(n, Integral) or n < 0:
+            raise ValueError("n must be a nonnegative integer")
+        with self._lock, self.ctx.activate():
+            self._require("ready")
+            kw = dict(block=self._block, stream=self._stream)
+            try:
+                self._status = "advancing"
+                for _ in range(n):
+                    self._copy(self._current, self._dst, self._src,
+                               np.uint32(len(self.layout.halo_dst)), grid=self._copy_grid, **kw)
+                    self._step(self._current, self._spare, np.uint32(self.layout.halo),
+                               grid=self._grid, **kw)
+                    self._stream.synchronize()
+                    self._current, self._spare = self._spare, self._current
+                    self._step_count += 1
+                self._status = "ready"
+            except BaseException:
+                self._status = "failed"
+                raise
+
+    def download(self):
+        """Return dense (9, ny, nx) FP32 core populations; no host cache."""
+        with self._lock, self.ctx.activate():
+            self._require("ready")
+            blocked = np.empty(self.layout.blocked_shape, np.float32)
+            try:
+                cuda.memcpy_dtoh(blocked, self._current)
+            except BaseException:
+                self._status = "failed"
+                raise
+            return self.layout.to_dense(blocked)
+
+    def close(self):
+        """Release owned resources; retain them for retry if context activation fails."""
+        with self._lock:
+            if self._status == "closed":
+                return
+            self._status = "failed"
+            with self.ctx.activate():
+                self._status = "closed"
+                error = None
+                try:
+                    if self._stream is not None:
+                        self._stream.synchronize()
+                except Exception as exc:
+                    error = exc
+                self._copy = self._step = self._module = None
+                for pointer in self._allocations:
+                    try:
+                        pointer.free()
+                    except Exception as exc:
+                        error = error or exc
+                self._allocations.clear()
+                self._stream = None
+                if error is not None:
+                    raise error
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc):
+        try:
+            self.close()
+        except Exception:
+            if exc[0] is None:
+                raise

src/CelerisLab/lbm/kernels/blocks.cu

@@ -0,0 +1,49 @@
+// CelerisLab/lbm/kernels/blocks.cu
+// One periodic level of fixed-size blocks (lbm/blocks.py layout). Separate
+// compilation entry: NX, NY are the block window including halo, so index_f
+// and linear_index address one block; blockIdx.z selects the block, whose
+// populations start at b*NQ*TOTAL_CELLS (block-contiguous, Q-major inside).
+#include <stdint.h>
+#include "config.h"
+#include "core/precision.cuh"
+#include "core/descriptors.cuh"
+#include "core/layout.cuh"
+#include "core/params.cuh"
+#include "operators/macro.cuh"
+#include "operators/collision_srt.cuh"
+#include "streaming/pull_double_buffer.cuh"
+
+#if DIM != 2 || NQ != 9 || COLLISION_MODEL != 0 || STORE_PRECISION != 0 || USE_DDF_SHIFTING || USE_LES || STREAMING_MODEL != 0
+#error "Blocks require unshifted FP32 D2Q9 SRT double buffers without LES"
+#endif
+
+// Halo refresh: block cell ids b*TOTAL_CELLS+c, sources are always core cells.
+extern "C" __global__ void BlockHaloCopy(
+    fpxx* f, const uint64_t* dst, const uint64_t* src, unsigned int count)
+{
+    unsigned int row = blockIdx.x * blockDim.x + threadIdx.x;
+    if (row >= count) return;
+    const uint64_t d = dst[row], s = src[row];
+    fpxx* to = f + (d / TOTAL_CELLS) * NQ * TOTAL_CELLS;
+    const fpxx* from = f + (s / TOTAL_CELLS) * NQ * TOTAL_CELLS;
+    for (unsigned int q = 0; q < NQ; ++q)
+        to[index_f(d % TOTAL_CELLS, q)] = from[index_f(s % TOTAL_CELLS, q)];
+}
+
+// Pull/collide/store core cells only; halo cells of out are left untouched.
+extern "C" __global__ void BlockStep(const fpxx* in, fpxx* out, unsigned int halo) {
+    unsigned int x, y; unsigned long k;
+    index_from_thread(x, y, k);
+    if (x < halo || y < halo || x >= NX - halo || y >= NY - halo) return;
+    const unsigned long base = (unsigned long)blockIdx.z * NQ * TOTAL_CELLS;
+    in += base;
+    out += base;
+    unsigned long j[NQ];
+    compute_neighbors(x, y, j);  // Core neighbors stay inside the window: no wrap.
+    float f[NQ], rho, ux, uy, feq[NQ];
+    stream_pull_load(k, f, in, j);
+    compute_rho_u(f, rho, ux, uy);
+    compute_feq(rho, ux, uy, feq);
+    collide_srt_no_force(f, feq, d_params.omega);
+    stream_pull_store(k, f, out);
+}

tests/validation/multigrid/test_block_stepper.py

@@ -0,0 +1,216 @@
+"""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._step_count = 0
+        backend._stream = Mock()
+        backend._module = 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, module = backend._allocations[:], backend._stream, backend._module
+        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)
+        self.assertIs(backend._module, module)
+        backend.close()
+        self.assertEqual(backend._status, "closed")
+        self.assertIsNone(backend._module)
+        self.assertIsNone(backend._copy)
+        self.assertIsNone(backend._step)
+        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()