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

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>

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))

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)

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,

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):

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()

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"))

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()

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)

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')

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:

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

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)

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()
 
 

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, {}))