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:
co-authored by
Claude Opus 5.5
parent
d27d344f43
commit
4029a0b27e
@@ -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()
|
||||
Reference in New Issue
Block a user