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

This commit is contained in:
Frank14f
2026-10-15 11:08:38 +08:00
3 changed files with 466 additions and 0 deletions
+201
View File
@@ -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
+49
View File
@@ -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);
}
@@ -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()