Frank14f/CelerisLab
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: opt-in MRT collision for blocked/refined stepping

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
Claude
2026-10-18 10:26:55 +00:00
co-authored by Claude Opus 5.5
parent 9805383f9c
commit 00d7a54829
8 changed files with 168 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/development/refinement_mrt.md
+34
View File
@@ -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.
src/CelerisLab/lbm/channel.py
+2 -2
View File
@@ -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
src/CelerisLab/lbm/kernels/blocks.cu
+13 -2
View File
@@ -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);
src/CelerisLab/lbm/kernels/refinement.cu
+12 -2
View File
@@ -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
src/CelerisLab/lbm/kernels/step/block_channel.cu
+8
View File
@@ -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);
src/CelerisLab/lbm/refinement.py
+4 -4
View File
@@ -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")
tests/unit/test_refinement_mrt.py
+40
View File
@@ -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
tests/validation/multigrid/test_refinement_mrt_gpu.py
+55
View File
@@ -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()