diff --git a/src/CelerisLab/lbm/block_stepper.py b/src/CelerisLab/lbm/block_stepper.py index 13c1c9c..8389bd8 100644 --- a/src/CelerisLab/lbm/block_stepper.py +++ b/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._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" diff --git a/tests/unit/test_block_step_sync.py b/tests/unit/test_block_step_sync.py new file mode 100644 index 0000000..5779a70 --- /dev/null +++ b/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()