ClaudeandClaude Opus 5.5
050d12b76c
Dense root + blocked fine level is level_blocks=[None, block]; bodies go to level_objects[1]. Both before-CUDA rejections (non-dividing and malformed fine block) keep their inputs. Assertions unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
374 lines
21 KiB
Python
374 lines
21 KiB
Python
"""Opt-in blocked fine storage for FixedPatchStepper: CPU tables plus GPU gate.
|
|
|
|
Frozen GPU gate (same stream, 8 coarse steps, nonequilibrium start), in both
|
|
periodic and slot_reservoir modes: dense fine storage vs 1x1, 2x2 and 4x2
|
|
fine blocks give bit-identical masked downloads AND bit-identical raw coarse
|
|
and unpacked owner-core fine arrays, including patch-ghost values/NaN masks.
|
|
"""
|
|
from contextlib import nullcontext
|
|
import unittest
|
|
from unittest.mock import Mock, patch
|
|
|
|
import numpy as np
|
|
import pycuda.driver as cuda
|
|
|
|
from CelerisLab.body.geometry.circle import CircleGeometry
|
|
from CelerisLab.body.objects import SimObject, SurfaceActuation
|
|
from CelerisLab.config import LBMConfig
|
|
from CelerisLab.cuda import compiler_v2 as compiler
|
|
from CelerisLab.lbm import curved_links
|
|
from CelerisLab.lbm.blocks import NO_DONOR, BlockLayout2D, blocked_link_tables
|
|
from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY
|
|
from CelerisLab.lbm.descriptors import SOLID
|
|
from CelerisLab.lbm.refinement import (ACTIVE, INNER_GHOST, OUTER_GHOST, FixedPatch2D,
|
|
prepare_fine_boundary)
|
|
from CelerisLab.lbm.refinement_stepper import FixedPatchStepper, blocked_fine_tables
|
|
|
|
NX, NY, BOUNDS, STEPS = 32, 24, (8, 6, 24, 18), 8
|
|
FINE_BLOCKS = ((32, 40), (16, 20), (16, 10)) # fine window 32x40: 1x1, 2x2, 4x2
|
|
W = np.array([4 / 9] + [1 / 9] * 4 + [1 / 36] * 4)
|
|
CX, CY = np.array(D2Q9_EX, float), np.array(D2Q9_EY, float)
|
|
|
|
|
|
def config(mode):
|
|
if mode == "periodic":
|
|
return LBMConfig(nx=NX, ny=NY, viscosity=.1)
|
|
return LBMConfig(nx=NX, ny=NY, viscosity=.1, velocity=0, wall_normal_velocity=.001,
|
|
wall_slot_width=8, y_wall_bc="on_site_velocity", top_wall_bc="reservoir")
|
|
|
|
|
|
def state(ny, nx, seed):
|
|
"""Positive low-Mach equilibrium plus a seeded nonequilibrium part."""
|
|
y, x = np.mgrid[0:ny, 0:nx]
|
|
rho = 1 + 1e-3 * np.cos(2 * np.pi * x / nx)
|
|
ux, uy = .02 * np.sin(2 * np.pi * y / ny), .01 * np.cos(2 * np.pi * x / nx)
|
|
cu = CX[:, None, None] * ux + CY[:, None, None] * uy
|
|
feq = W[:, None, None] * rho * (1 + 3 * cu + 4.5 * cu * cu - 1.5 * (ux * ux + uy * uy))
|
|
noise = np.random.default_rng(seed).uniform(-1, 1, feq.shape) * 1e-3 * W[:, None, None]
|
|
return (feq + noise).astype(np.float32)
|
|
|
|
|
|
def circle(spec=None):
|
|
"""Astra fixture: coarse cell-edge C=(16,12), R=2.15 -> fine C=(19.5,15.5), R=4.3."""
|
|
return SimObject(0, geometry=CircleGeometry(16.0, 12.0, 2.15), center=(16.0, 12.0),
|
|
radius=2.15, surface_actuation=spec)
|
|
|
|
|
|
def actuation(speed):
|
|
"""Astra's seam fixture: 21 aperture links; in each decomposed layout 3 of them
|
|
write a halo image and 2 read a q<.5 donor from a halo. Coarse lattice units:
|
|
prepare_fine_boundary does the (acoustic) fine scaling, never scale twice."""
|
|
return SurfaceActuation(0.7, 0.85, speed)
|
|
|
|
|
|
def fine_circle(mode, spec=None):
|
|
cfg = config(mode)
|
|
patch = FixedPatch2D.from_config(cfg, BOUNDS)
|
|
return patch, prepare_fine_boundary(cfg, patch, (circle(spec),),
|
|
wall_aware=mode == "slot_reservoir")
|
|
|
|
|
|
def fine_cfg(mode, patch):
|
|
from dataclasses import replace
|
|
cfg = replace(config(mode), nx=patch.fine_shape[1], ny=patch.fine_shape[0], viscosity=.2)
|
|
if mode == "slot_reservoir":
|
|
cfg = replace(cfg, y_wall_bc="bounce_back", top_wall_bc="velocity",
|
|
wall_slot_width=0, wall_normal_velocity=0.0, velocity=0.0)
|
|
cfg.validate() # as the stepper does before surface_link_velocity
|
|
return cfg
|
|
|
|
|
|
def seam_cells(layout, mapped):
|
|
"""Distinct non-core cells written / read as q<.5 donors (receiver-local halos)."""
|
|
L = int(np.prod(layout.local_shape))
|
|
core = set(layout.owner.ravel().tolist())
|
|
writes = {int(w) // (9 * L) * L + int(w) % L for w in mapped.write}
|
|
donors = {int(d) for d in mapped.donor if d != NO_DONOR}
|
|
return len(writes - core), len(donors - core)
|
|
|
|
|
|
def raw(backend):
|
|
"""Unmasked device arrays; blocked fine storage unpacked from owner cores."""
|
|
coarse, fine = backend._levels
|
|
out = [np.empty(coarse.shape, np.float32)]
|
|
out.append(np.empty(fine.shape if backend.fine_layout is None
|
|
else backend.fine_layout.blocked_shape, np.float32))
|
|
cuda.memcpy_dtoh(out[0], coarse.current)
|
|
cuda.memcpy_dtoh(out[1], fine.current)
|
|
if backend.fine_layout is not None:
|
|
out[1] = backend.fine_layout.to_dense(out[1])
|
|
return out
|
|
|
|
|
|
class BlockedFineCPU(unittest.TestCase):
|
|
def forbidden_context(self):
|
|
return Mock(activate=Mock(side_effect=AssertionError("CUDA reached")), sm_arch="sm_70")
|
|
|
|
def test_rejections_before_cuda(self):
|
|
for kw in ({"level_blocks": [None, (7, 7)]}, {"level_blocks": [None, (16,)]}):
|
|
with self.subTest(kw=kw), self.assertRaises(ValueError):
|
|
FixedPatchStepper(config("periodic"), context=self.forbidden_context(),
|
|
refinements=[BOUNDS], **kw)
|
|
|
|
def test_single_block_layout_allocates_no_empty_halo_tables(self):
|
|
# A 1x1 non-periodic fine layout has zero halo pairs; CUDA rejects mem_alloc(0).
|
|
context = Mock(activate=nullcontext, sm_arch="sm_70")
|
|
module = Mock()
|
|
module.get_global.return_value = (Mock(), 20)
|
|
sizes = []
|
|
|
|
def alloc(size):
|
|
if size <= 0:
|
|
raise AssertionError("zero-byte device allocation")
|
|
sizes.append(size)
|
|
return Mock()
|
|
with patch.object(cuda, "Stream"), patch.object(compiler, "generate_config"), \
|
|
patch.object(compiler, "compile_kernel"), patch.object(compiler, "cleanup_build_dir"), \
|
|
patch.object(compiler, "load_module", return_value=module), \
|
|
patch.object(cuda, "mem_alloc", side_effect=alloc), patch.object(cuda, "memcpy_htod"):
|
|
backend = FixedPatchStepper(config("periodic"), context=context, refinements=[BOUNDS],
|
|
level_blocks=[None, FINE_BLOCKS[0]])
|
|
self.assertEqual(len(backend.fine_layout.halo_dst), 0)
|
|
self.assertIsNone(backend._levels[1].halo_dst)
|
|
backend._gather_halo(backend._levels[1], Mock()) # no launch without pairs
|
|
module.get_function.return_value.assert_not_called()
|
|
backend.close()
|
|
self.assertTrue(sizes)
|
|
|
|
def test_circle_fixture_counts_and_seam_staging(self):
|
|
for mode in ("periodic", "slot_reservoir"):
|
|
patch, boundary = fine_circle(mode)
|
|
fluid_idx, dirs, q, _, _, _, _, fallback = boundary.curved_host
|
|
self.assertEqual(len(fluid_idx), 76, mode)
|
|
self.assertEqual(int(np.count_nonzero(q < .5)), 68, mode)
|
|
self.assertTrue((fallback[q < .5] == 0).all(), mode) # every q<.5 link uses its donor
|
|
for by, bx in FINE_BLOCKS:
|
|
layout = BlockLayout2D.build(*patch.fine_shape, by, bx, periodic=(False, False))
|
|
_, mapped = blocked_link_tables(layout, boundary.flags, fluid_idx, dirs, q, fallback)
|
|
single = layout.nblocks == 1
|
|
self.assertEqual(seam_cells(layout, mapped), (0, 0) if single else (8, 8), (mode, by, bx))
|
|
self.assertEqual(len(set(mapped.block.tolist())), 1 if single else 4, (mode, by, bx))
|
|
|
|
def test_storage_adapter_flags_sizing(self):
|
|
patch, boundary = fine_circle("periodic")
|
|
columns = [boundary.curved_host[i] for i in (0, 1, 2, 7)]
|
|
layout = BlockLayout2D.build(*patch.fine_shape, 16, 20, periodic=(False, False))
|
|
tables, mapped = blocked_link_tables(layout, boundary.flags, *columns)
|
|
L = int(np.prod(layout.local_shape))
|
|
self.assertEqual(tables["flags"].shape, (layout.nblocks * L,))
|
|
np.testing.assert_array_equal(tables["flags"][layout.owner.ravel().astype(np.intp)],
|
|
boundary.flags)
|
|
np.testing.assert_array_equal(tables["cell"], mapped.receiver % np.uint64(L))
|
|
for bad in (boundary.flags[:-1], boundary.flags.astype(np.uint32)):
|
|
with self.assertRaises(ValueError):
|
|
blocked_link_tables(layout, bad, *columns)
|
|
|
|
def mocked_backend(self, objects):
|
|
"""Real constructor/initialize with driver doubles; kernels record (name, args)."""
|
|
context = Mock(activate=nullcontext, sm_arch="sm_70")
|
|
calls, functions = [], {}
|
|
|
|
def get_function(name):
|
|
return functions.setdefault(name, Mock(side_effect=lambda *a, **k: calls.append((name, a))))
|
|
module = Mock(get_function=Mock(side_effect=get_function))
|
|
module.get_global.return_value = (Mock(), 20)
|
|
with patch.object(cuda, "Stream"), patch.object(compiler, "generate_config"), \
|
|
patch.object(compiler, "compile_kernel"), patch.object(compiler, "cleanup_build_dir"), \
|
|
patch.object(compiler, "load_module", return_value=module), \
|
|
patch.object(cuda, "mem_alloc", side_effect=lambda size: Mock()), \
|
|
patch.object(cuda, "memcpy_htod"), \
|
|
patch.object(curved_links, "CurvedLinkSoA", side_effect=AssertionError("dense SoA built")):
|
|
backend = FixedPatchStepper(config("periodic"), context=context, refinements=[BOUNDS],
|
|
level_blocks=[None, FINE_BLOCKS[1]],
|
|
level_objects=[(), tuple(objects)])
|
|
backend.initialize(state(NY, NX, 1), state(*backend.patch.fine_shape, 2))
|
|
return backend, calls
|
|
|
|
def test_same_buffer_halo_bfl_step_launch_order(self):
|
|
for spec in (None, actuation(.002)):
|
|
backend, calls = self.mocked_backend([circle(spec)])
|
|
c, f = backend._levels
|
|
coarse, cur, spare = c.current, f.current, f.spare
|
|
flags, kind = backend._fine_flags, backend._levels[1].kind
|
|
velocity = backend._link_tables.get("surface_velocity")
|
|
self.assertEqual(velocity is None, spec is None)
|
|
backend.step(1)
|
|
backend.close()
|
|
order = [(name, a[0]) for name, a in calls]
|
|
self.assertEqual(order, [("BlockExplode", coarse),
|
|
("BlockHaloCopy", cur), ("BlockCurvedBoundary", cur), ("BlockStep", cur),
|
|
("BlockHaloCopy", spare), ("BlockCurvedBoundary", spare),
|
|
("BlockStep", spare), ("PatchCoarseStep", coarse)])
|
|
steps = [a for name, a in calls if name == "BlockStep"]
|
|
self.assertEqual([a[1] for a in steps], [spare, cur])
|
|
for a in steps:
|
|
self.assertIs(a[3], flags) # blocked flags, not null
|
|
self.assertIs(a[4], kind)
|
|
for a in (a for name, a in calls if name == "BlockCurvedBoundary"):
|
|
if spec is None: # explicit null at the raw launch ABI
|
|
self.assertIsInstance(a[7], np.uintp)
|
|
self.assertEqual(int(a[7]), 0)
|
|
else:
|
|
self.assertIs(a[7], velocity)
|
|
|
|
def test_actuation_without_velocity_buffer_fails_before_bfl(self):
|
|
for speed in (0.0, .002): # any spec, even zero amplitude
|
|
backend, calls = self.mocked_backend([circle(actuation(speed))])
|
|
self.assertIsNotNone(backend._fine_tables["surface_velocity"])
|
|
del backend._link_tables["surface_velocity"]
|
|
with self.assertRaisesRegex(RuntimeError, "Surface velocity device buffer is missing"):
|
|
backend.step(1)
|
|
names = [name for name, _ in calls]
|
|
self.assertNotIn("BlockCurvedBoundary", names)
|
|
self.assertNotIn("BlockStep", names)
|
|
self.assertEqual(backend._status, "failed")
|
|
backend.close()
|
|
|
|
def test_surface_velocity_helper_matches_dense_soa(self):
|
|
for mode in ("periodic", "slot_reservoir"):
|
|
patch, boundary = fine_circle(mode, actuation(-.002))
|
|
cfg = fine_cfg(mode, patch)
|
|
soa = curved_links.CurvedLinkSoA()
|
|
soa.assign_host(*boundary.curved_host)
|
|
soa.assign_surface(boundary.objects, cfg, boundary.flags)
|
|
fluid_idx, dirs, _, body_id, rx, ry, _, _ = boundary.curved_host
|
|
velocity, aperture = curved_links.surface_link_velocity(
|
|
boundary.objects, cfg, boundary.flags, fluid_idx, dirs, body_id, rx, ry)
|
|
np.testing.assert_array_equal(velocity, soa.surface_velocity)
|
|
np.testing.assert_array_equal(aperture, soa.aperture)
|
|
self.assertEqual(curved_links.surface_link_velocity(
|
|
(circle(),), cfg, boundary.flags, fluid_idx, dirs, body_id, rx, ry), (None, None))
|
|
# Actuated (not just geometric) seam coverage in both decomposed layouts.
|
|
self.assertEqual(int(aperture.sum()), 21, mode)
|
|
self.assertTrue(np.all(np.any(velocity[aperture] != 0, axis=1)), mode)
|
|
for by, bx in FINE_BLOCKS[1:]:
|
|
layout = BlockLayout2D.build(*patch.fine_shape, by, bx, periodic=(False, False))
|
|
_, mapped = blocked_link_tables(layout, boundary.flags, fluid_idx, dirs,
|
|
boundary.curved_host[2], boundary.curved_host[7])
|
|
L = int(np.prod(layout.local_shape))
|
|
core = set(layout.owner.ravel().tolist())
|
|
halo_write = np.array([int(w) // (9 * L) * L + int(w) % L not in core
|
|
for w in mapped.write])
|
|
halo_donor = np.array([d != NO_DONOR and int(d) not in core for d in mapped.donor])
|
|
self.assertEqual((int((halo_write & aperture).sum()),
|
|
int((halo_donor & aperture).sum())), (3, 2), (mode, by, bx))
|
|
|
|
def test_tables_map_back_to_dense_slots_in_order(self):
|
|
patch = FixedPatch2D.from_config(config("periodic"), BOUNDS)
|
|
fy, fx = patch.fine_shape
|
|
cells = fy * fx
|
|
order = np.argsort(patch.coalesce_coarse)
|
|
for by, bx in FINE_BLOCKS:
|
|
layout = BlockLayout2D.build(fy, fx, by, bx, periodic=(False, False))
|
|
core = set(layout.owner.ravel().tolist())
|
|
targets, children, kind = blocked_fine_tables(patch, layout, order)
|
|
# Explosion targets: owner cores of exactly the non-ACTIVE fine cells.
|
|
np.testing.assert_array_equal(targets, layout.owner.ravel()[patch.ghost_fine])
|
|
self.assertEqual(sorted(targets.tolist()),
|
|
sorted(layout.owner[patch.fine_kind != ACTIVE].tolist()))
|
|
# Children: same rows/child order, same (q, dense cell), owner cores.
|
|
L = int(np.prod(layout.local_shape))
|
|
b, rest = np.divmod(children, np.uint64(9 * L))
|
|
q, c = np.divmod(rest, np.uint64(L))
|
|
dense = patch.coalesce_fine[order]
|
|
np.testing.assert_array_equal(q, dense // cells)
|
|
np.testing.assert_array_equal(b * np.uint64(L) + c,
|
|
layout.owner.ravel()[(dense % cells).astype(np.intp)])
|
|
self.assertTrue(set((b * np.uint64(L) + c).ravel().tolist()) <= core)
|
|
# Kinds follow owner cores; unused halo entries are OUTER.
|
|
np.testing.assert_array_equal(kind[layout.owner.ravel().astype(np.intp)],
|
|
patch.fine_kind.ravel())
|
|
self.assertEqual(set(np.unique(kind).tolist()) - {ACTIVE, INNER_GHOST, OUTER_GHOST}, set())
|
|
|
|
|
|
class BlockedFineGPU(unittest.TestCase):
|
|
def test_dense_and_blocked_fine_storage_are_bit_identical(self):
|
|
from CelerisLab.cuda.context import CudaContext
|
|
for mode in ("periodic", "slot_reservoir"):
|
|
cfg = config(mode)
|
|
patch = FixedPatch2D.from_config(cfg, BOUNDS)
|
|
g_coarse, g_fine = state(NY, NX, 1), state(*patch.fine_shape, 2)
|
|
results = []
|
|
with CudaContext(0) as context:
|
|
for fine_block in (None, *FINE_BLOCKS):
|
|
with FixedPatchStepper(cfg, context=context, coarse_boundary=mode,
|
|
refinements=[BOUNDS],
|
|
level_blocks=[None, fine_block]) as backend:
|
|
backend.initialize(g_coarse, g_fine)
|
|
backend.step(STEPS)
|
|
self.assertEqual(backend.step_count, STEPS)
|
|
results.append((backend.download(), raw(backend)))
|
|
(dense_public, dense_raw), blocked = results[0], results[1:]
|
|
kind = patch.fine_kind
|
|
self.assertTrue(np.isnan(dense_raw[1][:, kind == OUTER_GHOST]).all(), mode)
|
|
# INNER populations may be NaN after hop 2 (pulled from poisoned OUTER);
|
|
# their exact pattern is pinned by the dense-vs-blocked comparison below.
|
|
self.assertTrue(np.isfinite(dense_raw[1][:, kind == ACTIVE]).all(), mode)
|
|
for fine_block, (public, raw_arrays) in zip(FINE_BLOCKS, blocked):
|
|
msg = f"{mode} fine_block={fine_block}"
|
|
for a, b in zip(public + tuple(raw_arrays), dense_public + tuple(dense_raw)):
|
|
np.testing.assert_array_equal(a, b, err_msg=msg) # values and NaN masks
|
|
print(f"{mode}: dense vs {len(blocked)} fine-block layouts bit-identical "
|
|
f"after {STEPS} coarse steps", flush=True)
|
|
|
|
def test_passive_circle_dense_and_blocked_fine_storage_are_bit_identical(self):
|
|
from CelerisLab.cuda.context import CudaContext
|
|
for mode in ("periodic", "slot_reservoir"):
|
|
cfg = config(mode)
|
|
patch, boundary = fine_circle(mode)
|
|
solid = (boundary.flags.reshape(patch.fine_shape) & SOLID) != 0
|
|
g_coarse, g_fine = state(NY, NX, 1), state(*patch.fine_shape, 2)
|
|
results = []
|
|
with CudaContext(0) as context:
|
|
for fine_block in (None, *FINE_BLOCKS):
|
|
with FixedPatchStepper(cfg, context=context, coarse_boundary=mode,
|
|
refinements=[BOUNDS], level_blocks=[None, fine_block],
|
|
level_objects=[(), (circle(),)]) as backend:
|
|
backend.initialize(g_coarse, g_fine)
|
|
backend.step(STEPS)
|
|
results.append((backend.download(), raw(backend)))
|
|
(dense_public, dense_raw), blocked = results[0], results[1:]
|
|
fine = dense_raw[1]
|
|
self.assertTrue(np.isfinite(fine[:, (patch.fine_kind == ACTIVE) & ~solid]).all(), mode)
|
|
self.assertTrue(np.isnan(fine[:, solid]).all(), mode)
|
|
self.assertTrue(np.isnan(fine[:, patch.fine_kind == OUTER_GHOST]).all(), mode)
|
|
for fine_block, (public, raw_arrays) in zip(FINE_BLOCKS, blocked):
|
|
msg = f"circle {mode} fine_block={fine_block}"
|
|
for a, b in zip(public + tuple(raw_arrays), dense_public + tuple(dense_raw)):
|
|
np.testing.assert_array_equal(a, b, err_msg=msg) # values and NaN masks
|
|
print(f"circle {mode}: dense vs {len(blocked)} fine-block layouts bit-identical "
|
|
f"after {STEPS} coarse steps", flush=True)
|
|
|
|
def test_actuated_circle_dense_and_blocked_fine_storage_are_bit_identical(self):
|
|
from CelerisLab.cuda.context import CudaContext
|
|
for mode in ("periodic", "slot_reservoir"):
|
|
for speed in (.002, -.002):
|
|
cfg = config(mode)
|
|
patch, boundary = fine_circle(mode, actuation(speed))
|
|
solid = (boundary.flags.reshape(patch.fine_shape) & SOLID) != 0
|
|
g_coarse, g_fine = state(NY, NX, 1), state(*patch.fine_shape, 2)
|
|
results = []
|
|
with CudaContext(0) as context:
|
|
for fine_block in (None, *FINE_BLOCKS):
|
|
with FixedPatchStepper(cfg, context=context, coarse_boundary=mode,
|
|
refinements=[BOUNDS], level_blocks=[None, fine_block],
|
|
level_objects=[(), (circle(actuation(speed)),)]) as backend:
|
|
backend.initialize(g_coarse, g_fine)
|
|
backend.step(STEPS)
|
|
results.append((backend.download(), raw(backend)))
|
|
(dense_public, dense_raw), blocked = results[0], results[1:]
|
|
fine = dense_raw[1]
|
|
self.assertTrue(np.isfinite(fine[:, (patch.fine_kind == ACTIVE) & ~solid]).all())
|
|
self.assertTrue(np.isnan(fine[:, solid]).all())
|
|
self.assertTrue(np.isnan(fine[:, patch.fine_kind == OUTER_GHOST]).all())
|
|
for fine_block, (public, raw_arrays) in zip(FINE_BLOCKS, blocked):
|
|
msg = f"actuated {speed} {mode} fine_block={fine_block}"
|
|
for a, b in zip(public + tuple(raw_arrays), dense_public + tuple(dense_raw)):
|
|
np.testing.assert_array_equal(a, b, err_msg=msg) # values and NaN masks
|
|
print(f"actuated {speed:+.3f} {mode}: dense vs {len(blocked)} fine-block layouts "
|
|
f"bit-identical after {STEPS} coarse steps", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|