diff --git a/src/CelerisLab/lbm/block_stepper.py b/src/CelerisLab/lbm/block_stepper.py index 76f2195..4d6a203 100644 --- a/src/CelerisLab/lbm/block_stepper.py +++ b/src/CelerisLab/lbm/block_stepper.py @@ -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 diff --git a/src/CelerisLab/lbm/kernels/blocks.cu b/src/CelerisLab/lbm/kernels/blocks.cu index da3f6ad..543f07d 100644 --- a/src/CelerisLab/lbm/kernels/blocks.cu +++ b/src/CelerisLab/lbm/kernels/blocks.cu @@ -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); } diff --git a/src/CelerisLab/lbm/refinement_stepper.py b/src/CelerisLab/lbm/refinement_stepper.py index f38aaa0..56aabda 100644 --- a/src/CelerisLab/lbm/refinement_stepper.py +++ b/src/CelerisLab/lbm/refinement_stepper.py @@ -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() diff --git a/tests/validation/multigrid/test_refinement_blocked_fine.py b/tests/validation/multigrid/test_refinement_blocked_fine.py new file mode 100644 index 0000000..9525ff9 --- /dev/null +++ b/tests/validation/multigrid/test_refinement_blocked_fine.py @@ -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()