Add opt-in row-owned raw boundary exchange for blocked refinement

tests/unit/test_raw_exchange.py

@@ -0,0 +1,171 @@
+"""CPU reduction and real host scheduling with CUDA doubles, not GPU physics."""
+from contextlib import nullcontext
+from dataclasses import replace
+from unittest.mock import Mock, patch
+
+import numpy as np
+import unittest
+
+from CelerisLab.body.geometry.circle import CircleGeometry
+from CelerisLab.body.objects import SimObject, SurfaceActuation
+from CelerisLab.config import LBMConfig
+from CelerisLab.lbm.blocks import BlockLayout2D
+from CelerisLab.lbm.refinement import FixedPatch2D, prepare_fine_boundary
+from CelerisLab.lbm.refinement_tables import _prepare_blocked_boundary
+from CelerisLab.lbm.raw_exchange import summarize_exchange
+from test_refinement_three_level_multipatch import host_module
+
+
+def body(center=(16., 16.), radius=.5, speed=0.):
+    return SimObject(91, geometry=CircleGeometry(*center, radius), center=center, radius=radius,
+                     surface_actuation=SurfaceActuation(.7, .85, speed))
+
+
+def check_weights_split_and_window(level):
+    h = 2. ** -level
+    rows = np.array([[2, -3, 5, 7], [-1, 4, -2, -6], [99, 99, 99, 99]], np.float32)
+    kw = dict(body_id=0, level=level, h=h, time_end=3, reference_center=(7, 8))
+    last = summarize_exchange(rows, [0, 0, 1], [False, True, False], window='last_substep', **kw)
+    integral = rows.astype(float)*[h*h, h*h, h**3, h*h]/h
+    full = summarize_exchange(integral, [0, 0, 1], [False, True, False], window='last_base_step', **kw)
+    np.testing.assert_array_equal(last['total']['force'], np.array([1, 1])*h)
+    assert last['total']['torque'] == 3*h*h
+    assert last['total']['mass_into_fluid'] == h*h
+    for name in ('force', 'torque', 'mass_rate'):
+        np.testing.assert_array_equal(last['total'][name], full['total'][name])
+    for name in ('force', 'impulse', 'torque', 'angular_impulse', 'mass_into_fluid'):
+        np.testing.assert_array_equal(last['total'][name],
+                                     np.asarray(last['wall'][name])+last['aperture'][name])
+    assert last['time_start'] == 3-h and full['time_start'] == 2
+    assert last['sample_count'] == 1 and full['sample_count'] == 2**level
+    assert full['reference_center'] == (7, 8)
+
+
+def test_final_substep_is_not_interval_mean():
+    kw = dict(body_id=0, level=1, h=.5, time_end=1, reference_center=(0, 0))
+    last = summarize_exchange([[4, 0, 0, 0]], [0], [False], window='last_substep', **kw)
+    full = summarize_exchange([[1.5, 0, 0, 0]], [0], [False], window='last_base_step', **kw)
+    assert last['total']['force'][0] == 2 and full['total']['force'][0] == 1.5
+
+
+def test_host_metadata_zero_jet_and_rows():
+    cfg = LBMConfig(nx=32, ny=32, viscosity=.1)
+    pair = FixedPatch2D.from_config(cfg, (4, 4, 28, 28))
+    boundary = prepare_fine_boundary(cfg, pair, [body(radius=2.15)])
+    local = replace(cfg, nx=pair.fine_shape[1], ny=pair.fine_shape[0], viscosity=.2)
+    local.validate()
+    layout = BlockLayout2D.build(*pair.fine_shape, 8, 8, periodic=(False, False))
+    b = _prepare_blocked_boundary(boundary, layout, local, ())
+    for key, idx in [('body_id', 3), ('rx', 4), ('ry', 5)]:
+        np.testing.assert_array_equal(b.tables[key], boundary.curved_host[idx])
+    assert b.tables['aperture'].any() and (~b.tables['aperture']).any()
+    assert not b.tables['surface_velocity'].any()
+    # Classification does not disappear when prescribed velocity is zero.
+    assert len(b.tables['cell']) == len(boundary.curved_host[0])
+
+
+def backend(observe=True, obj=None):
+    host = host_module()
+    calls = []
+    module = Mock()
+    module.get_global.return_value = (Mock(), 20)
+    module.get_function.side_effect = lambda name: Mock(side_effect=lambda *a, **k: calls.append((name, a, k)))
+    host.compiler.load_module.return_value = module
+    host.cuda.mem_alloc.side_effect = lambda n: Mock(size=n)
+    host.cuda.memcpy_dtoh.side_effect = lambda array, pointer: array.fill(0)
+    ctx = Mock(activate=nullcontext, sm_arch='sm_70')
+    run = host.FixedPatchStepper(LBMConfig(nx=48, ny=48, viscosity=.02), context=ctx,
+        refinements=[[(6, 6, 26, 26)], [(20, 20, 44, 44)], [(56, 56, 72, 72)]],
+        level_blocks=[(4, 4)]*4, level_objects=[(), (), (), (body() if obj is None else obj,)],
+        observe_raw_exchange=observe)
+    arrays = [np.full(l.shape, 1/9, np.float32) for l in run._levels]
+    run.initialize(*arrays)
+    return run, host, calls, arrays
+
+
+def test_four_level_schedule_publication_and_fixed_storage():
+    run, host, calls, arrays = backend()
+    b = run._level_boundaries[3]
+    rows = len(b.tables['cell'])
+    assert b.last_raw.size == rows*16 and b.interval_integral.size == rows*32
+    with unittest.TestCase().assertRaisesRegex(RuntimeError, 'completed'):
+        run.read_raw_exchange(0, level=3)
+    run.step()
+    names = [name for name, _, _ in calls]
+    assert names.count('BlockAccumulateExchange') == 8
+    for i, (name, args, kw) in enumerate(calls):
+        if name == 'BlockAccumulateExchange':
+            assert names[i-1] == 'BlockCurvedBoundary' and names[i+1] == 'BlockStep'
+            assert args[0] is b.last_raw and args[1] is b.interval_integral
+            assert args[3:] == (.125, .125)
+            assert kw['stream'] is run._stream
+    assert host.cuda.memset_d8_async.call_count == 1
+    value = run.read_raw_exchange(0, level=3)
+    assert value['sample_count'] == 8 and value['time_start'] == 0
+    run.step(0)
+    assert run.read_raw_exchange(0, level=3)['time_end'] == 1
+    assert run.read_raw_exchange(0, level=3, window='last_substep')['time_start'] == .875
+    run.step(2)
+    assert run.read_raw_exchange(0, level=3)['time_start'] == 2
+    run.initialize(*arrays)
+    with unittest.TestCase().assertRaisesRegex(RuntimeError, 'completed'):
+        run.read_raw_exchange(0, level=3)
+    run.close()
+
+
+
+def test_torque_reference_uses_geometry_center_not_optional_object_center():
+    for center in (None, (3., 4.)):
+        obj = SimObject(91, geometry=CircleGeometry(16., 16., .5), center=center, radius=.5)
+        run, _, _, _ = backend(obj=obj)
+        run.step()
+        result = run.read_raw_exchange(0, level=3)
+        assert result['reference_center'] == (16., 16.)
+        assert result['reference_center'] != tuple(obj.center)
+        run.close()
+
+def test_disabled_has_nulls_and_no_observation_allocation_or_launch():
+    run, host, calls, _ = backend(False)
+    b = run._level_boundaries[3]
+    assert b.last_raw is b.interval_integral is None
+    run.step()
+    for name, args, _ in calls:
+        if name == 'BlockCurvedBoundary':
+            assert tuple(args[8:]) == (0, 0, 0)
+    assert all(n != 'BlockAccumulateExchange' for n, _, _ in calls)
+    host.cuda.memset_d8_async.assert_not_called()
+    with unittest.TestCase().assertRaisesRegex(RuntimeError, 'disabled'):
+        run.read_raw_exchange(0, level=3)
+    run.close()
+
+
+def test_failed_sync_never_publishes_partial_interval():
+    run, host, calls, _ = backend()
+    run.step()
+    run._stream.synchronize.side_effect = RuntimeError('device failure')
+    with unittest.TestCase().assertRaisesRegex(RuntimeError, 'device failure'):
+        run.step()
+    assert not run._exchange_valid
+    with unittest.TestCase().assertRaisesRegex(RuntimeError, 'failed'):
+        run.read_raw_exchange(0, level=3)
+    run._stream.synchronize.side_effect = None
+    run.close()
+
+
+def test_dense_observation_rejected_before_device_access():
+    host = host_module()
+    ctx = Mock()
+    with unittest.TestCase().assertRaisesRegex(ValueError, 'blocked body'):
+        host.FixedPatchStepper(LBMConfig(nx=32, ny=32, viscosity=.1), (4, 4, 28, 28),
+                               context=ctx, fine_objects=[body(radius=2.15)], observe_raw_exchange=True)
+    ctx.activate.assert_not_called()
+
+
+def test_all_depth_weights():
+    for level in (1, 2, 3, 7):
+        check_weights_split_and_window(level)
+
+
+def load_tests(loader, tests, pattern):
+    return unittest.TestSuite(unittest.FunctionTestCase(value) for name, value in globals().items()
+                              if name.startswith("test_") and callable(value))

