BlockLayout2D: explicit selected block origins with hole owners

build(..., origins=None|[(y0, x0), ...]) stores only grid-aligned selected
blocks (unique, strictly row-major, no silent reordering). Cells of absent
blocks own layout.hole (= nblocks*L, one past the last block cell) so an
unguarded index fails rather than wrapping like -1; layout.owned is the
dense mask. Halo pairs never read holes, to_blocked ignores hole values,
to_dense returns NaN there, map_cut_links rejects receivers, solid
neighbors or q<.5 donors in absent blocks, blocked_link_tables stores
owned flags only. origins=None or every origin reproduces the regular
layout exactly. Storage only: no CUDA or stepper changes.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Claude
2026-10-16 03:09:15 +00:00
co-authored by Claude Opus 5.5
parent 99ddb5fc27
commit c397e08565
2 changed files with 226 additions and 21 deletions
+90 -20
View File
@@ -3,8 +3,14 @@
Dense arrays are Q-major (q, y, x) with cell y*nx+x, as in core/layout.cuh.
Blocked storage is block-contiguous and Q-major inside each block:
(nblocks, q, by+2*halo, bx+2*halo). Block b = jy*nbx + jx has core origin
(jy*by, jx*bx). Every dense cell, fluid or solid, has exactly one core owner.
(nblocks, q, by+2*halo, bx+2*halo). A full layout stores every grid block,
block b = jy*nbx + jx with core origin (jy*by, jx*bx), so every dense cell,
fluid or solid, has exactly one core owner. A sparse layout stores only the
selected grid-aligned blocks (row-major order, b = position in that list);
cells of absent blocks are HOLES: owner = layout.hole (one past the last
block cell, so any unguarded use fails instead of wrapping like -1), no halo
pair reads them, packing ignores them and unpacking returns NaN there. This
is physical-block storage, not arbitrary cell refinement.
A "block cell" id is b*L + c with L=(by+2h)*(bx+2h) and c the local row-major
cell; its population q lives at flat slot b*NQ*L + q*L + c. Halo pairs copy
@@ -57,6 +63,24 @@ def _positive_int(value, name):
return int(value)
def _selected_blocks(origins, ny, nx, by, bx):
"""Row-major grid block indices of explicit (y0, x0) origins, validated."""
array = np.asarray(origins)
if array.dtype.kind not in "iu" or array.ndim != 2 or array.shape[1] != 2 or not len(array):
raise ValueError("origins must be a non-empty integer (n, 2) array of (y0, x0)")
y0, x0 = array[:, 0].astype(np.int64), array[:, 1].astype(np.int64)
if (y0 % by).any() or (x0 % bx).any():
raise ValueError("origins must be multiples of the block size")
if (y0 < 0).any() or (y0 >= ny).any() or (x0 < 0).any() or (x0 >= nx).any():
raise ValueError("origins must lie inside the grid")
index = (y0 // by) * (nx // bx) + x0 // bx
if (np.diff(index) == 0).any():
raise ValueError("origins must be unique")
if (np.diff(index) < 0).any():
raise ValueError("origins must be in row-major order")
return index
@dataclass(frozen=True, eq=False)
class BlockCutLinks:
"""One row per input link, in input order (so force is counted once).
@@ -79,8 +103,9 @@ class BlockLayout2D:
"""Fixed-size block partition of an (ny, nx) grid; construct with build().
``origins`` rows are (y0, x0) core origins. ``owner`` maps each dense cell
to its owning block cell id. ``halo_dst``/``halo_src`` are block cell ids
(destination halo cell, source owner core cell) sorted by destination;
to its owning block cell id, or ``hole`` if its block is absent.
``halo_dst``/``halo_src`` are block cell ids (destination halo cell,
source owner core cell) sorted by destination;
``halo_kind`` marks FACE or CORNER pairs. Sources are never halo cells, so
one gather in any order fills every valid halo cell.
"""
@@ -95,8 +120,14 @@ class BlockLayout2D:
halo_kind: np.ndarray
@classmethod
def build(cls, ny, nx, by, bx, *, halo=1, periodic=(True, True)) -> "BlockLayout2D":
"""``periodic`` is (periodic_x, periodic_y); block sizes must divide the grid."""
def build(cls, ny, nx, by, bx, *, halo=1, periodic=(True, True), origins=None) -> "BlockLayout2D":
"""``periodic`` is (periodic_x, periodic_y); block sizes must divide the grid.
``origins`` (optional) selects stored blocks by (y0, x0) core origin:
integer multiples of (by, bx) inside the grid, at least one, unique
and strictly increasing in row-major order (no silent reordering).
None stores all blocks; passing every origin gives the same layout.
"""
ny, nx, by, bx, halo = (_positive_int(v, n) for v, n in
((ny, "ny"), (nx, "nx"), (by, "by"), (bx, "bx"), (halo, "halo")))
if ny % by or nx % bx:
@@ -109,14 +140,23 @@ class BlockLayout2D:
nby, nbx = ny // by, nx // bx
ly, lx = by + 2 * halo, bx + 2 * halo
cells = ly * lx
if nby * nbx * NQ * cells > np.iinfo(np.int64).max // 2:
if origins is None:
jy, jx = np.divmod(np.arange(nby * nbx), nbx)
origins = np.stack([jy * by, jx * bx], axis=1)
selected = _selected_blocks(origins, ny, nx, by, bx)
origins = np.asarray(origins, dtype=np.int64).reshape(-1, 2)
if (len(origins) + 1) * NQ * cells > np.iinfo(np.int64).max // 2:
raise ValueError("blocked storage exceeds 64-bit indexing")
jy, jx = np.divmod(np.arange(nby * nbx), nbx)
origins = np.stack([jy * by, jx * bx], axis=1)
# Grid block -> stored block, or len(origins) for an absent block, so
# absent cells land exactly on hole = nblocks * cells.
stored = np.full(nby * nbx, len(origins), dtype=np.int64)
stored[selected] = np.arange(len(origins))
gy, gx = np.mgrid[0:ny, 0:nx]
owner = ((gy // by) * nbx + gx // bx) * cells + (gy % by + halo) * lx + gx % bx + halo
gb = stored[(gy // by) * nbx + gx // bx]
hole = len(origins) * cells
owner = np.where(gb < len(origins),
gb * cells + (gy % by + halo) * lx + gx % bx + halo, hole)
# Halo cells of one block in local row-major order, then all blocks.
hy, hx = np.mgrid[0:ly, 0:lx]
@@ -132,6 +172,7 @@ class BlockLayout2D:
if px:
sx %= nx
keep = (sy >= 0) & (sy < ny) & (sx >= 0) & (sx < nx)
keep[keep] = owner[sy[keep], sx[keep]] != hole
dst.append(b * cells + hy[keep] * lx + hx[keep])
src.append(owner[sy[keep], sx[keep]])
kinds.append(kind[keep])
@@ -145,6 +186,16 @@ class BlockLayout2D:
def nblocks(self) -> int:
return len(self.origins)
@property
def hole(self) -> int:
"""Owner value of cells in absent blocks: one past the last block cell."""
return self.nblocks * int(np.prod(self.local_shape))
@property
def owned(self) -> np.ndarray:
"""Dense (ny, nx) mask of cells stored in some block core."""
return self.owner != self.hole
@property
def local_shape(self) -> tuple[int, int]:
return self.block[0] + 2 * self.halo, self.block[1] + 2 * self.halo
@@ -168,6 +219,8 @@ class BlockLayout2D:
if two images fit, the core-side one is returned (use local offsets,
not this lookup, when a stencil must follow a specific image).
Raises if the cell is not represented there or lies across an exterior.
In a sparse layout a represented halo cell whose source is a hole has
no pair and stays invalid (NaN after packing).
"""
by, bx = self.block
ly, lx = self.local_shape
@@ -193,6 +246,8 @@ class BlockLayout2D:
x %= nx
if not (0 <= y < ny and 0 <= x < nx):
raise ValueError("cut-link stencil crosses a non-periodic exterior")
if self.owner[y, x] == self.hole:
raise ValueError("cut-link stencil touches an absent block")
return y, x
def map_cut_links(self, links) -> BlockCutLinks:
@@ -200,7 +255,9 @@ class BlockLayout2D:
Rejects repeated physical links and repeated local writers. Addresses
are geometry-free: q, fallback and BFL math stay as built on the dense
grid. Halo must be at least one cell (always true here).
grid. Halo must be at least one cell (always true here). In a sparse
layout the receiver, its solid neighbor and any q<.5 donor must all lie
in stored blocks.
"""
ny, nx = self.shape
lx = self.local_shape[1]
@@ -214,6 +271,8 @@ class BlockLayout2D:
seen.add((r, q))
y, x = divmod(r, nx)
receiver = int(self.owner[y, x])
if receiver == self.hole:
raise ValueError("cut-link receiver lies in an absent block")
block = receiver // cells
# Local stencil offsets from the core receiver: the window image the
# next local pull reads, even where a periodic seam wraps physically.
@@ -244,23 +303,32 @@ class BlockLayout2D:
flat[db, :, dc] = flat[sb, :, sc]
def to_blocked(self, dense):
"""Pack (NQ, ny, nx) floats; halo cells without a pair stay NaN."""
"""Pack (NQ, ny, nx) floats; halo cells without a pair stay NaN.
Values at holes (absent blocks) are ignored.
"""
dense = np.asarray(dense)
if dense.shape != (NQ, *self.shape) or not np.issubdtype(dense.dtype, np.floating):
raise ValueError(f"expected floating (NQ, ny, nx) = {(NQ, *self.shape)}")
blocked = np.full(self.blocked_shape, np.nan, dtype=dense.dtype)
flat = blocked.reshape(self.nblocks, NQ, -1)
ob, oc = np.divmod(self.owner.ravel().astype(np.intp), flat.shape[2])
flat[ob, :, oc] = dense.reshape(NQ, -1).T
owned = self.owned.ravel()
ob, oc = np.divmod(self.owner.ravel()[owned].astype(np.intp), flat.shape[2])
flat[ob, :, oc] = dense.reshape(NQ, -1)[:, owned].T
self.exchange(blocked)
return blocked
def to_dense(self, blocked):
"""Unpack owner cores only; halo contents are ignored."""
"""Unpack owner cores only; halo contents are ignored, holes are NaN."""
blocked = self._check_blocked(blocked)
flat = blocked.reshape(self.nblocks, NQ, -1)
ob, oc = np.divmod(self.owner.ravel().astype(np.intp), flat.shape[2])
return np.ascontiguousarray(flat[ob, :, oc].T).reshape(NQ, *self.shape)
owned = self.owned.ravel()
ob, oc = np.divmod(self.owner.ravel()[owned].astype(np.intp), flat.shape[2])
dense = np.empty((NQ, owned.size), dtype=blocked.dtype)
if not owned.all(): # full layouts keep any dtype; holes need a float NaN
dense[:, ~owned] = np.nan
dense[:, owned] = flat[ob, :, oc].T
return dense.reshape(NQ, *self.shape)
def _check_blocked(self, blocked):
if not isinstance(blocked, np.ndarray) or blocked.shape != self.blocked_shape \
@@ -276,7 +344,8 @@ def blocked_link_tables(layout, flags, fluid_idx, dir, q, fallback):
own builder; this only maps storage via map_cut_links (unique receiver
owner and local writer, input row order kept). Returns (tables, mapped):
tables holds device-ready arrays, with flags sized nblocks*L (unused halo
entries 0) and cell = the receiver's local core cell.
entries 0) and cell = the receiver's local core cell. Flags at holes
(absent blocks) are not stored.
"""
flags = np.asarray(flags)
if flags.shape != (int(np.prod(layout.shape)),) or flags.dtype != np.uint16:
@@ -288,7 +357,8 @@ def blocked_link_tables(layout, flags, fluid_idx, dir, q, fallback):
for r, d, v, f in zip(*columns))
cells = int(np.prod(layout.local_shape))
blocked_flags = np.zeros(layout.nblocks * cells, np.uint16)
blocked_flags[layout.owner.ravel().astype(np.intp)] = flags
owned = layout.owned.ravel()
blocked_flags[layout.owner.ravel()[owned].astype(np.intp)] = flags[owned]
tables = dict(flags=blocked_flags, block=mapped.block,
cell=(mapped.receiver % cells).astype(np.uint32),
dir=columns[1].astype(np.uint8), q=columns[2].astype(np.float32),