205 lines
10 KiB
Python
205 lines
10 KiB
Python
"""CPU stencil and actual closed-interface packet budgets; no CUDA claims."""
|
|
import numpy as np
|
|
import unittest
|
|
from pathlib import Path
|
|
import shutil
|
|
import subprocess
|
|
import tempfile
|
|
|
|
from CelerisLab.config import LBMConfig
|
|
from CelerisLab.lbm.blocks import BlockLayout2D
|
|
from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY
|
|
from CelerisLab.lbm.refinement import FixedPatch2D, FixedPatchSet2D, ACTIVE, OUTER_GHOST
|
|
from CelerisLab.lbm.refinement_tables import projected_explosion_tables
|
|
|
|
E = np.array([D2Q9_EX, D2Q9_EY]).T
|
|
W = np.array([4/9] + [1/9]*4 + [1/36]*4)
|
|
|
|
|
|
def setup(union=False, bounds=(4, 4, 12, 12)):
|
|
cfg = LBMConfig(nx=16, ny=16, viscosity=.1)
|
|
p = (FixedPatchSet2D.from_config(cfg, [bounds]) if union else
|
|
FixedPatch2D.from_config(cfg, bounds))
|
|
layout = BlockLayout2D.build(16, 16, 4, 4, periodic=(True, True))
|
|
neighbors, offsets = projected_explosion_tables(p, layout)
|
|
inverse = np.full(layout.nblocks * np.prod(layout.local_shape), -1, dtype=int)
|
|
inverse[layout.owner.ravel()] = np.arange(256)
|
|
return p, inverse[neighbors], offsets
|
|
|
|
|
|
def reconstruct(c, p, neighbors, offsets, project=True):
|
|
# Independent NumPy evaluation of the documented stencil; neighbors have
|
|
# been decoded from actual blocked owner IDs, not recomputed geometry.
|
|
values = c.reshape(9, -1)[:, neighbors]
|
|
g = .5 * (values[:, :, 1::2] - values[:, :, ::2])
|
|
if project:
|
|
g[1:] -= (np.einsum('qra,qa->qr', g[1:], E[1:]) /
|
|
np.sum(E[1:]**2, axis=1)[:, None])[:, :, None] * E[1:, None, :]
|
|
out = c.reshape(9, -1)[:, p.ghost_coarse] + np.einsum('qra,ra->qr', g, offsets)
|
|
return out, g
|
|
|
|
|
|
def check_legal_footprint_mean_and_characteristic(union, bounds):
|
|
p, neighbors, offsets = setup(union, bounds)
|
|
assert set(neighbors.ravel()) <= set(p.ghost_coarse)
|
|
assert not p.covered.ravel()[neighbors].any()
|
|
assert (neighbors >= 0).all()
|
|
c = np.random.default_rng(72).normal(size=(9,16,16))
|
|
c[:, p.covered] = np.nan # poison every forbidden donor
|
|
out, g = reconstruct(c, p, neighbors, offsets)
|
|
assert np.isfinite(out).all()
|
|
np.testing.assert_allclose(np.einsum('qra,qa->qr', g, E), 0, atol=5e-16)
|
|
for d in np.unique(p.ghost_coarse):
|
|
rows = p.ghost_coarse == d
|
|
assert rows.sum() == 4
|
|
np.testing.assert_allclose(out[:, rows].mean(axis=1), c.reshape(9,-1)[:,d], atol=5e-16)
|
|
for q in range(1,9):
|
|
delta = offsets[rows, None, :] - offsets[None, rows, :]
|
|
# Parallel displacements have zero 2D cross product.
|
|
parallel = np.isclose(delta[...,0]*E[q,1]-delta[...,1]*E[q,0], 0)
|
|
diff = out[q, rows][:,None]-out[q, rows][None,:]
|
|
np.testing.assert_allclose(diff[parallel], 0, atol=5e-16)
|
|
const = np.broadcast_to(W[:,None,None], (9,16,16)).copy()
|
|
constant, _ = reconstruct(const,p,neighbors,offsets)
|
|
np.testing.assert_array_equal(constant, np.broadcast_to(W[:,None], constant.shape))
|
|
|
|
|
|
def collide(f, tau):
|
|
rho = f.sum(axis=0)
|
|
u = np.einsum('qa,q...->a...', E, f) / rho
|
|
cu = np.einsum('qa,a...->q...', E, u)
|
|
eq = W[:,None] * rho * (1 + 3*cu + 4.5*cu**2 - 1.5*np.sum(u*u,axis=0))
|
|
return f + (eq-f)/tau
|
|
|
|
|
|
def totals(c, f, p):
|
|
q = c[:,~p.covered].sum(axis=1)+.25*f[:,p.fine_kind==ACTIVE].sum(axis=1)
|
|
return np.r_[q.sum(), E.T@q]
|
|
|
|
|
|
def check_full_two_hop_closed_corner_budget(union, collision, bounds):
|
|
p, neighbors, offsets = setup(union,bounds)
|
|
rng=np.random.default_rng(17)
|
|
c=W[:,None,None]*(1+rng.normal(0,.002,(9,16,16)))
|
|
f=W[:,None,None]*(1+rng.normal(0,.002,(9,*p.fine_shape)))
|
|
start=totals(c,f,p)
|
|
# Count exactly the eight coarse diagonal links whose source is also
|
|
# uncovered, but whose midpoint lies on the closed rectangular interface.
|
|
q, k=np.divmod(p.coalesce_coarse.astype(np.int64),256)
|
|
y,x=np.divmod(k,16)
|
|
grazing=~p.covered[y-E[q,1],x-E[q,0]]
|
|
assert grazing.sum()==8
|
|
c[:,p.covered]=np.nan
|
|
for _ in range(8):
|
|
f.reshape(9,-1)[:,p.ghost_fine]=reconstruct(c,p,neighbors,offsets)[0]
|
|
for _ in range(2):
|
|
f=np.array([np.roll(f[q],(ey,ex),(0,1)) for q,(ex,ey) in enumerate(E)])
|
|
f[:,p.fine_kind==OUTER_GHOST]=np.nan
|
|
if collision:
|
|
f[:,p.fine_kind==ACTIVE]=collide(f[:,p.fine_kind==ACTIVE],p.tau_fine)
|
|
out=np.array([np.roll(c[q],(ey,ex),(0,1)) for q,(ex,ey) in enumerate(E)])
|
|
packets=f.reshape(-1)[p.coalesce_fine].mean(axis=1)
|
|
assert np.isfinite(packets[grazing]).all()
|
|
out.reshape(-1)[p.coalesce_coarse]=packets
|
|
if collision:out[:,~p.covered]=collide(out[:,~p.covered],p.tau_coarse)
|
|
out[:,p.covered]=np.nan
|
|
c=out
|
|
np.testing.assert_allclose(totals(c,f,p),start,rtol=0,atol=2e-12)
|
|
|
|
|
|
def check_projection_is_distinguished_by_characteristic_not_global_mass():
|
|
p,n,o=setup()
|
|
c=np.random.default_rng(8).normal(size=(9,16,16))
|
|
_,g=reconstruct(c,p,n,o,project=False)
|
|
assert np.max(abs(np.einsum('qra,qa->qr',g,E))) > .1
|
|
|
|
|
|
def check_multiple_rectangles_rejected():
|
|
cfg=LBMConfig(nx=24,ny=16,viscosity=.1)
|
|
p=FixedPatchSet2D.from_config(cfg, [(4,4,8,12),(12,4,16,12)])
|
|
layout=BlockLayout2D.build(16,24,4,4,periodic=(True,True))
|
|
with unittest.TestCase().assertRaisesRegex(ValueError, 'one rectangle'):
|
|
projected_explosion_tables(p,layout)
|
|
|
|
|
|
class ProjectedExplosionTests(unittest.TestCase):
|
|
def test_stencil(self):
|
|
for union in (False, True):
|
|
for bounds in ((4,4,12,12),(2,3,9,10),(2,2,3,3)):
|
|
with self.subTest(union=union,bounds=bounds):
|
|
check_legal_footprint_mean_and_characteristic(union,bounds)
|
|
|
|
def test_closed_interface(self):
|
|
for union in (False,True):
|
|
for collision in (False,True):
|
|
for bounds in ((4,4,12,12),(2,3,14,12),(2,2,3,3)):
|
|
with self.subTest(union=union,collision=collision,bounds=bounds):
|
|
check_full_two_hop_closed_corner_budget(union,collision,bounds)
|
|
|
|
def test_projection(self):
|
|
check_projection_is_distinguished_by_characteristic_not_global_mass()
|
|
|
|
def test_affine_and_shift_invariance(self):
|
|
p,n,o=setup()
|
|
y,x=np.indices((16,16))
|
|
c=np.array([.7+q*.01+x*.03+y*.05 for q in range(9)])
|
|
out,g=reconstruct(c,p,n,o)
|
|
available=n[:,::2]!=n[:,1::2]
|
|
gradient=np.array([.03,.05])*available
|
|
expected=np.broadcast_to(gradient,(9,*gradient.shape)).copy()
|
|
expected[1:]-=(np.einsum('qra,qa->qr',expected[1:],E[1:])/
|
|
np.sum(E[1:]**2,axis=1)[:,None])[:,:,None]*E[1:,None,:]
|
|
np.testing.assert_allclose(g,expected,atol=2e-16)
|
|
shifted,_=reconstruct(c-W[:,None,None],p,n,o)
|
|
np.testing.assert_allclose(out-W[:,None],shifted,atol=5e-16)
|
|
|
|
def test_reject(self):
|
|
check_multiple_rectangles_rejected()
|
|
|
|
|
|
class ProjectedKernelTests(unittest.TestCase):
|
|
@unittest.skipUnless(shutil.which("g++"), "g++ is required for actual-kernel CPU arithmetic")
|
|
def test_actual_kernel_fp32_shifted_and_unshifted(self):
|
|
p, neighbors, offsets = setup(True)
|
|
parent = BlockLayout2D.build(16,16,4,4,periodic=(True,True))
|
|
child = p.fine_layout((8,4)) # deliberately unequal parent/child strides
|
|
ids, offsets = projected_explosion_tables(p,parent)
|
|
cells = int(np.prod(parent.local_shape))
|
|
targets = child.owner.ravel()[p.ghost_fine]
|
|
donors = parent.owner.ravel()[p.ghost_coarse]
|
|
source = (Path(__file__).resolve().parents[2] /
|
|
"src/CelerisLab/lbm/kernels/blocks.cu").read_text()
|
|
kernel = source.split('extern "C" __global__ void BlockExplodeProjected(',1)[1]
|
|
kernel = 'void BlockExplodeProjected(' + kernel.split('// Targets are receiver-local SOURCE slots',1)[0]
|
|
def array(name, values, kind):
|
|
v = np.asarray(values).ravel()
|
|
text = ','.join(str(int(x)) if kind=='uint64_t' else float(x).hex()+'f' for x in v)
|
|
return f"{kind} {name}[] = {{{text}}};\n"
|
|
for shifted in (False,True):
|
|
c=np.random.default_rng(11).normal(0,.001,(9,16,16))
|
|
if not shifted:c+=W[:,None,None]
|
|
c=c.astype(np.float32)
|
|
packed=np.full(parent.nblocks*9*cells,np.nan,np.float32)
|
|
for q in range(9):packed[parent.slot(parent.owner.ravel(),q)]=c[q].ravel()
|
|
# All halo/covered inputs are poisoned; kernel must use legal cores.
|
|
for q in range(9):packed[parent.slot(parent.owner.ravel()[p.covered.ravel()],q)]=np.nan
|
|
# Hex initializer needs a portable expression for poisoned NaNs.
|
|
data=array('coarse',packed,'float').replace('nanf','NAN')
|
|
code = '#include <cstdint>\n#include <cmath>\n#include <cstdio>\nusing fpxx=float;\n'
|
|
code += f'constexpr unsigned NQ=9, TOTAL_CELLS={int(np.prod(child.local_shape))};\n'
|
|
code += 'struct Id {unsigned x;}; Id blockIdx{0},blockDim{1},threadIdx{0};\n'
|
|
code += 'unsigned index_f(unsigned k,unsigned q){return q*TOTAL_CELLS+k;}\n'
|
|
code += 'int d_cx[]={0,1,-1,0,0,1,-1,1,-1},d_cy[]={0,0,0,1,-1,1,-1,-1,1};\n'+kernel
|
|
code += data+array('target',targets,'uint64_t')+array('donor',donors,'uint64_t')
|
|
code += array('neighbor',ids,'uint64_t')+array('offset',offsets,'float')
|
|
code += f'float fine[{child.nblocks*9*int(np.prod(child.local_shape))}]={{}};\n'
|
|
code += f'int main(){{for(unsigned r=0;r<{len(targets)};++r){{threadIdx.x=r;BlockExplodeProjected(coarse,fine,target,donor,{len(targets)},{cells},neighbor,offset);}}'
|
|
code += f'for(unsigned q=0;q<9;++q)for(unsigned r=0;r<{len(targets)};++r){{auto t=target[r];printf("%.9g\\n",fine[(t/TOTAL_CELLS)*9*TOTAL_CELLS+q*TOTAL_CELLS+t%TOTAL_CELLS]);}}}}'
|
|
with tempfile.TemporaryDirectory() as folder:
|
|
src=Path(folder)/'kernel.cpp'; binary=Path(folder)/'kernel'
|
|
src.write_text(code)
|
|
subprocess.run(['g++','-std=c++17','-O2',str(src),'-o',str(binary)],check=True,capture_output=True)
|
|
result=np.fromstring(subprocess.check_output([str(binary)],text=True),sep=' ').reshape(9,-1)
|
|
expected,_=reconstruct(c.astype(float),p,neighbors,offsets)
|
|
np.testing.assert_allclose(result,expected,atol=6e-8,rtol=2e-7)
|