Add CPU BlockLayout2D: core ownership, halo pairs, blocked packing
Fixed-size 2D blocks, block-contiguous Q-major storage, explicit per-axis periodicity; non-periodic exterior halo cells stay unpaired (NaN). Provides owner/halo tables for a later periodic GPU consumer. Cut-link mapping deliberately omitted until the BFL write policy is settled. Tests: ownership, pair invariants, pack round trip, tagged pull vs dense. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Claude
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
3a7967f4cc
commit
a64af5bece
@@ -0,0 +1,173 @@
|
||||
"""CPU layout for fixed-size 2D blocks: core ownership, halo pairs and packing.
|
||||
|
||||
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.
|
||||
|
||||
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
|
||||
all NQ populations from an owner core cell into another block's halo cell.
|
||||
Periodicity is explicit per axis. Across a non-periodic exterior, halo cells
|
||||
have no pair and stay invalid (NaN after packing): no wrap, no implied BC.
|
||||
|
||||
Preparation only: no CUDA, no boundary conditions, no refinement levels.
|
||||
"""
|
||||
from dataclasses import dataclass
|
||||
from numbers import Integral
|
||||
|
||||
import numpy as np
|
||||
|
||||
NQ = 9
|
||||
FACE, CORNER = 0, 1
|
||||
|
||||
|
||||
def _readonly(values, dtype):
|
||||
array = np.ascontiguousarray(values, dtype=dtype)
|
||||
array.setflags(write=False)
|
||||
return array
|
||||
|
||||
|
||||
def _positive_int(value, name):
|
||||
if isinstance(value, (bool, np.bool_)) or not isinstance(value, Integral) or value <= 0:
|
||||
raise ValueError(f"{name} must be a positive integer")
|
||||
return int(value)
|
||||
|
||||
|
||||
@dataclass(frozen=True, eq=False)
|
||||
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;
|
||||
``halo_kind`` marks FACE or CORNER pairs. Sources are never halo cells, so
|
||||
one gather in any order fills every valid halo cell.
|
||||
"""
|
||||
shape: tuple[int, int]
|
||||
block: tuple[int, int]
|
||||
halo: int
|
||||
periodic: tuple[bool, bool]
|
||||
origins: np.ndarray
|
||||
owner: np.ndarray
|
||||
halo_dst: np.ndarray
|
||||
halo_src: np.ndarray
|
||||
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."""
|
||||
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:
|
||||
raise ValueError("block sizes must divide the grid")
|
||||
if halo > min(by, bx):
|
||||
raise ValueError("halo may not exceed the block size")
|
||||
if len(periodic) != 2 or not all(isinstance(p, (bool, np.bool_)) for p in periodic):
|
||||
raise ValueError("periodic must be two booleans (x, y)")
|
||||
px, py = map(bool, periodic)
|
||||
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:
|
||||
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)
|
||||
|
||||
gy, gx = np.mgrid[0:ny, 0:nx]
|
||||
owner = ((gy // by) * nbx + gx // bx) * cells + (gy % by + halo) * lx + gx % bx + halo
|
||||
|
||||
# Halo cells of one block in local row-major order, then all blocks.
|
||||
hy, hx = np.mgrid[0:ly, 0:lx]
|
||||
out_y = (hy < halo) | (hy >= by + halo)
|
||||
out_x = (hx < halo) | (hx >= bx + halo)
|
||||
ring = out_y | out_x
|
||||
hy, hx, kind = hy[ring], hx[ring], np.where(out_y[ring] & out_x[ring], CORNER, FACE)
|
||||
dst, src, kinds = [], [], []
|
||||
for b, (y0, x0) in enumerate(origins):
|
||||
sy, sx = y0 + hy - halo, x0 + hx - halo
|
||||
if py:
|
||||
sy %= ny
|
||||
if px:
|
||||
sx %= nx
|
||||
keep = (sy >= 0) & (sy < ny) & (sx >= 0) & (sx < nx)
|
||||
dst.append(b * cells + hy[keep] * lx + hx[keep])
|
||||
src.append(owner[sy[keep], sx[keep]])
|
||||
kinds.append(kind[keep])
|
||||
return cls((ny, nx), (by, bx), halo, (px, py),
|
||||
_readonly(origins, np.int64), _readonly(owner, np.uint64),
|
||||
_readonly(np.concatenate(dst), np.uint64),
|
||||
_readonly(np.concatenate(src), np.uint64),
|
||||
_readonly(np.concatenate(kinds), np.uint8))
|
||||
|
||||
@property
|
||||
def nblocks(self) -> int:
|
||||
return len(self.origins)
|
||||
|
||||
@property
|
||||
def local_shape(self) -> tuple[int, int]:
|
||||
return self.block[0] + 2 * self.halo, self.block[1] + 2 * self.halo
|
||||
|
||||
@property
|
||||
def blocked_shape(self) -> tuple[int, int, int, int]:
|
||||
return (self.nblocks, NQ, *self.local_shape)
|
||||
|
||||
def slot(self, block_cell, q):
|
||||
"""Flat blocked population slot of block cell id(s) for direction q."""
|
||||
cells = int(np.prod(self.local_shape))
|
||||
b, c = np.divmod(np.asarray(block_cell, dtype=np.uint64), np.uint64(cells))
|
||||
return (b * np.uint64(NQ) + np.asarray(q, dtype=np.uint64)) * np.uint64(cells) + c
|
||||
|
||||
def block_cell(self, block, y, x):
|
||||
"""Address of in-grid cell (y, x) in ``block``'s own storage, core or halo.
|
||||
|
||||
Periodic axes resolve to the image inside the block's halo window.
|
||||
Raises if the cell is not represented there or lies across an exterior.
|
||||
"""
|
||||
by, bx = self.block
|
||||
ly, lx = self.local_shape
|
||||
ny, nx = self.shape
|
||||
y0, x0 = self.origins[block]
|
||||
ry, rx = int(y) - int(y0) + self.halo, int(x) - int(x0) + self.halo
|
||||
if self.periodic[1] and not 0 <= ry < ly:
|
||||
ry = next((r for r in (ry - ny, ry + ny) if 0 <= r < ly), ry)
|
||||
if self.periodic[0] and not 0 <= rx < lx:
|
||||
rx = next((r for r in (rx - nx, rx + nx) if 0 <= r < lx), rx)
|
||||
if not (0 <= ry < ly and 0 <= rx < lx and 0 <= y < ny and 0 <= x < nx):
|
||||
raise ValueError("cell is not represented in this block's storage")
|
||||
return int(block) * ly * lx + ry * lx + rx
|
||||
|
||||
def exchange(self, blocked) -> None:
|
||||
"""Fill every valid halo cell in place from its owner's core."""
|
||||
blocked = self._check_blocked(blocked)
|
||||
flat = blocked.reshape(self.nblocks, NQ, -1)
|
||||
cells = flat.shape[2]
|
||||
db, dc = np.divmod(self.halo_dst.astype(np.intp), cells)
|
||||
sb, sc = np.divmod(self.halo_src.astype(np.intp), cells)
|
||||
flat[db, :, dc] = flat[sb, :, sc]
|
||||
|
||||
def to_blocked(self, dense):
|
||||
"""Pack (NQ, ny, nx) floats; halo cells without a pair stay NaN."""
|
||||
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
|
||||
self.exchange(blocked)
|
||||
return blocked
|
||||
|
||||
def to_dense(self, blocked):
|
||||
"""Unpack owner cores only; halo contents are ignored."""
|
||||
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)
|
||||
|
||||
def _check_blocked(self, blocked):
|
||||
if not isinstance(blocked, np.ndarray) or blocked.shape != self.blocked_shape \
|
||||
or not blocked.flags.c_contiguous:
|
||||
raise ValueError(f"expected C-contiguous blocked array {self.blocked_shape}")
|
||||
return blocked
|
||||
Reference in New Issue
Block a user