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>
84 lines
3.4 KiB
Python
84 lines
3.4 KiB
Python
"""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()
|