feat(cuda): scoped device-memory arena (withCudaArena) + use-after-free detector - #22
Open
NicolasRouquette wants to merge 1 commit into
Open
feat(cuda): scoped device-memory arena (withCudaArena) + use-after-free detector#22NicolasRouquette wants to merge 1 commit into
NicolasRouquette wants to merge 1 commit into
Conversation
NicolasRouquette
added a commit
to NicolasRouquette/TorchLean
that referenced
this pull request
Jul 23, 2026
…n-dojo#22) # Conflicts: # NN/Tests/Runtime/Cuda/Stress.lean # NN/Tests/Suite.lean
NicolasRouquette
force-pushed
the
cuda-arena-scope
branch
2 times, most recently
from
August 3, 2026 18:01
1267903 to
2f12a15
Compare
…ee detector A long *pure* eager loop — a `foldl` of `Buffer → Buffer` ops with no IO sequencing — keeps every intermediate `Buffer` GC-reachable until the final readback, so the reference-counting finalizers that would free the device memory never run and the working set grows with the loop length. Explicit `release` cannot reach those buffers: a pure carrier has no IO point at which to call it. This adds a scoped arena that sidesteps GC reachability. `Buffer.arenaEnter` opens an allocation epoch; every device buffer allocated while it is open is registered to it. `Buffer.arenaExit keep` frees the device data of *all* of them — reachable or not — except the `keep` results, which are promoted to an enclosing arena (or left to ordinary finalization at the outermost level). `Buffer.withCudaArena keep body` is the bracketed form (arena still closed if `body` raises). This matches a training loop's natural phase boundary (one step / one fold). Implementation: - `csrc/cuda/common/torchlean_cuda_arena.h`: a mutex-guarded epoch-stack registry. Each tracked buffer points at a heap `arena_reg` (and back); a buffer freed mid-scope flips its reg's `alive` flag so the exit walk skips it (no dangling/double-free). No-arena and untracked paths take no lock. - Two `arena_reg`/`arena_freed_depth` fields on the buffer struct; register in `buffer_alloc`, unlink in `finalize`/`drop_unboxed`; `arena_enter`/`arena_exit` IO exports — mirrored in the CUDA `.cu` and the portable stub `.c`. It also adds an opt-in use-after-free detector (`TORCHLEAN_ARENA_DEBUG=1`): a reclaimed buffer records the epoch that freed it in `arena_freed_depth`, and the `require_same_size2/3` choke point (every binary/ternary op) asserts liveness before the size compare — turning a stale operand into a panic that names the freeing epoch instead of a silent launch on freed memory, and catching the both-operands-freed case the bare `0 == 0` size check misses. When the flag is off it is one predicted branch on a cached int. Tests (`NN/Tests/Runtime/Cuda/Stress.lean`, run by `nn_tests_suite`): - `runArenaStress`: k buffers held live (never released) across a scope exit are reclaimed anyway (the case `release` cannot reach), and a promoted buffer survives and stays usable. - `runArenaDetectorDeathTest`: forks the suite via `/proc/self/exe` — positive (detector on + planted UAF ⇒ aborts naming the hazard), negative (detector on + a valid promotion reused in a binary op ⇒ clean, no false positive), control (detector off ⇒ the UAF slips through).
NicolasRouquette
force-pushed
the
cuda-arena-scope
branch
from
August 4, 2026 19:48
2f12a15 to
0e350d6
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
A long pure eager loop — a
foldlofBuffer → Bufferops with no IO sequencing — keeps everyintermediate
BufferGC-reachable until the final readback, so the reference-counting finalizers thatwould free the device memory never run and the working set grows with the loop length. Explicit
Buffer.releasecannot reach those buffers: a pure carrier has no IO point at which to call it.This adds a scoped arena that sidesteps GC reachability at a training loop's natural phase boundary:
Buffer.arenaEnteropens an allocation epoch; every device buffer allocated while it is open is registered to it.Buffer.arenaExit keepfrees the device data of all of them — reachable or not — except thekeepresults, which are promoted to an enclosing arena (or left to ordinary finalization at the outermost level).Buffer.withCudaArena keep bodyis the bracketed form (the arena is still closed ifbodyraises).Implementation
csrc/cuda/common/torchlean_cuda_arena.h: a mutex-guarded epoch-stack registry. Each tracked buffer points at a heaparena_reg(and back); a buffer freed mid-scope flips its reg'saliveflag so the exit walk skips it (no dangling / double-free). The no-arena and untracked paths take no lock.arena_reg,arena_freed_depth) on the buffer struct; register inbuffer_alloc, unlink infinalize/drop_unboxed;arena_enter/arena_exitIO exports — mirrored in the CUDA.cuand the portable stub.c.Use-after-free detector (opt-in)
Under
TORCHLEAN_ARENA_DEBUG=1, a reclaimed buffer records the epoch that freed it inarena_freed_depth, and therequire_same_size2/3choke point (every binary/ternary op) asserts liveness before the size compare — turning a stale operand into a panic that names the freeing epoch instead of a silent launch on freed memory, and catching the both-operands-freed case the bare0 == 0size check misses. With the flag off it is one predicted branch on a cached int.Tests (
nn_tests_suite)runArenaStress:kbuffers held live (never released) across a scope exit are reclaimed anyway (the casereleasecannot reach), and a promoted buffer survives and stays usable.runArenaDetectorDeathTest: forks the suite via/proc/self/exe— positive (detector on + planted UAF ⇒ aborts naming the hazard), negative (detector on + a valid promotion reused in a binary op ⇒ clean, no false positive), control (detector off ⇒ the UAF slips through).Rebased onto the current
mainafter the numerics/backend refactor; the CUDA build and the fullnn_tests_suite(including the arena stress and the detector fork death test) pass on an RTX A4500.