Frank14f/CelerisLab
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add opt-in blocked fine storage to FixedPatchStepper

Browse Source 
fine_block=(by, bx) partitions the ENTIRE fine_shape (incl. the four-cell
patch ghost band) with a non-periodic one-cell-halo BlockLayout2D; body-
free only (fine_objects rejected in blocked mode; dense fine-body and
actuation path unchanged). Per coarse step on the same stream:
BlockExplode into ghost owner cores -> gather -> BlockStep(kind) ->
gather -> BlockStep(kind) -> unchanged PatchCoarseStep reading remapped
owner-core child slots in the original row/child order. Public arrays
keep the dense fine shape; only the private allocation is blocked.
blocks.cu: explicit BlockExplode; BlockStep gains nullable kind
(OUTER skip pull + NaN, INNER pull only, ACTIVE pull + SRT).
PeriodicBlockStepper passes an explicit null kind (raw launch ABI).
refinement.cu/refinement.py and coarse numerics unchanged.
Tests: CPU rejections and table remap (dense (q,cell) and child order
preserved, owner cores, kinds); GPU gate: periodic and slot_reservoir,
8 coarse steps, dense vs 1x1/2x2/4x2 fine blocks bit-identical in
masked downloads and raw coarse/unpacked fine incl. ghost NaN masks.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
ClaudeandClaude Opus 5.5 committed 2026-10-15 03:48:54 +00:00
1 parent 3e44e90ca5
commit dc5850c198
4 files changed
+257 -27
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/CelerisLab/lbm/block_stepper.py
+2 -1
View File
@@ -215,7 +215,8 @@ class PeriodicBlockStepper:
self._curved(self._current, t["block"], t["cell"], t["dir"], t["q"],
t["fallback"], np.uint32(self._links), grid=self._link_grid, **kw)
self._step(self._current, self._spare, np.uint32(self.layout.halo),
self._tables.get("flags", np.uintp(0)), grid=self._grid, **kw)
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
src/CelerisLab/lbm/kernels/blocks.cu
+29 -6
View File
@@ -21,6 +21,9 @@
#error "Blocks require unshifted FP32 D2Q9 SRT double buffers without LES"
#endif
// Fixed-patch cell kinds, matching refinement.py ACTIVE/INNER_GHOST/OUTER_GHOST.
constexpr unsigned char KIND_ACTIVE = 0u, KIND_INNER = 1u, KIND_OUTER = 2u;
// 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)
@@ -34,6 +37,20 @@ extern "C" __global__ void BlockHaloCopy(
to[index_f(d % TOTAL_CELLS, q)] = from[index_f(s % TOTAL_CELLS, q)];
}
// Fixed-patch explosion into blocked fine storage: target = owner block cell
// id of a patch-ghost cell, donor = dense coarse cell (coarse is Q-major dense).
extern "C" __global__ void BlockExplode(
const fpxx* coarse, fpxx* fine, const uint64_t* target, const uint32_t* donor,
unsigned int count, unsigned long coarse_cells)
{
unsigned int row = blockIdx.x * blockDim.x + threadIdx.x;
if (row >= count) return;
const uint64_t t = target[row];
fpxx* to = fine + (t / TOTAL_CELLS) * NQ * TOTAL_CELLS;
for (unsigned int q = 0; q < NQ; ++q)
to[index_f(t % TOTAL_CELLS, q)] = coarse[q * coarse_cells + donor[row]];
}
// Static cut links, receiver-local staging: cell is the receiver's local core
// cell, so the existing BFL writes the solid-source window image that this
// block's next pull reads (no halo copy between this kernel and BlockStep).
@@ -49,23 +66,29 @@ extern "C" __global__ void BlockCurvedBoundary(
// Pull/collide/store core cells only; halo cells of out are left untouched.
// flags (block cell ids, nullable): solid cores store NaN and never collide.
// kind (block cell ids, nullable = all ACTIVE): OUTER skips the pull and
// stores NaN, INNER pulls without collision, ACTIVE pulls and collides.
extern "C" __global__ void BlockStep(const fpxx* in, fpxx* out, unsigned int halo,
const uint16_t* flags) {
const uint16_t* flags, const unsigned char* kind) {
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;
const unsigned long cell = (unsigned long)blockIdx.z * TOTAL_CELLS + k;
const unsigned char cell_kind = kind != nullptr ? kind[cell] : KIND_ACTIVE;
float f[NQ];
if (flags != nullptr && is_solid(flags[(unsigned long)blockIdx.z * TOTAL_CELLS + k])) {
if (cell_kind == KIND_OUTER || (flags != nullptr && is_solid(flags[cell]))) {
for (int q = 0; q < NQ; ++q) f[q] = NAN;
} else {
unsigned long j[NQ];
compute_neighbors(x, y, j); // Core neighbors stay inside the window: no wrap.
float rho, ux, uy, feq[NQ];
stream_pull_load(k, f, in + base, j);
compute_rho_u(f, rho, ux, uy);
compute_feq(rho, ux, uy, feq);
collide_srt_no_force(f, feq, d_params.omega);
if (cell_kind == KIND_ACTIVE) {
float rho, ux, uy, feq[NQ];
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 + base);
}
src/CelerisLab/lbm/refinement_stepper.py
+99 -20
View File
@@ -14,10 +14,28 @@ import numpy as np
import pycuda.driver as cuda
from ..cuda import compiler_v2 as compiler
from .refinement import ACTIVE, HALO, FixedPatch2D, prepare_fine_boundary
from .blocks import BlockLayout2D
from .block_stepper import MAX_GRID_DIM
from .refinement import ACTIVE, HALO, OUTER_GHOST, FixedPatch2D, prepare_fine_boundary
from .descriptors import SOLID
def blocked_fine_tables(patch, layout, order):
"""Remap fixed-patch fine tables to owner-core blocked storage (host only).
Returns (explosion target block cells, coalescence child slots in the same
rows/child order as ``patch.coalesce_fine[order]``, per-block-cell kind with
unused halo entries OUTER). Coarse donors/targets are unchanged.
"""
owner = layout.owner.ravel().astype(np.intp)
cells = patch.fine_shape[0] * patch.fine_shape[1]
q, cell = np.divmod(patch.coalesce_fine[order], np.uint64(cells))
kind = np.full(layout.nblocks * int(np.prod(layout.local_shape)), OUTER_GHOST, np.uint8)
kind[owner] = patch.fine_kind.ravel()
return (owner[patch.ghost_fine.astype(np.intp)].astype(np.uint64),
layout.slot(owner[cell.astype(np.intp)], q), kind)
@dataclass
class _Level:
shape: tuple
@@ -46,8 +64,18 @@ class FixedPatchStepper:
Every ``step()`` publishes only after both levels complete on one stream.
On a device/launch failure the instance becomes unusable: close and recreate
it, rather than pretending the partially advanced fine buffer rolled back.
``fine_block=(by, bx)`` (opt-in, body-free only) stores the fine level in
blocks.cu storage: the ENTIRE fine_shape, including the four-cell patch
ghost band, is partitioned by a non-periodic BlockLayout2D with one-cell
same-level halos. Patch kinds stay per cell (OUTER invalid, INNER stream
only, ACTIVE collide); halos are gathered before each fine hop. Public
arrays keep the dense (9, ny, nx) fine shape; only the allocation differs.
"""
def __init__(self, cfg, bounds, *, context, coarse_boundary="periodic", fine_objects=()):
fine_layout = None # dense fine storage unless fine_block is given
def __init__(self, cfg, bounds, *, context, coarse_boundary="periodic", fine_objects=(),
fine_block=None):
if HALO != 4:
raise RuntimeError("refinement.cu requires a four-cell fine halo")
if coarse_boundary not in ("periodic", "slot_reservoir"):
@@ -61,8 +89,20 @@ class FixedPatchStepper:
"velocity=0 and 2 <= wall_slot_width < nx")
if self.patch.bounds[1] < 4 or self.patch.bounds[3] > cfg.ny - 4:
raise ValueError("slot_reservoir requires y0 >= 4 and y1 <= ny-4")
fine_objects = tuple(fine_objects)
self.fine_layout = None
if fine_block is not None:
if fine_objects:
raise ValueError("fine_block storage is body-free; fine_objects need the dense fine level")
if len(fine_block) != 2:
raise ValueError("fine_block must be (by, bx)")
fy, fx = self.patch.fine_shape
self.fine_layout = BlockLayout2D.build(fy, fx, *fine_block, periodic=(False, False))
if (self.fine_layout.nblocks > MAX_GRID_DIM
or self.fine_layout.local_shape[0] > MAX_GRID_DIM):
raise ValueError("fine blocks exceed launch limits")
self.fine_boundary = prepare_fine_boundary(
cfg, self.patch, tuple(fine_objects), wall_aware=self._wall_aware)
cfg, self.patch, fine_objects, wall_aware=self._wall_aware)
self._fine_leaf = self.patch.fine_kind == ACTIVE
if self.fine_boundary is not None:
self._fine_leaf &= (self.fine_boundary.flags.reshape(self.patch.fine_shape) & SOLID) == 0
@@ -74,6 +114,7 @@ class FixedPatchStepper:
self._levels = []
self._stream = None
self._fine_flags = self._fine_curved = self._curved = None
self._kind = self._halo_copy = self._halo_dst = self._halo_src = None
self._block = (cfg.threads_per_block, 1, 1)
fine_cfg = replace(cfg, nx=self.patch.fine_shape[1], ny=self.patch.fine_shape[0],
viscosity=2 * cfg.viscosity)
@@ -83,8 +124,14 @@ class FixedPatchStepper:
with self.ctx.activate():
try:
self._stream = cuda.Stream()
for config, name, tau in ((replace(cfg), "PatchCoarseStep", self.patch.tau_coarse),
(fine_cfg, "PatchFineStep", self.patch.tau_fine)):
layout = self.fine_layout
fine = ((fine_cfg, "PatchFineStep", "refinement.cu") if layout is None else
(replace(fine_cfg, nx=layout.local_shape[1], ny=layout.local_shape[0],
y_wall_bc="bounce_back", top_wall_bc="velocity", wall_slot_width=0,
wall_normal_velocity=0.0, velocity=0.0), "BlockStep", "blocks.cu"))
for config, name, tau, source in (
(replace(cfg), "PatchCoarseStep", self.patch.tau_coarse, "refinement.cu"),
(*fine[:2], self.patch.tau_fine, fine[2])):
config.validate()
if not np.isclose(config.omega, 1 / tau, rtol=1e-14, atol=0):
raise RuntimeError("level omega/tau mismatch")
@@ -92,7 +139,7 @@ class FixedPatchStepper:
if name == "PatchFineStep" and self.fine_boundary is not None else 0)
build = compiler.generate_config(config, n_objects=n_objects)
try:
ptx = compiler.compile_kernel(self.ctx.sm_arch, source="refinement.cu",
ptx = compiler.compile_kernel(self.ctx.sm_arch, source=source,
build_dir=build)
module = compiler.load_module(ptx)
finally:
@@ -103,17 +150,30 @@ class FixedPatchStepper:
cuda.memcpy_htod(ptr, struct.pack("=ffffI", config.omega, 0, 0, 0, n_objects))
shape = (9, config.ny, config.nx)
size = int(np.prod(shape)) * 4
grid = ((config.nx + self._block[0] - 1) // self._block[0], config.ny, 1)
if name == "BlockStep": # public dense shape, private blocked allocation
shape = (9, *self.patch.fine_shape)
size = int(np.prod(layout.blocked_shape)) * 4
grid = grid[:2] + (layout.nblocks,)
self._levels.append(_Level(shape, module, self._allocate(size),
self._allocate(size), module.get_function(name),
((config.nx + self._block[0] - 1) // self._block[0],
config.ny, 1)))
self._allocate(size), module.get_function(name), grid))
p = self.patch
order = np.argsort(p.coalesce_coarse)
self._ghosts = self._upload(p.ghost_fine)
tables = None if layout is None else blocked_fine_tables(p, layout, order)
self._ghosts = self._upload(p.ghost_fine if layout is None else tables[0])
self._donors = self._upload(p.ghost_coarse)
self._targets = self._upload(p.coalesce_coarse[order])
self._children = self._upload(p.coalesce_fine[order])
self._explode = self._levels[1].module.get_function("PatchExplode")
if layout is None:
self._children = self._upload(p.coalesce_fine[order])
self._explode = self._levels[1].module.get_function("PatchExplode")
else:
self._children = self._upload(tables[1])
self._kind = self._upload(tables[2])
self._halo_dst = self._upload(layout.halo_dst)
self._halo_src = self._upload(layout.halo_src)
module = self._levels[1].module
self._explode = module.get_function("BlockExplode")
self._halo_copy = module.get_function("BlockHaloCopy")
if self.fine_boundary is not None:
from .curved_links import CurvedLinkSoA
boundary = self.fine_boundary
@@ -174,6 +234,8 @@ class FixedPatchStepper:
if not np.isfinite(array[:, mask]).all() or (array[:, mask].sum(axis=0) <= 0).any():
raise ValueError("Leaf populations must be finite with positive density")
self._mask_nonleaves(*arrays)
if self.fine_layout is not None:
arrays[1] = self.fine_layout.to_blocked(arrays[1])
try:
for level, array in zip(self._levels, arrays):
cuda.memcpy_htod(level.current, array)
@@ -205,6 +267,12 @@ class FixedPatchStepper:
cl.link_exchange_gpu if cl.link_exchange_gpu is not None else np.uintp(0),
grid=((cl.count + self._block[0] - 1) // self._block[0], 1), **kw)
def _fine_halo(self, buffer, **kw):
pairs = len(self.fine_layout.halo_dst)
if pairs:
self._halo_copy(buffer, self._halo_dst, self._halo_src, np.uint32(pairs),
grid=((pairs + self._block[0] - 1) // self._block[0], 1), **kw)
def step(self, n=1):
"""Advance G→G: one explosion, two fine collisions, one coarse collision."""
if isinstance(n, bool) or not isinstance(n, Integral) or n < 0:
@@ -221,13 +289,20 @@ class FixedPatchStepper:
np.uint32(len(p.ghost_fine)), np.uint64(np.prod(p.coarse_shape)),
grid=((len(p.ghost_fine) + self._block[0] - 1) // self._block[0], 1),
**kw)
flags = self._fine_flags if self._fine_flags is not None else np.uintp(0)
if self.fine_boundary is not None and self._fine_flags is None:
raise RuntimeError("Fine boundary flags are missing")
self._apply_fine_boundary(f.current, **kw)
f.kernel(f.current, f.spare, flags, grid=f.grid, **kw)
self._apply_fine_boundary(f.spare, **kw)
f.kernel(f.spare, f.current, flags, grid=f.grid, **kw)
if self.fine_layout is not None:
# Gather one-cell block halos before EACH fine hop.
for src, dst in ((f.current, f.spare), (f.spare, f.current)):
self._fine_halo(src, **kw)
f.kernel(src, dst, np.uint32(self.fine_layout.halo), np.uintp(0),
self._kind, grid=f.grid, **kw)
else:
flags = self._fine_flags if self._fine_flags is not None else np.uintp(0)
if self.fine_boundary is not None and self._fine_flags is None:
raise RuntimeError("Fine boundary flags are missing")
self._apply_fine_boundary(f.current, **kw)
f.kernel(f.current, f.spare, flags, grid=f.grid, **kw)
self._apply_fine_boundary(f.spare, **kw)
f.kernel(f.spare, f.current, flags, grid=f.grid, **kw)
c.kernel(c.current, c.spare, f.current, self._targets, self._children,
np.uint32(len(p.coalesce_coarse)),
*(np.uint32(v) for v in p.bounds), np.uint32(self._wall_aware),
@@ -245,12 +320,16 @@ class FixedPatchStepper:
with self._lock, self.ctx.activate():
self._require("ready")
arrays = [np.empty(level.shape, np.float32) for level in self._levels]
if self.fine_layout is not None:
arrays[1] = np.empty(self.fine_layout.blocked_shape, np.float32)
try:
for level, array in zip(self._levels, arrays):
cuda.memcpy_dtoh(array, level.current)
except BaseException:
self._status = "failed"
raise
if self.fine_layout is not None:
arrays[1] = self.fine_layout.to_dense(arrays[1])
self._mask_nonleaves(*arrays)
return tuple(arrays)
@@ -269,7 +348,7 @@ class FixedPatchStepper:
except Exception as exc:
error = exc
self._levels.clear()
self._explode = None
self._explode = self._halo_copy = None
if getattr(self, "_fine_curved", None) is not None:
try:
self._fine_curved.free()