refinement: free-memory preflight for every level before any device work

Restores the PeriodicBlockStepper safeguard dropped in 8e1feff, for all
FixedPatchStepper levels: both G buffers, halo pairs and boundary columns
of every level are compared with cuda.mem_get_info() before the stream,
any module or any buffer exists, raising MemoryError('level buffers exceed
free device memory'). Transfer tables are not counted (lower bound).

Tests: the single-level MemoryError case is back beside the param-layout
check (no stream, no allocation); host test_single_level covers depth 1
and depth 2 (no stream, build or allocation) and the wait failure after a
good step (count stays 1). Driver-mocking fixtures report ample free memory.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Claude
2026-10-27 01:10:15 +00:00
co-authored by Claude Opus 5.5
parent 8c776309c9
commit 19db7d362f
9 changed files with 57 additions and 8 deletions
+18
View File
@@ -324,6 +324,9 @@ class FixedPatchStepper:
self._block = (cfg.threads_per_block, 1, 1)
with self.ctx.activate():
try:
free, _ = cuda.mem_get_info()
if self._planned_bytes() > free: # before any stream, module or buffer
raise MemoryError("level buffers exceed free device memory")
self._stream = cuda.Stream()
for level in range(depth):
self._levels.append(self._build_level(level))
@@ -369,6 +372,21 @@ class FixedPatchStepper:
pass # Preserve the original construction failure.
raise
def _planned_bytes(self):
"""Both G buffers, halo pairs and boundary columns of every level (transfer tables excluded)."""
total = 0
for index, layout in enumerate(self._layouts):
cells = int(np.prod(layout.blocked_shape if layout is not None
else (9, *self.hierarchy.kinds[index].shape)))
total += 2 * 4 * cells
if layout is not None:
total += layout.halo_dst.nbytes + layout.halo_src.nbytes
b = self._level_boundaries[index]
if b is not None:
total += sum(np.asarray(v).nbytes for k, v in b.tables.items()
if k == "flags" or k in BLOCKED_BOUNDARY_DEVICE_COLUMNS)
return total
def _build_level(self, index):
"""Compile one level's step module and allocate its two G buffers."""
layout, config = self._layouts[index], self._level_cfgs[index]