From 00d7a54829175f565027adb3a9f6104c542ec90a Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 18 Oct 2026 10:26:55 +0000 Subject: [PATCH] feat: opt-in MRT collision for blocked/refined stepping Additive compile-time COLLISION_MODEL == 2 branches in BlockStep, BlockChannelStep and patch_collide repeat the uniform collide_dispatch MRT sequence (zero forcing, omega clamp, collide_mrt_no_force). SRT statements, including the shifted retained-delta/direct-momentum path, are unchanged under the alternate branch. Guards accept SRT or MRT only; TRT, LES, non-FP32, non-double-buffer and shifted dense/mixed refinement stay rejected. Transfers, subcycling, BFL, raw ledgers, channel closures, shifted-Jet SRT guard and PeriodicBlockStepper are untouched. Existing MRT operator, shifted representation and fixed kinetic rates are reused. Co-Authored-By: Claude Opus 5.5 --- docs/development/refinement_mrt.md | 34 ++++++++++++ src/CelerisLab/lbm/channel.py | 4 +- src/CelerisLab/lbm/kernels/blocks.cu | 15 ++++- src/CelerisLab/lbm/kernels/refinement.cu | 14 ++++- .../lbm/kernels/step/block_channel.cu | 8 +++ src/CelerisLab/lbm/refinement.py | 8 +-- tests/unit/test_refinement_mrt.py | 40 ++++++++++++++ .../multigrid/test_refinement_mrt_gpu.py | 55 +++++++++++++++++++ 8 files changed, 168 insertions(+), 10 deletions(-) create mode 100644 docs/development/refinement_mrt.md create mode 100644 tests/unit/test_refinement_mrt.py create mode 100644 tests/validation/multigrid/test_refinement_mrt_gpu.py diff --git a/docs/development/refinement_mrt.md b/docs/development/refinement_mrt.md new file mode 100644 index 0000000..2aa1554 --- /dev/null +++ b/docs/development/refinement_mrt.md @@ -0,0 +1,34 @@ +# MRT reuse in blocked/refined stepping (candidate) + +Base `9805383`. Opt-in by `LBMConfig(collision="MRT")`; SRT remains the default and its +kernel statements, including the shifted retained-delta/direct-momentum path, are unchanged +under the alternate compile branch. + +## What changes +- `blocks.cu` (`BlockStep`), `step/block_channel.cu` (`BlockChannelStep`) and `refinement.cu` + (`patch_collide`) gain a compile-time `COLLISION_MODEL == 2` branch that repeats the uniform + `collide_dispatch` sequence with zero forcing: `compute_rho_u`, omega clamp to + `[OMEGA_COLLISION_MIN, OMEGA_COLLISION_MAX]`, `collide_mrt_no_force`. +- Kernel and host guards accept SRT or MRT only. TRT, LES, non-FP32 storage, non-double-buffer + streaming, and shifted dense/mixed refinement remain rejected. +- Transfers, subcycling, BFL, raw exchange ledgers and channel closures are untouched. The + existing pressure-NEQ outlet selects its MRT branch (unknowns-only NEQ extrapolation) from + `COLLISION_MODEL`, so MRT channel parity must use uniform MRT as reference. + +## Contracts (existing operator, not redefined here) +- Shifted storage uses the existing MRT representation (equilibrium moments from `rho - 1`), + not the SRT retained-delta/direct-momentum precision. Mass and momentum are conserved + moments (relaxation rate 0). +- Kinetic rates are the operator's fixed `s_e = s_eps = s_q = 1.2` on every level; the shear + rate is the existing per-level omega. Under acoustic 2:1 scaling the physical bulk viscosity + therefore halves on each finer level. No rate rescaling is applied. +- No LES, no external forcing (`d_params` force fields stay zero). +- Shifted Jets stay SRT-only: `body/jets.py` keeps its existing guard. This extension does not + provide shifted MRT Jets. +- `PeriodicBlockStepper` (uniform blocked stepper) remains SRT-only. + +## Evidence status +CPU: `tests/unit/test_refinement_mrt.py` (host guards, source contract). GPU: +`tests/validation/multigrid/test_refinement_mrt_gpu.py` (uniform-MRT channel parity, AMR MRT +smoke) and unchanged SRT GPU tests. Strict parity may be subject to cross-kernel FMA +contraction differences already documented for SRT. Not a DFG/benchmark validation. diff --git a/src/CelerisLab/lbm/channel.py b/src/CelerisLab/lbm/channel.py index 8915f0c..146dc59 100644 --- a/src/CelerisLab/lbm/channel.py +++ b/src/CelerisLab/lbm/channel.py @@ -1,4 +1,4 @@ -"""Host contract for the blocked root's uniform-SRT channel boundary. +"""Host contract for the blocked root's uniform SRT/MRT channel boundary. Boundary ghosts are evolved state, but are not physical fluid leaves. Halo ownership remains periodic to reproduce the uniform kernel's wrapped pulls. @@ -29,7 +29,7 @@ def validate_channel_config(cfg, blocks): def validate_channel_transfers(pair, layout, targets): """Reject actual exchange/closure overlap, not a padded geometric envelope. - Explosion donors and restriction receivers must use ordinary bulk SRT + Explosion donors and restriction receivers must use ordinary bulk collision semantics. Restriction writes receiver-local source slots; these may not overwrite the channel's extra immutable-input reads (pressure donor and same-node halfway-wall corrections). The slot check deliberately respects diff --git a/src/CelerisLab/lbm/kernels/blocks.cu b/src/CelerisLab/lbm/kernels/blocks.cu index add72db..746a4b8 100644 --- a/src/CelerisLab/lbm/kernels/blocks.cu +++ b/src/CelerisLab/lbm/kernels/blocks.cu @@ -14,12 +14,15 @@ #include "core/obs.cuh" #include "operators/macro.cuh" #include "operators/collision_srt.cuh" +#if COLLISION_MODEL == 2 +#include "operators/collision_mrt.cuh" +#endif #include "streaming/pull_double_buffer.cuh" #include "boundary/curved_boundary.cuh" #include "boundary/planar_velocity.cuh" -#if DIM != 2 || NQ != 9 || COLLISION_MODEL != 0 || STORE_PRECISION != 0 || USE_LES || STREAMING_MODEL != 0 -#error "Blocks require FP32 D2Q9 SRT double buffers without LES" +#if DIM != 2 || NQ != 9 || (COLLISION_MODEL != 0 && COLLISION_MODEL != 2) || STORE_PRECISION != 0 || USE_LES || STREAMING_MODEL != 0 +#error "Blocks require FP32 D2Q9 SRT or MRT double buffers without LES" #endif // Fixed-patch cell kinds, matching refinement.py ACTIVE/INNER_GHOST/OUTER_GHOST. @@ -141,6 +144,13 @@ extern "C" __global__ void BlockStep(const fpxx* in, fpxx* out, unsigned int hal gx, d_block_wall.nx, d_block_wall.slot_width, d_block_wall.amplitude)); else if (gy == d_block_wall.ny - 2) reconstruct_top_reservoir(f, RHO); } +#if COLLISION_MODEL == 2 + // Uniform MRT sequence (collide_dispatch, zero forcing); SRT path below unchanged. + float rho, ux, uy; + compute_rho_u(f, rho, ux, uy); + collide_mrt_no_force(f, rho, ux, uy, + fminf(OMEGA_COLLISION_MAX, fmaxf(OMEGA_COLLISION_MIN, d_params.omega))); +#else float rho, ux, uy, feq[NQ]; compute_rho_u(f, rho, ux, uy); #if USE_DDF_SHIFTING @@ -153,6 +163,7 @@ extern "C" __global__ void BlockStep(const fpxx* in, fpxx* out, unsigned int hal compute_feq(rho, ux, uy, feq); #endif collide_srt_no_force(f, feq, d_params.omega); +#endif } } stream_pull_store(k, f, out + base); diff --git a/src/CelerisLab/lbm/kernels/refinement.cu b/src/CelerisLab/lbm/kernels/refinement.cu index ab4bf20..93c79b9 100644 --- a/src/CelerisLab/lbm/kernels/refinement.cu +++ b/src/CelerisLab/lbm/kernels/refinement.cu @@ -11,6 +11,9 @@ #include "core/obs.cuh" #include "operators/macro.cuh" #include "operators/collision_srt.cuh" +#if COLLISION_MODEL == 2 +#include "operators/collision_mrt.cuh" +#endif #include "streaming/pull_double_buffer.cuh" #include "boundary/planar_velocity.cuh" #include "boundary/curved_boundary.cuh" @@ -19,8 +22,8 @@ extern "C" { #include "step/curved_double.cu" } -#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" +#if DIM != 2 || NQ != 9 || (COLLISION_MODEL != 0 && COLLISION_MODEL != 2) || STORE_PRECISION != 0 || USE_DDF_SHIFTING || USE_LES || STREAMING_MODEL != 0 +#error "Refinement requires unshifted FP32 D2Q9 SRT or MRT double buffers without LES" #endif // Enabled only in the validation build; no diagnostic arrays in normal builds. @@ -33,10 +36,17 @@ __device__ unsigned int incoming_counts[NQ * TOTAL_CELLS]; #endif __device__ __forceinline__ void patch_collide(float* f, unsigned long k) { +#if COLLISION_MODEL == 2 + float rho, ux, uy; + compute_rho_u(f, rho, ux, uy); + collide_mrt_no_force(f, rho, ux, uy, + fminf(OMEGA_COLLISION_MAX, fmaxf(OMEGA_COLLISION_MIN, d_params.omega))); +#else 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); +#endif #ifdef CELERIS_REFINEMENT_AUDIT ++collision_counts[k]; #endif diff --git a/src/CelerisLab/lbm/kernels/step/block_channel.cu b/src/CelerisLab/lbm/kernels/step/block_channel.cu index 4e95674..17e90fb 100644 --- a/src/CelerisLab/lbm/kernels/step/block_channel.cu +++ b/src/CelerisLab/lbm/kernels/step/block_channel.cu @@ -54,12 +54,20 @@ extern "C" __global__ void BlockChannelStep( f[7] = load_ddf(in + base, index_f(k, 8)); } if (!wall && !outlet) { +#if COLLISION_MODEL == 2 + // Match uniform MRT arithmetic (collide_dispatch, zero forcing). + float rho, ux, uy; + compute_rho_u(f, rho, ux, uy); + collide_mrt_no_force(f, rho, ux, uy, + fminf(OMEGA_COLLISION_MAX, fmaxf(OMEGA_COLLISION_MIN, d_params.omega))); +#else // Match uniform SRT arithmetic, including its shifted equilibrium path. float rho, ux, uy, feq[NQ], forcing[NQ] = {}; compute_rho_u(f, rho, ux, uy); compute_feq(rho, ux, uy, feq); const float omega = fminf(OMEGA_COLLISION_MAX, fmaxf(OMEGA_COLLISION_MIN, d_params.omega)); collide_srt(f, feq, forcing, omega); +#endif } } stream_pull_store(k, f, out + base); diff --git a/src/CelerisLab/lbm/refinement.py b/src/CelerisLab/lbm/refinement.py index ab9a61c..37d04b9 100644 --- a/src/CelerisLab/lbm/refinement.py +++ b/src/CelerisLab/lbm/refinement.py @@ -52,17 +52,17 @@ class FixedPatch2D: def from_config(cls, cfg: LBMConfig, bounds) -> "FixedPatch2D": """Prepare maps without changing cfg or allocating coarse-sized tables. - The initial backend contract is D2Q9/SRT/FP32/double_buffer, + The backend contract is D2Q9/SRT or MRT/FP32/double_buffer, without LES. Two coarse cells of margin keep all halo donors in bounds; this does NOT certify that those donors are fluid or periodic. """ checked = replace(cfg) checked.validate() - if (not checked.is_d2q9 or checked.collision != "SRT" + if (not checked.is_d2q9 or checked.collision not in ("SRT", "MRT") or checked.streaming != "double_buffer" or checked.store_precision != "FP32" or checked.les_enabled): - raise ValueError("FixedPatch2D requires D2Q9/SRT/FP32/double_buffer without LES") + raise ValueError("FixedPatch2D requires D2Q9/SRT or MRT/FP32/double_buffer without LES") bounds = tuple(bounds) if len(bounds) != 4 or any(isinstance(v, (bool, np.bool_)) or not isinstance(v, Integral) for v in bounds): @@ -451,7 +451,7 @@ def prepare_fine_boundary(cfg, patch, objects, *, wall_aware=False, jets=(), checked = replace(cfg) checked.validate() if (patch.coarse_shape != (cfg.ny, cfg.nx) or not checked.is_d2q9 - or checked.collision != "SRT" or checked.streaming != "double_buffer" + or checked.collision not in ("SRT", "MRT") or checked.streaming != "double_buffer" or checked.store_precision != "FP32" or checked.les_enabled or patch.tau_coarse != 3*checked.viscosity+.5): raise ValueError("Fine boundary requires the matching fixed-patch configuration") diff --git a/tests/unit/test_refinement_mrt.py b/tests/unit/test_refinement_mrt.py new file mode 100644 index 0000000..07f0db8 --- /dev/null +++ b/tests/unit/test_refinement_mrt.py @@ -0,0 +1,40 @@ +"""Host guards and source contract for MRT reuse in refined/blocked stepping (no CUDA).""" +from pathlib import Path +import pytest + +from CelerisLab.config import LBMConfig +from CelerisLab.lbm.refinement import FixedPatch2D + +KERNELS = Path(__file__).resolve().parents[2] / "src/CelerisLab/lbm/kernels" + + +@pytest.mark.parametrize("collision", ["SRT", "MRT"]) +def test_patch_accepts_implemented_models(collision): + FixedPatch2D.from_config(LBMConfig(nx=24, ny=20, viscosity=.1, collision=collision), (8, 6, 16, 14)) + + +@pytest.mark.parametrize("options", [dict(collision="TRT"), dict(collision="MRT", les_enabled=True)]) +def test_patch_rejects_unimplemented(options): + with pytest.raises(ValueError): + FixedPatch2D.from_config(LBMConfig(nx=24, ny=20, viscosity=.1, **options), (8, 6, 16, 14)) + + +def test_srt_statements_and_guards_preserved(): + """MRT is an additive compile-time branch; SRT statements and other guards remain.""" + blocks = (KERNELS / "blocks.cu").read_text() + channel = (KERNELS / "step/block_channel.cu").read_text() + refinement = (KERNELS / "refinement.cu").read_text() + for text in ("compute_feq_from_momentum(drho, jx, jy, ux, uy, feq);", + "compute_feq(rho, ux, uy, feq);", + "collide_srt_no_force(f, feq, d_params.omega);"): + assert text in blocks + assert "collide_srt(f, feq, forcing, omega);" in channel + assert "collide_srt_no_force(f, feq, d_params.omega);" in refinement + for source in (blocks, refinement): + assert "(COLLISION_MODEL != 0 && COLLISION_MODEL != 2)" in source + assert "USE_LES" in source and "STORE_PRECISION != 0" in source + assert "USE_DDF_SHIFTING ||" in refinement # dense/mixed shifted stays rejected + for source in (blocks, channel, refinement): + assert source.count("collide_mrt_no_force(") == 1 # one MRT-only collision site each + for source in (blocks, refinement): + assert '#if COLLISION_MODEL == 2\n#include "operators/collision_mrt.cuh"\n#endif' in source diff --git a/tests/validation/multigrid/test_refinement_mrt_gpu.py b/tests/validation/multigrid/test_refinement_mrt_gpu.py new file mode 100644 index 0000000..196c03c --- /dev/null +++ b/tests/validation/multigrid/test_refinement_mrt_gpu.py @@ -0,0 +1,55 @@ +"""MRT reuse GPU gates. Caller owns GPU admission (CUDA_VISIBLE_DEVICES). + +Run: python -m unittest discover -s tests/validation/multigrid -p test_refinement_mrt_gpu.py +Parity compares the blocked MRT channel root with uniform MRT OneStep: the existing +pressure-NEQ outlet branch differs by collision model, so SRT is never the reference. +AMR checks are finite-evolution/split-invariance smoke tests, not benchmark validation. +""" +import unittest +import numpy as np + +from CelerisLab.config import LBMConfig +from CelerisLab.cuda.context import CudaContext +from CelerisLab.lbm.population_storage import _encode_physical +from CelerisLab.lbm.refinement_stepper import FixedPatchStepper +from test_channel_root_gpu import run_root, state + + +class RefinementMRTGPU(unittest.TestCase): + def test_uniform_matched_mrt_root(self): + with CudaContext(0) as context: + for shifted in (False, True): + for profile in ("uniform", "parabolic"): + cfg = LBMConfig(nx=24, ny=20, viscosity=.1, velocity=.025, collision="MRT", + ddf_shifting=shifted, inlet_profile=profile) + initial = _encode_physical(state(cfg.ny, cfg.nx), shifted) + expected = run_root(cfg, context, initial) + self.assertTrue(np.isfinite(expected).all()) + for block in ((20, 24), (5, 6), (4, 1)): + np.testing.assert_array_equal(run_root(cfg, context, initial, block), expected) + + def test_amr_mrt_smoke(self): + from CelerisLab.body.geometry.circle import CircleGeometry + from CelerisLab.body.objects import SimObject + center = (24., 16.) + body = SimObject(0, geometry=CircleGeometry(*center, 2.15), center=center, radius=2.15) + with CudaContext(0) as context: + for boundary in ("periodic", "channel"): + for shifted in (False, True): + cfg = LBMConfig(nx=48, ny=32, viscosity=.1, velocity=.015 if boundary == "channel" else 0., + collision="MRT", ddf_shifting=shifted) + reference = None + for block in ((32, 48), (8, 8), (4, 1)): + with FixedPatchStepper(cfg, (8, 6, 40, 26), context=context, coarse_boundary=boundary, + base_block=block, fine_block=(8, 8), fine_objects=(body,)) as backend: + backend.initialize(*[state(*level.shape[1:]) for level in backend._levels]) + backend.step(8) + raw = backend._download_raw() + for a, mask in zip(raw.values, backend._state): + self.assertTrue(np.isfinite(a[:, mask]).all()) + if reference is None: reference = raw.values + else: + for a, b in zip(reference, raw.values): np.testing.assert_array_equal(a, b) + + +if __name__ == "__main__": unittest.main()