Files
CelerisLab/tests/validation/multigrid/test_spacetime_gpu.py
T

347 lines
19 KiB
Python

"""Explicit host-only checks or admitted tiny GPU audit; never a long trajectory."""
import argparse
from contextlib import nullcontext
from dataclasses import replace
import hashlib
import json
import os
from pathlib import Path
import subprocess
import sys
import tempfile
import time
from types import SimpleNamespace
import unittest
from unittest.mock import Mock
import numpy as np
ROOT = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(ROOT/'src'))
from CelerisLab.config import LBMConfig
import spacetime_gpu as gpu
import run_spacetime_gpu_n64 as trajectory
ref = gpu.ref
RECORDS = []
REPORT = None
OUTPUT = None
GPU_ADMITTED = False
def record(value):
RECORDS.append(value)
if OUTPUT is not None:
ref.write_receipt(OUTPUT,REPORT)
def config(n):
return LBMConfig(nx=n,ny=n,viscosity=.0192,velocity=0.,ddf_shifting=False)
def initial(plan, tagged=True):
arrays = []
for n, mask in ((plan.n,plan.maps.coarse_active),(plan.nf,plan.maps.fine_active)):
a = np.broadcast_to(ref.cpu.W[:,None,None], (9,n,n)).copy()
if tagged:
a *= 1+1e-3*np.random.default_rng(n).standard_normal(a.shape)
a[:,~mask] = np.nan
arrays.append(a)
return ref.cpu.State(*arrays)
class Host(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.plan = ref.Plan(12,(3,3,9,8))
def test_pins_and_supported_config(self):
gpu.check_sources()
for n in (12,16,64):
c,f = gpu.validate_config(config(n))
self.assertAlmostEqual(c.omega,1/.5576)
self.assertAlmostEqual(f.omega,1/.6152)
for change in (dict(ddf_shifting=True),dict(collision='TRT'),dict(nx=128,ny=128),
dict(ny=16),dict(les_enabled=True),dict(velocity=.03)):
with self.subTest(change=change), self.assertRaises(ValueError):
gpu.validate_config(replace(config(12),**change))
def test_packed_weights_and_packet_ownership(self):
for plan in (self.plan,ref.Plan(16,(4,5,11,12))):
t = gpu.pack_plan(plan)
self.assertEqual(t['sources'].shape,(len(plan.receivers),4))
self.assertEqual(np.count_nonzero(t['times']==-1),24)
self.assertEqual(gpu.table_hash(t),gpu.table_hash(gpu.pack_plan(plan)))
for i,g in enumerate(plan.groups):
np.testing.assert_array_equal(t['weights'][i,:,:len(g.samples)],g.weights)
np.testing.assert_allclose(t['weights'][i].sum(axis=1),1,rtol=0,atol=2e-14)
np.testing.assert_array_equal(t['samples'][i,:len(g.samples)],g.q*plan.n**2+g.samples)
for m in (0,1):
self.assertEqual(len(t[f'incoming{m}']),len(plan.outgoing[m]))
def test_duplicate_writer_is_rejected(self):
plan = self.plan
groups = list(plan.groups)
a,b = next((i,k) for i,g in enumerate(groups) for k,e in enumerate(g.events) if e is not None)
event = groups[a].events[b]
other = next(i for i,g in enumerate(groups) if i!=a and g.q==groups[a].q and event[0] in [e[0] for e in g.events if e])
events = list(groups[other].events)
k = next(k for k,e in enumerate(events) if e is not None and e[0]==event[0])
events[k] = event; groups[other] = replace(groups[other],events=tuple(events))
fake = SimpleNamespace(n=plan.n,nf=plan.nf,groups=groups,receivers=plan.receivers,outgoing=plan.outgoing)
with self.assertRaisesRegex(ValueError,'duplicate incoming'):
gpu.pack_plan(fake)
def test_rest_ledger_and_missing_packet_discrimination(self):
p,t = self.plan,gpu.pack_plan(self.plan)
state = initial(p,False)
state = ref.cpu.State(state.coarse_f.astype(np.float32),state.fine_f.astype(np.float32))
g0 = np.broadcast_to(ref.cpu.W.astype(np.float32)[:,None,None],(9,p.n,p.n)).copy()
pre = np.repeat(g0.ravel()[t['origins'],None].astype(float),4,axis=1)
h = pre.astype(np.float32); registers = np.zeros(len(p.receivers))
phases = [(state.fine_f,state.fine_f)]*2
args = (p,t,state,state,g0,g0,phases,pre,h)
ledger = gpu.rounding_ledger(*args,registers)
self.assertEqual(max(map(abs,ledger['unexplained'])),0)
registers[0] = .01
with self.assertRaisesRegex(AssertionError,'register mismatch'):
gpu.rounding_ledger(*args,registers)
def test_second_substep_characteristic_donor(self):
p = ref.Plan(16,(4,4,12,12)); t = gpu.pack_plan(p)
row = t['incoming1'][t['incoming1'][:,0]==5*32*32+10*32+7][0]
self.assertEqual(t['origins'][row[1]//4],5*16*16+4*16+3)
def test_failed_free_remains_retryable(self):
run=gpu.SpaceTimeGPU.__new__(gpu.SpaceTimeGPU)
run.ctx=SimpleNamespace(activate=nullcontext); run.stream=Mock()
first,second=Mock(),Mock(); first.free.side_effect=[RuntimeError('free failed'),None]
run.allocations=[first,second]; run.modules=[]; run.status='ready'
with self.assertRaisesRegex(RuntimeError,'free failed'):
run.close()
self.assertEqual(run.allocations,[first]); self.assertEqual(run.status,'cleanup_failed')
run.close()
self.assertEqual(run.status,'closed'); self.assertEqual(run.allocations,[])
second.free.assert_called_once()
def test_unmeasured_schedule_and_health_failure(self):
run=gpu.SpaceTimeGPU.__new__(gpu.SpaceTimeGPU)
run.status='ready'; run.step_count=0; run.identity=False; run.mode='spacetime'
run.ctx=SimpleNamespace(activate=nullcontext); run.cuda=Mock(); run.stream=Mock()
run.plan=self.plan; run.bounds=(3,3,9,8)
run.configs=gpu.validate_config(config(12)); run.coarse=list('abcd'); run.fine=list('efg')
run.tables=gpu.pack_plan(self.plan); run.device_tables={key:key for key in run.tables}
run.active='active'; run.pre='pre'; run.h='h'; run.registers='registers'; run.counts='counts'
run.health_counts=np.zeros(6,np.uint64)
run.launch=Mock(); run.read=Mock(return_value=np.zeros(6,np.uint32)); run.download=Mock()
names=['CoarseStart','Pull','Collide','Interpolate','Collide','Register','Inject','Pull',
'Collide','Register','Inject','Pull','ApplyRegisters']
for _ in range(2):
result=run.step(audit=False)
self.assertEqual(result['boundary_budget_status'],'not_measured')
self.assertIsNone(result['unexplained'])
self.assertEqual([call.args[1] for call in run.launch.call_args_list],names*2)
self.assertEqual(run.read.call_count,2); run.download.assert_not_called()
self.assertEqual(run.step_count,2); self.assertEqual(run.coarse,list('abcd'))
run.read.return_value=np.array([0,0,0,0,1,0],np.uint32)
with self.assertRaisesRegex(AssertionError,'health failed'):
run.step(audit=False)
self.assertEqual(run.status,'failed'); self.assertEqual(run.step_count,2)
self.assertEqual(run.health_counts.tolist(),[0,0,0,0,1,0])
def test_warmup_resets_arrivals_and_counters(self):
class Fake:
def initialize(self,state):
self.state=ref.cpu.State(state.coarse_f.copy(),state.fine_f.copy())
self.step_count=0; self.health_counts=np.zeros(6,np.uint64)
def download(self):
return ref.cpu.State(self.state.coarse_f.copy(),self.state.fine_f.copy(),tick=2*self.step_count)
def step(self,*,audit):
assert audit is False
self.step_count+=1; self.state.coarse_f[:]+=1e-4
fake=Fake(); check=Mock(); state=initial(self.plan)
actual,warm=trajectory.warm_and_reset(fake,state,check)
np.testing.assert_array_equal(actual.coarse_f,state.coarse_f)
self.assertEqual(check.call_count,32); self.assertEqual(fake.step_count,0)
self.assertTrue(warm['reset_verified']); self.assertEqual(warm['intervals'],32)
def test_trajectory_screens_differences_and_json(self):
rows=[dict(tstar=t,velocity_l2=np.float64(.01),pressure_l2=.03,energy_relative=.01,
mass_relative=1e-7,momentum_change=[1e-6,-2e-6],minimum_population=.02,
minimum_boundary_pre=.02,minimum_boundary_stored=.02) for t in trajectory.TIMES]
cpu=dict(runs=[dict(worst=dict(velocity_l2=.02,pressure_l2=.04,energy_relative=.02)),
dict(rows=[dict(tstar=t,velocity_l2=.009,momentum_change=[0.,0.]) for t in trajectory.TIMES])])
uniform=dict(rows=[dict(tstar=.12,pressure_l2=.012253173806349701)])
result=trajectory.assess(rows,cpu,uniform,[0]*6)
self.assertTrue(result['campaign_pass']); self.assertTrue(result['experimental_pass'])
self.assertFalse(result['historical_pressure_gate']['passed'])
self.assertFalse(result['historical_pressure_gate']['matched_storage'])
self.assertAlmostEqual(result['gpu_minus_cpu']['rows'][0]['fields']['velocity_l2']['signed'],.001)
with tempfile.TemporaryDirectory() as directory:
path=Path(directory)/'receipt.json'
ref.write_receipt(path,dict(completed=True,**result))
self.assertTrue(json.loads(path.read_text())['completed'])
rows[0]['velocity_l2']=.06
failed=trajectory.assess(rows,cpu,uniform,[0,1,0,0,0,0])
ref.write_receipt(path,dict(completed=True,**failed))
saved=json.loads(path.read_text())
self.assertFalse(saved['campaign_pass']); self.assertFalse(saved['experimental_pass'])
for invalid in (float('nan'),-.01):
rows[0]['velocity_l2']=.01; rows[-1]['velocity_l2']=invalid
failed=trajectory.assess(rows,cpu,uniform,[0]*6)
self.assertFalse(failed['campaign_pass']); self.assertFalse(failed['experimental_pass'])
path.write_text('changed reference')
with self.assertRaisesRegex(ValueError,'pinned reference'):
trajectory.load_references(path,path)
class Device(unittest.TestCase):
@classmethod
def setUpClass(cls):
if not GPU_ADMITTED:
raise unittest.SkipTest('CUDA compilation/execution requires explicit tiny-job admission')
from CelerisLab.cuda.context import CudaContext
cls.context = CudaContext(0)
@classmethod
def tearDownClass(cls):
cls.context.close()
def test_tiny_equivalence_budgets_and_phase(self):
for mode,identity in (('spacetime',False),('spacetime',True),('frozen',False),('frozen',True)):
with self.subTest(mode=mode,identity=identity), gpu.SpaceTimeGPU(
config(12),(3,3,9,8),context=self.context,mode=mode,identity=identity) as run:
spec = ref.Spec(n_coarse=12,patch_lo=1,halo=2,tau_coarse=.5576,tau_fine=.6152,
collision='identity' if identity else 'srt')
for tagged,steps in ((False,2),(True,8)):
expected = initial(run.plan,tagged); run.initialize(expected)
before = run.download()
for k in range(1,steps+1):
expected,_ = ref.advance(spec,run.plan,expected,mode)
ledger = run.step(); actual = run.download()
self.assertEqual(actual.tick,2*k); self.assertEqual(run.step_count,k)
for a,b in zip((actual.coarse_f,actual.fine_f),(expected.coarse_f,expected.fine_f)):
np.testing.assert_array_equal(np.isnan(a),np.isnan(b))
np.testing.assert_allclose(a,b,rtol=0,atol=3e-6,equal_nan=True)
total = lambda s: ref.cpu.MOMENTS.T @ ref.directional(
ref.cpu.State(s.coarse_f.astype(float),s.fine_f.astype(float)),run.plan.maps)
np.testing.assert_allclose(total(actual),total(before),rtol=0,atol=2e-4)
if tagged and not identity and k==1:
for a,b,tau,mask in zip((actual.coarse_f,actual.fine_f),
(expected.coarse_f,expected.fine_f),(.5576,.6152),
(run.plan.maps.coarse_active,run.plan.maps.fine_active)):
wrong = ref.cpu._collide(spec,b,np.isfinite(b),tau,mask,np.zeros(mask.shape,int))
self.assertGreater(float(np.nanmax(abs(a-wrong))),3e-6)
if k in (1,2,8):
record(dict(case='tiny',mode=mode,identity=identity,tagged=tagged,
step=k,device_bytes=run.device_bytes,ledger=ledger))
def test_omitted_packets_and_ghosts_fail(self):
for mutation in (1,2,3,4):
with self.subTest(mutation=mutation), gpu.SpaceTimeGPU(
config(12),(3,3,9,8),context=self.context) as run:
run.initialize(initial(run.plan))
with self.assertRaises(AssertionError):
run.step(mutation=mutation)
self.assertEqual(run.step_count,0)
self.assertEqual(run.status,'failed')
with self.assertRaises(RuntimeError):
run.download()
def test_frozen_donor_trap(self):
with gpu.SpaceTimeGPU(config(16),(4,4,12,12),context=self.context,mode='frozen',identity=True) as run:
state = initial(run.plan,False)
state.coarse_f[5,4,3]=1.; state.coarse_f[5,5,3]=2.
run.initialize(state); run.step()
self.assertEqual(run.download().fine_f[5,11,8],1.)
# Deliberately negative boundary data are reported and remain unclipped.
with self.context.activate():
data=np.full((9,16,16),-.01,np.float32)
run.cuda.memcpy_htod(run.coarse[1],data); run.cuda.memcpy_htod(run.coarse[3],data)
run.cuda.memset_d32_async(run.counts,0,6,run.stream)
t=run.device_tables; ng=len(run.plan.groups)
run.launch(0,'Interpolate',run.coarse[1],run.coarse[3],t['samples'],t['origins'],
t['sizes'],t['weights'],run.pre,run.h,np.uint32(ng),np.uint32(0),run.counts,count=ng)
counts=run.read(run.counts,(6,),np.uint32)
self.assertEqual(counts[1],4*ng); self.assertEqual(counts[3],4*ng)
self.assertTrue((run.read(run.h,(ng,4))<0).all())
def test_n64_one_interval_only(self):
sys.path.insert(0,str(Path(__file__).parents[1]/'taylor_green'))
import oracle as tg
with gpu.SpaceTimeGPU(config(64),(8,8,56,48),context=self.context) as run:
arrays = [tg.initial((64,64),0,64,.0192),tg.initial((128,128),1,64,.0192)]
for a,mask in zip(arrays,(run.plan.maps.coarse_active,run.plan.maps.fine_active)):
a[:,~mask]=np.nan
state = ref.cpu.State(*arrays); run.initialize(state)
spec = ref.Spec(n_coarse=64,patch_lo=8,halo=16,tau_coarse=.5576,tau_fine=.6152,collision='srt')
expected,_ = ref.advance(spec,run.plan,state)
started=time.monotonic(); ledger=run.step(); seconds=time.monotonic()-started
actual=run.download()
errors=[]
for a,b in zip((actual.coarse_f,actual.fine_f),(expected.coarse_f,expected.fine_f)):
np.testing.assert_array_equal(np.isnan(a),np.isnan(b))
np.testing.assert_allclose(a,b,rtol=0,atol=3e-6,equal_nan=True)
errors.append(float(np.nanmax(abs(a-b))))
record(dict(case='N64_one_interval',errors=errors,ledger=ledger,seconds=seconds,
device_bytes=run.device_bytes,table_sha256=gpu.table_hash(run.tables),
build_metadata=run.build_metadata,device=self.context.name))
def test_unmeasured_path_and_warm_reset(self):
with gpu.SpaceTimeGPU(config(12),(3,3,9,8),context=self.context) as run:
state=initial(run.plan)
run.initialize(state)
for _ in range(2):
run.step()
expected=run.download()
run.initialize(state)
for _ in range(2):
result=run.step(audit=False)
self.assertEqual(result['boundary_budget_status'],'not_measured')
actual=run.download()
for a,b in zip((actual.coarse_f,actual.fine_f),(expected.coarse_f,expected.fine_f)):
np.testing.assert_array_equal(a,b)
reset,warm=trajectory.warm_and_reset(run,state,lambda:None)
for a,b in zip((reset.coarse_f,reset.fine_f),(state.coarse_f,state.fine_f)):
np.testing.assert_array_equal(a,b.astype(np.float32))
record(dict(case='unmeasured_bitwise_and_reset',warmup=warm,health_counts=run.health_counts.tolist()))
def main():
global REPORT,OUTPUT,GPU_ADMITTED
parser=argparse.ArgumentParser(description=__doc__)
mode=parser.add_mutually_exclusive_group(required=True)
mode.add_argument('--host-only',action='store_true')
mode.add_argument('--gpu-admitted',action='store_true')
parser.add_argument('--output',type=Path)
args=parser.parse_args()
GPU_ADMITTED=args.gpu_admitted
if args.gpu_admitted and (not os.environ.get('CUDA_VISIBLE_DEVICES') or args.output is None):
parser.error('GPU job requires explicit CUDA_VISIBLE_DEVICES and --output')
if args.output is not None and args.output.exists():
parser.error('output exists; retain evidence and choose a new path')
started=time.monotonic()
report=dict(scope='tiny GPU equivalence and rounding audit; no long trajectory',completed=False,
requested_mode='admitted_tiny_gpu' if args.gpu_admitted else 'host_only',
commit=subprocess.check_output(['git','rev-parse','HEAD'],cwd=ROOT,text=True).strip(),
dirty=bool(subprocess.check_output(['git','status','--porcelain'],cwd=ROOT,text=True)),
hashes={p.name:hashlib.sha256(p.read_bytes()).hexdigest() for p in
(Path(__file__),Path(gpu.__file__),Path(gpu.__file__).with_suffix('.cu'))},
reference_hashes=gpu.PINS,cuda_visible_devices=os.environ.get('CUDA_VISIBLE_DEVICES'),records=RECORDS)
REPORT,OUTPUT=report,args.output
if args.output:
ref.write_receipt(args.output,report)
suite=unittest.TestLoader().loadTestsFromTestCase(Host)
if args.gpu_admitted:
suite.addTests(unittest.TestLoader().loadTestsFromTestCase(Device))
result=unittest.TextTestRunner(verbosity=2).run(suite)
report.update(completed=True,passed=result.wasSuccessful(),tests=result.testsRun,
failures=[dict(test=str(test),error=error) for test,error in result.failures+result.errors],
seconds=time.monotonic()-started)
if args.output:
ref.write_receipt(args.output,report)
return 0 if result.wasSuccessful() else 1
if __name__=='__main__':
raise SystemExit(main())