197 lines
10 KiB
Python
197 lines
10 KiB
Python
"""Production host dispatch with CUDA doubles; no CUDA compilation or flow runs."""
|
|
from contextlib import nullcontext
|
|
from dataclasses import replace
|
|
from unittest.mock import Mock
|
|
import unittest
|
|
|
|
import numpy as np
|
|
|
|
from CelerisLab.config import LBMConfig
|
|
from tests.unit.test_hierarchy import CFG, FOUR
|
|
from tests.unit.test_refinement_three_level_multipatch import host_module
|
|
|
|
|
|
def backend(**options):
|
|
host = host_module()
|
|
calls, allocations, storage = [], [], {}
|
|
context = Mock(activate=nullcontext, sm_arch='sm_70')
|
|
host.cuda.mem_alloc.side_effect = lambda size: allocate(size)
|
|
|
|
def allocate(size):
|
|
pointer = Mock(name=f'buffer{len(allocations)}', size=size)
|
|
allocations.append(pointer)
|
|
return pointer
|
|
|
|
def load(_):
|
|
module = Mock()
|
|
module.get_global.return_value = (Mock(), 20)
|
|
module.get_function.side_effect = lambda name: Mock(
|
|
name=name, side_effect=lambda *a, **kw: calls.append((name, a, kw)))
|
|
return module
|
|
|
|
host.compiler.build_module.side_effect = lambda *a, **k: load(None)
|
|
host.cuda.memcpy_htod.side_effect = lambda ptr, a: storage.__setitem__(id(ptr), np.array(a, copy=True))
|
|
host.cuda.memcpy_dtoh.side_effect = lambda a, ptr: np.copyto(a, storage[id(ptr)])
|
|
cfg = options.pop('cfg', CFG)
|
|
refinements = options.pop('refinements', FOUR)
|
|
blocks = options.pop('level_blocks', [(4, 4)] * (len(refinements) + 1))
|
|
run = host.FixedPatchStepper(cfg, context=context, refinements=refinements,
|
|
level_blocks=blocks, **options)
|
|
run.initialize(*(np.full(level.shape, 1/9, np.float32) for level in run._levels))
|
|
return run, host, calls, allocations, storage
|
|
|
|
|
|
def trace(run, calls, allocations, storage):
|
|
"""Stable pointer identities expose allocations/arguments, not just kernel names."""
|
|
pointer = {id(value): index for index, value in enumerate(allocations)}
|
|
launches = [(name, tuple(('ptr', pointer[id(a)]) if id(a) in pointer else int(a)
|
|
for a in args), kw['grid'], kw['block'])
|
|
for name, args, kw in calls]
|
|
uploads = [(index, value.dtype.str, value.shape, value.tobytes())
|
|
for index, ptr in enumerate(allocations) if (value := storage.get(id(ptr))) is not None]
|
|
return [p.size for p in allocations], uploads, launches
|
|
|
|
|
|
class ProjectedPairDispatch(unittest.TestCase):
|
|
def test_omitted_and_constant_have_identical_allocations_uploads_and_launches(self):
|
|
results = []
|
|
for options in ({}, {'explosion': 'constant'}):
|
|
run, host, calls, allocations, storage = backend(**options)
|
|
self.assertEqual(run._explosion, 'constant')
|
|
self.assertTrue(all(p.neighbors is None and p.offsets is None for p in run._pairs))
|
|
run.step()
|
|
explosion = [args for name, args, _ in calls if name == 'BlockExplode']
|
|
self.assertEqual(len(explosion), 7)
|
|
self.assertTrue(all(len(args) == 6 for args in explosion))
|
|
self.assertFalse(any(name == 'BlockExplodeProjected' for name, _, _ in calls))
|
|
results.append(trace(run, calls, allocations, storage))
|
|
run.close()
|
|
self.assertEqual(*results)
|
|
|
|
def test_projected_adds_exactly_two_tables_per_adjacent_pair(self):
|
|
for depth in (2, 3, 4):
|
|
constant, _, _, ca, _ = backend(refinements=FOUR[:depth-1])
|
|
projected, _, calls, pa, storage = backend(
|
|
refinements=FOUR[:depth-1], explosion='projected_linear')
|
|
self.assertEqual(len(pa)-len(ca), 2*(depth-1))
|
|
for pair, geometry in zip(projected._pairs, projected.hierarchy.pairs):
|
|
self.assertEqual(storage[id(pair.neighbors)].shape, (len(geometry.ghost_fine), 4))
|
|
self.assertEqual(storage[id(pair.offsets)].shape, (len(geometry.ghost_fine), 2))
|
|
self.assertEqual(pair.neighbors.size, len(geometry.ghost_fine)*4*8)
|
|
self.assertEqual(pair.offsets.size, len(geometry.ghost_fine)*2*4)
|
|
projected.step()
|
|
explosions = [args for name, args, _ in calls if name == 'BlockExplodeProjected']
|
|
self.assertEqual(len(explosions), 2**(depth-1)-1)
|
|
self.assertTrue(all(len(args) == 8 for args in explosions))
|
|
self.assertFalse(any(name == 'BlockExplode' for name, _, _ in calls))
|
|
for args in explosions:
|
|
pair = next(p for p in projected._pairs if args[2] is p.ghosts)
|
|
self.assertIs(args[6], pair.neighbors)
|
|
self.assertIs(args[7], pair.offsets)
|
|
self.assertEqual(args[5], pair.parent_cells)
|
|
constant.close(); projected.close()
|
|
|
|
def test_three_and_four_level_schedule_has_fresh_midpoint_sources(self):
|
|
for depth in (3, 4):
|
|
for explosion in ('constant', 'projected_linear'):
|
|
run, host, calls, allocations, storage = backend(
|
|
refinements=FOUR[:depth-1], explosion=explosion)
|
|
before = len(allocations)
|
|
events = []
|
|
for index, pair in enumerate(run._pairs):
|
|
pair.explode.side_effect = lambda source, *a, index=index, **kw: events.append(('explode', index, source))
|
|
pair.coalesce.side_effect = lambda child, source, *a, index=index, **kw: events.append(('restrict', index, source))
|
|
for index, level in enumerate(run._levels):
|
|
level.kernel.side_effect = lambda source, *a, index=index, **kw: events.append(('step', index, source))
|
|
level.halo_copy.side_effect = lambda source, *a, index=index, **kw: events.append(('halo', index, source))
|
|
expected = []
|
|
|
|
def hop(index, source):
|
|
if index+1 < depth:
|
|
expected.append(('explode', index, source))
|
|
child = run._levels[index+1]
|
|
hop(index+1, child.current)
|
|
hop(index+1, child.spare)
|
|
expected.append(('halo', index, source))
|
|
if index+1 < depth:
|
|
expected.append(('restrict', index, source))
|
|
expected.append(('step', index, source))
|
|
|
|
hop(0, run._levels[0].current)
|
|
run.step()
|
|
self.assertEqual(events, expected)
|
|
self.assertEqual([level.kernel.call_count for level in run._levels], [2**i for i in range(depth)])
|
|
self.assertEqual(len(allocations), before)
|
|
self.assertEqual(run.step_count, 1)
|
|
for index in range(1, depth-1):
|
|
feeds = [source for name, pair, source in events if name == 'explode' and pair == index]
|
|
self.assertEqual(feeds, [run._levels[index].current, run._levels[index].spare] * (2**(index-1)))
|
|
run.close()
|
|
|
|
def test_scalar_windows_follow_the_same_pair_path(self):
|
|
run, _, _, _, _ = backend(cfg=LBMConfig(nx=16,ny=16,viscosity=.02),
|
|
refinements=[(4,4,12,12),(8,8,16,16),(8,8,16,16)],
|
|
level_blocks=[(4,4),(4,8),(8,4),(4,4)], explosion='projected_linear')
|
|
self.assertTrue(all(pair.neighbors is not None for pair in run._pairs))
|
|
self.assertEqual([pair.parent_cells for pair in run._pairs], [36,60,60])
|
|
run.step()
|
|
self.assertEqual([level.kernel.call_count for level in run._levels], [1,2,4,8])
|
|
run.close()
|
|
|
|
def test_physics_gate_precedes_context_and_allocations(self):
|
|
host = host_module()
|
|
context = Mock()
|
|
options = dict(refinements=FOUR[:1], level_blocks=[(4,4)]*2,
|
|
explosion='projected_linear', context=context)
|
|
for cfg, changes in [(CFG, {'explosion':'unknown'}),
|
|
(CFG, {'coarse_boundary':'slot_reservoir'}),
|
|
(CFG, {'coarse_boundary':'channel'}),
|
|
(CFG, {'level_objects':[(),(object(),)]}),
|
|
(CFG, {'level_jets':[(),(object(),)]}),
|
|
(replace(CFG,collision='MRT'), {}),
|
|
(replace(CFG,les_enabled=True), {})]:
|
|
with self.subTest(changes=changes, cfg=cfg), self.assertRaisesRegex(ValueError,'explosion'):
|
|
host.FixedPatchStepper(cfg, **dict(options, **changes))
|
|
context.activate.assert_not_called()
|
|
host.cuda.mem_alloc.assert_not_called()
|
|
# The constant path still allows MRT.
|
|
for options in ({'cfg':replace(CFG,collision='MRT')},):
|
|
run, _, _, _, _ = backend(**options)
|
|
run.close()
|
|
|
|
def test_remap_preserves_selector_tables_and_constant_default(self):
|
|
for explosion in ('constant','projected_linear'):
|
|
run, host, _, _, storage = backend(refinements=FOUR[:2], explosion=explosion)
|
|
run._step_count = 7
|
|
old = run.download()
|
|
clone = run.remap_finest(FOUR[1])
|
|
self.assertEqual(clone._explosion, explosion)
|
|
self.assertEqual(clone.step_count, 7)
|
|
for pair in clone._pairs:
|
|
self.assertEqual(pair.neighbors is None, explosion == 'constant')
|
|
for expected, actual in zip(old, clone.download()):
|
|
np.testing.assert_array_equal(actual,expected)
|
|
clone.step()
|
|
self.assertEqual((run.step_count,clone.step_count), (7,8))
|
|
clone.close(); run.close()
|
|
|
|
def test_changed_remap_geometry_keeps_projected_in_both_representations(self):
|
|
changed = [(22,22,42,42), (72,68,76,76)]
|
|
for shifted in (False, True):
|
|
run, _, _, _, _ = backend(cfg=replace(CFG,ddf_shifting=shifted),
|
|
refinements=FOUR[:2], explosion='projected_linear')
|
|
run._step_count = 3
|
|
before = run.download()
|
|
clone = run.remap_finest(changed)
|
|
self.assertEqual(clone._explosion, 'projected_linear')
|
|
self.assertEqual(clone.hierarchy.pairs[1].bounds, tuple(changed))
|
|
self.assertNotEqual(clone._pairs[1].n_ghosts, run._pairs[1].n_ghosts)
|
|
for old, new in zip(run._pairs, clone._pairs):
|
|
self.assertIsNot(old.neighbors, new.neighbors)
|
|
self.assertIsNotNone(new.neighbors)
|
|
for expected, actual in zip(before, run.download()):
|
|
np.testing.assert_array_equal(actual, expected)
|
|
clone.step()
|
|
self.assertEqual((run.step_count, clone.step_count), (3, 4))
|
|
clone.close(); run.close()
|