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>
This commit is contained in:
Claude
2026-10-26 19:12:11 +00:00
co-authored by Claude Opus 5.5
parent 21fe71d332
commit 18e8ef45f0
4 changed files with 54 additions and 34 deletions
+8 -17
View File
@@ -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)
+35 -1
View File
@@ -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)
+10 -15
View File
@@ -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.
@@ -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