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

Merge base-blocked FixedPatchStepper 906e0d9 into claude/jet

Browse Source 
Exact codex/refinement-blocked-base-20261015 head, cleared by Astra/root.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
ClaudeandClaude Opus 5.5 committed 2026-10-15 08:41:38 +00:00
2 parents 038a759ff1 906e0d9a0e
commit f26410f94a
6 files changed
+496 -66
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/CelerisLab/lbm/kernels/blocks.cu
+34 -3
View File
@@ -16,6 +16,7 @@
#include "operators/collision_srt.cuh"
#include "streaming/pull_double_buffer.cuh"
#include "boundary/curved_boundary.cuh"
#include "boundary/planar_velocity.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"
@@ -24,6 +25,11 @@
// Fixed-patch cell kinds, matching refinement.py ACTIVE/INNER_GHOST/OUTER_GHOST.
constexpr unsigned char KIND_ACTIVE = 0u, KIND_INNER = 1u, KIND_OUTER = 2u;
// Zero-initialized for ordinary blocks/fine levels. Set only for parent walls;
// NX/NY remain local storage strides, never physical wall coordinates.
struct BlockWall { unsigned int nx, ny, slot_width; float amplitude; };
extern "C" { __constant__ BlockWall d_block_wall; }
// 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)
@@ -38,17 +44,32 @@ extern "C" __global__ void BlockHaloCopy(
}
// 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).
// id of a patch-ghost cell; donors are parent block cell ids. Dense parent
// storage is the one-window case with coarse_cells = full domain size.
extern "C" __global__ void BlockExplode(
const fpxx* coarse, fpxx* fine, const uint64_t* target, const uint32_t* donor,
const fpxx* coarse, fpxx* fine, const uint64_t* target, const uint64_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;
const uint64_t d = donor[row];
const fpxx* from = coarse + (d / coarse_cells) * NQ * coarse_cells;
for (unsigned int q = 0; q < NQ; ++q)
to[index_f(t % TOTAL_CELLS, q)] = coarse[q * coarse_cells + donor[row]];
to[index_f(t % TOTAL_CELLS, q)] = from[q * coarse_cells + d % coarse_cells];
}
// Targets are receiver-local SOURCE slots for the immediately following pull.
// Keep child summation order identical to PatchCoarseStep; no collision here.
extern "C" __global__ void BlockCoalesce(
const fpxx* fine, fpxx* coarse, const uint64_t* target, const uint64_t* children,
unsigned int count)
{
unsigned int row = blockIdx.x * blockDim.x + threadIdx.x;
if (row >= count) return;
const uint64_t* s = children + 4ul * row;
coarse[target[row]] = 0.25f * (fine[s[0]] + fine[s[1]] + fine[s[2]] + fine[s[3]]);
}
// Static cut links, receiver-local staging: cell is the receiver's local core
@@ -92,6 +113,16 @@ extern "C" __global__ void BlockStep(const fpxx* in, fpxx* out, unsigned int hal
compute_neighbors(x, y, j); // Core neighbors stay inside the window: no wrap.
stream_pull_load(k, f, in + base, j);
if (cell_kind == KIND_ACTIVE) {
if (d_block_wall.nx != 0) {
const unsigned int bx = NX - 2 * halo, by = NY - 2 * halo;
const unsigned int nbx = d_block_wall.nx / bx;
const unsigned int gx = (blockIdx.z % nbx) * bx + x - halo;
const unsigned int gy = (blockIdx.z / nbx) * by + y - halo;
if (gy == 1)
reconstruct_bottom_velocity(f, 0.0f, bottom_wall_normal_velocity(
gx, d_block_wall.nx, d_block_wall.slot_width, d_block_wall.amplitude));
else if (gy == d_block_wall.ny - 2) reconstruct_top_reservoir(f, RHO);
}
float rho, ux, uy, feq[NQ];
compute_rho_u(f, rho, ux, uy);
compute_feq(rho, ux, uy, feq);
src/CelerisLab/lbm/kernels/boundary/planar_velocity.cuh
+15 -5
View File
@@ -3,16 +3,26 @@
// the caller owns collision/store. No force, shifting or off-lattice walls.
#ifndef CELERIS_BOUNDARY_PLANAR_VELOCITY_CUH
#define CELERIS_BOUNDARY_PLANAR_VELOCITY_CUH
#include <math.h>
__device__ __forceinline__ float bottom_wall_normal_velocity(
unsigned int x, unsigned int width, unsigned int slot_width, float amplitude)
{
// Width is the physical domain width, even when NX is a block window.
if (slot_width > 0 && slot_width < width) {
const float s = (2.0f*(float)x-(float)(width-1))/(float)slot_width;
const float w = fmaxf(0.0f, 1.0f-s*s);
return amplitude*w*w;
}
return amplitude; // Disabled or explicitly uniform full width.
}
__device__ __forceinline__ float bottom_wall_normal_velocity(unsigned int x)
{
#if WALL_SLOT_WIDTH > 0 && WALL_SLOT_WIDTH < NX
// C1 compact profile centered at (NX-1)/2; amplitude is its continuous peak.
const float s = (2.0f*(float)x-(float)(NX-1))/(float)WALL_SLOT_WIDTH;
const float w = fmaxf(0.0f, 1.0f-s*s);
return WALL_NORMAL_VELOCITY*w*w;
return bottom_wall_normal_velocity(x, NX, WALL_SLOT_WIDTH, WALL_NORMAL_VELOCITY);
#else
return WALL_NORMAL_VELOCITY; // Disabled or explicitly uniform full width.
return WALL_NORMAL_VELOCITY;
#endif
}
src/CelerisLab/lbm/refinement_stepper.py
+119 -54
View File
@@ -18,7 +18,7 @@ from .blocks import BlockLayout2D, blocked_link_tables
from .curved_links import surface_link_velocity
from .block_stepper import MAX_GRID_DIM
from .refinement import ACTIVE, HALO, OUTER_GHOST, FixedPatch2D, prepare_fine_boundary
from .descriptors import SOLID
from .descriptors import SOLID, D2Q9_EX, D2Q9_EY
def blocked_fine_tables(patch, layout, order):
@@ -37,6 +37,27 @@ def blocked_fine_tables(patch, layout, order):
layout.slot(owner[cell.astype(np.intp)], q), kind)
def blocked_base_tables(patch, layout, order, wall_aware=False):
"""Parent donor cells and receiver-local arrival slots, in transfer row order.
Stage arrivals only after child reads and parent halo gather, immediately
before parent pull. Source *owner* slots are wrong at block seams.
"""
owner = layout.owner.ravel()
q, receiver = np.divmod(patch.coalesce_coarse[order], np.uint64(owner.size))
offset = (np.asarray(D2Q9_EY)[q.astype(np.intp)] * layout.local_shape[1]
+ np.asarray(D2Q9_EX)[q.astype(np.intp)])
source = owner[receiver.astype(np.intp)].astype(np.int64) - offset
kind = np.full(layout.nblocks * int(np.prod(layout.local_shape)), OUTER_GHOST, np.uint8)
leaves = np.ones(patch.coarse_shape, bool)
x0, y0, x1, y1 = patch.bounds
leaves[y0:y1, x0:x1] = False
if wall_aware:
leaves[[0, -1]] = False
kind[owner[leaves.ravel()].astype(np.intp)] = ACTIVE
return (owner[patch.ghost_coarse].astype(np.uint64), layout.slot(source, q), kind)
@dataclass
class _Level:
shape: tuple
@@ -45,6 +66,11 @@ class _Level:
spare: object
kernel: object
grid: tuple
layout: object = None
kind: object = None
halo_copy: object = None
halo_dst: object = None
halo_src: object = None
class FixedPatchStepper:
@@ -75,11 +101,14 @@ class FixedPatchStepper:
Blocked fine_objects are stationary circles, optionally surface-actuated:
prepare_fine_boundary decides legality on the dense fine grid, then cut
links are staged receiver-local per hop: halo -> BFL -> step, same buffer.
``base_block=(by, bx)`` also blocks the parent; it requires ``fine_block``.
Parent halos precede receiver-local coalescence staging and pull/collision;
covered cells are explicitly poisoned on every step. Public shapes stay dense.
Blocked mode applies prescribed surface velocity only; its per-link
exchange diagnostic is unavailable (rows keep physical cut-link order).
"""
def __init__(self, cfg, bounds, *, context, coarse_boundary="periodic", fine_objects=(),
fine_block=None):
fine_block=None, base_block=None):
if HALO != 4:
raise RuntimeError("refinement.cu requires a four-cell fine halo")
if coarse_boundary not in ("periodic", "slot_reservoir"):
@@ -94,6 +123,17 @@ class FixedPatchStepper:
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)
if base_block is not None and fine_block is None:
raise ValueError("base_block requires fine_block")
self.base_layout = None
if base_block is not None:
if len(base_block) != 2:
raise ValueError("base_block must be (by, bx)")
self.base_layout = BlockLayout2D.build(
cfg.ny, cfg.nx, *base_block, periodic=(True, not self._wall_aware))
if (self.base_layout.nblocks > MAX_GRID_DIM
or self.base_layout.local_shape[0] > MAX_GRID_DIM):
raise ValueError("base blocks exceed launch limits")
self.fine_layout = None
if fine_block is not None:
if len(fine_block) != 2:
@@ -121,7 +161,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._coalesce = None
self._link_tables = {}
self._block = (cfg.threads_per_block, 1, 1)
fine_cfg = replace(cfg, nx=self.patch.fine_shape[1], ny=self.patch.fine_shape[0],
@@ -141,14 +181,15 @@ class FixedPatchStepper:
with self.ctx.activate():
try:
self._stream = cuda.Stream()
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])):
for index, (config, layout, tau, name) in enumerate((
(replace(cfg), self.base_layout, self.patch.tau_coarse, "PatchCoarseStep"),
(fine_cfg, self.fine_layout, self.patch.tau_fine, "PatchFineStep"))):
source = "refinement.cu"
if layout is not None:
config = replace(config, 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)
name, source = "BlockStep", "blocks.cu"
config.validate()
if not np.isclose(config.omega, 1 / tau, rtol=1e-14, atol=0):
raise RuntimeError("level omega/tau mismatch")
@@ -165,33 +206,45 @@ class FixedPatchStepper:
if size != struct.calcsize("=ffffI"):
raise RuntimeError("LBMParams host/device layout mismatch")
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), grid))
shape = (9, *(self.patch.coarse_shape if index == 0 else self.patch.fine_shape))
size = int(np.prod(shape if layout is None else layout.blocked_shape)) * 4
grid = ((config.nx + self._block[0] - 1) // self._block[0], config.ny,
1 if layout is None else layout.nblocks)
level = _Level(shape, module, self._allocate(size), self._allocate(size),
module.get_function(name), grid, layout)
self._levels.append(level)
if layout is not None:
level.halo_copy = module.get_function("BlockHaloCopy")
if len(layout.halo_dst):
level.halo_dst = self._upload(layout.halo_dst)
level.halo_src = self._upload(layout.halo_src)
if index == 0 and self._wall_aware:
ptr, size = module.get_global("d_block_wall")
if size != struct.calcsize("=IIIf"):
raise RuntimeError("Block wall host/device layout mismatch")
cuda.memcpy_htod(ptr, struct.pack("=IIIf", cfg.nx, cfg.ny,
cfg.wall_slot_width, cfg.wall_normal_velocity))
p = self.patch
c, f = self._levels
order = np.argsort(p.coalesce_coarse)
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])
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])
if len(layout.halo_dst): # a 1x1 non-periodic layout has no pairs
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")
fine_tables = (None if f.layout is None else
blocked_fine_tables(p, f.layout, order))
base_tables = (None if c.layout is None else
blocked_base_tables(p, c.layout, order, self._wall_aware))
self._ghosts = self._upload(p.ghost_fine if fine_tables is None else fine_tables[0])
donors = p.ghost_coarse if base_tables is None else base_tables[0]
self._donors = self._upload(donors if fine_tables is None else donors.astype(np.uint64))
self._targets = self._upload(p.coalesce_coarse[order] if base_tables is None
else base_tables[1])
self._children = self._upload(p.coalesce_fine[order] if fine_tables is None
else fine_tables[1])
self._explode = f.module.get_function("PatchExplode" if fine_tables is None
else "BlockExplode")
if fine_tables is not None:
f.kind = self._upload(fine_tables[2])
if base_tables is not None:
c.kind = self._upload(base_tables[2])
self._coalesce = c.module.get_function("BlockCoalesce")
if self._fine_tables is not None:
self._fine_flags = self._upload(self._fine_tables["flags"])
self._link_tables = {k: self._upload(v) for k, v in self._fine_tables.items()
@@ -257,8 +310,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])
arrays = [a if level.layout is None else level.layout.to_blocked(a)
for a, level in zip(arrays, self._levels)]
try:
for level, array in zip(self._levels, arrays):
cuda.memcpy_htod(level.current, array)
@@ -303,10 +356,10 @@ class FixedPatchStepper:
np.uint32(count), velocity if velocity is not None else np.uintp(0),
grid=((count + self._block[0] - 1) // self._block[0], 1), **kw)
def _fine_halo(self, buffer, **kw):
pairs = len(self.fine_layout.halo_dst)
def _gather_halo(self, level, buffer, **kw):
pairs = len(level.layout.halo_dst)
if pairs:
self._halo_copy(buffer, self._halo_dst, self._halo_src, np.uint32(pairs),
level.halo_copy(buffer, level.halo_dst, level.halo_src, np.uint32(pairs),
grid=((pairs + self._block[0] - 1) // self._block[0], 1), **kw)
def step(self, n=1):
@@ -322,7 +375,9 @@ class FixedPatchStepper:
self._status = "advancing"
for _ in range(n):
self._explode(c.current, f.current, self._ghosts, self._donors,
np.uint32(len(p.ghost_fine)), np.uint64(np.prod(p.coarse_shape)),
np.uint32(len(p.ghost_fine)),
np.uint64(np.prod(p.coarse_shape if c.layout is None
else c.layout.local_shape)),
grid=((len(p.ghost_fine) + self._block[0] - 1) // self._block[0], 1),
**kw)
if self.fine_layout is not None:
@@ -330,10 +385,10 @@ class FixedPatchStepper:
# Same buffer per hop: halo -> receiver-local BFL -> step.
flags = self._fine_flags if self._fine_flags is not None else np.uintp(0)
for src, dst in ((f.current, f.spare), (f.spare, f.current)):
self._fine_halo(src, **kw)
self._gather_halo(f, src, **kw)
self._blocked_fine_boundary(src, **kw)
f.kernel(src, dst, np.uint32(self.fine_layout.halo), flags,
self._kind, grid=f.grid, **kw)
f.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:
@@ -342,10 +397,21 @@ class FixedPatchStepper:
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),
grid=c.grid, **kw)
if c.layout is None:
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),
grid=c.grid, **kw)
else:
# Child donors must be read before staging mutates parent G.
# No halo gather may overwrite these receiver-local arrivals.
self._gather_halo(c, c.current, **kw)
self._coalesce(f.current, c.current, self._targets, self._children,
np.uint32(len(p.coalesce_coarse)),
grid=((len(p.coalesce_coarse) + self._block[0] - 1)
// self._block[0], 1), **kw)
c.kernel(c.current, c.spare, np.uint32(c.layout.halo), np.uintp(0),
c.kind, grid=c.grid, **kw)
self._stream.synchronize()
c.current, c.spare = c.spare, c.current
self._step_count += 1
@@ -358,17 +424,16 @@ class FixedPatchStepper:
"""Return two synchronized G arrays, NaN on all non-leaves; no host cache."""
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)
arrays = [np.empty(level.shape if level.layout is None else level.layout.blocked_shape,
np.float32) for level in self._levels]
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])
arrays = [a if level.layout is None else level.layout.to_dense(a)
for a, level in zip(arrays, self._levels)]
self._mask_nonleaves(*arrays)
return tuple(arrays)
@@ -387,7 +452,7 @@ class FixedPatchStepper:
except Exception as exc:
error = exc
self._levels.clear()
self._explode = self._halo_copy = None
self._explode = self._coalesce = None
if getattr(self, "_fine_curved", None) is not None:
try:
self._fine_curved.free()