328 lines
18 KiB
Python
328 lines
18 KiB
Python
"""All-blocked fixed patch: transfer ordering, failure ownership and CUDA parity.
|
|
|
|
Run CPU and GPU classes in separate processes. GPU callers must first check
|
|
resources/device availability; this file never selects or reserves a GPU.
|
|
"""
|
|
from contextlib import ExitStack, contextmanager, nullcontext
|
|
from dataclasses import replace
|
|
from pathlib import Path
|
|
import struct
|
|
import unittest
|
|
from unittest.mock import Mock, patch
|
|
|
|
import numpy as np
|
|
import pycuda.driver as cuda
|
|
|
|
from CelerisLab.cuda import compiler_v2 as compiler
|
|
from CelerisLab.lbm import refinement_stepper
|
|
from CelerisLab.lbm.blocks import BlockLayout2D
|
|
from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY
|
|
from CelerisLab.lbm.refinement import ACTIVE, OUTER_GHOST, FixedPatch2D
|
|
from CelerisLab.lbm.refinement_stepper import FixedPatchStepper, blocked_base_tables
|
|
from test_refinement_blocked_fine import NX, NY, BOUNDS, config, state, circle, actuation
|
|
|
|
BASE_BLOCKS = ((24, 32), (12, 8), (3, 4), (1, 1))
|
|
FINE_BLOCK = (16, 20)
|
|
|
|
|
|
def raw(backend):
|
|
arrays = []
|
|
for level in backend._levels:
|
|
array = np.empty(level.shape if level.layout is None else level.layout.blocked_shape,
|
|
np.float32)
|
|
cuda.memcpy_dtoh(array, level.current)
|
|
arrays.append(array if level.layout is None else level.layout.to_dense(array))
|
|
return tuple(arrays)
|
|
|
|
|
|
def exact(test, actual, expected, label):
|
|
"""Compare the complete invalid mask and every finite FP32 bit (incl. -0)."""
|
|
np.testing.assert_array_equal(np.isnan(actual), np.isnan(expected), err_msg=label)
|
|
test.assertFalse(np.isinf(actual).any(), label)
|
|
test.assertFalse(np.isinf(expected).any(), label)
|
|
finite = np.isfinite(expected)
|
|
np.testing.assert_array_equal(actual[finite].view(np.uint32),
|
|
expected[finite].view(np.uint32), err_msg=label)
|
|
|
|
|
|
@contextmanager
|
|
def mocked_backend(mode="periodic", fail_alloc=None):
|
|
"""Use the real constructor, uploads and launch logic, recording driver calls."""
|
|
calls, allocations, uploads, functions = [], [], [], {}
|
|
stream = Mock()
|
|
context = Mock(activate=nullcontext, sm_arch="sm_70")
|
|
|
|
def function(name):
|
|
if name not in functions:
|
|
functions[name] = Mock(side_effect=lambda *a, **kw: calls.append((name, a, kw)))
|
|
return functions[name]
|
|
|
|
def allocate(size):
|
|
if len(allocations) == fail_alloc:
|
|
raise MemoryError("injected allocation failure")
|
|
if size <= 0:
|
|
raise AssertionError("empty allocation")
|
|
pointer = Mock()
|
|
allocations.append(pointer)
|
|
return pointer
|
|
|
|
module = Mock(get_function=Mock(side_effect=function))
|
|
module.get_global.side_effect = lambda name: (name, 16 if name == "d_block_wall" else 20)
|
|
with ExitStack() as stack:
|
|
for name in ("generate_config", "compile_kernel", "cleanup_build_dir"):
|
|
stack.enter_context(patch.object(compiler, name))
|
|
stack.enter_context(patch.object(compiler, "load_module", return_value=module))
|
|
stack.enter_context(patch.object(cuda, "Stream", return_value=stream))
|
|
stack.enter_context(patch.object(cuda, "mem_alloc", side_effect=allocate))
|
|
stack.enter_context(patch.object(cuda, "memcpy_htod",
|
|
side_effect=lambda p, a: uploads.append((p, a))))
|
|
yield dict(context=context, calls=calls, allocations=allocations, uploads=uploads,
|
|
functions=functions, stream=stream)
|
|
|
|
|
|
class BlockedBaseCPU(unittest.TestCase):
|
|
def test_reject_before_cuda(self):
|
|
context = Mock(activate=Mock(side_effect=AssertionError("CUDA reached")))
|
|
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):
|
|
p = FixedPatch2D.from_config(config("periodic"), BOUNDS)
|
|
order = np.argsort(p.coalesce_coarse)
|
|
self.assertEqual(len(order), 172)
|
|
for mode in ("periodic", "slot_reservoir"):
|
|
leaves = np.ones((NY, NX), bool)
|
|
x0, y0, x1, y1 = BOUNDS
|
|
leaves[y0:y1, x0:x1] = False
|
|
if mode != "periodic":
|
|
leaves[[0, -1]] = False
|
|
for by, bx in BASE_BLOCKS:
|
|
layout = BlockLayout2D.build(NY, NX, by, bx, periodic=(True, mode == "periodic"))
|
|
donors, targets, kind = blocked_base_tables(p, layout, order, mode != "periodic")
|
|
ly, lx = layout.local_shape
|
|
L = ly * lx
|
|
expected, source_cells = [], []
|
|
for slot in p.coalesce_coarse[order]:
|
|
q, receiver = divmod(int(slot), NX * NY)
|
|
y, x = divmod(receiver, NX)
|
|
block = (y // by) * (NX // bx) + x // bx
|
|
# Scalar coordinate construction independent of layout.owner / slot.
|
|
sy, sx = y % by + 1 - D2Q9_EY[q], x % bx + 1 - D2Q9_EX[q]
|
|
expected.append((block * 9 + q) * L + sy * lx + sx)
|
|
source_cells.append((y - D2Q9_EY[q], x - D2Q9_EX[q]))
|
|
np.testing.assert_array_equal(targets, expected)
|
|
self.assertEqual(len(set(targets.tolist())), 172)
|
|
self.assertEqual(sum(not (y0 <= y < y1 and x0 <= x < x1)
|
|
for y, x in source_cells), 8)
|
|
expected_donors = []
|
|
for cell in p.ghost_coarse:
|
|
y, x = divmod(int(cell), NX)
|
|
block = (y // by) * (NX // bx) + x // bx
|
|
expected_donors.append(block * L + (y % by + 1) * lx + x % bx + 1)
|
|
np.testing.assert_array_equal(donors, expected_donors)
|
|
expected_kind = np.full(layout.nblocks * L, OUTER_GHOST, np.uint8)
|
|
expected_kind[layout.owner[leaves].astype(np.intp)] = ACTIVE
|
|
np.testing.assert_array_equal(kind, expected_kind)
|
|
|
|
def test_staging_negative_controls_detect_seam_and_donor_corruption(self):
|
|
p = FixedPatch2D.from_config(config("periodic"), BOUNDS)
|
|
order = np.argsort(p.coalesce_coarse)
|
|
layout = BlockLayout2D.build(NY, NX, 3, 4)
|
|
donors, targets, _ = blocked_base_tables(p, layout, order)
|
|
dense = np.arange(9 * NY * NX, dtype=np.float32).reshape(9, NY, NX)
|
|
initial = layout.to_blocked(dense)
|
|
arrivals = np.arange(len(targets), dtype=np.float32) + 10000
|
|
staged = initial.copy()
|
|
staged.ravel()[targets.astype(np.intp)] = arrivals
|
|
def pull(storage):
|
|
out = np.empty_like(dense)
|
|
by, bx = layout.block
|
|
ly, lx = layout.local_shape
|
|
for y in range(NY):
|
|
for x in range(NX):
|
|
b = (y // by) * (NX // bx) + x // bx
|
|
for q in range(9):
|
|
out[q, y, x] = storage[b, q, y % by + 1 - D2Q9_EY[q],
|
|
x % bx + 1 - D2Q9_EX[q]]
|
|
return out
|
|
|
|
expected = np.array([np.roll(dense[q], (D2Q9_EY[q], D2Q9_EX[q]), axis=(0, 1))
|
|
for q in range(9)])
|
|
expected.ravel()[p.coalesce_coarse[order].astype(np.intp)] = arrivals
|
|
leaves = np.ones((NY, NX), bool)
|
|
leaves[6:18, 8:24] = False
|
|
# Compare ALL parent arrivals, not only the slots deliberately written.
|
|
np.testing.assert_array_equal(pull(staged)[:, leaves], expected[:, leaves])
|
|
# A late halo gather erases receiver-local arrivals at seams.
|
|
late_gather = staged.copy()
|
|
layout.exchange(late_gather)
|
|
self.assertGreater(np.count_nonzero(pull(late_gather)[:, leaves] != expected[:, leaves]), 0)
|
|
# Canonical source owners are not the slots read by each receiver's pull.
|
|
q, receiver = np.divmod(p.coalesce_coarse[order], NX * NY)
|
|
y, x = np.divmod(receiver.astype(np.intp), NX)
|
|
owners = layout.owner[y - np.asarray(D2Q9_EY)[q.astype(np.intp)],
|
|
x - np.asarray(D2Q9_EX)[q.astype(np.intp)]]
|
|
wrong = initial.copy()
|
|
wrong.ravel()[layout.slot(owners, q).astype(np.intp)] = arrivals
|
|
self.assertGreater(np.count_nonzero(pull(wrong)[:, leaves] != expected[:, leaves]), 0)
|
|
# With one parent block, all eight grazing sources are owner cores.
|
|
# Each is read for four fine ghost children: late explosion corrupts32.
|
|
single = BlockLayout2D.build(NY, NX, NY, NX)
|
|
donors, targets, _ = blocked_base_tables(p, single, order)
|
|
initial = single.to_blocked(dense)
|
|
staged = initial.copy()
|
|
staged.ravel()[targets.astype(np.intp)] = arrivals
|
|
donor_slots = single.slot(donors[None, :], np.arange(9)[:, None]).astype(np.intp)
|
|
self.assertEqual(np.count_nonzero(staged.ravel()[donor_slots]
|
|
!= initial.ravel()[donor_slots]), 32)
|
|
|
|
def test_raw_mask_gate_detects_leak_hidden_by_download_mask(self):
|
|
valid = state(NY, NX, 1)
|
|
valid[:, 6:18, 8:24] = np.nan
|
|
leaked = valid.copy()
|
|
leaked[5, 6, 8] = 42 # skipped invalid-cell store leaves stale staging data
|
|
public = leaked.copy()
|
|
public[:, 6:18, 8:24] = np.nan
|
|
exact(self, public, valid, "masked download cannot detect stale staging")
|
|
with self.assertRaises(AssertionError):
|
|
exact(self, leaked, valid, "raw owner-core must detect stale staging")
|
|
|
|
def test_launch_order_stream_wall_constants_and_buffer_swaps(self):
|
|
with mocked_backend("slot_reservoir") as m:
|
|
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
|
|
current, spare, fine_current = coarse.current, coarse.spare, fine.current
|
|
backend.step(2)
|
|
names = [name for name, _, _ in m["calls"]]
|
|
interval = ["BlockExplode", "BlockHaloCopy", "BlockCurvedBoundary", "BlockStep",
|
|
"BlockHaloCopy", "BlockCurvedBoundary", "BlockStep",
|
|
"BlockHaloCopy", "BlockCoalesce", "BlockStep"]
|
|
self.assertEqual(names, interval * 2)
|
|
self.assertEqual(backend.step_count, 2)
|
|
self.assertIs(coarse.current, current)
|
|
self.assertIs(coarse.spare, spare)
|
|
self.assertIs(fine.current, fine_current)
|
|
for _, _, kw in m["calls"]:
|
|
self.assertIs(kw["stream"], m["stream"])
|
|
for offset, src, dst in ((0, current, spare), (10, spare, current)):
|
|
explode, gather, stage, collide = [m["calls"][offset + i][1] for i in (0, 7, 8, 9)]
|
|
self.assertIs(explode[0], src)
|
|
self.assertIs(gather[0], src)
|
|
self.assertIs(stage[0], fine_current)
|
|
self.assertIs(stage[1], src)
|
|
self.assertIs(collide[0], src)
|
|
self.assertIs(collide[1], dst)
|
|
self.assertIs(collide[4], coarse.kind)
|
|
walls = [a for ptr, a in m["uploads"] if isinstance(ptr, str) and ptr == "d_block_wall"]
|
|
self.assertEqual(len(walls), 1)
|
|
self.assertEqual([call.kwargs["block_root_bc"] for call in
|
|
compiler.generate_config.call_args_list], [1, 0])
|
|
width, height, slot, speed = struct.unpack("=IIIf", walls[0])
|
|
self.assertEqual((width, height, slot), (NX, NY, 8))
|
|
self.assertAlmostEqual(speed, .001)
|
|
finally:
|
|
backend.close()
|
|
for allocation in m["allocations"]:
|
|
allocation.free.assert_called_once_with()
|
|
|
|
def test_allocation_failure_releases_every_owned_buffer(self):
|
|
# Fail during the first level, second level and transfer-table uploads.
|
|
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"), 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"), 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
|
|
with self.assertRaisesRegex(RuntimeError, "injected parent stage"):
|
|
backend.step()
|
|
self.assertEqual(backend.step_count, 0)
|
|
self.assertEqual(backend._status, "failed")
|
|
self.assertIs(backend._levels[0].current, current)
|
|
with self.assertRaisesRegex(RuntimeError, "failed"):
|
|
backend.download()
|
|
with self.assertRaisesRegex(RuntimeError, "failed"):
|
|
backend.step()
|
|
backend.close()
|
|
for allocation in m["allocations"]:
|
|
allocation.free.assert_called_once_with()
|
|
|
|
|
|
class BlockedBaseGPU(unittest.TestCase):
|
|
def test_bodyfree_exact_parity_and_raw_invalid_slots(self):
|
|
self.compare([(None, ((24, 32), (3, 4), (1, 1)), 8)])
|
|
|
|
def test_actuated_circle_exact_parity_and_raw_invalid_slots(self):
|
|
self.compare([(.002, ((12, 8),), 8), (-.002, ((12, 8),), 8)])
|
|
|
|
def test_non_power_two_slot_width_exact_parity(self):
|
|
self.compare([(None, ((3, 4),), 7), (None, ((3, 4),), 10)], ("slot_reservoir",))
|
|
|
|
def compare(self, cases, modes=("periodic", "slot_reservoir")):
|
|
from CelerisLab.cuda.context import CudaContext
|
|
expected_source = Path(__file__).resolve().parents[3] / "src/CelerisLab/lbm/refinement_stepper.py"
|
|
self.assertEqual(Path(refinement_stepper.__file__).resolve(), expected_source)
|
|
with CudaContext(0) as context:
|
|
for mode in modes:
|
|
for speed, layouts, width in cases:
|
|
cfg = config(mode)
|
|
if mode == "slot_reservoir":
|
|
cfg = replace(cfg, wall_slot_width=width)
|
|
p = FixedPatch2D.from_config(cfg, BOUNDS)
|
|
arrays = state(NY, NX, 1), state(*p.fine_shape, 2)
|
|
reference = None
|
|
for base_block in (None, *layouts):
|
|
label = f"{mode} width={width} speed={speed} base={base_block}"
|
|
snapshots = []
|
|
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:
|
|
# Blocked staging dirties invalid input slots; its output
|
|
# must poison them. Dense parent does not stage there.
|
|
coarse = backend._levels[0]
|
|
spare = coarse.layout.to_blocked(np.full(coarse.shape, 42, np.float32))
|
|
cuda.memcpy_htod(coarse.spare, spare)
|
|
backend.step(step - backend.step_count)
|
|
public, owner = backend.download(), raw(backend)
|
|
self.assertEqual(backend.step_count, step)
|
|
self.assertTrue(np.isnan(owner[0][:, 6:18, 8:24]).all(), label)
|
|
if mode != "periodic":
|
|
self.assertTrue(np.isnan(owner[0][:, (0, -1), :]).all(), label)
|
|
snapshots.append(public + owner)
|
|
backend.initialize(*arrays)
|
|
self.assertEqual(backend.step_count, 0)
|
|
for actual, expected in zip(backend.download() + raw(backend), snapshots[0]):
|
|
exact(self, actual, expected, label + " reset")
|
|
if reference is None:
|
|
reference = snapshots
|
|
else:
|
|
for step, actual, expected in zip((0, 1, 2, 8), snapshots, reference):
|
|
for index, (a, b) in enumerate(zip(actual, expected)):
|
|
exact(self, a, b, f"{label} step={step} array={index}")
|
|
print(f"PASS {label}: full finite bits/NaN masks at 0,1,2,8 + reset", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|