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:
Claude
2026-10-15 03:48:54 +00:00
co-authored by Claude Opus 5.5
parent 3e44e90ca5
commit dc5850c198
4 changed files with 257 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
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()
tests/validation/multigrid/test_refinement_blocked_fine.py
+127
View File
@@ -0,0 +1,127 @@
"""Opt-in blocked fine storage for FixedPatchStepper: CPU tables plus GPU gate.
Frozen GPU gate (same stream, 8 coarse steps, nonequilibrium start), in both
periodic and slot_reservoir modes: dense fine storage vs 1x1, 2x2 and 4x2
fine blocks give bit-identical masked downloads AND bit-identical raw coarse
and unpacked owner-core fine arrays, including patch-ghost values/NaN masks.
"""
import unittest
from unittest.mock import Mock
import numpy as np
import pycuda.driver as cuda
from CelerisLab.body.geometry.circle import CircleGeometry
from CelerisLab.body.objects import SimObject
from CelerisLab.config import LBMConfig
from CelerisLab.lbm.blocks import BlockLayout2D
from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY
from CelerisLab.lbm.refinement import ACTIVE, INNER_GHOST, OUTER_GHOST, FixedPatch2D
from CelerisLab.lbm.refinement_stepper import FixedPatchStepper, blocked_fine_tables
NX, NY, BOUNDS, STEPS = 32, 24, (8, 6, 24, 18), 8
FINE_BLOCKS = ((32, 40), (16, 20), (16, 10)) # fine window 32x40: 1x1, 2x2, 4x2
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(mode):
if mode == "periodic":
return LBMConfig(nx=NX, ny=NY, viscosity=.1)
return LBMConfig(nx=NX, ny=NY, viscosity=.1, velocity=0, wall_normal_velocity=.001,
wall_slot_width=8, y_wall_bc="on_site_velocity", top_wall_bc="reservoir")
def state(ny, nx, seed):
"""Positive low-Mach equilibrium plus a seeded nonequilibrium part."""
y, x = np.mgrid[0:ny, 0:nx]
rho = 1 + 1e-3 * np.cos(2 * np.pi * x / nx)
ux, uy = .02 * np.sin(2 * np.pi * y / ny), .01 * np.cos(2 * np.pi * x / nx)
cu = CX[:, None, None] * ux + CY[:, None, None] * uy
feq = W[:, None, None] * rho * (1 + 3 * cu + 4.5 * cu * cu - 1.5 * (ux * ux + uy * uy))
noise = np.random.default_rng(seed).uniform(-1, 1, feq.shape) * 1e-3 * W[:, None, None]
return (feq + noise).astype(np.float32)
def raw(backend):
"""Unmasked device arrays; blocked fine storage unpacked from owner cores."""
coarse, fine = backend._levels
out = [np.empty(coarse.shape, np.float32)]
out.append(np.empty(fine.shape if backend.fine_layout is None
else backend.fine_layout.blocked_shape, np.float32))
cuda.memcpy_dtoh(out[0], coarse.current)
cuda.memcpy_dtoh(out[1], fine.current)
if backend.fine_layout is not None:
out[1] = backend.fine_layout.to_dense(out[1])
return out
class BlockedFineCPU(unittest.TestCase):
def forbidden_context(self):
return Mock(activate=Mock(side_effect=AssertionError("CUDA reached")), sm_arch="sm_70")
def test_rejections_before_cuda(self):
circle = SimObject(0, geometry=CircleGeometry(16.0, 12.0, 2.0), center=(16.0, 12.0), radius=2.0)
for kw in ({"fine_block": (7, 7)}, {"fine_block": (16,)},
{"fine_block": (16, 20), "fine_objects": [circle]}):
with self.subTest(kw=kw), self.assertRaises(ValueError):
FixedPatchStepper(config("periodic"), BOUNDS, context=self.forbidden_context(), **kw)
def test_tables_map_back_to_dense_slots_in_order(self):
patch = FixedPatch2D.from_config(config("periodic"), BOUNDS)
fy, fx = patch.fine_shape
cells = fy * fx
order = np.argsort(patch.coalesce_coarse)
for by, bx in FINE_BLOCKS:
layout = BlockLayout2D.build(fy, fx, by, bx, periodic=(False, False))
core = set(layout.owner.ravel().tolist())
targets, children, kind = blocked_fine_tables(patch, layout, order)
# Explosion targets: owner cores of exactly the non-ACTIVE fine cells.
np.testing.assert_array_equal(targets, layout.owner.ravel()[patch.ghost_fine])
self.assertEqual(sorted(targets.tolist()),
sorted(layout.owner[patch.fine_kind != ACTIVE].tolist()))
# Children: same rows/child order, same (q, dense cell), owner cores.
L = int(np.prod(layout.local_shape))
b, rest = np.divmod(children, np.uint64(9 * L))
q, c = np.divmod(rest, np.uint64(L))
dense = patch.coalesce_fine[order]
np.testing.assert_array_equal(q, dense // cells)
np.testing.assert_array_equal(b * np.uint64(L) + c,
layout.owner.ravel()[(dense % cells).astype(np.intp)])
self.assertTrue(set((b * np.uint64(L) + c).ravel().tolist()) <= core)
# Kinds follow owner cores; unused halo entries are OUTER.
np.testing.assert_array_equal(kind[layout.owner.ravel().astype(np.intp)],
patch.fine_kind.ravel())
self.assertEqual(set(np.unique(kind).tolist()) - {ACTIVE, INNER_GHOST, OUTER_GHOST}, set())
class BlockedFineGPU(unittest.TestCase):
def test_dense_and_blocked_fine_storage_are_bit_identical(self):
from CelerisLab.cuda.context import CudaContext
for mode in ("periodic", "slot_reservoir"):
cfg = config(mode)
patch = FixedPatch2D.from_config(cfg, BOUNDS)
g_coarse, g_fine = state(NY, NX, 1), state(*patch.fine_shape, 2)
results = []
with CudaContext(0) as context:
for fine_block in (None, *FINE_BLOCKS):
with FixedPatchStepper(cfg, BOUNDS, context=context, coarse_boundary=mode,
fine_block=fine_block) as backend:
backend.initialize(g_coarse, g_fine)
backend.step(STEPS)
self.assertEqual(backend.step_count, STEPS)
results.append((backend.download(), raw(backend)))
(dense_public, dense_raw), blocked = results[0], results[1:]
kind = patch.fine_kind
self.assertTrue(np.isnan(dense_raw[1][:, kind == OUTER_GHOST]).all(), mode)
self.assertTrue(np.isfinite(dense_raw[1][:, kind != OUTER_GHOST]).all(), mode)
for fine_block, (public, raw_arrays) in zip(FINE_BLOCKS, blocked):
msg = f"{mode} fine_block={fine_block}"
for a, b in zip(public + tuple(raw_arrays), dense_public + tuple(dense_raw)):
np.testing.assert_array_equal(a, b, err_msg=msg) # values and NaN masks
print(f"{mode}: dense vs {len(blocked)} fine-block layouts bit-identical "
f"after {STEPS} coarse steps", flush=True)
if __name__ == "__main__":
unittest.main()