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