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:
ClaudeandClaude Opus 5.5 committed
1 parent
d27d344f43
commit
4029a0b27e
3 files changed
+459
No files matched your search
@@ -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._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 = 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
|
||||
@@ -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);
|
||||
}
|
||||
Reference in new issue
Block a user