Add body-free periodic GPU consumer for BlockLayout2D

blocks.cu: separate compilation entry with NX/NY = block window (halo
included), blockIdx.z selects the block; BlockHaloCopy refreshes halos
from owner cores, BlockStep pulls/collides core cells with the existing
compute_rho_u/compute_feq/collide_srt_no_force. Existing kernels and
refinement numerics unchanged.
block_stepper.py: PeriodicBlockStepper (D2Q9/SRT/FP32 double buffers,
doubly periodic layout); launch-limit and free-memory checks before any
allocation; uploads only halo pairs; spare buffer starts NaN; lifecycle
mirrors FixedPatchStepper.
Tests: CPU contract doubles (rejections before CUDA/allocation, initialize
guards, failure/close/retry) and one GPU case: 32x24, 8 steps, 1x1/2x2/4x2
splits bit-identical, FP64 oracle max abs <= 3e-6, drift reported.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Claude
2026-10-15 02:49:09 +00:00
co-authored by Claude Opus 5.5
parent d27d344f43
commit 4029a0b27e
3 changed files with 459 additions and 0 deletions
@@ -0,0 +1,209 @@
"""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._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()