tests/validation/multigrid/test_raw_exchange_gpu.py

@@ -0,0 +1,264 @@
+"""GPU gates for raw exchange. Run explicitly; requires PyCUDA and CUDA.
+
+Independent census reads actual staged incoming/outgoing populations, including
+non-body ACTIVE-boundary flux, then checks postcollision fluid momentum/mass.
+This is a discrete implementation gate, NOT a physical nozzle-load validation.
+"""
+from dataclasses import replace
+import unittest
+import numpy as np
+import pycuda.driver as cuda
+
+from CelerisLab.body.geometry.circle import CircleGeometry
+from CelerisLab.body.objects import SimObject
+from CelerisLab.body.jets import Jet, Tube, Parallel
+from CelerisLab.config import LBMConfig
+from CelerisLab.cuda.context import CudaContext
+from CelerisLab.lbm.refinement_stepper import FixedPatchStepper
+from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY, SOLID
+
+CX, CY = np.asarray(D2Q9_EX), np.asarray(D2Q9_EY)
+OPP = np.array([np.flatnonzero((CX == -x) & (CY == -y))[0] for x, y in zip(CX, CY)])
+W = np.array([4/9] + [1/9]*4 + [1/36]*4, np.float32)
+CFG = LBMConfig(nx=32, ny=24, viscosity=.1)
+BOUNDS = (8, 6, 24, 18)
+
+
+def circle():
+    return SimObject(17, geometry=CircleGeometry(16., 12., 2.15), center=(16., 12.), radius=2.15)
+
+
+def jet():
+    # Asymmetric aperture avoids hiding exchange errors by paired-jet cancellation.
+    return Jet(Tube((16., 12.), 0., 3., .7), Parallel(0.), amplitude=0.)
+
+
+def state(shape, ux=.013):
+    cu = CX[:, None, None]*ux
+    values = W[:, None, None]*(1+3*cu+4.5*cu*cu-1.5*ux*ux)
+    return np.broadcast_to(values, shape).astype(np.float32).copy()
+
+
+def download(pointer, shape, dtype=np.float32):
+    result = np.empty(shape, dtype)
+    cuda.memcpy_dtoh(result, pointer)
+    return result
+
+
+def totals(g, leaf):
+    values = g[:, leaf].astype(np.float64)
+    return np.array([(values*CX[:, None]).sum(), (values*CY[:, None]).sum(), values.sum()])
+
+
+def census(run, index, buffer):
+    """Independent incoming/outgoing face census on ACTIVE fluid cells.
+
+    Reads the receiver's actual block window, not a freshly gathered owner image
+    (that would erase BFL/restriction staging). Classifies external flux solely
+    from geometry, without consulting raw observation arrays.
+    """
+    level = run._levels[index]
+    layout, b = level.layout, level.boundary
+    g = download(buffer, layout.blocked_shape)
+    dense = layout.to_dense(g)
+    leaf = run._leaf[index]
+    ny, nx = leaf.shape
+    solid = (b.host.flags.reshape(ny, nx) & SOLID) != 0
+    y, x = np.nonzero(leaf)
+    owner = layout.owner[y, x].astype(np.intp)
+    cells = int(np.prod(layout.local_shape))
+    block, cell = np.divmod(owner, cells)
+    flat = g.reshape(layout.nblocks, 9, cells)
+    external = np.zeros(3, np.float64)
+    for q, (cx, cy) in enumerate(zip(CX, CY)):
+        sy, sx = y-cy, x-cx
+        incoming = flat[block, q, cell-cy*layout.local_shape[1]-cx].astype(np.float64)
+        entering = ~leaf[sy, sx] & ~solid[sy, sx]
+        external += np.array([cx, cy, 1])*incoming[entering].sum()
+        dy, dx = y+cy, x+cx
+        leaving = ~leaf[dy, dx] & ~solid[dy, dx]
+        outgoing = flat[block, q, cell].astype(np.float64)
+        external -= np.array([cx, cy, 1])*outgoing[leaving].sum()
+    receiver, direction, _, _, rx, ry, _, _ = b.host.curved_host
+    owner = layout.owner.ravel()[receiver].astype(np.intp)
+    block, cell = np.divmod(owner, cells)
+    a = flat[block, direction, cell].astype(np.float64)
+    source = cell+CY[direction]*layout.local_shape[1]+CX[direction]
+    incoming = flat[block, OPP[direction], source].astype(np.float64)
+    jx, jy = CX[direction]*(a+incoming), CY[direction]*(a+incoming)
+    independent = np.column_stack((jx, jy, rx*jy-ry*jx, incoming-a))
+    return totals(dense, leaf), external, independent
+
+
+class RawExchangeGPU(unittest.TestCase):
+    def test_row_ledger_interval_jet_switch_and_no_population_effect(self):
+        with CudaContext(0) as ctx:
+            for block in ((32, 40), (16, 20), (8, 10)):
+                with FixedPatchStepper(CFG, BOUNDS, context=ctx, fine_block=block,
+                        fine_objects=[circle()], jets=[jet()], observe_raw_exchange=True) as measured, \
+                     FixedPatchStepper(CFG, BOUNDS, context=ctx, fine_block=block,
+                        fine_objects=[circle()], jets=[jet()]) as plain:
+                    arrays = [state(l.shape) for l in measured._levels]
+                    measured.initialize(*arrays)
+                    plain.initialize(*arrays)
+                    b = measured._level_boundaries[1]
+                    count = len(b.tables['cell'])
+                    kernel = measured._levels[1].kernel
+                    samples = []
+
+                    def audited_step(src, dst, *args, **kw):
+                        measured._stream.synchronize()
+                        before, external, independent = census(measured, 1, src)
+                        raw = download(b.last_raw, (count, 4))
+                        np.testing.assert_allclose(raw, independent, rtol=2e-6, atol=2e-7)
+                        kernel(src, dst, *args, **kw)
+                        measured._stream.synchronize()
+                        after = measured._levels[1].layout.to_dense(
+                            download(dst, measured._levels[1].layout.blocked_shape))
+                        change = totals(after, measured._leaf[1])-before
+                        expected = external + np.array([-raw[:, 0].astype(float).sum(),
+                                                        -raw[:, 1].astype(float).sum(),
+                                                         raw[:, 3].astype(float).sum()])
+                        np.testing.assert_allclose(change, expected, rtol=2e-4, atol=8e-5)
+                        samples.append(raw.copy())
+                    measured._levels[1].kernel = audited_step
+                    for speed in (0., .002, -.001, 0.):
+                        measured.set_jet_amplitude(0, speed)
+                        plain.set_jet_amplitude(0, speed)
+                        samples.clear()
+                        measured.step()
+                        plain.step()
+                        for a, e in zip(measured.download(), plain.download()):
+                            np.testing.assert_array_equal(a, e)
+                        integral = download(b.interval_integral, (count, 4), np.float64)
+                        expected = sum(s.astype(float)*[.25, .25, .125, .25] for s in samples)
+                        np.testing.assert_array_equal(integral, expected)
+                        last = download(b.last_raw, (count, 4))
+                        np.testing.assert_array_equal(last, samples[-1])
+                        obs = measured.read_raw_exchange(0)
+                        assert obs['sample_count'] == 2
+                        assert b.tables['aperture'].any()
+                        measured.step(0)
+                        np.testing.assert_array_equal(measured.read_raw_exchange(0)['total']['impulse'],
+                                                      obs['total']['impulse'])
+
+    def test_dense_fine_storage_reference(self):
+        """Same fine resolution/initial data with dense storage: a numerical
+        reference for cut-link rows, not a uniform-whole-domain convergence study."""
+        with CudaContext(0) as ctx:
+            with FixedPatchStepper(CFG, BOUNDS, context=ctx, fine_objects=[circle()], jets=[jet()]) as dense, \
+                 FixedPatchStepper(CFG, BOUNDS, context=ctx, fine_block=(16, 20),
+                    fine_objects=[circle()], jets=[jet()], observe_raw_exchange=True) as blocked:
+                arrays = [state(l.shape) for l in blocked._levels]
+                dense.initialize(*arrays)
+                blocked.initialize(*arrays)
+                for speed in (0., .002, 0.):
+                    dense.set_jet_amplitude(0, speed)
+                    blocked.set_jet_amplitude(0, speed)
+                    dense.step(); blocked.step()
+                    for a, e in zip(blocked.download(), dense.download()):
+                        np.testing.assert_array_equal(a, e)
+                    count = dense._fine_curved.count
+                    expected = download(dense._fine_curved.link_exchange_gpu, (count, 3))
+                    actual = download(blocked._level_boundaries[1].last_raw, (count, 4))
+                    np.testing.assert_array_equal(actual[:, :3], expected)
+
+    def test_depth_four_interval_storage_and_sampling(self):
+        cfg = LBMConfig(nx=48, ny=48, viscosity=.02)
+        obj = SimObject(0, geometry=CircleGeometry(16., 16., .5), center=(16., 16.), radius=.5)
+        drive = Jet(Tube((16., 16.), 0., .75, .2), Parallel(0.), amplitude=0.)
+        settings = dict(refinements=[[(6, 6, 26, 26)], [(20, 20, 44, 44)], [(56, 56, 72, 72)]],
+                        level_blocks=[(4, 4)]*4, level_objects=[(), (), (), (obj,)],
+                        level_jets=[(), (), (), (drive,)])
+        with CudaContext(0) as ctx:
+            with FixedPatchStepper(cfg, context=ctx, observe_raw_exchange=True, **settings) as run, \
+                 FixedPatchStepper(cfg, context=ctx, **settings) as plain:
+                arrays = [state(l.shape, ux=0.) for l in run._levels]
+                run.initialize(*arrays); plain.initialize(*arrays)
+                b = run._level_boundaries[3]
+                count = len(b.tables['cell'])
+                kernel = run._levels[3].kernel
+                samples, independent_samples = [], []
+                def audited_step(src, dst, *args, **kw):
+                    run._stream.synchronize()
+                    before, external, independent = census(run, 3, src)
+                    raw = download(b.last_raw, (count, 4))
+                    np.testing.assert_allclose(raw, independent, rtol=2e-6, atol=2e-7)
+                    kernel(src, dst, *args, **kw)
+                    run._stream.synchronize()
+                    after = run._levels[3].layout.to_dense(
+                        download(dst, run._levels[3].layout.blocked_shape))
+                    change = totals(after, run._leaf[3])-before
+                    expected = external + np.array([-independent[:, 0].sum(),
+                                                    -independent[:, 1].sum(), independent[:, 3].sum()])
+                    np.testing.assert_allclose(change, expected, rtol=2e-4, atol=8e-5)
+                    samples.append(raw)
+                    independent_samples.append(independent)
+                run._levels[3].kernel = audited_step
+                for speed in (0., .002, 0.):
+                    run.set_jet_amplitude(0, speed, level=3)
+                    plain.set_jet_amplitude(0, speed, level=3)
+                    samples.clear(); independent_samples.clear()
+                    run.step(); plain.step()
+                    self.assertEqual(len(samples), 8)
+                    weights = np.array([1/64, 1/64, 1/512, 1/64])
+                    integral = download(b.interval_integral, (count, 4), np.float64)
+                    np.testing.assert_array_equal(integral, sum(v.astype(float)*weights for v in samples))
+                    np.testing.assert_allclose(integral, sum(v*weights for v in independent_samples),
+                                               rtol=2e-6, atol=2e-7)
+                    np.testing.assert_array_equal(download(b.last_raw, (count, 4)), samples[-1])
+                    for actual, expected in zip(run.download(), plain.download()):
+                        np.testing.assert_array_equal(actual, expected)
+                    result = run.read_raw_exchange(0, level=3)
+                    self.assertEqual(result['sample_count'], 8)
+                    self.assertEqual(result['time_start'], run.step_count-1)
+                    if run.step_count == 1:
+                        np.testing.assert_allclose(result['total']['force'], 0, atol=2e-7)
+                        self.assertAlmostEqual(result['total']['torque'], 0, delta=2e-7)
+
+    def test_uniform_whole_domain_fine_short_time_reference(self):
+        """Same root geometry, nu and Jet speed on a whole-domain fine mesh.
+        Eight local hops are a short-time equivalence check, not convergence.
+        The test adapter only enables the standalone kernel's optional inputs.
+        """
+        from CelerisLab.lbm.blocks import BlockLayout2D
+        from CelerisLab.lbm.block_stepper import PeriodicBlockStepper
+        fine_cfg = replace(CFG, nx=64, ny=48, viscosity=.2)
+        obj = SimObject(0, geometry=CircleGeometry(31.5, 23.5, 4.3),
+                        center=(31.5, 23.5), radius=4.3)
+        layout = BlockLayout2D.build(48, 64, 8, 8)
+        with CudaContext(0) as ctx:
+            with PeriodicBlockStepper(fine_cfg, layout, context=ctx, objects=[obj]) as uniform, \
+                 FixedPatchStepper(CFG, BOUNDS, context=ctx, fine_block=(16, 20),
+                    fine_objects=[circle()], jets=[jet()], observe_raw_exchange=True) as amr:
+                amr.initialize(*(state(l.shape) for l in amr._levels))
+                uniform.initialize(state((9, 48, 64)))
+                b = amr._level_boundaries[1]
+                count = len(b.tables['cell'])
+                links = obj.geometry.build_cut_links(64, 48)
+                np.testing.assert_array_equal([v.dir for v in links], b.tables['dir'])
+                np.testing.assert_allclose([v.q for v in links], b.tables['q'], atol=2e-7, rtol=0)
+                # Same fine-grid force arms after translation to the same center.
+                rx = uniform._upload(b.tables['rx'])
+                ry = uniform._upload(b.tables['ry'])
+                velocity = uniform._upload(b.tables['surface_velocity'])
+                raw = uniform._allocate(count*16)
+                original = uniform._curved
+                samples = []
+                def measured_boundary(*args, **kw):
+                    original(*args[:7], velocity, rx, ry, raw, **kw)
+                    uniform._stream.synchronize()
+                    samples.append(download(raw, (count, 4)))
+                uniform._curved = measured_boundary
+                for speed in (0., .002, -.001, 0.):
+                    amr.set_jet_amplitude(0, speed)
+                    cuda.memcpy_htod(velocity, b.tables['surface_velocity'])
+                    samples.clear()
+                    amr.step(); uniform.step(2)
+                    expected = sum(v.astype(float)*[.25, .25, .125, .25] for v in samples)
+                    actual = download(b.interval_integral, (count, 4), np.float64)
+                    np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=2e-6)
+
+
+if __name__ == '__main__':
+    unittest.main()