Files
CelerisLab/tests/validation/multigrid/test_refinement_blocked_fine.py
T
ClaudeandClaude Opus 5.5 59921a0181 Blocked fine: fix INNER assertion, skip empty halo uploads, explicit fixture
- GPU test: after hop 2 some INNER_GHOST populations are NaN (pulled from
  OUTER poisoned on hop 1; 236 slots in pinned CPUREF); assert ACTIVE
  finite and OUTER NaN, keep the full dense-vs-blocked raw value/NaN-mask
  comparison.
- A 1x1 non-periodic fine layout has zero halo pairs: skip the halo table
  uploads (CUDA rejects mem_alloc(0)); _fine_halo already skips launch.
  CPU guard test with driver doubles.
- Drop the class-level fine_layout default; the followup fake backend now
  sets fine_layout = None explicitly.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-15 03:52:42 +00:00

157 lines
8.0 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
from CelerisLab.config import LBMConfig
from CelerisLab.cuda import compiler_v2 as compiler
from CelerisLab.lbm.blocks import BlockLayout2D
from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY
from CelerisLab.lbm.refinement import ACTIVE, INNER_GHOST, OUTER_GHOST, FixedPatch2D
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 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):
circle = SimObject(0, geometry=CircleGeometry(16.0, 12.0, 2.0), center=(16.0, 12.0), radius=2.0)
for kw in ({"fine_block": (7, 7)}, {"fine_block": (16,)},
{"fine_block": (16, 20), "fine_objects": [circle]}):
with self.subTest(kw=kw), self.assertRaises(ValueError):
FixedPatchStepper(config("periodic"), BOUNDS, context=self.forbidden_context(), **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"), BOUNDS, context=context,
fine_block=FINE_BLOCKS[0])
self.assertEqual(len(backend.fine_layout.halo_dst), 0)
self.assertIsNone(backend._halo_dst)
backend._fine_halo(Mock()) # no launch without pairs
module.get_function.return_value.assert_not_called()
backend.close()
self.assertTrue(sizes)
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, BOUNDS, context=context, coarse_boundary=mode,
fine_block=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)
if __name__ == "__main__":
unittest.main()