ClaudeandClaude Opus 5.5
7281d91fd6
close() now clears _module (initialized to None) inside the active context together with the function references; lifecycle test asserts it is retained on failed activation and cleared after close. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
202 lines
8.7 KiB
Python
202 lines
8.7 KiB
Python
# 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
|