test_block_step_sync: locate block_stepper without importing it

The module-level import of CelerisLab.lbm.block_stepper loaded
pycuda.driver into the shared test process, unlike the other CPU host
tests (which load mocked copies via AST). find_spec gives the source path
without the import. Test content and assertions unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

src/CelerisLab/lbm/block_stepper.py

@@ -204,11 +204,21 @@ class PeriodicBlockStepper:
                 raise
 
     def step(self, n=1):
-        """Advance n steps: halo refresh, then core pull/collide, per step."""
+        """Advance n steps: halo refresh, then core pull/collide, per step.
+
+        All n steps are enqueued in order on the one stream and the host waits
+        once before returning (same launches and bits as n step(1) calls).
+        Asynchronous device errors surface at that single wait, so they are not
+        attributed to the step that caused them; call step(1) repeatedly to
+        localize one. On any failure the instance is failed and step_count keeps
+        its entry value.
+        """
         if isinstance(n, bool) or not isinstance(n, Integral) or n < 0:
             raise ValueError("n must be a nonnegative integer")
         with self._lock, self.ctx.activate():
             self._require("ready")
+            if n == 0:
+                return
             kw = dict(block=self._block, stream=self._stream)
             try:
                 self._status = "advancing"
@@ -224,9 +234,9 @@ class PeriodicBlockStepper:
                     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)
-                    self._stream.synchronize()
-                    self._current, self._spare = self._spare, self._current
-                    self._step_count += 1
+                    self._current, self._spare = self._spare, self._current  # host handles only
+                self._stream.synchronize()
+                self._step_count += n
                 self._status = "ready"
             except BaseException:
                 self._status = "failed"

src/CelerisLab/lbm/kernels/blocks.cu

@@ -178,7 +178,11 @@ extern "C" __global__ void BlockStep(const fpxx* in, fpxx* out, unsigned int hal
                 else if (gy == d_block_wall.ny - 2) reconstruct_top_reservoir(f, RHO);
             }
 #if COLLISION_MODEL == 2
-            collide_mrt_unforced(f);  // SRT path below unchanged.
+            // Uniform MRT sequence (collide_dispatch, zero forcing); SRT path below unchanged.
+            float rho, ux, uy;
+            compute_rho_u(f, rho, ux, uy);
+            collide_mrt_no_force(f, rho, ux, uy,
+                fminf(OMEGA_COLLISION_MAX, fmaxf(OMEGA_COLLISION_MIN, d_params.omega)));
 #else
             float rho, ux, uy, feq[NQ];
             compute_rho_u(f, rho, ux, uy);

src/CelerisLab/lbm/kernels/operators/collision_mrt.cuh

@@ -114,16 +114,6 @@ __device__ __forceinline__ void collide_mrt_no_force(float* __restrict__ g,
     collide_mrt(g, rho, ux, uy, Fin, omega);
 }
 
