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

compiler: one build_module/set_params path for all three module owners

Browse Source 
Simulation, PeriodicBlockStepper and FixedPatchStepper levels each repeated
generate_config -> compile_kernel -> load_module -> cleanup and the
d_params pack/upload. They now call compiler.build_module(cfg, entry,
arch=..., n_objects=..., block_root_bc=...) and compiler.set_params.
Simulation no longer keeps its build directory alive until close (nothing
read it after loading). Removes the unused compile_kernel_v2 alias.
Host fakes and lifecycle tests follow the new seam; the close-retry case
for build-directory cleanup goes with the removed cleanup step.
No kernel, numerical or launch change.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Claude
2026-10-26 13:59:39 +00:00
co-authored by Claude Opus 5.5
parent 96c5decf69
commit 0d62c3cfb8
14 changed files with 51 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/CelerisLab/cuda/compiler_v2.py
+28 -8
View File
@@ -2,10 +2,13 @@
"""
Kernel configuration, compilation, and module loading.
Workflow:
Owners call build_module(cfg, entry, arch=...), then set_params(module, omega, ...).
build_module runs, in a private directory it removes afterwards:
1. build_dir = generate_config(lbm_cfg, n_objects) — isolated source snapshot
2. compile_kernel(arch, build_dir=build_dir) — nvcc -ptx kernel_v2.cu
2. compile_kernel(arch, source=entry, build_dir=build_dir) — nvcc -ptx
3. load_module(ptx_path) — PyCUDA module from PTX
Entries: kernel_v2.cu (uniform grid), blocks.cu (blocked level), refinement.cu
(dense refinement level).
"""
from __future__ import annotations
@@ -293,10 +296,27 @@ def load_module(ptx_path: Optional[str] = None):
return cuda.module_from_file(ptx_path)
# ---------------------------------------------------------------------------
# Back-compat helpers (used by test scripts)
# ---------------------------------------------------------------------------
def build_module(cfg: LBMConfig, source: str, *, arch: str, n_objects: int = 0,
block_root_bc: int = 0):
"""Generate, compile and load one private module; the build directory is removed.
def compile_kernel_v2(arch: str = "sm_70") -> str:
"""Alias for compile_kernel() kept for test-script compatibility."""
return compile_kernel(arch=arch)
The single build path of every owner (Simulation, PeriodicBlockStepper,
FixedPatchStepper levels). Each call returns a NEW module: callers must not
share it, because its globals (d_params, ...) are mutable per module.
"""
build_dir = generate_config(cfg, n_objects=n_objects, block_root_bc=block_root_bc)
try:
return load_module(compile_kernel(arch, source=source, build_dir=build_dir))
finally:
cleanup_build_dir(build_dir)
def set_params(module, omega: float, fx: float = 0.0, fy: float = 0.0, fz: float = 0.0,
n_objects: int = 0) -> None:
"""Write the module's __constant__ LBMParams (core/params.cuh layout)."""
import struct
import pycuda.driver as cuda
ptr, size = module.get_global("d_params")
if size != struct.calcsize("=ffffI"):
raise RuntimeError("LBMParams host/device layout mismatch")
cuda.memcpy_htod(ptr, struct.pack("=ffffI", omega, fx, fy, fz, n_objects))
src/CelerisLab/lbm/block_stepper.py
+2 -11
View File
@@ -12,7 +12,6 @@ parameter mutation; wall/slot/inlet fields in cfg are ignored.
"""
from dataclasses import replace
from numbers import Integral
import struct
from threading import RLock
import numpy as np
@@ -129,16 +128,8 @@ class PeriodicBlockStepper:
config.validate()
if not np.isclose(config.omega, 1 / (3 * checked.viscosity + .5), rtol=1e-14, atol=0):
raise RuntimeError("block omega/tau mismatch")
build = compiler.generate_config(config)
try:
ptx = compiler.compile_kernel(self.ctx.sm_arch, source="blocks.cu", build_dir=build)
module = compiler.load_module(ptx)
finally:
compiler.cleanup_build_dir(build)
ptr, size = module.get_global("d_params")
if size != struct.calcsize("=ffffI"):
raise RuntimeError("LBMParams host/device layout mismatch")
cuda.memcpy_htod(ptr, struct.pack("=ffffI", config.omega, 0, 0, 0, 0))
module = compiler.build_module(config, "blocks.cu", arch=self.ctx.sm_arch)
compiler.set_params(module, config.omega)
self._module = module
self._current = self._allocate(self._bytes)
self._spare = self._allocate(self._bytes)
src/CelerisLab/lbm/refinement_stepper.py
+3 -10
View File
@@ -365,16 +365,9 @@ class FixedPatchStepper:
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)
finally:
compiler.cleanup_build_dir(build)
ptr, size = module.get_global("d_params")
if size != struct.calcsize("=ffffI"):
raise RuntimeError("LBMParams host/device layout mismatch")
cuda.memcpy_htod(ptr, struct.pack("=ffffI", config.omega, 0, 0, 0, n_objects))
module = compiler.build_module(config, source, arch=self.ctx.sm_arch, n_objects=n_objects,
block_root_bc=block_root_bc)
compiler.set_params(module, config.omega, n_objects=n_objects)
shape = (9, *self.hierarchy.kinds[index].shape)
size = int(np.prod(shape if layout is None else layout.blocked_shape)) * 4
grid = ((config.nx + self._block[0] - 1) // self._block[0], config.ny,
src/CelerisLab/simulation.py
+3 -16
View File
@@ -46,7 +46,6 @@ class Simulation:
device_id: int = 0):
# Internal stream for async DTOH coordination
self._stream: Optional[cuda.Stream] = None
self._build_dir = None
self._closed = False
# Load configs
@@ -59,9 +58,7 @@ class Simulation:
try:
with self.ctx.activate():
arch = self._resolve_compile_arch()
self._build_dir = compiler.generate_config(self.lbm_cfg, n_objects=0)
self._ptx_path = compiler.compile_kernel(arch=arch, build_dir=self._build_dir)
self._module = compiler.load_module(self._ptx_path)
self._module = compiler.build_module(self.lbm_cfg, "kernel_v2.cu", arch=arch)
self._compiled_n_objects = 0
self._allocation_signature = self.lbm_cfg.allocation_signature()
self._compiled_viscosity = self.lbm_cfg.viscosity
@@ -412,15 +409,8 @@ class Simulation:
self.lbm_cfg.validate_runtime()
self._assert_allocation_contract()
arch = self._resolve_compile_arch()
next_dir = compiler.generate_config(self.lbm_cfg, n_objects=expected_n_objects)
try:
next_ptx = compiler.compile_kernel(arch=arch, build_dir=next_dir)
next_module = compiler.load_module(next_ptx)
except Exception:
compiler.cleanup_build_dir(next_dir)
raise
old_dir = self._build_dir
self._build_dir, self._ptx_path, self._module = next_dir, next_ptx, next_module
self._module = compiler.build_module(self.lbm_cfg, "kernel_v2.cu", arch=arch,
n_objects=expected_n_objects)
self._compiled_n_objects = int(expected_n_objects)
# Reconnect field and stepper to new module
@@ -438,7 +428,6 @@ class Simulation:
self.stepper._step_count = _prev_step_count
self._assert_object_count_contract(expected_count=expected_n_objects)
compiler.cleanup_build_dir(old_dir)
def recompile(self):
"""Re-generate config headers and recompile kernel.
@@ -715,8 +704,6 @@ class Simulation:
self._module = None
self._stream = None
ctx.close()
compiler.cleanup_build_dir(getattr(self, "_build_dir", None))
self._build_dir = None
self._closed = True
def __enter__(self):
tests/integration/test_runtime_contracts_gpu.py
+1 -1
View File
@@ -21,7 +21,7 @@ def test_two_contexts_interleaved_and_non_lifo_close():
first = make_sim(nx=32, ny=16)
second = make_sim(nx=48, ny=24, streaming="esopull", shifted=True)
try:
assert first._build_dir != second._build_dir
assert first._module is not second._module
assert cuda.Context.get_current() == before
first.initialize()
second.initialize()
tests/unit/test_block_step_sync.py
+1 -1
View File
@@ -37,7 +37,7 @@ def backend():
module.get_global.return_value = (Mock(), 20)
module.get_function.side_effect = lambda name: Mock(
side_effect=lambda *a, **k: log.append((name, a[:2])))
host.compiler.load_module.return_value = module
host.compiler.build_module.return_value = module
cfg = LBMConfig(nx=16, ny=16, viscosity=.1)
run = host.PeriodicBlockStepper(cfg, BlockLayout2D.build(16, 16, 8, 8),
context=Mock(activate=nullcontext, sm_arch="sm_70"))
tests/unit/test_channel_root.py
+3 -3
View File
@@ -74,7 +74,7 @@ def test_evolved_ghost_snapshot_and_physical_download():
from .test_refinement_three_level_multipatch import host_module
host = host_module()
storage = {}
host.compiler.load_module.return_value.get_global.side_effect = lambda name: (name, 12 if name == "d_block_channel" else 20)
host.compiler.build_module.return_value.get_global.side_effect = lambda name: (name, 12 if name == "d_block_channel" else 20)
host.cuda.mem_alloc.side_effect = lambda size: Mock()
host.cuda.memcpy_htod.side_effect = lambda ptr, a: storage.__setitem__(id(ptr), np.array(a, copy=True))
host.cuda.memcpy_dtoh.side_effect = lambda a, ptr: np.copyto(a, storage[id(ptr)])
@@ -96,7 +96,7 @@ def test_evolved_ghost_snapshot_and_physical_download():
arrays[0][0, 0, 0] = np.nan
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]
assert backend._levels[0].kernel == host.compiler.build_module.return_value.get_function("BlockChannelStep")
assert [call.kwargs["block_root_bc"] for call in host.compiler.build_module.call_args_list] == [2, 0]
finally:
backend.close()
tests/unit/test_projected_pair_stepper.py
+1 -1
View File
@@ -29,7 +29,7 @@ def backend(**options):
name=name, side_effect=lambda *a, **kw: calls.append((name, a, kw)))
return module
host.compiler.load_module.side_effect = load
host.compiler.build_module.side_effect = lambda *a, **k: load(None)
host.cuda.memcpy_htod.side_effect = lambda ptr, a: storage.__setitem__(id(ptr), np.array(a, copy=True))
host.cuda.memcpy_dtoh.side_effect = lambda a, ptr: np.copyto(a, storage[id(ptr)])
cfg = options.pop('cfg', CFG)
tests/unit/test_raw_exchange.py
+1 -1
View File
@@ -25,7 +25,7 @@ def backend(observe=True, obj=None):
module = Mock()
module.get_global.return_value = (Mock(), 20)
module.get_function.side_effect = lambda name: Mock(side_effect=lambda *a, **k: calls.append((name, a, k)))
host.compiler.load_module.return_value = module
host.compiler.build_module.return_value = module
host.cuda.mem_alloc.side_effect = lambda n: Mock(size=n)
host.cuda.memcpy_dtoh.side_effect = lambda array, pointer: array.fill(0)
ctx = Mock(activate=nullcontext, sm_arch='sm_70')
tests/unit/test_refinement_three_level_multipatch.py
+2 -2
View File
@@ -193,9 +193,9 @@ def mock_backend(*, children=CHILDREN, bodies=False, fail_compile=False, fail_up
gpu.get_function.side_effect=lambda name: Mock(name=name)
gpu.get_global.side_effect=lambda name: (Mock(),20 if name=='d_params' else 16)
return gpu
module.compiler.load_module.side_effect=load
module.compiler.build_module.side_effect=lambda *a,**k: load(None)
if fail_compile:
module.compiler.compile_kernel.side_effect=['first','second',RuntimeError('L2 compile failed')]
module.compiler.build_module.side_effect=[load(None),load(None),RuntimeError('L2 compile failed')]
if fail_upload:
def upload(*args):
if len(allocations)>20:
tests/unit/test_shifted_jets.py
+1 -1
View File
@@ -62,7 +62,7 @@ class ShiftedJets(unittest.TestCase):
level_blocks=list(blocks),level_objects=[(),(),(body(),)],
level_jets=[(),(),(jet(),)])
module.cuda.mem_alloc.assert_not_called()
module.compiler.compile_kernel.assert_not_called()
module.compiler.build_module.assert_not_called()
def test_switch_upload_only_then_raw_remap_keeps_live_velocity(self):
run,module,allocations,storage,objects=ready();clone=None
tests/unit/test_shifted_storage.py
+2 -2
View File
@@ -135,7 +135,7 @@ class ShiftedStorage(unittest.TestCase):
with self.assertRaisesRegex(ValueError,'all-blocked|blocked refined levels below'):
module.FixedPatchStepper(cfg,context=None,refinements=[(4,4,20,28)],
level_blocks=blocks)
module.compiler.compile_kernel.assert_not_called()
module.compiler.build_module.assert_not_called()
def test_uniform_block_codec_and_raw_snapshot(self):
path = Path(backend.r.__file__).with_name('block_stepper.py')
@@ -151,7 +151,7 @@ class ShiftedStorage(unittest.TestCase):
cuda.mem_get_info.return_value=(10**9,10**9)
cuda.mem_alloc.side_effect=lambda size: Mock()
gpu=Mock();gpu.get_global.return_value=(Mock(),20)
compiler.load_module.return_value=gpu
compiler.build_module.return_value=gpu
context=Mock(activate=nullcontext,sm_arch='sm_70')
layout=module.BlockLayout2D.build(16,16,4,4)
cfg=replace(backend.CFG,nx=16,ny=16,ddf_shifting=True)
tests/unit/test_simulation_lifecycle.py
+1 -8
View File
@@ -36,12 +36,10 @@ def make_simulation(monkeypatch):
sim = module.Simulation.__new__(module.Simulation)
sim.ctx = Owner()
sim._closed = False
sim._build_dir = "private-build"
for name in ("stepper", "bodies", "field", "_module", "_stream"):
setattr(sim, name, object())
monkeypatch.setattr(module, "cuda", SimpleNamespace(
Context=SimpleNamespace(synchronize=lambda: None)))
monkeypatch.setattr(module.compiler, "cleanup_build_dir", lambda path: None)
return sim
@@ -57,7 +55,7 @@ def test_close_inside_activation_preserves_resources(monkeypatch):
assert sim._closed and sim.ctx.is_closed
@pytest.mark.parametrize("failure", ["synchronize", "detach", "directory"])
@pytest.mark.parametrize("failure", ["synchronize", "detach"])
def test_close_failure_is_retryable(monkeypatch, failure):
sim = make_simulation(monkeypatch)
if failure == "detach":
@@ -68,17 +66,12 @@ def test_close_failure_is_retryable(monkeypatch, failure):
calls.append(args)
if len(calls) == 1:
raise RuntimeError(failure + " failed")
if failure == "synchronize":
monkeypatch.setattr(module.cuda.Context, "synchronize", fail_once)
else:
monkeypatch.setattr(module.compiler, "cleanup_build_dir", fail_once)
with pytest.raises(RuntimeError, match="failed"):
sim.close()
assert not sim._closed
assert sim._build_dir == "private-build"
sim.close()
assert sim._closed and sim.ctx.is_closed
assert sim._build_dir is None
sim.close()
tests/unit/test_step_sync.py
+1 -1
View File
@@ -31,7 +31,7 @@ def backend(window=False):
module.get_global.return_value = (Mock(), 20)
module.get_function.side_effect = lambda name: Mock(
side_effect=lambda *a, **k: log.append((name, a, k)))
host.compiler.load_module.return_value = module
host.compiler.build_module.return_value = module
host.cuda.mem_alloc.side_effect = allocate
host.cuda.memcpy_dtoh.side_effect = lambda array, pointer: array.fill(0)
host.cuda.memset_d8_async.side_effect = lambda *a: log.append(("memset", a, {}))