From 18e8ef45f0cf7e534668588d187836ec42ab3635 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 26 Oct 2026 19:12:11 +0000 Subject: [PATCH] blocks: one definition of a blocked level's config and kernel launches (C1a) PeriodicBlockStepper and FixedPatchStepper each built the blocked level's kernel config (block window NX/NY, neutralized wall/slot/inlet fields, omega/tau check) and each spelled out the raw BlockHaloCopy and BlockCurvedBoundary argument lists. These now live once in lbm/blocks.py (level_config, launch_halo_copy, launch_curved_boundary), so a kernel ABI change has one host site. Both owners keep their own lifecycle; no new class, no behaviour or default change. PeriodicBlockStepper now checks its level config before any device work (same errors). Dense refinement levels keep their own config check. Co-Authored-By: Claude Opus 5.5 --- src/CelerisLab/lbm/block_stepper.py | 25 +++++-------- src/CelerisLab/lbm/blocks.py | 36 ++++++++++++++++++- src/CelerisLab/lbm/refinement_stepper.py | 25 ++++++------- .../multigrid/test_block_stepper.py | 2 +- 4 files changed, 54 insertions(+), 34 deletions(-) diff --git a/src/CelerisLab/lbm/block_stepper.py b/src/CelerisLab/lbm/block_stepper.py index 1b7ce9e..257256e 100644 --- a/src/CelerisLab/lbm/block_stepper.py +++ b/src/CelerisLab/lbm/block_stepper.py @@ -18,7 +18,8 @@ import numpy as np import pycuda.driver as cuda from ..cuda import compiler_v2 as compiler -from .blocks import NQ, BlockLayout2D, blocked_link_tables +from .blocks import (NQ, BlockLayout2D, blocked_link_tables, launch_curved_boundary, + launch_halo_copy, level_config) from .descriptors import D2Q9_EX, D2Q9_EY, FLUID, SOLID from .population_storage import _StoredPopulations, _encode_physical, _decode_stored @@ -103,13 +104,9 @@ class PeriodicBlockStepper: self.layout = layout self.ctx = context self._links = links - self._link_grid = ((links + tpb - 1) // tpb, 1) - # Wall/slot fields are not applied; neutralize them like the fine patch level. - self._local_cfg = replace(checked, nx=lx, ny=ly, y_wall_bc="bounce_back", top_wall_bc="velocity", - wall_slot_width=0, wall_normal_velocity=0.0, velocity=0.0) + self._local_cfg = level_config(checked, layout, 3 * checked.viscosity + .5) self._block = (tpb, 1, 1) self._grid = ((lx + tpb - 1) // tpb, ly, layout.nblocks) - self._copy_grid = ((pairs + tpb - 1) // tpb, 1) self._bytes = int(np.prod(layout.blocked_shape)) * 4 self._lock = RLock() self._status = "new" @@ -125,9 +122,6 @@ class PeriodicBlockStepper: raise MemoryError("blocked buffers exceed free device memory") self._stream = cuda.Stream() config = self._local_cfg - config.validate() - if not np.isclose(config.omega, 1 / (3 * checked.viscosity + .5), rtol=1e-14, atol=0): - raise RuntimeError("block omega/tau mismatch") module = compiler.build_module(config, "blocked.cu", arch=self.ctx.sm_arch) compiler.set_params(module, config.omega) self._module = module @@ -214,14 +208,11 @@ class PeriodicBlockStepper: try: self._status = "advancing" for _ in range(n): - self._copy(self._current, self._dst, self._src, - np.uint32(len(self.layout.halo_dst)), grid=self._copy_grid, **kw) - if self._curved is not None: - t = self._tables - self._curved(self._current, t["block"], t["cell"], t["dir"], t["q"], - t["fallback"], np.uint32(self._links), np.uintp(0), # no actuation - np.uintp(0), np.uintp(0), np.uintp(0), # no observation - grid=self._link_grid, **kw) + launch_halo_copy(self._copy, self._current, self._dst, self._src, + len(self.layout.halo_dst), **kw) + if self._curved is not None: # stationary: no actuation or observation + launch_curved_boundary(self._curved, self._current, self._tables, + self._links, **kw) self._step(self._current, self._spare, np.uint32(self.layout.halo), self._tables.get("flags", np.uintp(0)), np.uintp(0), # no kind grid=self._grid, **kw) diff --git a/src/CelerisLab/lbm/blocks.py b/src/CelerisLab/lbm/blocks.py index 5fe109f..3626926 100644 --- a/src/CelerisLab/lbm/blocks.py +++ b/src/CelerisLab/lbm/blocks.py @@ -31,7 +31,7 @@ receiver-local boundary values: the solid's owner-core copy of that slot may be stale, which is distinct from the unique core ownership of stored state. No moving geometry is implied. """ -from dataclasses import dataclass +from dataclasses import dataclass, replace from numbers import Integral from types import SimpleNamespace @@ -373,3 +373,37 @@ def blocked_link_tables(layout, flags, fluid_idx, dir, q, fallback): dir=columns[1].astype(np.uint8), q=columns[2].astype(np.float32), fallback=columns[3].astype(np.uint8)) return tables, mapped + + +# One blocked level as both owners (PeriodicBlockStepper, FixedPatchStepper) build and +# launch it. Raw argument orders of the blocked.cu kernels live only here. + +def level_config(cfg, layout, tau): + """Kernel config of one blocked level: NX/NY = block window incl. halo. + + Physical wall/slot/inlet fields are neutralized; a blocked root gets them + from its module globals (d_block_wall/d_block_channel) instead. + """ + ly, lx = layout.local_shape + local = replace(cfg, nx=lx, ny=ly, y_wall_bc="bounce_back", top_wall_bc="velocity", + wall_slot_width=0, wall_normal_velocity=0.0, velocity=0.0) + local.validate() + if not np.isclose(local.omega, 1 / tau, rtol=1e-14, atol=0): + raise RuntimeError("block level omega/tau mismatch") + return local + + +def launch_halo_copy(kernel, buffer, dst, src, count, *, block, stream): + """BlockHaloCopy over ``count`` uploaded (dst, src) block-cell pairs.""" + kernel(buffer, dst, src, np.uint32(count), + grid=((count + block[0] - 1) // block[0], 1), block=block, stream=stream) + + +def launch_curved_boundary(kernel, buffer, tables, count, *, block, stream, + velocity=None, rx=None, ry=None, raw=None): + """BlockCurvedBoundary on uploaded link tables; None = no actuation/observation.""" + none = np.uintp(0) + kernel(buffer, tables["block"], tables["cell"], tables["dir"], tables["q"], tables["fallback"], + np.uint32(count), none if velocity is None else velocity, + none if rx is None else rx, none if ry is None else ry, none if raw is None else raw, + grid=((count + block[0] - 1) // block[0], 1), block=block, stream=stream) diff --git a/src/CelerisLab/lbm/refinement_stepper.py b/src/CelerisLab/lbm/refinement_stepper.py index 45e87da..8be028a 100644 --- a/src/CelerisLab/lbm/refinement_stepper.py +++ b/src/CelerisLab/lbm/refinement_stepper.py @@ -15,7 +15,7 @@ import numpy as np import pycuda.driver as cuda from ..cuda import compiler_v2 as compiler -from .blocks import BlockLayout2D +from .blocks import BlockLayout2D, launch_curved_boundary, launch_halo_copy, level_config from .channel import channel_masks, validate_channel_config, validate_channel_transfers from ..body.jets import jet_link_velocity from .block_stepper import MAX_GRID_DIM @@ -353,13 +353,12 @@ class FixedPatchStepper: tau = pairs[0].tau_coarse if index == 0 else pairs[index - 1].tau_fine name, source = ("PatchCoarseStep" if index == 0 else "PatchFineStep"), "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) + config = level_config(config, layout, tau) name, source = ("BlockChannelStep" if index == 0 and self._channel else "BlockStep"), "blocked.cu" - config.validate() - if not np.isclose(config.omega, 1 / tau, rtol=1e-14, atol=0): - raise RuntimeError("level omega/tau mismatch") + else: + config.validate() + if not np.isclose(config.omega, 1 / tau, rtol=1e-14, atol=0): + raise RuntimeError("level omega/tau mismatch") host = self._boundary_hosts[index] n_objects = len(host.objects) if name == "PatchFineStep" and host is not None else 0 block_root_bc = 0 @@ -676,15 +675,12 @@ class FixedPatchStepper: return if b.kernel is None or b.flags is None: raise RuntimeError("Boundary device buffers are missing") - t, count = b.device, len(b.tables["cell"]) + t = b.device velocity = t.get("surface_velocity") if b.actuated and velocity is None: raise RuntimeError("Surface velocity device buffer is missing") - b.kernel(buffer, t["block"], t["cell"], t["dir"], t["q"], t["fallback"], - np.uint32(count), velocity if velocity is not None else np.uintp(0), - t.get("rx", np.uintp(0)), t.get("ry", np.uintp(0)), - b.last_raw if b.last_raw is not None else np.uintp(0), - grid=((count + self._block[0] - 1) // self._block[0], 1), **kw) + launch_curved_boundary(b.kernel, buffer, t, len(b.tables["cell"]), velocity=velocity, + rx=t.get("rx"), ry=t.get("ry"), raw=b.last_raw, **kw) def _advance(self, index, parent_buffer, **kw): """Explode from the given parent buffer, then two hops of level ``index``.""" @@ -741,8 +737,7 @@ class FixedPatchStepper: def _gather_halo(self, level, buffer, **kw): pairs = len(level.layout.halo_dst) if 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) + launch_halo_copy(level.halo_copy, buffer, level.halo_dst, level.halo_src, pairs, **kw) def step(self, n=1): """Advance G→G by n base intervals: level l takes 2**l hops per interval. diff --git a/tests/validation/multigrid/test_block_stepper.py b/tests/validation/multigrid/test_block_stepper.py index c6eb987..c03ccad 100644 --- a/tests/validation/multigrid/test_block_stepper.py +++ b/tests/validation/multigrid/test_block_stepper.py @@ -198,7 +198,7 @@ class CpuContracts(unittest.TestCase): backend._copy = backend._step = Mock() backend.boundary, backend._curved, backend._tables, backend._links = None, None, {}, 0 backend._dst, backend._src = Mock(), Mock() - backend._block, backend._grid, backend._copy_grid = (32, 1, 1), (1, 1, 1), (1, 1) + backend._block, backend._grid = (32, 1, 1), (1, 1, 1) backend._allocations = [Mock(), Mock()] backend._current, backend._spare = backend._allocations return backend