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);
+}