From a64af5beceece3f2a1415d37671d95bf7a33ae65 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 15 Oct 2026 02:30:40 +0000 Subject: [PATCH] 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 --- src/CelerisLab/lbm/blocks.py | 173 ++++++++++++++++++++++++++++++++ tests/unit/test_block_layout.py | 139 +++++++++++++++++++++++++ 2 files changed, 312 insertions(+) create mode 100644 src/CelerisLab/lbm/blocks.py create mode 100644 tests/unit/test_block_layout.py diff --git a/src/CelerisLab/lbm/blocks.py b/src/CelerisLab/lbm/blocks.py new file mode 100644 index 0000000..b7885d8 --- /dev/null +++ b/src/CelerisLab/lbm/blocks.py @@ -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 diff --git a/tests/unit/test_block_layout.py b/tests/unit/test_block_layout.py new file mode 100644 index 0000000..39d8623 --- /dev/null +++ b/tests/unit/test_block_layout.py @@ -0,0 +1,139 @@ +"""CPU block layout contracts: ownership, halo pairs, packing and tagged pull.""" +import sys +import unittest + +import numpy as np + +from CelerisLab.lbm.blocks import CORNER, FACE, NQ, BlockLayout2D +from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY + +CASES = ( # (ny, nx, by, bx, halo) + (6, 8, 3, 4, 1), # 2x2 split + (6, 12, 3, 4, 1), # 2x3 split + (8, 8, 4, 4, 2), # wider halo + (4, 8, 4, 4, 1), # one block along y +) +PERIODIC = ((True, True), (False, False), (True, False), (False, True)) + + +def tags(ny, nx): + """Unique, exactly representable value per dense population slot.""" + return np.arange(NQ * ny * nx, dtype=np.float64).reshape(NQ, ny, nx) + + +def dense_pull(a, periodic): + """f[q](x) = a[q](x - c_q); NaN where a non-periodic exterior is crossed.""" + px, py = periodic + _, ny, nx = a.shape + gy, gx = np.mgrid[0:ny, 0:nx] + out = np.empty_like(a) + for q in range(NQ): + sy, sx = gy - D2Q9_EY[q], gx - D2Q9_EX[q] + valid = ((py | ((sy >= 0) & (sy < ny))) & (px | ((sx >= 0) & (sx < nx)))) + out[q] = np.where(valid, a[q, sy % ny, sx % nx], np.nan) + return out + + +def blocked_pull(layout, blocked): + """Exchange, then pull every core cell from its own block storage only.""" + layout.exchange(blocked) + out = np.full_like(blocked, np.nan) + (by, bx), h = layout.block, layout.halo + for q in range(NQ): + ey, ex = D2Q9_EY[q], D2Q9_EX[q] + out[:, q, h:h + by, h:h + bx] = blocked[:, q, h - ey:h - ey + by, h - ex:h - ex + bx] + return out + + +class BlockLayoutTests(unittest.TestCase): + def test_cpu_import_boundary(self): + self.assertFalse(any(name == "pycuda" or name.startswith("pycuda.") for name in sys.modules)) + + def test_origins_and_unique_core_ownership(self): + layout = BlockLayout2D.build(6, 8, 3, 4) + np.testing.assert_array_equal(layout.origins, [[0, 0], [0, 4], [3, 0], [3, 4]]) + self.assertEqual(layout.blocked_shape, (4, NQ, 5, 6)) + for ny, nx, by, bx, h in CASES: + layout = BlockLayout2D.build(ny, nx, by, bx, halo=h) + ly, lx = layout.local_shape + core = [b * ly * lx + (y + h) * lx + x + h + for b in range(layout.nblocks) for y in range(by) for x in range(bx)] + self.assertEqual(sorted(layout.owner.ravel().tolist()), sorted(core)) + for y in range(ny): + for x in range(nx): + b = (y // by) * (nx // bx) + x // bx + self.assertEqual(layout.block_cell(b, y, x), layout.owner[y, x]) + + def test_halo_pairs_are_owner_core_to_unique_halo(self): + for ny, nx, by, bx, h in CASES: + for periodic in PERIODIC: + layout = BlockLayout2D.build(ny, nx, by, bx, halo=h, periodic=periodic) + core = set(layout.owner.ravel().tolist()) + dst = layout.halo_dst.tolist() + self.assertEqual(dst, sorted(dst)) + self.assertEqual(len(dst), len(set(dst))) + self.assertTrue(core.isdisjoint(dst)) + self.assertTrue(set(layout.halo_src.tolist()) <= core) + self.assertTrue(set(layout.halo_kind.tolist()) <= {FACE, CORNER}) + full = BlockLayout2D.build(6, 8, 3, 4) + per_block = np.bincount(full.halo_kind, minlength=2) // full.nblocks + np.testing.assert_array_equal(per_block, [2 * (3 + 4), 4]) + closed = BlockLayout2D.build(6, 8, 3, 4, periodic=(False, False)) + # Each 2x2 block keeps one inner face column, one face row and one corner. + self.assertEqual(len(closed.halo_dst), 4 * (3 + 4 + 1)) + + def test_pack_round_trip_and_exterior_stays_invalid(self): + for ny, nx, by, bx, h in CASES: + for periodic in PERIODIC: + layout = BlockLayout2D.build(ny, nx, by, bx, halo=h, periodic=periodic) + a = tags(ny, nx) + blocked = layout.to_blocked(a) + np.testing.assert_array_equal(layout.to_dense(blocked), a) + filled = np.isfinite(blocked[:, 0]).sum() + self.assertEqual(filled, ny * nx + len(layout.halo_dst)) + flat = blocked.ravel() + for q in (0, 5): + np.testing.assert_array_equal(flat[layout.slot(layout.halo_dst, q).astype(np.intp)], + flat[layout.slot(layout.halo_src, q).astype(np.intp)]) + + def test_tagged_pull_matches_dense_pull(self): + for ny, nx, by, bx, h in CASES: + for periodic in PERIODIC: + with self.subTest(case=(ny, nx, by, bx, h), periodic=periodic): + layout = BlockLayout2D.build(ny, nx, by, bx, halo=h, periodic=periodic) + a = tags(ny, nx) + got = layout.to_dense(blocked_pull(layout, layout.to_blocked(a))) + np.testing.assert_array_equal(got, dense_pull(a, periodic)) + + def test_block_cell_resolves_periodic_halo_images(self): + layout = BlockLayout2D.build(6, 8, 3, 4) + ly, lx = layout.local_shape + # Block 0 sees x=7 (wrapped) in its west halo and y=5 in its south halo. + self.assertEqual(layout.block_cell(0, 0, 7), 1 * lx + 0) + self.assertEqual(layout.block_cell(0, 5, 0), 0 * lx + 1) + closed = BlockLayout2D.build(6, 8, 3, 4, periodic=(False, False)) + with self.assertRaises(ValueError): + closed.block_cell(0, 0, 7) + with self.assertRaises(ValueError): + layout.block_cell(0, 4, 0) # two rows past the one-cell halo + + def test_invalid_inputs_rejected(self): + for args, kw in (((6, 8, 4, 4), {}), ((6, 8, 3, 4), {"halo": 0}), + ((6, 8, 3, 4), {"halo": 4}), ((6, 8, 3, 4), {"periodic": (1, 1)}), + ((6, 8, 3, True), {}), ((6, 8.0, 3, 4), {})): + with self.assertRaises(ValueError): + BlockLayout2D.build(*args, **kw) + layout = BlockLayout2D.build(6, 8, 3, 4) + with self.assertRaises(ValueError): + layout.to_blocked(np.zeros((NQ, 6, 8), dtype=np.int64)) + with self.assertRaises(ValueError): + layout.to_dense(np.zeros((4, NQ, 5, 5))) + + def test_tables_are_readonly(self): + layout = BlockLayout2D.build(6, 8, 3, 4) + for array in (layout.origins, layout.owner, layout.halo_dst, layout.halo_src, layout.halo_kind): + self.assertFalse(array.flags.writeable) + + +if __name__ == "__main__": + unittest.main()