222 lines
12 KiB
Python
222 lines
12 KiB
Python
"""Bounded FP32 GPU / pinned FP64 CPUREF-0.3 equivalence, not accuracy validation.
|
|
|
|
Run explicitly in pycuda_3_10 with PYTHONPATH=<this repo>/src, one BLAS thread,
|
|
and CUDA_VISIBLE_DEVICES=0 (or 1) only after checking the physical GPU is idle.
|
|
At most eight coarse steps per case. No files, plots or flow arrays are saved.
|
|
"""
|
|
import hashlib
|
|
import importlib.util
|
|
from contextlib import nullcontext
|
|
from pathlib import Path
|
|
import resource
|
|
import sys
|
|
from threading import RLock
|
|
import unittest
|
|
from unittest.mock import Mock, patch
|
|
|
|
import numpy as np
|
|
import pycuda.driver as cuda
|
|
|
|
import CelerisLab
|
|
from CelerisLab.config import LBMConfig
|
|
from CelerisLab.cuda import compiler_v2 as compiler
|
|
from CelerisLab.cuda.context import CudaContext
|
|
from CelerisLab.lbm.refinement_stepper import PeriodicPatchStepper
|
|
|
|
ROOT = Path(__file__).resolve().parents[3]
|
|
POP_ATOL = 3e-6
|
|
MOMENT_ATOL = 2e-4
|
|
REFERENCE_SHA = "bba2a06a9f6517b32079ffff15e4060d8c3b7ed35f16ca134e1877226c9e4e1e"
|
|
REF_PATH = Path(__file__).with_name("cpu_reference.py")
|
|
assert hashlib.sha256(REF_PATH.read_bytes()).hexdigest() == REFERENCE_SHA
|
|
loader = importlib.util.spec_from_file_location("pinned_patch_reference", REF_PATH)
|
|
ref = importlib.util.module_from_spec(loader)
|
|
sys.modules[loader.name] = ref
|
|
loader.loader.exec_module(ref)
|
|
|
|
|
|
def postcollision(state, spec, maps):
|
|
# Observation conversion only. Never applied to production G.
|
|
return tuple(ref._collide(spec, a, np.isfinite(a), tau, mask,
|
|
np.zeros(mask.shape, dtype=np.int64))
|
|
for a, tau, mask in ((state.coarse_f, spec.tau_coarse, maps.coarse_active),
|
|
(state.fine_f, spec.tau_fine, maps.fine_active)))
|
|
|
|
|
|
def totals(arrays, maps):
|
|
# Leaf sum independent of the reference's flux/budget bookkeeping helpers.
|
|
c, f = arrays
|
|
populations = c[:, maps.coarse_active].astype(np.float64).sum(axis=1)
|
|
populations += .25 * f[:, maps.fine_active].astype(np.float64).sum(axis=1)
|
|
return np.array([populations.sum(), populations @ ref.CX, populations @ ref.CY])
|
|
|
|
|
|
def global_array(module, name, shape):
|
|
out = np.empty(shape, dtype=np.uint32)
|
|
ptr, size = module.get_global(name)
|
|
assert size == out.nbytes
|
|
cuda.memcpy_dtoh(out, ptr)
|
|
return out
|
|
|
|
|
|
class CoupledGPU(unittest.TestCase):
|
|
def test_cleanup_attempts_every_allocation(self):
|
|
backend = PeriodicPatchStepper.__new__(PeriodicPatchStepper)
|
|
backend.ctx = Mock(activate=nullcontext)
|
|
backend._lock = RLock()
|
|
backend._status = "failed"
|
|
backend._stream = Mock()
|
|
backend._stream.synchronize.side_effect = RuntimeError("original stream failure")
|
|
first, second = Mock(), Mock()
|
|
first.free.side_effect = RuntimeError("secondary free failure")
|
|
backend._allocations = [first, second]
|
|
backend._levels = [Mock()]
|
|
with self.assertRaisesRegex(RuntimeError, "original stream failure"):
|
|
backend.close()
|
|
first.free.assert_called_once()
|
|
second.free.assert_called_once()
|
|
self.assertEqual(backend._allocations, [])
|
|
self.assertEqual(backend._levels, [])
|
|
backend.close()
|
|
|
|
def test_partial_advance_cannot_be_published(self):
|
|
spec = ref.make_spec(collision="srt")
|
|
maps = ref.build_maps(spec)
|
|
with CudaContext(0) as context:
|
|
with PeriodicPatchStepper(LBMConfig(nx=16, ny=16, viscosity=.1),
|
|
(4, 4, 12, 12), context=context) as backend:
|
|
backend.initialize(*postcollision(ref.make_state(spec), spec, maps))
|
|
coarse, fine = backend._levels
|
|
|
|
def fail_after_fine(*args, **kwargs):
|
|
# Ensure both queued fine hops actually completed before failing.
|
|
backend._stream.synchronize()
|
|
raise RuntimeError("coarse launch failed after two fine hops")
|
|
|
|
with patch.object(coarse, "kernel", side_effect=fail_after_fine):
|
|
with self.assertRaisesRegex(RuntimeError, "after two fine hops"):
|
|
backend.step()
|
|
self.assertEqual(backend.step_count, 0)
|
|
for operation in (backend.download, backend.step):
|
|
with self.assertRaisesRegex(RuntimeError, "failed"):
|
|
operation()
|
|
raw = np.empty(fine.shape, np.float32)
|
|
cuda.memcpy_dtoh(raw, fine.current)
|
|
self.assertTrue(np.isfinite(raw.ravel()[backend.patch.coalesce_fine]).all())
|
|
self.assertFalse(context.is_closed)
|
|
print("partial failure: two fine hops completed, no coarse step published; cleanup PASS", flush=True)
|
|
|
|
def test_bounded_equivalence_and_ownership(self):
|
|
self.assertEqual(Path(CelerisLab.__file__).resolve(), ROOT / "src/CelerisLab/__init__.py")
|
|
self.assertEqual(Path(compiler._KERNEL_DIR).resolve(), ROOT / "src/CelerisLab/lbm/kernels")
|
|
spec = ref.make_spec(collision="srt")
|
|
maps = ref.build_maps(spec)
|
|
cfg = LBMConfig(nx=16, ny=16, viscosity=.1)
|
|
compile_original = compiler.compile_kernel
|
|
|
|
def audit_compile(*args, **kwargs):
|
|
# Only this private build's entry differs. Production source unchanged.
|
|
entry = Path(kwargs["build_dir"]) / "refinement_audit.cu"
|
|
entry.write_text('#define CELERIS_REFINEMENT_AUDIT\n#include "refinement.cu"\n')
|
|
kwargs["source"] = entry.name
|
|
return compile_original(*args, **kwargs)
|
|
|
|
with CudaContext(0) as context:
|
|
with patch.object(compiler, "compile_kernel", audit_compile):
|
|
backend = PeriodicPatchStepper(cfg, (4, 4, 12, 12), context=context)
|
|
with backend:
|
|
# Production allocations: four DDF arrays and four compact tables.
|
|
device_bytes = 2 * 9 * (16**2 + 24**2) * 4
|
|
p = backend.patch
|
|
device_bytes += sum(a.nbytes for a in (p.ghost_fine, p.ghost_coarse,
|
|
p.coalesce_coarse, p.coalesce_fine))
|
|
audit_bytes = sum((2 + 9) * np.prod(level.shape[1:]) * 4
|
|
for level in backend._levels)
|
|
self.assertLess(device_bytes + audit_bytes, 128 * 1024)
|
|
print(f"device_arrays={device_bytes} audit_arrays={audit_bytes} bytes", flush=True)
|
|
for case in ("rest", "nonequilibrium_seed1"):
|
|
with self.subTest(case=case):
|
|
state = ref.make_state(spec, velocity=(0, 0) if case == "rest" else (.03, -.02))
|
|
if case != "rest":
|
|
rng = np.random.default_rng(1)
|
|
state = ref.State(*(a * (1 + 1e-3 * rng.standard_normal(a.shape))
|
|
for a in (state.coarse_f, state.fine_f)))
|
|
initial = postcollision(state, spec, maps)
|
|
backend.initialize(*initial)
|
|
for level in backend._levels:
|
|
for name in ("collision_counts", "explosion_counts", "incoming_counts"):
|
|
ptr, size = level.module.get_global(name)
|
|
cuda.memset_d8(ptr, 0, size)
|
|
before = totals(backend.download(), maps)
|
|
max_error = max_drift = max_moment_error = max_corner = 0.0
|
|
for step in range(1, 9):
|
|
state, diagnostic = ref.advance_coarse_once(spec, maps, state)
|
|
backend.step()
|
|
actual = backend.download()
|
|
expected = postcollision(state, spec, maps)
|
|
self.assertEqual(backend.step_count, step)
|
|
p.check_synchronization(step, state.tick, phase="postcollision_g")
|
|
for got, want in zip(actual, expected):
|
|
np.testing.assert_array_equal(np.isnan(got), np.isnan(want))
|
|
error = float(np.nanmax(np.abs(got - want)))
|
|
self.assertLessEqual(error, POP_ATOL)
|
|
max_error = max(max_error, error)
|
|
drift = float(np.max(np.abs(totals(actual, maps) - before)))
|
|
moment_error = float(np.max(np.abs(totals(actual, maps) - totals(expected, maps))))
|
|
self.assertLessEqual(drift, MOMENT_ATOL)
|
|
self.assertLessEqual(moment_error, MOMENT_ATOL)
|
|
max_drift, max_moment_error = max(max_drift, drift), max(max_moment_error, moment_error)
|
|
c, f = backend._levels
|
|
np.testing.assert_array_equal(global_array(c.module, "collision_counts", (16, 16)),
|
|
step * maps.coarse_active)
|
|
np.testing.assert_array_equal(global_array(f.module, "collision_counts", (24, 24)),
|
|
2 * step * maps.fine_active)
|
|
np.testing.assert_array_equal(global_array(f.module, "explosion_counts", (24, 24)),
|
|
step * ~maps.fine_active)
|
|
owners = np.zeros((9, 16, 16), np.uint32)
|
|
owners.ravel()[p.coalesce_coarse] = step
|
|
np.testing.assert_array_equal(global_array(c.module, "incoming_counts", owners.shape), owners)
|
|
raw = np.empty(f.shape, np.float32)
|
|
cuda.memcpy_dtoh(raw, f.current)
|
|
np.testing.assert_array_equal(np.isfinite(raw), diagnostic["hop2_valid"])
|
|
self.assertTrue(np.isfinite(raw.ravel()[p.coalesce_fine]).all())
|
|
self.assertEqual(len(diagnostic["corner_packets"]), 8)
|
|
for packet in diagnostic["corner_packets"].values():
|
|
sx, sy = packet["source_coarse_xy"]
|
|
self.assertTrue(maps.coarse_active[sy, sx])
|
|
values = [raw[packet["direction"], y, x]
|
|
for y, x in packet["children_fine_yx"]]
|
|
error = float(np.max(np.abs(np.array(values) - packet["packets"])))
|
|
self.assertLessEqual(error, POP_ATOL)
|
|
max_corner = max(max_corner, error)
|
|
print(f"{case}: steps=8 max_population={max_error:.9g} "
|
|
f"max_leaf_drift={max_drift:.9g} max_leaf_error={max_moment_error:.9g} "
|
|
f"max_corner_packet={max_corner:.9g}; ownership/counts PASS", flush=True)
|
|
|
|
# Fail before any additional time step; no stale synchronized export.
|
|
with patch.object(backend, "_explode", side_effect=RuntimeError("injected launch failure")):
|
|
with self.assertRaisesRegex(RuntimeError, "injected"):
|
|
backend.step()
|
|
with self.assertRaisesRegex(RuntimeError, "failed"):
|
|
backend.download()
|
|
with self.assertRaisesRegex(RuntimeError, "failed"):
|
|
backend.step()
|
|
self.assertFalse(context.is_closed)
|
|
with self.assertRaisesRegex(RuntimeError, "closed"):
|
|
backend.download()
|
|
|
|
# Compile and execute the actual non-audited entry as well: one step.
|
|
with PeriodicPatchStepper(cfg, (4, 4, 12, 12), context=context) as plain:
|
|
rest = postcollision(ref.make_state(spec), spec, maps)
|
|
plain.initialize(*rest)
|
|
plain.step()
|
|
for got, expected in zip(plain.download(), rest):
|
|
np.testing.assert_allclose(got, expected, rtol=0, atol=POP_ATOL, equal_nan=True)
|
|
print("normal-build rest: steps=1 PASS; source/kernel paths pinned", flush=True)
|
|
rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024
|
|
print(f"host_peak_rss={rss:.3f} MiB", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main(verbosity=2)
|