From 3a374faace34c837981746aeff193f1a89db74f1 Mon Sep 17 00:00:00 2001 From: Frank14f <1515444314@qq.com> Date: Wed, 21 Oct 2026 18:35:27 +0800 Subject: [PATCH] Specialize blocked root boundary mode at compile time --- src/CelerisLab/cuda/compiler_v2.py | 16 ++++-- src/CelerisLab/lbm/kernels/blocks.cu | 54 ++++++++----------- .../lbm/kernels/boundary/block_wall.cuh | 20 +++++++ .../lbm/kernels/config/config_method.h | 3 ++ .../lbm/kernels/operators/block_collision.cuh | 24 +++++++++ .../lbm/kernels/step/block_channel.cu | 4 ++ src/CelerisLab/lbm/refinement_stepper.py | 5 +- tests/unit/cpp/channel_root.cpp | 1 + tests/unit/test_build_contracts.py | 26 +++++++++ tests/unit/test_channel_root.py | 1 + tests/unit/test_refinement_mrt.py | 8 ++- .../multigrid/test_channel_root_gpu.py | 2 +- .../multigrid/test_refinement_blocked_base.py | 2 + 13 files changed, 127 insertions(+), 39 deletions(-) create mode 100644 src/CelerisLab/lbm/kernels/boundary/block_wall.cuh create mode 100644 src/CelerisLab/lbm/kernels/operators/block_collision.cuh diff --git a/src/CelerisLab/cuda/compiler_v2.py b/src/CelerisLab/cuda/compiler_v2.py index 48dff26..cb7a43e 100644 --- a/src/CelerisLab/cuda/compiler_v2.py +++ b/src/CelerisLab/cuda/compiler_v2.py @@ -54,14 +54,18 @@ def _write(path: str, content: str): def generate_config(cfg: LBMConfig, n_objects: int = 0, - *, build_root: Optional[str] = None) -> str: + *, build_root: Optional[str] = None, block_root_bc: int = 0) -> str: """Snapshot kernels and render headers in a fresh private directory. Returns its path. Pass it explicitly to compile_kernel for interleaved builds. Legacy generate_config(); compile_kernel() remains context-local. The caller owns cleanup_build_dir() once the loaded module is no longer needed for recompilation/debugging. No package files are changed. + block_root_bc specializes a blocked root module; fine/dense builds use 0. """ + if (isinstance(block_root_bc, bool) or not isinstance(block_root_bc, int) + or block_root_bc not in (0, 1, 2)): + raise ValueError("block_root_bc must be an integer mode 0, 1 or 2") macros = cfg.to_macros() if isinstance(n_objects, bool) or not isinstance(n_objects, int) or n_objects < 0: raise ValueError("n_objects must be a nonnegative integer") @@ -71,14 +75,15 @@ def generate_config(cfg: LBMConfig, n_objects: int = 0, try: shutil.copytree(_KERNEL_DIR, build_dir, dirs_exist_ok=True, ignore=shutil.ignore_patterns("*.ptx", "*.cubin", "*.tmp")) - _render_config(macros, n_objects, os.path.join(build_dir, "config")) + _render_config(macros, n_objects, os.path.join(build_dir, "config"), + block_root_bc=block_root_bc) digest = hashlib.sha256() for path in sorted(Path(build_dir).rglob("*")): if path.is_file(): digest.update(str(path.relative_to(build_dir)).encode()) digest.update(path.read_bytes()) metadata = {"source_sha256": digest.hexdigest(), "config": asdict(cfg), - "n_objects": n_objects} + "n_objects": n_objects, "block_root_bc": block_root_bc} _write(os.path.join(build_dir, ".celeris-build"), json.dumps(metadata, sort_keys=True)) except BaseException: shutil.rmtree(build_dir) @@ -101,7 +106,7 @@ def cleanup_build_dir(build_dir: Optional[str]) -> None: _last_ptx_path.set(None) -def _render_config(m: dict, n_objects: int, cdir: str): +def _render_config(m: dict, n_objects: int, cdir: str, *, block_root_bc: int = 0): _write(os.path.join(cdir, "config_grid.h"), f"""{_HEADER}\ // Layer: Global / Grid #ifndef CELERIS_CONFIG_GRID_H @@ -165,6 +170,9 @@ def _render_config(m: dict, n_objects: int, cdir: str): #define STORE_PRECISION {m['STORE_PRECISION']} #define USE_DDF_SHIFTING {m['USE_DDF_SHIFTING']} +// Internal blocked root mode: 0=periodic/fine, 1=slot/reservoir, 2=channel. +#define BLOCK_ROOT_BC {block_root_bc} + #define USE_LES {m['USE_LES']} #define LES_CS {m['LES_CS']} #define LES_CLOSED_FORM {m['LES_CLOSED_FORM']} diff --git a/src/CelerisLab/lbm/kernels/blocks.cu b/src/CelerisLab/lbm/kernels/blocks.cu index b8c7ef4..2a9ef94 100644 --- a/src/CelerisLab/lbm/kernels/blocks.cu +++ b/src/CelerisLab/lbm/kernels/blocks.cu @@ -1,5 +1,6 @@ // CelerisLab/lbm/kernels/blocks.cu -// One periodic level of fixed-size blocks (lbm/blocks.py layout). Separate +// One fixed-size blocked level (lbm/blocks.py layout), specialized at compile +// time for periodic/fine, slot/reservoir or channel stepping. Separate // compilation entry: NX, NY are the block window including halo, so index_f // and linear_index address one block; blockIdx.z selects the block, whose // populations start at b*NQ*TOTAL_CELLS (block-contiguous, Q-major inside). @@ -19,7 +20,18 @@ #endif #include "streaming/pull_double_buffer.cuh" #include "boundary/curved_boundary.cuh" -#include "boundary/planar_velocity.cuh" +#ifndef BLOCK_ROOT_BC +#error "Blocks require a generated BLOCK_ROOT_BC mode" +#elif BLOCK_ROOT_BC < 0 || BLOCK_ROOT_BC > 2 +#error "BLOCK_ROOT_BC must be 0, 1 or 2" +#endif + +#if BLOCK_ROOT_BC == 1 +#include "boundary/block_wall.cuh" +#endif +#if BLOCK_ROOT_BC != 2 +#include "operators/block_collision.cuh" +#endif #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" @@ -28,11 +40,6 @@ // Fixed-patch cell kinds, matching refinement.py ACTIVE/INNER_GHOST/OUTER_GHOST. constexpr unsigned char KIND_ACTIVE = 0u, KIND_INNER = 1u, KIND_OUTER = 2u; -// Zero-initialized for ordinary blocks/fine levels. Set only for parent walls; -// NX/NY remain local storage strides, never physical wall coordinates. -struct BlockWall { unsigned int nx, ny, slot_width; float amplitude; }; -extern "C" { __constant__ BlockWall d_block_wall; } - // Halo refresh: block cell ids b*TOTAL_CELLS+c, sources are always core cells. extern "C" __global__ void BlockHaloCopy( fpxx* f, const uint64_t* dst, const uint64_t* src, unsigned int count) @@ -147,6 +154,7 @@ extern "C" __global__ void BlockAccumulateExchange( integral[k + 3] += h * h * (double)raw[k + 3]; } +#if BLOCK_ROOT_BC != 2 // Pull/collide/store core cells only; halo cells of out are left untouched. // flags (block cell ids, nullable): solid cores store NaN and never collide. // kind (block cell ids, nullable = all ACTIVE): OUTER skips the pull and @@ -167,34 +175,18 @@ extern "C" __global__ void BlockStep(const fpxx* in, fpxx* out, unsigned int hal compute_neighbors(x, y, j); // Core neighbors stay inside the window: no wrap. stream_pull_load(k, f, in + base, j); if (cell_kind == KIND_ACTIVE) { - if (d_block_wall.nx != 0) { - unsigned int gx, gy; - block_global_xy(x, y, halo, d_block_wall.nx, gx, gy); - if (gy == 1) - reconstruct_bottom_velocity(f, 0.0f, bottom_wall_normal_velocity( - 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 - collide_mrt_unforced(f); // SRT path below unchanged. -#else - float rho, ux, uy, feq[NQ]; - compute_rho_u(f, rho, ux, uy); -#if USE_DDF_SHIFTING - // Preserve the density perturbation before adding the unit reference. - const float drho = f[0]+f[1]+f[2]+f[3]+f[4]+f[5]+f[6]+f[7]+f[8]; - const float jx = f[1]-f[2]+f[5]-f[6]+f[7]-f[8]; - const float jy = f[3]-f[4]+f[5]-f[6]-f[7]+f[8]; - compute_feq_from_momentum(drho, jx, jy, ux, uy, feq); -#else - compute_feq(rho, ux, uy, feq); -#endif - collide_srt_no_force(f, feq, d_params.omega); +#if BLOCK_ROOT_BC == 1 + block_wall(f, x, y, halo); #endif + block_collide(f); } } stream_pull_store(k, f, out + base); } -// Root channel closure; fine and periodic levels keep BlockStep. +#endif + +// Root channel is a replacement full step, selected once by the host. +#if BLOCK_ROOT_BC == 2 #include "step/block_channel.cu" +#endif diff --git a/src/CelerisLab/lbm/kernels/boundary/block_wall.cuh b/src/CelerisLab/lbm/kernels/boundary/block_wall.cuh new file mode 100644 index 0000000..beee159 --- /dev/null +++ b/src/CelerisLab/lbm/kernels/boundary/block_wall.cuh @@ -0,0 +1,20 @@ +// Physical parent walls; NX/NY remain local storage strides. +#ifndef CELERIS_BOUNDARY_BLOCK_WALL_CUH +#define CELERIS_BOUNDARY_BLOCK_WALL_CUH + +#include "planar_velocity.cuh" + +struct BlockWall { unsigned int nx, ny, slot_width; float amplitude; }; +extern "C" { __constant__ BlockWall d_block_wall; } + +__device__ __forceinline__ void block_wall(float* f, unsigned int x, unsigned int y, + unsigned int halo) { + unsigned int gx, gy; + block_global_xy(x, y, halo, d_block_wall.nx, gx, gy); + if (gy == 1) + reconstruct_bottom_velocity(f, 0.0f, bottom_wall_normal_velocity( + 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); +} + +#endif diff --git a/src/CelerisLab/lbm/kernels/config/config_method.h b/src/CelerisLab/lbm/kernels/config/config_method.h index 225be90..4de8e0c 100644 --- a/src/CelerisLab/lbm/kernels/config/config_method.h +++ b/src/CelerisLab/lbm/kernels/config/config_method.h @@ -8,6 +8,9 @@ #define STORE_PRECISION 0 #define USE_DDF_SHIFTING 0 +// Internal blocked root mode: 0=periodic/fine, 1=slot/reservoir, 2=channel. +#define BLOCK_ROOT_BC 0 + #define USE_LES 0 #define LES_CS 0.160000f #define LES_CLOSED_FORM 1 diff --git a/src/CelerisLab/lbm/kernels/operators/block_collision.cuh b/src/CelerisLab/lbm/kernels/operators/block_collision.cuh new file mode 100644 index 0000000..56e0d26 --- /dev/null +++ b/src/CelerisLab/lbm/kernels/operators/block_collision.cuh @@ -0,0 +1,24 @@ +// Collision used by periodic/fine and slot/reservoir BlockStep. +#ifndef CELERIS_OPERATORS_BLOCK_COLLISION_CUH +#define CELERIS_OPERATORS_BLOCK_COLLISION_CUH + +__device__ __forceinline__ void block_collide(float* f) { +#if COLLISION_MODEL == 2 + collide_mrt_unforced(f); // SRT path below unchanged. +#else + float rho, ux, uy, feq[NQ]; + compute_rho_u(f, rho, ux, uy); +#if USE_DDF_SHIFTING + // Preserve the density perturbation before adding the unit reference. + const float drho = f[0]+f[1]+f[2]+f[3]+f[4]+f[5]+f[6]+f[7]+f[8]; + const float jx = f[1]-f[2]+f[5]-f[6]+f[7]-f[8]; + const float jy = f[3]-f[4]+f[5]-f[6]-f[7]+f[8]; + compute_feq_from_momentum(drho, jx, jy, ux, uy, feq); +#else + compute_feq(rho, ux, uy, feq); +#endif + collide_srt_no_force(f, feq, d_params.omega); +#endif +} + +#endif diff --git a/src/CelerisLab/lbm/kernels/step/block_channel.cu b/src/CelerisLab/lbm/kernels/step/block_channel.cu index d31bd0d..1ff1cd2 100644 --- a/src/CelerisLab/lbm/kernels/step/block_channel.cu +++ b/src/CelerisLab/lbm/kernels/step/block_channel.cu @@ -1,3 +1,7 @@ +#if BLOCK_ROOT_BC != 2 +#error "BlockChannelStep requires BLOCK_ROOT_BC == 2" +#endif + // Root-only blocked counterpart of OneStep's D2Q9/SRT channel closure. // NX/NY are LOCAL strides; all physical boundary decisions use global coordinates. #include "../boundary/inlet/common.cuh" diff --git a/src/CelerisLab/lbm/refinement_stepper.py b/src/CelerisLab/lbm/refinement_stepper.py index 4d719aa..50883ba 100644 --- a/src/CelerisLab/lbm/refinement_stepper.py +++ b/src/CelerisLab/lbm/refinement_stepper.py @@ -362,7 +362,10 @@ class FixedPatchStepper: raise RuntimeError("level omega/tau mismatch") host = self._boundary_hosts[index] n_objects = len(host.objects) if name == "PatchFineStep" and host is not None else 0 - build = compiler.generate_config(config, n_objects=n_objects) + block_root_bc = 0 + if index == 0 and layout is not None: + block_root_bc = 2 if self._channel else 1 if self._wall_aware else 0 + build = compiler.generate_config(config, n_objects=n_objects, block_root_bc=block_root_bc) try: ptx = compiler.compile_kernel(self.ctx.sm_arch, source=source, build_dir=build) module = compiler.load_module(ptx) diff --git a/tests/unit/cpp/channel_root.cpp b/tests/unit/cpp/channel_root.cpp index 3d7cd21..6385d28 100644 --- a/tests/unit/cpp/channel_root.cpp +++ b/tests/unit/cpp/channel_root.cpp @@ -16,6 +16,7 @@ #define NZ 1 #define NT 1 #define COLLISION_MODEL 0 +#define BLOCK_ROOT_BC 2 #define USE_LES 0 #define RHO 1.f #define U0 .025f diff --git a/tests/unit/test_build_contracts.py b/tests/unit/test_build_contracts.py index a23bb9f..4d8f11a 100644 --- a/tests/unit/test_build_contracts.py +++ b/tests/unit/test_build_contracts.py @@ -83,6 +83,32 @@ assert abs(sum(compute_feq_d2q9(1.,0.,0.)) - 1.) < 1e-6 class BuildContracts(unittest.TestCase): + def test_block_root_modes_are_rendered_and_recorded(self): + with tempfile.TemporaryDirectory() as root: + for mode in (0, 1, 2): + with self.subTest(mode=mode): + build = compiler.generate_config(LBMConfig(), build_root=root, block_root_bc=mode) + try: + method = Path(build, "config/config_method.h").read_text() + self.assertIn(f"#define BLOCK_ROOT_BC {mode}\n", method) + metadata = json.loads(Path(build, ".celeris-build").read_text()) + self.assertEqual(metadata["block_root_bc"], mode) + self.assertNotIn("block_root_bc", metadata["config"]) + finally: + compiler.cleanup_build_dir(build) + build = compiler.generate_config(LBMConfig(), build_root=root) + try: + self.assertEqual(json.loads(Path(build, ".celeris-build").read_text())["block_root_bc"], 0) + finally: + compiler.cleanup_build_dir(build) + + def test_invalid_block_root_mode_fails_before_snapshot(self): + for mode in (-1, 3, True, False, 1.0, "1", None): + with self.subTest(mode=mode), patch.object(compiler.tempfile, "mkdtemp") as allocate: + with self.assertRaisesRegex(ValueError, "block_root_bc"): + compiler.generate_config(LBMConfig(), block_root_bc=mode) + allocate.assert_not_called() + @staticmethod def fake_compile(command, check): source = Path(command[-3]) diff --git a/tests/unit/test_channel_root.py b/tests/unit/test_channel_root.py index 4d59830..780315b 100644 --- a/tests/unit/test_channel_root.py +++ b/tests/unit/test_channel_root.py @@ -97,5 +97,6 @@ def test_evolved_ghost_snapshot_and_physical_download(): with pytest.raises(ValueError, match="finite"): backend.initialize(*arrays) assert backend._levels[0].kernel == host.compiler.load_module.return_value.get_function("BlockChannelStep") + assert [call.kwargs["block_root_bc"] for call in host.compiler.generate_config.call_args_list] == [2, 0] finally: backend.close() diff --git a/tests/unit/test_refinement_mrt.py b/tests/unit/test_refinement_mrt.py index be4d6f5..9302b72 100644 --- a/tests/unit/test_refinement_mrt.py +++ b/tests/unit/test_refinement_mrt.py @@ -24,11 +24,14 @@ class RefinementMRTHost(unittest.TestCase): """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() + block_collision = (KERNELS / "operators/block_collision.cuh").read_text() + self.assertIn('#include "operators/block_collision.cuh"', blocks) + self.assertEqual(blocks.count("block_collide(f);"), 1) 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);"): - self.assertIn(text, blocks) + self.assertIn(text, block_collision) self.assertIn("collide_srt(f, feq, forcing, omega);", channel) self.assertIn("collide_srt_no_force(f, feq, d_params.omega);", refinement) for source in (blocks, refinement): @@ -40,10 +43,11 @@ class RefinementMRTHost(unittest.TestCase): # One shared unforced MRT body; each blocked/patch path calls it exactly once. mrt = (KERNELS / "operators/collision_mrt.cuh").read_text() self.assertEqual(mrt.count("void collide_mrt_unforced("), 1) - for source in (blocks, channel, refinement): + for source in (blocks + "\n" + block_collision, channel, refinement): self.assertEqual(source.count("collide_mrt_unforced(f);"), 1) self.assertNotIn("collide_mrt_no_force(", source) self.assertNotIn("compute_rho_u(f, rho, ux, uy);\n collide_mrt", source) + self.assertNotRegex(source, r"compute_rho_u\(f, rho, ux, uy\);\s*collide_mrt") if __name__ == "__main__": unittest.main() diff --git a/tests/validation/multigrid/test_channel_root_gpu.py b/tests/validation/multigrid/test_channel_root_gpu.py index cf0593e..905b720 100644 --- a/tests/validation/multigrid/test_channel_root_gpu.py +++ b/tests/validation/multigrid/test_channel_root_gpu.py @@ -27,7 +27,7 @@ def state(ny, nx): def compile_module(cfg, context, source): cfg.validate() - build = compiler.generate_config(cfg) + build = compiler.generate_config(cfg, block_root_bc=2 if source == "blocks.cu" else 0) try: module = compiler.load_module(compiler.compile_kernel(context.sm_arch, source=source, build_dir=build)) finally: diff --git a/tests/validation/multigrid/test_refinement_blocked_base.py b/tests/validation/multigrid/test_refinement_blocked_base.py index bb54da5..5af68d5 100644 --- a/tests/validation/multigrid/test_refinement_blocked_base.py +++ b/tests/validation/multigrid/test_refinement_blocked_base.py @@ -223,6 +223,8 @@ class BlockedBaseCPU(unittest.TestCase): self.assertIs(collide[4], coarse.kind) walls = [a for ptr, a in m["uploads"] if isinstance(ptr, str) and ptr == "d_block_wall"] self.assertEqual(len(walls), 1) + self.assertEqual([call.kwargs["block_root_bc"] for call in + compiler.generate_config.call_args_list], [1, 0]) width, height, slot, speed = struct.unpack("=IIIf", walls[0]) self.assertEqual((width, height, slot), (NX, NY, 8)) self.assertAlmostEqual(speed, .001)