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 <noreply@anthropic.com>

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)
+        self._local_cfg = level_config(checked, layout, 3 * checked.viscosity + .5)
-        # 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._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,
+                    launch_halo_copy(self._copy, self._current, self._dst, self._src,
-                               np.uint32(len(self.layout.halo_dst)), grid=self._copy_grid, **kw)
+                                     len(self.layout.halo_dst), **kw)
-                    if self._curved is not None:
+                    if self._curved is not None:  # stationary: no actuation or observation
-                        t = self._tables
+                        launch_curved_boundary(self._curved, self._current, self._tables,
-                        self._curved(self._current, t["block"], t["cell"], t["dir"], t["q"],
+                                               self._links, **kw)
-                                     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)
                     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)

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)

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],
+            config = level_config(config, layout, tau)
-                             y_wall_bc="bounce_back", top_wall_bc="velocity",
-                             wall_slot_width=0, wall_normal_velocity=0.0, velocity=0.0)
             name, source = ("BlockChannelStep" if index == 0 and self._channel else "BlockStep"), "blocked.cu"
-        config.validate()
+        else:
-        if not np.isclose(config.omega, 1 / tau, rtol=1e-14, atol=0):
+            config.validate()
-            raise RuntimeError("level omega/tau mismatch")
+            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"],
+        launch_curved_boundary(b.kernel, buffer, t, len(b.tables["cell"]), velocity=velocity,
-                 np.uint32(count), velocity if velocity is not None else np.uintp(0),
+                               rx=t.get("rx"), ry=t.get("ry"), raw=b.last_raw, **kw)
-                 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)
 
     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),
+            launch_halo_copy(level.halo_copy, buffer, level.halo_dst, level.halo_src, pairs, **kw)
-                            grid=((pairs + self._block[0] - 1) // self._block[0], 1), **kw)
 
     def step(self, n=1):
         """Advance G→G by n base intervals: level l takes 2**l hops per interval.

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