185 lines
9.6 KiB
Python
185 lines
9.6 KiB
Python
"""CPU host contracts; mocked CUDA executes real lifecycle, not numerical kernels."""
|
|
import ast
|
|
import types
|
|
from pathlib import Path
|
|
from contextlib import nullcontext
|
|
from dataclasses import replace
|
|
import unittest
|
|
from unittest.mock import Mock, patch
|
|
import numpy as np
|
|
from CelerisLab.lbm.population_storage import (_StoredPopulations, _WEIGHTS,
|
|
_encode_physical, _decode_stored)
|
|
from CelerisLab.lbm.synchronized_remap import _remap_stored, remap_synchronized
|
|
from tests.unit.test_synchronized_remap import topology, A, B
|
|
from tests.unit import test_refinement_three_level_multipatch as backend
|
|
|
|
|
|
def shifted_fields(mesh):
|
|
rng = np.random.default_rng(129)
|
|
arrays = []
|
|
for leaf in mesh.leaf_masks:
|
|
a = rng.normal(0, 1e-7, (9, *leaf.shape)).astype(np.float32)
|
|
a[:, ~leaf] = np.nan
|
|
y, x = np.nonzero(leaf)
|
|
a[0, y[:6], x[:6]] = np.array([0., -0., 2**-80, -2**-80, 2**-140, -2**-140], np.float32)
|
|
arrays.append(a)
|
|
return tuple(arrays)
|
|
|
|
|
|
def remap(old, new, arrays, shifted=True):
|
|
return _remap_stored(old, new, _StoredPopulations(arrays, shifted), shifted=True,
|
|
base_count=7, level_steps=(7, 14, 28), phase='postcollision_g')
|
|
|
|
|
|
def ready():
|
|
with patch.object(backend, 'CFG', replace(backend.CFG, ddf_shifting=True)):
|
|
run, module, allocations = backend.mock_backend(children=A)
|
|
storage = {}
|
|
module.cuda.memcpy_htod.side_effect = lambda p, a: storage.__setitem__(id(p), np.array(a, copy=True))
|
|
module.cuda.memcpy_dtoh.side_effect = lambda a, p: np.copyto(a, storage[id(p)])
|
|
run._initialize_raw(_StoredPopulations(shifted_fields(topology(A)), True))
|
|
run._step_count = 7
|
|
return run, module, allocations, storage
|
|
|
|
|
|
class ShiftedStorage(unittest.TestCase):
|
|
def test_codec_preserves_sub_fp32_physical_perturbation(self):
|
|
physical = np.broadcast_to(_WEIGHTS, (9, 2, 2)).copy()
|
|
physical[1] += 2**-35
|
|
encoded = _encode_physical(physical, True)
|
|
self.assertTrue((encoded[1] == np.float32(2**-35)).all())
|
|
np.testing.assert_array_equal(_decode_stored(encoded, True), physical)
|
|
self.assertEqual(_decode_stored(encoded, True).dtype, np.float64)
|
|
np.testing.assert_array_equal(_encode_physical(_WEIGHTS, True), np.zeros((9,1,1), np.float32))
|
|
self.assertEqual(_encode_physical(physical, False).dtype, np.float32)
|
|
|
|
def test_actual_fill_loop_bitwise_copy_and_single_rounding_restriction(self):
|
|
old, new = topology(A), topology(B)
|
|
arrays = shifted_fields(old)
|
|
saved = [a.tobytes() for a in arrays]
|
|
out = remap(old, new, arrays)
|
|
self.assertTrue(out.shifted)
|
|
for a, b, om, nm, before in zip(arrays, out.g, old.leaf_masks, new.leaf_masks, saved):
|
|
self.assertEqual(a[:, om & nm].tobytes(), b[:, om & nm].tobytes())
|
|
self.assertTrue(np.isnan(b[:, ~nm]).all())
|
|
self.assertEqual(a.tobytes(), before)
|
|
self.assertFalse(np.shares_memory(a,b))
|
|
y,x = np.nonzero(new.leaf_masks[2] & ~old.leaf_masks[2])
|
|
self.assertEqual(out.g[2][:,y,x].tobytes(), arrays[1][:,y//2,x//2].tobytes())
|
|
y,x = np.nonzero(new.leaf_masks[1] & ~old.leaf_masks[1])
|
|
mean = sum(arrays[2][:,2*y+dy,2*x+dx].astype(np.float64)*.25
|
|
for dy in (0,1) for dx in (0,1))
|
|
self.assertEqual(out.g[1][:,y,x].tobytes(), mean.astype(np.float32).tobytes())
|
|
def integral(mesh, values):
|
|
return sum(a[:,m].sum(axis=1,dtype=np.float64)/4**l
|
|
for l,(a,m) in enumerate(zip(values,mesh.leaf_masks)))
|
|
cast_error = abs(mean.astype(np.float32).astype(np.float64)-mean).sum(axis=1)/4
|
|
self.assertTrue(np.all(abs(integral(old, arrays)-integral(new,out.g)) <= cast_error+2e-18))
|
|
|
|
def test_shifted_density_and_tag_validation(self):
|
|
old,new = topology(A),topology(B)
|
|
a = tuple(np.where(m[None], np.float32(-.001), np.nan).repeat(9,axis=0) for m in old.leaf_masks)
|
|
remap(old,new,a)
|
|
with self.assertRaisesRegex(ValueError,'representation'): remap(old,new,a,False)
|
|
with self.assertRaisesRegex(ValueError,'FP32'): remap(old,new,tuple(x.astype(np.float64) for x in a))
|
|
a[0][:,old.leaf_masks[0]] = -1
|
|
with self.assertRaisesRegex(ValueError,'positive density'): remap(old,new,a)
|
|
rest = tuple(np.where(m[None],np.float32(0),np.nan).repeat(9,axis=0) for m in old.leaf_masks)
|
|
remap(old,new,rest)
|
|
with self.assertRaisesRegex(ValueError,'positive density'):
|
|
remap_synchronized(old,new,rest,base_count=7,level_steps=(7,14,28),phase='postcollision_g')
|
|
|
|
def test_runtime_remap_bypasses_public_codec_and_preserves_both_buffers(self):
|
|
run,module,allocations,storage = ready()
|
|
clone = None
|
|
try:
|
|
before = run._download_raw()
|
|
expected = remap(topology(A),topology(B),before.values)
|
|
with patch.object(run,'download',side_effect=AssertionError('public decode used')):
|
|
clone = run.remap_child(B)
|
|
self.assertEqual((run.step_count,clone.step_count),(7,7))
|
|
self.assertFalse(clone._exchange_valid)
|
|
for level, want, got in zip(clone._levels,expected.g,clone._download_raw().values):
|
|
self.assertEqual(want.tobytes(),got.tobytes())
|
|
for pointer in (level.current,level.spare):
|
|
np.testing.assert_array_equal(storage[id(pointer)],level.layout.to_blocked(want))
|
|
for a,b in zip(before.values,run._download_raw().values): self.assertEqual(a.tobytes(),b.tobytes())
|
|
for a in clone.download(): self.assertEqual(a.dtype,np.float64)
|
|
finally:
|
|
if clone is not None: clone.close()
|
|
run.close()
|
|
|
|
def test_runtime_public_encode_and_raw_mismatch_fail_before_upload(self):
|
|
run,module,_,_ = ready()
|
|
try:
|
|
physical = tuple(np.broadcast_to(_WEIGHTS, level.shape).copy() for level in run._levels)
|
|
for a in physical: a[1] += 2**-35
|
|
run.initialize(*physical)
|
|
for a,mask in zip(run._download_raw().values,run._leaf):
|
|
self.assertTrue((a[1,mask] == np.float32(2**-35)).all())
|
|
self.assertEqual(run.step_count,0)
|
|
module.cuda.memcpy_htod.reset_mock()
|
|
with self.assertRaisesRegex(ValueError,'representation'):
|
|
run._initialize_raw(_StoredPopulations(run._download_raw().values,False))
|
|
module.cuda.memcpy_htod.assert_not_called()
|
|
run._level_cfgs[0] = replace(run._level_cfgs[0],ddf_shifting=False)
|
|
with self.assertRaisesRegex(ValueError,'mixed'): run._download_raw()
|
|
finally: run.close()
|
|
|
|
def test_dense_or_mixed_shifted_rejected_before_compilation(self):
|
|
module=backend.host_module()
|
|
cfg=replace(backend.CFG,ddf_shifting=True)
|
|
for kwargs in ({},dict(base_block=(4,4)),dict(fine_block=(4,4))):
|
|
with self.assertRaisesRegex(ValueError,'all-blocked|base_block requires fine_block'):
|
|
module.FixedPatchStepper(cfg,(4,4,20,28),context=None,**kwargs)
|
|
module.compiler.compile_kernel.assert_not_called()
|
|
|
|
def test_uniform_block_codec_and_raw_snapshot(self):
|
|
path = Path(backend.r.__file__).with_name('block_stepper.py')
|
|
tree = ast.parse(path.read_text())
|
|
tree.body = [n for n in tree.body if not (
|
|
isinstance(n, ast.Import) and any(a.name == 'pycuda.driver' for a in n.names)
|
|
or isinstance(n, ast.ImportFrom) and n.module == 'cuda')]
|
|
module = types.ModuleType('CelerisLab.lbm._test_block_host')
|
|
module.__package__ = 'CelerisLab.lbm'
|
|
cuda, compiler = Mock(), Mock()
|
|
module.__dict__.update(cuda=cuda,compiler=compiler)
|
|
exec(compile(tree,str(path),'exec'),module.__dict__)
|
|
cuda.mem_get_info.return_value=(10**9,10**9)
|
|
cuda.mem_alloc.side_effect=lambda size: Mock()
|
|
gpu=Mock();gpu.get_global.return_value=(Mock(),20)
|
|
compiler.load_module.return_value=gpu
|
|
context=Mock(activate=nullcontext,sm_arch='sm_70')
|
|
layout=module.BlockLayout2D.build(16,16,4,4)
|
|
cfg=replace(backend.CFG,nx=16,ny=16,ddf_shifting=True)
|
|
run=module.PeriodicBlockStepper(cfg,layout,context=context)
|
|
storage={}
|
|
cuda.memcpy_htod.side_effect=lambda p,a: storage.__setitem__(id(p),np.array(a,copy=True))
|
|
cuda.memcpy_dtoh.side_effect=lambda a,p: np.copyto(a,storage[id(p)])
|
|
try:
|
|
physical=np.broadcast_to(_WEIGHTS,(9,16,16)).copy();physical[1]+=2**-35
|
|
run.initialize(physical)
|
|
raw=run._download_raw()
|
|
self.assertTrue(raw.shifted)
|
|
self.assertEqual(len(raw.values),1)
|
|
self.assertTrue((raw.values[0][1]==np.float32(2**-35)).all())
|
|
np.testing.assert_array_equal(run.download(),physical)
|
|
finally: run.close()
|
|
|
|
def test_passive_shifted_body_raw_remap_preserves_support(self):
|
|
from tests.unit.test_pinned_body_remap import mesh, body, A as rects, B as moved
|
|
old,new=mesh(rects,(body(),)),mesh(moved,(body(),))
|
|
result=remap(old,new,shifted_fields(old))
|
|
self.assertTrue(result.shifted)
|
|
self.assertTrue(np.isnan(result.g[2][:,~new.leaf_masks[2]]).all())
|
|
changed=mesh(moved,(body(center=(12.1,14.)),))
|
|
with self.assertRaisesRegex(ValueError,'boundary'):
|
|
remap(old,changed,shifted_fields(old))
|
|
|
|
def test_fp64_physical_roundtrip_cannot_replace_raw(self):
|
|
tiny=np.full((9,1,1),np.float32(2**-80))
|
|
self.assertFalse(np.array_equal(tiny,_encode_physical(_decode_stored(tiny,True),True)))
|
|
|
|
|
|
if __name__=='__main__': unittest.main()
|