diff --git a/src/CelerisLab/lbm/refinement_stepper.py b/src/CelerisLab/lbm/refinement_stepper.py new file mode 100644 --- /dev/null +++ b/src/CelerisLab/lbm/refinement_stepper.py @@ -0,0 +1,207 @@ +# CelerisLab/lbm/refinement_stepper.py +"""Opt-in fixed periodic patch backend; no Simulation/channel/objects adapter. + +The caller owns the CudaContext and must close this backend before that context. +Only synchronized postcollision_g leaf populations are public. No checkpoint, +snapshot, forcing, runtime parameter mutation or dynamic refinement is offered. +""" +from dataclasses import dataclass, replace +from numbers import Integral +import struct +from threading import RLock + +import numpy as np +import pycuda.driver as cuda + +from ..cuda import compiler_v2 as compiler +from .refinement import ACTIVE, FixedPatch2D + + +@dataclass +class _Level: + shape: tuple + module: object + current: object + spare: object + kernel: object + grid: tuple + + +class PeriodicPatchStepper: + """D2Q9/SRT/FP32, one unforced 2:1 patch in a doubly periodic coarse box. + + ``initialize(coarse_g, fine_g)`` accepts physical Q-major (9,ny,nx) arrays + already in postcollision phase. It does not collide or build channel flags. + Every ``step()`` publishes only after both levels complete on one stream. + On a device/launch failure the instance becomes unusable: close and recreate + it, rather than pretending the partially advanced fine buffer rolled back. + """ + def __init__(self, cfg, bounds, *, context): + self.patch = FixedPatch2D.from_config(cfg, bounds) + self.ctx = context + self._lock = RLock() + self._status = "new" + self._step_count = 0 + self._allocations = [] + self._levels = [] + self._stream = None + self._block = (cfg.threads_per_block, 1, 1) + fine_cfg = replace(cfg, nx=self.patch.fine_shape[1], ny=self.patch.fine_shape[0], + viscosity=2 * cfg.viscosity) + with self.ctx.activate(): + try: + self._stream = cuda.Stream() + for config, name in ((replace(cfg), "PatchCoarseStep"), + (fine_cfg, "PatchFineStep")): + config.validate() + build = compiler.generate_config(config, n_objects=0) + try: + ptx = compiler.compile_kernel(self.ctx.sm_arch, source="refinement.cu", + build_dir=build) + module = compiler.load_module(ptx) + finally: + compiler.cleanup_build_dir(build) + ptr, _ = module.get_global("d_params") + cuda.memcpy_htod(ptr, struct.pack("=ffffI", config.omega, 0, 0, 0, 0)) + shape = (9, config.ny, config.nx) + size = int(np.prod(shape)) * 4 + self._levels.append(_Level(shape, module, self._allocate(size), + self._allocate(size), module.get_function(name), + ((config.nx + self._block[0] - 1) // self._block[0], + config.ny, 1))) + p = self.patch + order = np.argsort(p.coalesce_coarse) + self._ghosts = self._upload(p.ghost_fine) + self._donors = self._upload(p.ghost_coarse) + self._targets = self._upload(p.coalesce_coarse[order]) + self._children = self._upload(p.coalesce_fine[order]) + self._explode = self._levels[1].module.get_function("PatchExplode") + except BaseException: + try: + self.close() + except Exception: + pass # Preserve the original construction failure. + raise + + def _allocate(self, size): + pointer = cuda.mem_alloc(size) + self._allocations.append(pointer) + return pointer + + def _upload(self, array): + array = np.ascontiguousarray(array) + pointer = self._allocate(array.nbytes) + cuda.memcpy_htod(pointer, array) + return pointer + + def _require(self, *states): + if self._status not in states: + raise RuntimeError(f"Periodic patch is {self._status}; expected {' or '.join(states)}") + + @property + def step_count(self): + return self._step_count + + def _mask_nonleaves(self, coarse, fine): + x0, y0, x1, y1 = self.patch.bounds + coarse[:, y0:y1, x0:x1] = np.nan + fine[:, self.patch.fine_kind != ACTIVE] = np.nan + + def initialize(self, coarse_g, fine_g): + """Upload an explicit periodic G_0; invalid non-leaves are discarded.""" + with self._lock, self.ctx.activate(): + self._require("new", "ready") + arrays = [np.array(a, dtype=np.float32, order="C", copy=True) + for a in (coarse_g, fine_g)] + if any(a.shape != level.shape for a, level in zip(arrays, self._levels)): + raise ValueError("Expected Q-major (9,ny,nx) arrays for both levels") + x0, y0, x1, y1 = self.patch.bounds + leaf = np.ones(self.patch.coarse_shape, dtype=bool) + leaf[y0:y1, x0:x1] = False + for array, mask in zip(arrays, (leaf, self.patch.fine_kind == ACTIVE)): + if not np.isfinite(array[:, mask]).all() or (array[:, mask].sum(axis=0) <= 0).any(): + raise ValueError("Leaf populations must be finite with positive density") + self._mask_nonleaves(*arrays) + try: + for level, array in zip(self._levels, arrays): + cuda.memcpy_htod(level.current, array) + cuda.memcpy_htod(level.spare, array) + self._step_count = 0 + self._status = "ready" + except BaseException: + self._status = "failed" + raise + + def step(self, n=1): + """Advance G→G: one explosion, two fine collisions, one coarse collision.""" + if isinstance(n, bool) or not isinstance(n, Integral) or n < 0: + raise ValueError("n must be a nonnegative integer") + with self._lock, self.ctx.activate(): + self._require("ready") + c, f = self._levels + p = self.patch + kw = dict(block=self._block, stream=self._stream) + try: + self._status = "advancing" + for _ in range(n): + self._explode(c.current, f.current, self._ghosts, self._donors, + np.uint32(len(p.ghost_fine)), np.uint64(np.prod(p.coarse_shape)), + grid=((len(p.ghost_fine) + self._block[0] - 1) // self._block[0], 1), + **kw) + f.kernel(f.current, f.spare, grid=f.grid, **kw) + f.kernel(f.spare, f.current, grid=f.grid, **kw) + c.kernel(c.current, c.spare, f.current, self._targets, self._children, + np.uint32(len(p.coalesce_coarse)), + *(np.uint32(v) for v in p.bounds), grid=c.grid, **kw) + self._stream.synchronize() + c.current, c.spare = c.spare, c.current + self._step_count += 1 + self._status = "ready" + except BaseException: + self._status = "failed" + raise + + def download(self): + """Return two synchronized G arrays, NaN on all non-leaves; no host cache.""" + with self._lock, self.ctx.activate(): + self._require("ready") + arrays = [np.empty(level.shape, np.float32) for level in self._levels] + for level, array in zip(self._levels, arrays): + cuda.memcpy_dtoh(array, level.current) + self._mask_nonleaves(*arrays) + return tuple(arrays) + + def close(self): + """Release owned allocations/modules; the caller's context stays open.""" + with self._lock: + if self._status == "closed": + return + self._status = "closed" + with self.ctx.activate(): + error = None + try: + if self._stream is not None: + self._stream.synchronize() + except Exception as exc: + error = exc + self._levels.clear() + self._explode = None + for pointer in self._allocations: + try: + pointer.free() + except Exception as exc: + error = error or exc + self._allocations.clear() + self._stream = None + if error is not None: + raise error + + def __enter__(self): + return self + + def __exit__(self, *exc): + try: + self.close() + except Exception: + if exc[0] is None: + raise diff --git a/src/CelerisLab/lbm/kernels/refinement.cu b/src/CelerisLab/lbm/kernels/refinement.cu new file mode 100644 --- /dev/null +++ b/src/CelerisLab/lbm/kernels/refinement.cu @@ -0,0 +1,105 @@ +// Fixed, unforced, doubly periodic 2:1 patch. Separate compilation entry; +// default OneStep/channel kernels are deliberately not included. +#include +#include +#include "config.h" +#include "core/precision.cuh" +#include "core/descriptors.cuh" +#include "core/layout.cuh" +#include "core/params.cuh" +#include "operators/macro.cuh" +#include "operators/collision_srt.cuh" +#include "streaming/pull_double_buffer.cuh" + +#if DIM != 2 || NQ != 9 || COLLISION_MODEL != 0 || STORE_PRECISION != 0 || USE_DDF_SHIFTING || USE_LES || STREAMING_MODEL != 0 +#error "Refinement requires unshifted FP32 D2Q9 SRT double buffers without LES" +#endif + +// Enabled only in the validation build; no diagnostic arrays in normal builds. +#ifdef CELERIS_REFINEMENT_AUDIT +extern "C" { +__device__ unsigned int collision_counts[TOTAL_CELLS]; +__device__ unsigned int explosion_counts[TOTAL_CELLS]; +__device__ unsigned int incoming_counts[NQ * TOTAL_CELLS]; +} +#endif + +__device__ __forceinline__ void patch_collide(float* f, unsigned long k) { + float rho, ux, uy, feq[NQ]; + compute_rho_u(f, rho, ux, uy); + compute_feq(rho, ux, uy, feq); + collide_srt_no_force(f, feq, d_params.omega); +#ifdef CELERIS_REFINEMENT_AUDIT + ++collision_counts[k]; +#endif +} + +extern "C" __global__ void PatchExplode( + const float* coarse, float* fine, const uint32_t* targets, + const uint32_t* donors, unsigned int count, unsigned long coarse_cells) +{ + unsigned int row = blockIdx.x * blockDim.x + threadIdx.x; + if (row >= count) return; + for (unsigned int q = 0; q < NQ; ++q) + fine[index_f(targets[row], q)] = coarse[q * coarse_cells + donors[row]]; +#ifdef CELERIS_REFINEMENT_AUDIT + ++explosion_counts[targets[row]]; +#endif +} + +extern "C" __global__ void PatchFineStep(const float* in, float* out) { + unsigned int x, y; unsigned long k; + index_from_thread(x, y, k); + if (x >= NX || y >= NY) return; + float f[NQ]; + // Outer halo is invalid after EACH hop, including when reusing the buffer + // that held explosion data. Inner halo only streams; it never collides. + if (x < 2 || y < 2 || x >= NX - 2 || y >= NY - 2) { + for (int q = 0; q < NQ; ++q) f[q] = NAN; + } else { + unsigned long j[NQ]; + compute_neighbors(x, y, j); // These interior targets never wrap. + stream_pull_load(k, f, in, j); + if (x >= 4 && y >= 4 && x < NX - 4 && y < NY - 4) + patch_collide(f, k); + } + stream_pull_store(k, f, out); +} + +extern "C" __global__ void PatchCoarseStep( + const float* in, float* out, const float* fine, + const uint64_t* targets, const uint64_t* children, unsigned int count, + unsigned int x0, unsigned int y0, unsigned int x1, unsigned int y1) +{ + unsigned int x, y; unsigned long k; + index_from_thread(x, y, k); + if (x >= NX || y >= NY) return; + // Both buffers were poisoned at initialization; covered cells stay invalid. + if (x >= x0 && x < x1 && y >= y0 && y < y1) return; + unsigned long j[NQ]; + compute_neighbors(x, y, j); + float f[NQ]; + stream_pull_load(k, f, in, j); + if (x >= x0 - 1 && x <= x1 && y >= y0 - 1 && y <= y1) { + for (unsigned int q = 1; q < NQ; ++q) { + uint64_t slot = index_f(k, q); + unsigned int lo = 0, hi = count; + while (lo < hi) { + unsigned int mid = lo + (hi - lo) / 2; + if (targets[mid] < slot) lo = mid + 1; + else hi = mid; + } + if (lo < count && targets[lo] == slot) { + const uint64_t* s = children + 4ul * lo; + f[q] = 0.25f * (fine[s[0]] + fine[s[1]] + fine[s[2]] + fine[s[3]]); +#ifdef CELERIS_REFINEMENT_AUDIT + ++incoming_counts[slot]; +#endif + } + } + } + // Includes the eight grazing-corner links with UNCOVERED coarse donors. + // All replacement arrivals must be in place before computing rho/u. + patch_collide(f, k); + stream_pull_store(k, f, out); +} diff --git a/tests/validation/multigrid/test_refinement_gpu.py b/tests/validation/multigrid/test_refinement_gpu.py new file mode 100644 --- /dev/null +++ b/tests/validation/multigrid/test_refinement_gpu.py @@ -0,0 +1,222 @@ +"""Bounded FP32 GPU / pinned FP64 CPUREF-0.3 equivalence, not accuracy validation. + +Run explicitly in pycuda_3_10 with PYTHONPATH=/src, one BLAS thread, +and CUDA_VISIBLE_DEVICES=0 (or 1) only after checking the physical GPU is idle. +At most eight coarse steps per case. No files, plots or flow arrays are saved. +""" +import hashlib +import importlib.util +from contextlib import nullcontext +from pathlib import Path +import resource +import sys +from threading import RLock +import unittest +from unittest.mock import Mock, patch + +import numpy as np +import pycuda.driver as cuda + +import CelerisLab +from CelerisLab.config import LBMConfig +from CelerisLab.cuda import compiler_v2 as compiler +from CelerisLab.cuda.context import CudaContext +from CelerisLab.lbm.refinement_stepper import PeriodicPatchStepper + +ROOT = Path(__file__).resolve().parents[3] +POP_ATOL = 3e-6 +MOMENT_ATOL = 2e-4 +REFERENCE_SHA = "bba2a06a9f6517b32079ffff15e4060d8c3b7ed35f16ca134e1877226c9e4e1e" +REF_PATH = Path(__file__).with_name("cpu_reference.py") +assert hashlib.sha256(REF_PATH.read_bytes()).hexdigest() == REFERENCE_SHA +loader = importlib.util.spec_from_file_location("pinned_patch_reference", REF_PATH) +ref = importlib.util.module_from_spec(loader) +sys.modules[loader.name] = ref +loader.loader.exec_module(ref) + + +def postcollision(state, spec, maps): + # Observation conversion only. Never applied to production G. + return tuple(ref._collide(spec, a, np.isfinite(a), tau, mask, + np.zeros(mask.shape, dtype=np.int64)) + for a, tau, mask in ((state.coarse_f, spec.tau_coarse, maps.coarse_active), + (state.fine_f, spec.tau_fine, maps.fine_active))) + + +def totals(arrays, maps): + # Leaf sum independent of the reference's flux/budget bookkeeping helpers. + c, f = arrays + populations = c[:, maps.coarse_active].astype(np.float64).sum(axis=1) + populations += .25 * f[:, maps.fine_active].astype(np.float64).sum(axis=1) + return np.array([populations.sum(), populations @ ref.CX, populations @ ref.CY]) + + +def global_array(module, name, shape): + out = np.empty(shape, dtype=np.uint32) + ptr, size = module.get_global(name) + assert size == out.nbytes + cuda.memcpy_dtoh(out, ptr) + return out + + +class CoupledGPU(unittest.TestCase): + def test_cleanup_attempts_every_allocation(self): + backend = PeriodicPatchStepper.__new__(PeriodicPatchStepper) + backend.ctx = Mock(activate=nullcontext) + backend._lock = RLock() + backend._status = "failed" + backend._stream = Mock() + backend._stream.synchronize.side_effect = RuntimeError("original stream failure") + first, second = Mock(), Mock() + first.free.side_effect = RuntimeError("secondary free failure") + backend._allocations = [first, second] + backend._levels = [Mock()] + with self.assertRaisesRegex(RuntimeError, "original stream failure"): + backend.close() + first.free.assert_called_once() + second.free.assert_called_once() + self.assertEqual(backend._allocations, []) + self.assertEqual(backend._levels, []) + backend.close() + + def test_partial_advance_cannot_be_published(self): + spec = ref.make_spec(collision="srt") + maps = ref.build_maps(spec) + with CudaContext(0) as context: + with PeriodicPatchStepper(LBMConfig(nx=16, ny=16, viscosity=.1), + (4, 4, 12, 12), context=context) as backend: + backend.initialize(*postcollision(ref.make_state(spec), spec, maps)) + coarse, fine = backend._levels + + def fail_after_fine(*args, **kwargs): + # Ensure both queued fine hops actually completed before failing. + backend._stream.synchronize() + raise RuntimeError("coarse launch failed after two fine hops") + + with patch.object(coarse, "kernel", side_effect=fail_after_fine): + with self.assertRaisesRegex(RuntimeError, "after two fine hops"): + backend.step() + self.assertEqual(backend.step_count, 0) + for operation in (backend.download, backend.step): + with self.assertRaisesRegex(RuntimeError, "failed"): + operation() + raw = np.empty(fine.shape, np.float32) + cuda.memcpy_dtoh(raw, fine.current) + self.assertTrue(np.isfinite(raw.ravel()[backend.patch.coalesce_fine]).all()) + self.assertFalse(context.is_closed) + print("partial failure: two fine hops completed, no coarse step published; cleanup PASS", flush=True) + + def test_bounded_equivalence_and_ownership(self): + self.assertEqual(Path(CelerisLab.__file__).resolve(), ROOT / "src/CelerisLab/__init__.py") + self.assertEqual(Path(compiler._KERNEL_DIR).resolve(), ROOT / "src/CelerisLab/lbm/kernels") + spec = ref.make_spec(collision="srt") + maps = ref.build_maps(spec) + cfg = LBMConfig(nx=16, ny=16, viscosity=.1) + compile_original = compiler.compile_kernel + + def audit_compile(*args, **kwargs): + # Only this private build's entry differs. Production source unchanged. + entry = Path(kwargs["build_dir"]) / "refinement_audit.cu" + entry.write_text('#define CELERIS_REFINEMENT_AUDIT\n#include "refinement.cu"\n') + kwargs["source"] = entry.name + return compile_original(*args, **kwargs) + + with CudaContext(0) as context: + with patch.object(compiler, "compile_kernel", audit_compile): + backend = PeriodicPatchStepper(cfg, (4, 4, 12, 12), context=context) + with backend: + # Production allocations: four DDF arrays and four compact tables. + device_bytes = 2 * 9 * (16**2 + 24**2) * 4 + p = backend.patch + device_bytes += sum(a.nbytes for a in (p.ghost_fine, p.ghost_coarse, + p.coalesce_coarse, p.coalesce_fine)) + audit_bytes = sum((2 + 9) * np.prod(level.shape[1:]) * 4 + for level in backend._levels) + self.assertLess(device_bytes + audit_bytes, 128 * 1024) + print(f"device_arrays={device_bytes} audit_arrays={audit_bytes} bytes", flush=True) + for case in ("rest", "nonequilibrium_seed1"): + with self.subTest(case=case): + state = ref.make_state(spec, velocity=(0, 0) if case == "rest" else (.03, -.02)) + if case != "rest": + rng = np.random.default_rng(1) + state = ref.State(*(a * (1 + 1e-3 * rng.standard_normal(a.shape)) + for a in (state.coarse_f, state.fine_f))) + initial = postcollision(state, spec, maps) + backend.initialize(*initial) + for level in backend._levels: + for name in ("collision_counts", "explosion_counts", "incoming_counts"): + ptr, size = level.module.get_global(name) + cuda.memset_d8(ptr, 0, size) + before = totals(backend.download(), maps) + max_error = max_drift = max_moment_error = max_corner = 0.0 + for step in range(1, 9): + state, diagnostic = ref.advance_coarse_once(spec, maps, state) + backend.step() + actual = backend.download() + expected = postcollision(state, spec, maps) + self.assertEqual(backend.step_count, step) + p.check_synchronization(step, state.tick, phase="postcollision_g") + for got, want in zip(actual, expected): + np.testing.assert_array_equal(np.isnan(got), np.isnan(want)) + error = float(np.nanmax(np.abs(got - want))) + self.assertLessEqual(error, POP_ATOL) + max_error = max(max_error, error) + drift = float(np.max(np.abs(totals(actual, maps) - before))) + moment_error = float(np.max(np.abs(totals(actual, maps) - totals(expected, maps)))) + self.assertLessEqual(drift, MOMENT_ATOL) + self.assertLessEqual(moment_error, MOMENT_ATOL) + max_drift, max_moment_error = max(max_drift, drift), max(max_moment_error, moment_error) + c, f = backend._levels + np.testing.assert_array_equal(global_array(c.module, "collision_counts", (16, 16)), + step * maps.coarse_active) + np.testing.assert_array_equal(global_array(f.module, "collision_counts", (24, 24)), + 2 * step * maps.fine_active) + np.testing.assert_array_equal(global_array(f.module, "explosion_counts", (24, 24)), + step * ~maps.fine_active) + owners = np.zeros((9, 16, 16), np.uint32) + owners.ravel()[p.coalesce_coarse] = step + np.testing.assert_array_equal(global_array(c.module, "incoming_counts", owners.shape), owners) + raw = np.empty(f.shape, np.float32) + cuda.memcpy_dtoh(raw, f.current) + np.testing.assert_array_equal(np.isfinite(raw), diagnostic["hop2_valid"]) + self.assertTrue(np.isfinite(raw.ravel()[p.coalesce_fine]).all()) + self.assertEqual(len(diagnostic["corner_packets"]), 8) + for packet in diagnostic["corner_packets"].values(): + sx, sy = packet["source_coarse_xy"] + self.assertTrue(maps.coarse_active[sy, sx]) + values = [raw[packet["direction"], y, x] + for y, x in packet["children_fine_yx"]] + error = float(np.max(np.abs(np.array(values) - packet["packets"]))) + self.assertLessEqual(error, POP_ATOL) + max_corner = max(max_corner, error) + print(f"{case}: steps=8 max_population={max_error:.9g} " + f"max_leaf_drift={max_drift:.9g} max_leaf_error={max_moment_error:.9g} " + f"max_corner_packet={max_corner:.9g}; ownership/counts PASS", flush=True) + + # Fail before any additional time step; no stale synchronized export. + with patch.object(backend, "_explode", side_effect=RuntimeError("injected launch failure")): + with self.assertRaisesRegex(RuntimeError, "injected"): + backend.step() + with self.assertRaisesRegex(RuntimeError, "failed"): + backend.download() + with self.assertRaisesRegex(RuntimeError, "failed"): + backend.step() + self.assertFalse(context.is_closed) + with self.assertRaisesRegex(RuntimeError, "closed"): + backend.download() + + # Compile and execute the actual non-audited entry as well: one step. + with PeriodicPatchStepper(cfg, (4, 4, 12, 12), context=context) as plain: + rest = postcollision(ref.make_state(spec), spec, maps) + plain.initialize(*rest) + plain.step() + for got, expected in zip(plain.download(), rest): + np.testing.assert_allclose(got, expected, rtol=0, atol=POP_ATOL, equal_nan=True) + print("normal-build rest: steps=1 PASS; source/kernel paths pinned", flush=True) + rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 + self.assertLess(rss, 512) + print(f"host_peak_rss={rss:.3f} MiB", flush=True) + + +if __name__ == "__main__": + unittest.main(verbosity=2)