PeriodicBlockStepper.step(n): wait on the stream once per call

Same change as the reviewed FixedPatchStepper end-sync (cherry-picked
ad1d82c+512dcc9 above): the n steps are enqueued in order and the host
waits once; step(0) launches nothing; step_count advances after the wait
and keeps its entry value on failure. Launches, buffer roles and bits are
those of n step(1) calls. Mocked-driver host test covers the single
wait, launch/role equality with n step(1), step(0) and failure.

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

src/CelerisLab/lbm/block_stepper.py

@@ -204,11 +204,18 @@ 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). On a
+        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 +231,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._current, self._spare = self._spare, self._current  # host handles only
                 self._stream.synchronize()
-                    self._current, self._spare = self._spare, self._current
-                    self._step_count += 1
+                self._step_count += n
                 self._status = "ready"
             except BaseException:
                 self._status = "failed"

tests/unit/test_block_step_sync.py

@@ -0,0 +1,83 @@
+"""Host contract of PeriodicBlockStepper.step(n): one stream wait per call (mocked driver)."""
+import ast
+from contextlib import nullcontext
+from pathlib import Path
+import types
+from unittest.mock import Mock
+import unittest
+
+import numpy as np
+
+import CelerisLab.lbm.block_stepper as b
+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."""
+    path = Path(b.__file__)
+    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"))
+
+
+if __name__ == "__main__":
+    unittest.main()