Skip to content

feat(cuda): scoped device-memory arena (withCudaArena) + use-after-free detector - #22

Open
NicolasRouquette wants to merge 1 commit into
lean-dojo:mainfrom
NicolasRouquette:cuda-arena-scope
Open

feat(cuda): scoped device-memory arena (withCudaArena) + use-after-free detector#22
NicolasRouquette wants to merge 1 commit into
lean-dojo:mainfrom
NicolasRouquette:cuda-arena-scope

Conversation

@NicolasRouquette

Copy link
Copy Markdown
Contributor

Summary

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
Buffer.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 at a training loop's natural phase boundary:

  • 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 (the arena is still closed if body raises).

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). The no-arena and untracked paths take no lock.
  • Two fields (arena_reg, arena_freed_depth) 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.

Use-after-free detector (opt-in)

Under 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. With the flag off it is one predicted branch on a cached int.

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

Rebased onto the current main after the numerics/backend refactor; the CUDA build and the full nn_tests_suite (including the arena stress and the detector fork death test) pass on an RTX A4500.

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
NicolasRouquette force-pushed the cuda-arena-scope branch 2 times, most recently from 1267903 to 2f12a15 Compare August 3, 2026 18:01
…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).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant