P1 (7/n): MRT, shifted-jet and blocked-base GPU tests migrated

test_refinement_mrt_gpu, test_shifted_jets_gpu and
test_refinement_blocked_base pass refinements/level_* lists. The three
blocked-base rejection cases keep their inputs per level; base-only
blocks now meet the level-list rule (blocked refined levels below) where
the legacy keyword check fired before. Assertions unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Claude
2026-10-21 01:58:57 +00:00
co-authored by Claude Opus 5.5
parent af0566db98
commit 0fcd4ac4cb
3 changed files with 30 additions and 26 deletions
@@ -83,11 +83,10 @@ def mocked_backend(mode="periodic", fail_alloc=None):
class BlockedBaseCPU(unittest.TestCase):
def test_reject_before_cuda(self):
context = Mock(activate=Mock(side_effect=AssertionError("CUDA reached")))
for kwargs in ({"base_block": (12, 8)},
{"base_block": (7, 8), "fine_block": FINE_BLOCK},
{"base_block": (12,), "fine_block": FINE_BLOCK}):
with self.subTest(kwargs=kwargs), self.assertRaises(ValueError):
FixedPatchStepper(config("periodic"), BOUNDS, context=context, **kwargs)
for blocks in ([(12, 8), None], [(7, 8), FINE_BLOCK], [(12,), FINE_BLOCK]):
with self.subTest(blocks=blocks), self.assertRaises(ValueError):
FixedPatchStepper(config("periodic"), context=context, refinements=[BOUNDS],
level_blocks=blocks)
context.activate.assert_not_called()
def test_all_arrival_rows_donors_and_kind_masks(self):
@@ -193,9 +192,10 @@ class BlockedBaseCPU(unittest.TestCase):
def test_launch_order_stream_wall_constants_and_buffer_swaps(self):
with mocked_backend("slot_reservoir") as m:
backend = FixedPatchStepper(config("slot_reservoir"), BOUNDS, context=m["context"],
coarse_boundary="slot_reservoir", fine_block=FINE_BLOCK,
base_block=(3, 4), fine_objects=[circle(actuation(.002))])
backend = FixedPatchStepper(config("slot_reservoir"), context=m["context"],
coarse_boundary="slot_reservoir", refinements=[BOUNDS],
level_blocks=[(3, 4), FINE_BLOCK],
level_objects=[(), (circle(actuation(.002)),)])
try:
backend.initialize(state(NY, NX, 1), state(*backend.patch.fine_shape, 2))
coarse, fine = backend._levels
@@ -236,16 +236,16 @@ class BlockedBaseCPU(unittest.TestCase):
for fail_at in (1, 6, 10):
with self.subTest(fail_at=fail_at), mocked_backend(fail_alloc=fail_at) as m:
with self.assertRaisesRegex(MemoryError, "injected allocation"):
FixedPatchStepper(config("periodic"), BOUNDS, context=m["context"],
fine_block=FINE_BLOCK, base_block=(3, 4))
FixedPatchStepper(config("periodic"), context=m["context"], refinements=[BOUNDS],
level_blocks=[(3, 4), FINE_BLOCK])
self.assertEqual(len(m["allocations"]), fail_at)
for allocation in m["allocations"]:
allocation.free.assert_called_once_with()
def test_failed_parent_stage_cannot_publish_or_resume(self):
with mocked_backend() as m:
backend = FixedPatchStepper(config("periodic"), BOUNDS, context=m["context"],
fine_block=FINE_BLOCK, base_block=(3, 4))
backend = FixedPatchStepper(config("periodic"), context=m["context"], refinements=[BOUNDS],
level_blocks=[(3, 4), FINE_BLOCK])
backend.initialize(state(NY, NX, 1), state(*backend.patch.fine_shape, 2))
m["functions"]["BlockCoalesce"].side_effect = RuntimeError("injected parent stage failure")
current = backend._levels[0].current
@@ -289,9 +289,10 @@ class BlockedBaseGPU(unittest.TestCase):
for base_block in (None, *layouts):
label = f"{mode} width={width} speed={speed} base={base_block}"
snapshots = []
with FixedPatchStepper(cfg, BOUNDS, context=context, coarse_boundary=mode,
fine_block=FINE_BLOCK, base_block=base_block,
fine_objects=() if speed is None else [circle(actuation(speed))]) as backend:
with FixedPatchStepper(cfg, context=context, coarse_boundary=mode,
refinements=[BOUNDS], level_blocks=[base_block, FINE_BLOCK],
level_objects=[(), () if speed is None
else (circle(actuation(speed)),)]) as backend:
backend.initialize(*arrays)
for step in (0, 1, 2, 8):
if step == 1 and base_block is not None: