287 lines
14 KiB
Python
287 lines
14 KiB
Python
# CelerisLab/lbm/block_stepper.py
|
|
"""Opt-in periodic backend for one BlockLayout2D level, optional static circles.
|
|
|
|
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.
|
|
Optional ``objects`` are stationary passive circles (static linear BFL via
|
|
receiver-local staging: halo refresh -> cut links -> pull, no copy between).
|
|
No force/actuation readout, motion, walls, refinement, checkpoint or runtime
|
|
parameter mutation; wall/slot/inlet fields in cfg are ignored.
|
|
"""
|
|
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, blocked_link_tables
|
|
from .descriptors import D2Q9_EX, D2Q9_EY, FLUID, SOLID
|
|
from .population_storage import _StoredPopulations, _encode_physical, _decode_stored
|
|
|
|
MAX_GRID_DIM = 65535 # CUDA gridDim.y/z limit (rows, blocks)
|
|
|
|
|
|
def _static_circles(layout, objects):
|
|
"""Host tables for stationary circles; rejects before any CUDA work.
|
|
|
|
Every periodic neighbor of every fluid core is inspected, including domain
|
|
edges that CircleGeometry does not enumerate, so clipped or incomplete
|
|
geometry cannot pass. Each physical link has one receiver-owning block.
|
|
"""
|
|
from ..body.geometry.circle import CircleGeometry
|
|
|
|
ny, nx = layout.shape
|
|
flags = np.full(ny * nx, FLUID, np.uint16)
|
|
for obj in objects:
|
|
if (type(obj.geometry) is not CircleGeometry or obj.is_sensor or obj.is_force_region
|
|
or obj.surface_actuation is not None
|
|
or any(getattr(obj.state, name) != 0 for name in ("vx", "vy", "omega"))
|
|
or any(vars(obj.control).values())):
|
|
raise ValueError("Block bodies must be stationary passive circles")
|
|
cx, cy = obj.geometry.center
|
|
radius = obj.geometry._radius
|
|
if not np.isfinite([cx, cy, radius]).all() or radius <= 0:
|
|
raise ValueError("Block circle center/radius must be finite with radius > 0")
|
|
mask = obj.geometry.build_flag_mask(nx, ny)
|
|
solid = (mask & SOLID) != 0
|
|
if not solid.any() or np.any(flags[solid] & SOLID):
|
|
raise ValueError("Block circles must be resolved and disjoint")
|
|
flags[solid] = mask[solid]
|
|
solid = ((flags & SOLID) != 0).reshape(ny, nx)
|
|
if solid[[0, -1]].any() or solid[:, [0, -1]].any():
|
|
raise ValueError("Block circles must not touch the domain edge")
|
|
links = [link for obj in objects
|
|
for link in obj.geometry.build_cut_links(nx, ny, domain_flags=flags)]
|
|
def source(r, q):
|
|
y, x = divmod(r, nx)
|
|
return (y + D2Q9_EY[q]) % ny, (x + D2Q9_EX[q]) % nx
|
|
receivers = {(int(link.fluid_idx), int(link.dir)) for link in links}
|
|
required = {(int(r), q) for r in np.flatnonzero(~solid.ravel())
|
|
for q in range(1, NQ) if solid[source(r, q)]}
|
|
if len(receivers) != len(links) or receivers != required:
|
|
raise ValueError("Cut links must cover every fluid pull from a solid exactly once")
|
|
tables, _ = blocked_link_tables(
|
|
layout, flags, np.array([link.fluid_idx for link in links], np.uint32),
|
|
np.array([link.dir for link in links], np.uint8),
|
|
np.array([link.q for link in links], np.float32),
|
|
np.array([link.fallback for link in links], np.uint8))
|
|
return dict(solid=solid, **tables)
|
|
|
|
|
|
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, objects=()):
|
|
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.les_enabled
|
|
or (checked.ny, checked.nx) != layout.shape):
|
|
raise ValueError("PeriodicBlockStepper requires matching D2Q9/SRT/FP32/"
|
|
"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")
|
|
objects = tuple(objects) # an empty iterable takes the body-free path
|
|
self.boundary = _static_circles(layout, objects) if objects else None
|
|
links = 0 if self.boundary is None else len(self.boundary["cell"])
|
|
if links >= 2**31:
|
|
raise ValueError("cut links exceed launch limits")
|
|
self._ddf_shifting = bool(checked.ddf_shifting)
|
|
self.layout = layout
|
|
self.ctx = context
|
|
self._links = links
|
|
self._link_grid = ((links + tpb - 1) // tpb, 1)
|
|
# 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 = self._curved = None
|
|
self._tables = {}
|
|
with self.ctx.activate():
|
|
try:
|
|
free, _ = cuda.mem_get_info()
|
|
tables = sum(a.nbytes for k, a in (self.boundary or {}).items() if k != "solid")
|
|
if 2 * self._bytes + layout.halo_dst.nbytes + layout.halo_src.nbytes + tables > 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")
|
|
if self.boundary is not None:
|
|
self._tables = {k: self._upload(a) for k, a in self.boundary.items() if k != "solid"}
|
|
self._curved = module.get_function("BlockCurvedBoundary")
|
|
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.float64 if self._ddf_shifting else np.float32,
|
|
order="C", copy=True)
|
|
if dense.shape != (NQ, *self.layout.shape):
|
|
raise ValueError("Expected a Q-major (9, ny, nx) array")
|
|
fluid = np.ones(self.layout.shape, bool) if self.boundary is None else ~self.boundary["solid"]
|
|
if not np.isfinite(dense[:, fluid]).all() or (dense[:, fluid].sum(axis=0) <= 0).any():
|
|
raise ValueError("Fluid populations must be finite with positive density")
|
|
dense[:, ~fluid] = np.nan # Solid cores are deliberately invalid.
|
|
encoded = _encode_physical(dense, self._ddf_shifting)
|
|
values = encoded[:, fluid]
|
|
if (not np.isfinite(values).all()
|
|
or (values.sum(axis=0, dtype=np.float64) + int(self._ddf_shifting) <= 0).any()):
|
|
raise ValueError("Stored fluid populations must be finite with positive density")
|
|
blocked = self.layout.to_blocked(encoded)
|
|
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)
|
|
if self._curved is not None:
|
|
t = self._tables
|
|
self._curved(self._current, t["block"], t["cell"], t["dir"], t["q"],
|
|
t["fallback"], np.uint32(self._links), np.uintp(0), # no actuation
|
|
np.uintp(0), np.uintp(0), np.uintp(0), # no observation
|
|
grid=self._link_grid, **kw)
|
|
self._step(self._current, self._spare, np.uint32(self.layout.halo),
|
|
self._tables.get("flags", np.uintp(0)), np.uintp(0), # no kind
|
|
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 physical core populations (FP64 when shifted); no host cache."""
|
|
snapshot = self._download_raw()
|
|
return _decode_stored(snapshot.values[0], snapshot.shifted)
|
|
|
|
def _download_raw(self):
|
|
"""Return one tagged dense FP32 population array without conversion."""
|
|
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 _StoredPopulations((self.layout.to_dense(blocked),), self._ddf_shifting)
|
|
|
|
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._curved = self._module = None
|
|
self._tables = {}
|
|
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
|