diff --git a/src/CelerisLab/cuda/compiler_v2.py b/src/CelerisLab/cuda/compiler_v2.py index cb7a43e..0a687fa 100644 --- a/src/CelerisLab/cuda/compiler_v2.py +++ b/src/CelerisLab/cuda/compiler_v2.py @@ -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)) diff --git a/src/CelerisLab/lbm/block_stepper.py b/src/CelerisLab/lbm/block_stepper.py index 4c32bad..6260aa1 100644 --- a/src/CelerisLab/lbm/block_stepper.py +++ b/src/CelerisLab/lbm/block_stepper.py @@ -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) diff --git a/src/CelerisLab/lbm/refinement_stepper.py b/src/CelerisLab/lbm/refinement_stepper.py index f9d561a..2a387fc 100644 --- a/src/CelerisLab/lbm/refinement_stepper.py +++ b/src/CelerisLab/lbm/refinement_stepper.py @@ -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, diff --git a/src/CelerisLab/simulation.py b/src/CelerisLab/simulation.py index b4f60f4..c2cfc95 100644 --- a/src/CelerisLab/simulation.py +++ b/src/CelerisLab/simulation.py @@ -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): diff --git a/tests/integration/test_runtime_contracts_gpu.py b/tests/integration/test_runtime_contracts_gpu.py index 0182509..e39d38e 100644 --- a/tests/integration/test_runtime_contracts_gpu.py +++ b/tests/integration/test_runtime_contracts_gpu.py @@ -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() diff --git a/tests/unit/test_block_step_sync.py b/tests/unit/test_block_step_sync.py index 3590fde..0e25587 100644 --- a/tests/unit/test_block_step_sync.py +++ b/tests/unit/test_block_step_sync.py @@ -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")) diff --git a/tests/unit/test_channel_root.py b/tests/unit/test_channel_root.py index 780315b..c932fbd 100644 --- a/tests/unit/test_channel_root.py +++ b/tests/unit/test_channel_root.py @@ -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() diff --git a/tests/unit/test_projected_pair_stepper.py b/tests/unit/test_projected_pair_stepper.py index 4dfcbee..b1cf084 100644 --- a/tests/unit/test_projected_pair_stepper.py +++ b/tests/unit/test_projected_pair_stepper.py @@ -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) diff --git a/tests/unit/test_raw_exchange.py b/tests/unit/test_raw_exchange.py index a413afd..9bfb065 100644 --- a/tests/unit/test_raw_exchange.py +++ b/tests/unit/test_raw_exchange.py @@ -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') diff --git a/tests/unit/test_refinement_three_level_multipatch.py b/tests/unit/test_refinement_three_level_multipatch.py index 021f5bb..d94a1b9 100644 --- a/tests/unit/test_refinement_three_level_multipatch.py +++ b/tests/unit/test_refinement_three_level_multipatch.py @@ -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: diff --git a/tests/unit/test_shifted_jets.py b/tests/unit/test_shifted_jets.py index a83aa9d..a753c24 100644 --- a/tests/unit/test_shifted_jets.py +++ b/tests/unit/test_shifted_jets.py @@ -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 diff --git a/tests/unit/test_shifted_storage.py b/tests/unit/test_shifted_storage.py index 8e5ea07..d3f3658 100644 --- a/tests/unit/test_shifted_storage.py +++ b/tests/unit/test_shifted_storage.py @@ -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) diff --git a/tests/unit/test_simulation_lifecycle.py b/tests/unit/test_simulation_lifecycle.py index 2405c32..69307c2 100644 --- a/tests/unit/test_simulation_lifecycle.py +++ b/tests/unit/test_simulation_lifecycle.py @@ -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) + monkeypatch.setattr(module.cuda.Context, "synchronize", 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() diff --git a/tests/unit/test_step_sync.py b/tests/unit/test_step_sync.py index 2cc04a5..255b16b 100644 --- a/tests/unit/test_step_sync.py +++ b/tests/unit/test_step_sync.py @@ -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, {}))