-// Unforced MRT with the clamped device omega, as in collide_dispatch with zero
-// forcing. Shared by the blocked, channel-root and dense-patch collisions.
-__device__ __forceinline__ void collide_mrt_unforced(float* __restrict__ g)
-{
-    float rho, ux, uy;
-    compute_rho_u(g, rho, ux, uy);
-    collide_mrt_no_force(g, rho, ux, uy,
-        fminf(OMEGA_COLLISION_MAX, fmaxf(OMEGA_COLLISION_MIN, d_params.omega)));
-}
-
 #elif LATTICE_MODEL == LATTICE_D3Q19
 
 __device__ __forceinline__ void collide_mrt(float* __restrict__ g,

src/CelerisLab/lbm/kernels/refinement.cu

@@ -37,7 +37,10 @@ __device__ unsigned int incoming_counts[NQ * TOTAL_CELLS];
 
 __device__ __forceinline__ void patch_collide(float* f, unsigned long k) {
 #if COLLISION_MODEL == 2
-    collide_mrt_unforced(f);
+    float rho, ux, uy;
+    compute_rho_u(f, rho, ux, uy);
+    collide_mrt_no_force(f, rho, ux, uy,
+        fminf(OMEGA_COLLISION_MAX, fmaxf(OMEGA_COLLISION_MIN, d_params.omega)));
 #else
     float rho, ux, uy, feq[NQ];
     compute_rho_u(f, rho, ux, uy);

src/CelerisLab/lbm/kernels/step/block_channel.cu

@@ -55,7 +55,11 @@ extern "C" __global__ void BlockChannelStep(
         }
         if (!wall && !outlet) {
 #if COLLISION_MODEL == 2
-            collide_mrt_unforced(f);
+            // Match uniform MRT arithmetic (collide_dispatch, zero forcing).
+            float rho, ux, uy;
+            compute_rho_u(f, rho, ux, uy);
+            collide_mrt_no_force(f, rho, ux, uy,
+                fminf(OMEGA_COLLISION_MAX, fmaxf(OMEGA_COLLISION_MIN, d_params.omega)));
 #else
             // Match uniform SRT arithmetic, including its shifted equilibrium path.
             float rho, ux, uy, feq[NQ], forcing[NQ] = {};

src/CelerisLab/lbm/refinement_stepper.py

@@ -822,27 +822,41 @@ class FixedPatchStepper:
                             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."""
+        """Advance G→G by n base intervals: level l takes 2**l hops per interval.
+
+        All n intervals are enqueued in order on the one stream, and the host waits
+        once, before returning; step(0) launches and waits for nothing. On return
+        the device work is complete, step_count and an open window cover all n
+        intervals, and the raw-exchange ledger holds the final interval only.
+
+        Asynchronous device errors surface at that single wait, so they are not
+        attributed to the interval that caused them; call step(1) repeatedly to
+        localize one (same launches, same bits). On any failure the instance is
+        failed, the exchange ledger and window are invalid, and step_count keeps
+        its value from entry to the failing call.
+        """
         if isinstance(n, bool) or not isinstance(n, Integral) or n < 0:
             raise ValueError("n must be a nonnegative integer")
         with self._lock, self.ctx.activate():
             self._require("ready")
+            if n == 0:
+                return
             root = self._levels[0]
             kw = dict(block=self._block, stream=self._stream)
             try:
                 self._status = "advancing"
+                self._exchange_valid = False
                 for _ in range(n):
-                    if self._observe_raw_exchange:
-                        self._exchange_valid = False
+                    if self._observe_raw_exchange:  # stream-ordered: holds the last interval only
                         for b in self._level_boundaries:
                             if b is not None:
                                 cuda.memset_d8_async(b.interval_integral, 0,
                                                      len(b.tables["cell"]) * 32, self._stream)
                     self._hop(0, root.current, root.spare, **kw)
-                    self._stream.synchronize()
-                    root.current, root.spare = root.spare, root.current
-                    self._step_count += 1
-                    self._exchange_valid = self._observe_raw_exchange
+                    root.current, root.spare = root.spare, root.current  # host handles only
+                self._stream.synchronize()
+                self._step_count += n
+                self._exchange_valid = self._observe_raw_exchange
                 self._status = "ready"
             except BaseException:
                 self._window_valid = False

tests/unit/test_block_step_sync.py

@@ -0,0 +1,104 @@
+"""Host contract of PeriodicBlockStepper.step(n): one stream wait per call (mocked driver)."""
+import ast
+from contextlib import nullcontext
+import importlib.util
+from pathlib import Path
+import types
+from unittest.mock import Mock
+import unittest
+
+import numpy as np
+
+from CelerisLab.config import LBMConfig
+from CelerisLab.lbm.blocks import BlockLayout2D
+
+
+def host_module():
+    """block_stepper with pycuda/compiler replaced by mocks; no GPU context."""
+    # Locate the source without importing it: importing block_stepper loads pycuda.driver,
+    # which CPU-only import-boundary tests (test_build_contracts) must never see.
+    path = Path(importlib.util.find_spec("CelerisLab.lbm.block_stepper").origin)
+    tree = ast.parse(path.read_text())
+    tree.body = [node for node in tree.body if not (
+        isinstance(node, ast.Import) and any(a.name == "pycuda.driver" for a in node.names)
+        or isinstance(node, ast.ImportFrom) and node.module == "cuda")]
+    module = types.ModuleType("CelerisLab.lbm._test_host_block_stepper")
+    module.__package__ = "CelerisLab.lbm"
+    module.__dict__.update(cuda=Mock(), compiler=Mock())
+    exec(compile(tree, str(path), "exec"), module.__dict__)
+    return module
+
+
+def backend():
+    host, log = host_module(), []
+    host.cuda.mem_get_info.return_value = (1 << 40, 1 << 40)
+    host.cuda.mem_alloc.side_effect = lambda size: Mock(size=size)  # distinct buffers
+    module = Mock()
+    module.get_global.return_value = (Mock(), 20)
+    module.get_function.side_effect = lambda name: Mock(
+        side_effect=lambda *a, **k: log.append((name, a[:2])))
+    host.compiler.load_module.return_value = module
+    cfg = LBMConfig(nx=16, ny=16, viscosity=.1)
+    run = host.PeriodicBlockStepper(cfg, BlockLayout2D.build(16, 16, 8, 8),
+                                    context=Mock(activate=nullcontext, sm_arch="sm_70"))
+    run._stream.synchronize.side_effect = lambda: log.append(("sync", ()))
+    run.initialize(np.full((9, 16, 16), 1 / 9))
+    log.clear()
+    return run, log
+
+
+class BlockStepSync(unittest.TestCase):
+    def test_step_n_waits_once_and_matches_n_single_steps(self):
+        many, log_many = backend()
+        many.step(3)
+        single, log_single = backend()
+        for _ in range(3):
+            single.step(1)
+        names = [name for name, _ in log_many]
+        self.assertEqual(names.count("sync"), 1)
+        self.assertEqual(names[-1], "sync")
+        self.assertEqual(names.count("BlockStep"), 3)
+        # Same launch order and buffer roles as three step(1) calls.
+        strip = lambda log: [(n, tuple(id(x) for x in a)) for n, a in log if n != "sync"]
+        roles = lambda run, log: [(n, tuple({id(run._current): "c", id(run._spare): "s"}.get(x, "-")
+                                            for x in ids)) for n, ids in strip(log)]
+        self.assertEqual([n for n, _ in strip(log_many)], [n for n, _ in strip(log_single)])
+        self.assertEqual(roles(many, log_many), roles(single, log_single))
+        self.assertIsNot(many._current, many._spare)
+        self.assertEqual((many.step_count, single.step_count), (3, 3))
+
+    def test_step_zero_launches_and_waits_for_nothing(self):
+        run, log = backend()
+        run.step(0)
+        self.assertEqual((log, run.step_count), ([], 0))
+
+    def test_failure_keeps_entry_step_count(self):
+        run, log = backend()
+        run.step(1)
+        run._stream.synchronize.side_effect = RuntimeError("device fault")
+        with self.assertRaises(RuntimeError):
+            run.step(2)
+        self.assertEqual((run.step_count, run._status), (1, "failed"))
+
+    def test_launch_failure_mid_batch_fails_closed_without_waiting(self):
+        run, log = backend()
+        calls = []
+        step = run._step
+
+        def failing(*args, **kwargs):
+            calls.append(1)
+            if len(calls) == 2:
+                raise RuntimeError("launch failure")
+            return step(*args, **kwargs)
+
+        run._step = failing
+        with self.assertRaises(RuntimeError):
+            run.step(3)
+        self.assertNotIn("sync", [name for name, _ in log])
+        self.assertEqual((run.step_count, run._status), (0, "failed"))
+        with self.assertRaises(RuntimeError):
+            run.download()
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/test_refinement_mrt.py

@@ -37,13 +37,8 @@ class RefinementMRTHost(unittest.TestCase):
             self.assertIn("STORE_PRECISION != 0", source)
             self.assertIn('#if COLLISION_MODEL == 2\n#include "operators/collision_mrt.cuh"\n#endif', source)
         self.assertIn("USE_DDF_SHIFTING ||", refinement)  # dense/mixed shifted stays rejected
-        # One shared unforced MRT body; each blocked/patch path calls it exactly once.
-        mrt = (KERNELS / "operators/collision_mrt.cuh").read_text()
-        self.assertEqual(mrt.count("void collide_mrt_unforced("), 1)
         for source in (blocks, channel, refinement):
-            self.assertEqual(source.count("collide_mrt_unforced(f);"), 1)
-            self.assertNotIn("collide_mrt_no_force(", source)
-            self.assertNotIn("compute_rho_u(f, rho, ux, uy);\n            collide_mrt", source)
+            self.assertEqual(source.count("collide_mrt_no_force("), 1)
 
 
 if __name__ == "__main__": unittest.main()

tests/unit/test_step_sync.py

@@ -0,0 +1,170 @@
+"""Host contract of FixedPatchStepper.step(n): one stream wait per call (mocked driver).
+
+All n base intervals are enqueued in order on the backend stream and the host
+waits once before step returns. Launch order and buffer arguments must equal n
+step(1) calls; completion metadata, raw-exchange and window ledgers, and failure
+invalidation are published only after that wait.
+"""
+from contextlib import nullcontext
+from unittest.mock import Mock
+import unittest
+
+import numpy as np
+
+from CelerisLab.config import LBMConfig
+from tests.unit.test_raw_exchange import body
+from tests.unit.test_refinement_three_level_multipatch import host_module
+
+BODY_LEVEL, HOPS = 3, 8  # body on L3: 2**3 hops per base interval
+
+
+def backend(window=False):
+    """Four-level host backend; every launch, memset and stream wait goes to one log."""
+    host = host_module()
+    log, allocations = [], []
+
+    def allocate(size):
+        allocations.append(Mock(size=size))
+        return allocations[-1]
+
+    module = Mock()
+    module.get_global.return_value = (Mock(), 20)
+    module.get_function.side_effect = lambda name: Mock(
+        side_effect=lambda *a, **k: log.append((name, a, k)))
+    host.compiler.load_module.return_value = module
+    host.cuda.mem_alloc.side_effect = allocate
+    host.cuda.memcpy_dtoh.side_effect = lambda array, pointer: array.fill(0)
+    host.cuda.memset_d8_async.side_effect = lambda *a: log.append(("memset", a, {}))
+    run = host.FixedPatchStepper(
+        LBMConfig(nx=48, ny=48, viscosity=.02), context=Mock(activate=nullcontext, sm_arch="sm_70"),
+        refinements=[[(6, 6, 26, 26)], [(20, 20, 44, 44)], [(56, 56, 72, 72)]],
+        level_blocks=[(4, 4)] * 4, level_objects=[(), (), (), (body(),)],
+        observe_raw_exchange=True, observe_window_exchange=window)
+    run._stream.synchronize.side_effect = lambda: log.append(("sync", (), {}))
+    run.initialize(*[np.full(level.shape, 1 / 9, np.float32) for level in run._levels])
+    log.clear()
+    return run, log, allocations
+
+
+def signature(log, allocations):
+    """Event names with device-pointer arguments replaced by allocation index."""
+    index = {id(a): i for i, a in enumerate(allocations)}
+    return [(name, tuple(index.get(id(a), "-") for a in args)) for name, args, _ in log]
+
+
+def names(log):
+    return [name for name, _, _ in log]
+
+
+class StepSync(unittest.TestCase):
+    def test_step1_waits_once_after_all_launches(self):
+        run, log, _ = backend()
+        run.step(1)
+        self.assertEqual(names(log).count("sync"), 1)
+        self.assertEqual(names(log)[-1], "sync")
+        self.assertEqual(names(log)[0], "memset")
+        self.assertEqual(names(log).count("BlockStep"), 1 + 2 + 4 + 8)
+        self.assertEqual(run.step_count, 1)
+        run.close()
+
+    def test_step_n_equals_n_step1_launches_with_one_wait(self):
+        single, single_log, single_alloc = backend()
+        for _ in range(3):
+            single.step(1)
+        batched, batched_log, batched_alloc = backend()
+        batched.step(3)
+        expected = [e for e in signature(single_log, single_alloc) if e[0] != "sync"]
+        got = signature(batched_log, batched_alloc)
+        self.assertEqual(got[-1][0], "sync")
+        self.assertEqual(names(batched_log).count("sync"), 1)
+        self.assertEqual(got[:-1], expected)  # same order, same ping-pong buffers
+        self.assertEqual(names(batched_log).count("memset"), 3)
+        self.assertEqual(batched.step_count, single.step_count)
+        index = [{id(a): i for i, a in enumerate(alloc)} for alloc in (single_alloc, batched_alloc)]
+        self.assertEqual(index[1][id(batched._levels[0].current)],
+                         index[0][id(single._levels[0].current)])
+        single.close()
+        batched.close()
+
+    def test_step0_launches_and_waits_for_nothing(self):
+        run, log, _ = backend()
+        run.step(2)
+        log.clear()
+        run.step(0)
+        self.assertEqual(log, [])
+        self.assertEqual(run.step_count, 2)
+        value = run.read_raw_exchange(0, level=BODY_LEVEL)  # step(0) keeps the ledger valid
+        self.assertEqual((value["time_start"], value["time_end"]), (1, 2))
+        run.close()
+
+    def test_completion_metadata_covers_all_intervals(self):
+        run, _, _ = backend()
+        with self.assertRaisesRegex(RuntimeError, "completed"):
+            run.read_raw_exchange(0, level=BODY_LEVEL)
+        run.step(3)
+        value = run.read_raw_exchange(0, level=BODY_LEVEL)
+        self.assertEqual((value["time_start"], value["time_end"], value["sample_count"]),
+                         (2, 3, HOPS))  # step(n) retains only the final base interval
+        run.step(2)
+        self.assertEqual(run.read_raw_exchange(0, level=BODY_LEVEL)["time_end"], 5)
+        run.close()
+
+    def test_window_accounting_across_batched_calls(self):
+        windows = []
+        for calls in ((1, 3), (4,), (1, 1, 1, 1)):
+            run, log, _ = backend(window=True)
+            run.step(1)
+            run.start_exchange_window()
+            log.clear()
+            for n in calls:
+                run.step(n)
+            self.assertEqual(names(log).count("BlockAccumulateExchange"), 2 * 4 * HOPS)
+            self.assertEqual(names(log).count("sync"), len(calls))
+            out = run.read_exchange_window(0, level=BODY_LEVEL)
+            windows.append((out["time_start"], out["time_end"], out["sample_count"]))
+            run.close()
+        self.assertEqual(windows, [(1., 5, 4 * HOPS)] * 3)
+
+    def test_failed_wait_publishes_nothing(self):
+        run, _, _ = backend(window=True)
+        run.step(2)
+        run.start_exchange_window()
+        run._stream.synchronize.side_effect = RuntimeError("device failure")
+        with self.assertRaisesRegex(RuntimeError, "device failure"):
+            run.step(3)
+        self.assertEqual(run.step_count, 2)  # value at entry to the failing call
+        self.assertFalse(run._exchange_valid)
+        self.assertFalse(run._window_valid)
+        for read in (lambda: run.read_raw_exchange(0, level=BODY_LEVEL),
+                     lambda: run.read_exchange_window(0, level=BODY_LEVEL), run.download,
+                     lambda: run.step(1)):
+            with self.assertRaisesRegex(RuntimeError, "failed"):
+                read()
+        run._stream.synchronize.side_effect = None
+        run.close()
+
+    def test_launch_error_mid_batch_skips_wait_and_fails(self):
+        run, log, _ = backend()
+        run.step(1)
+        log.clear()
+        hop, calls = run._hop, []
+
+        def failing(index, *args, **kw):
+            if index == 0:
+                calls.append(1)
+                if len(calls) == 2:
+                    raise RuntimeError("launch failure")
+            return hop(index, *args, **kw)
+        run._hop = failing
+        with self.assertRaisesRegex(RuntimeError, "launch failure"):
+            run.step(3)
+        self.assertNotIn("sync", names(log))
+        self.assertEqual(run.step_count, 1)
+        self.assertFalse(run._exchange_valid)
+        with self.assertRaisesRegex(RuntimeError, "failed"):
+            run.read_raw_exchange(0, level=BODY_LEVEL)
+        run.close()
+
+
+if __name__ == "__main__":
+    unittest.main()