From 16d6778b29254c576e2208501ad50a114e5122f7 Mon Sep 17 00:00:00 2001 From: Nicolas Rouquette Date: Tue, 23 Jun 2026 07:59:07 -0700 Subject: [PATCH 1/2] feat(cuda): add withCudaArena scoped device-memory reclamation Long pure eager loops (a foldl of Buffer -> Buffer ops with no IO sequencing) keep every intermediate Buffer GC-reachable until the final readback, so the reference-counting finalizers that free device memory never run and the working set grows with the loop length. Explicit release cannot help from inside a pure carrier: there is no IO point at which to call it. withCudaArena opens an allocation epoch; every device buffer allocated while it is open is tracked. On exit it frees the device data of all of them -- reachable or not -- except a small `keep` set (the step's results), which are promoted to the parent epoch (or left to ordinary finalization at the outermost level). This matches a training loop's natural phase boundary (one step / one fold) and keeps the pure carrier unchanged; only the driver wraps each step. Mechanism (csrc/cuda/common/torchlean_cuda_arena.h, shared by the CUDA backend and the CPU stub): each tracked buffer points at a heap reg, and back; a buffer finalized mid-scope flips its reg's alive flag instead of leaving a dangling pointer, so the exit walk skips it. Registry mutation is mutex-guarded; the no-arena fast paths take no lock, so allocation and finalization outside an arena are unchanged. Wiring: buffer_alloc registers, finalize/drop_unboxed unlink, plus two IO export wrappers; the buffer struct gains one pointer field. Lean API: Buffer.arenaEnter / arenaExit / withCudaArena (Buffer.lean). Test: runArenaStress (NN/Tests/Runtime/Cuda/Stress.lean, wired into Stress.run) drives k live, never-released buffers through a scope -- the case release cannot reach -- holding them alive across the exit so only the arena can free them, and asserts via the allocator telemetry that (1) they are live in-scope, (2) the arena reclaims exactly them at exit with live bytes restored, and (3) a kept buffer is promoted and stays usable. The assertions use only the shared allocator counters, so the test is backend-agnostic; verified on the real CUDA backend, and the CPU stub compiles cleanly for GPU-less CI. --- NN/Runtime/Autograd/Engine/Cuda/Buffer.lean | 43 +++++ NN/Tests/Runtime/Cuda/Stress.lean | 86 +++++++++ csrc/cuda/common/torchlean_cuda_arena.h | 181 ++++++++++++++++++ csrc/cuda/common/torchlean_cuda_buffer.h | 11 +- csrc/cuda/tensor/torchlean_cuda_tensor.cu | 30 +++ csrc/cuda/tensor/torchlean_cuda_tensor_stub.c | 29 +++ 6 files changed, 378 insertions(+), 2 deletions(-) create mode 100644 csrc/cuda/common/torchlean_cuda_arena.h diff --git a/NN/Runtime/Autograd/Engine/Cuda/Buffer.lean b/NN/Runtime/Autograd/Engine/Cuda/Buffer.lean index 26648ad..5effb91 100644 --- a/NN/Runtime/Autograd/Engine/Cuda/Buffer.lean +++ b/NN/Runtime/Autograd/Engine/Cuda/Buffer.lean @@ -239,6 +239,49 @@ opaque collectAllocatorRaw (force : UInt32) : UInt32 def collectAllocator (force : Bool := true) : UInt32 := collectAllocatorRaw (if force then 1 else 0) +/-! +### Scoped device-memory arena (`withCudaArena`) + +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 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. + +A scoped arena fixes this. `arenaEnter` opens an allocation epoch; every device buffer allocated while +it is open is tracked. `arenaExit keep` frees the device data of *all* of them — reachable or not — +except the `keep` results, which survive: promoted to an enclosing arena if there is one, or left to +ordinary reference-counted finalization at the outermost level. This matches a training loop's natural +phase boundary (one LM step / one fold). + +**Contract.** Every buffer that must outlive the scope MUST appear in `keep` (or be reachable only +through a kept buffer). Touching any other in-scope buffer after `arenaExit` is a use-after-free, the +same hazard as reading a buffer after `release`. +-/ + +@[extern "torchlean_cuda_arena_enter"] +opaque arenaEnter : IO Unit + +@[extern "torchlean_cuda_arena_exit"] +opaque arenaExit (keep : @& Array Buffer) : IO Unit + +/-- +Run `body` inside a fresh device-memory arena, then reclaim every buffer it allocated except those +named by `keep result`. + +`keep` extracts the buffers that must survive the scope (typically the step's result parameters). If +`body` raises, the arena is still closed and every in-scope buffer is reclaimed before the exception +propagates. +-/ +def withCudaArena {α : Type} (keep : α → Array Buffer) (body : IO α) : IO α := do + arenaEnter + try + let r ← body + arenaExit (keep r) + return r + catch e => + arenaExit #[] + throw e + /-- Allocate a length-`n` buffer filled with zeros. -/ @[extern "torchlean_cuda_buffer_zeros"] opaque zeros (n : UInt32) : Buffer diff --git a/NN/Tests/Runtime/Cuda/Stress.lean b/NN/Tests/Runtime/Cuda/Stress.lean index 79f2d8c..b02bd88 100644 --- a/NN/Tests/Runtime/Cuda/Stress.lean +++ b/NN/Tests/Runtime/Cuda/Stress.lean @@ -119,6 +119,91 @@ def runReleaseStress : IO Unit := do if Buffer.size b != 0 then throw <| IO.userError s!"release size reset: expected 0, got {Buffer.size b}" +/-- +Build `k` distinct length-`n` buffers and force their allocation immediately. + +The fill value varies per `(salt, index)` for two reasons: within a call it stops the compiler from +hoisting one shared `full` out of the inner loop, and across calls a distinct `salt` keeps the whole +expression from being treated as loop-invariant (and hoisted out of the *caller's* loop) or CSE'd with +another call site — either of which would allocate the buffers once, outside the arena under test. The +returned element-count total is a forcing witness the caller checks. +-/ +def buildArenaScratch (n : UInt32) (k : Nat) (salt : Nat) : Array Buffer × Nat := + Id.run do + let mut held : Array Buffer := Array.mkEmpty k + for i in [0:k] do + held := held.push (Buffer.full n (1.0 + Float.ofNat (salt * k + i))) + let mut touched : Nat := 0 + for b in held do + touched := touched + (Buffer.size b).toNat + return (held, touched) + +def runArenaStress : IO Unit := do + IO.println "== cuda arena scope stress ==" + + let n : UInt32 := 4096 + let k : Nat := 64 + let blocks : Nat := 4 + let expectTouched := k * n.toNat + let base ← Buffer.allocatorStats + IO.println s!" baseline: {base.format}" + + -- Reclaim path: `k` buffers are built and held live (never released) inside an arena, kept alive + -- across the scope exit, and reclaimed anyway. This is the case explicit `release` cannot reach: the + -- buffers stay GC-reachable for the whole scope, so only the arena can free them. + for blockIdx in [0:blocks] do + let before ← Buffer.allocatorStatsWithToken (UInt32.ofNat blockIdx) + Buffer.arenaEnter + let (held, touched) := buildArenaScratch n k blockIdx + if touched != expectTouched then + throw <| IO.userError s!"arena: scratch build under-allocated ({touched} vs {expectTouched})" + let inside ← Buffer.allocatorStatsWithToken (UInt32.ofNat (blockIdx + 100)) + if inside.allocCount != before.allocCount + UInt64.ofNat k then + throw <| IO.userError + s!"arena: expected {k} in-scope allocations ({before.allocCount} → {inside.allocCount})" + if inside.freeCount != before.freeCount then + throw <| IO.userError + s!"arena: buffers freed before scope exit ({before.freeCount} → {inside.freeCount})" + -- Keep nothing: every in-scope buffer is reclaimed even though `held` still references them all. + Buffer.arenaExit #[] + -- Touch `held` *after* the exit so it stays live across it: this proves the ARENA did the freeing, + -- not reference-counted finalization (which cannot run while `held` is still referenced). + let heldGuard := held.size + let after ← Buffer.allocatorStatsWithToken (UInt32.ofNat (blockIdx + 200)) + if heldGuard != k then + throw <| IO.userError s!"arena: held guard mismatch ({heldGuard} vs {k})" + if after.freeCount != before.freeCount + UInt64.ofNat k then + throw <| IO.userError + s!"arena: scope exit freed {after.freeCount - before.freeCount}, expected {k} live buffers" + if after.liveBytes != before.liveBytes then + throw <| IO.userError + s!"arena: live bytes not restored at scope exit ({before.liveBytes} → {after.liveBytes})" + IO.println s!" reclaimed {k} live in-scope buffers across {blocks} arenas" + + -- Promotion path: a kept buffer survives the scope and stays usable; the rest are reclaimed. + let before ← Buffer.allocatorStats + Buffer.arenaEnter + let keeper := Buffer.full n 2.0 + -- Force `keeper`'s allocation inside the scope (so it is registered and then promoted on exit). + if Buffer.size keeper != n then + throw <| IO.userError "arena: keep-path keeper not allocated" + let (scratch, scratchTouched) := buildArenaScratch n k blocks + if scratchTouched != expectTouched then + throw <| IO.userError "arena: keep-path scratch build under-allocated" + Buffer.arenaExit #[keeper] + let scratchGuard := scratch.size + let after ← Buffer.allocatorStats + if scratchGuard != k then + throw <| IO.userError "arena: keep-path scratch guard mismatch" + if after.freeCount != before.freeCount + UInt64.ofNat k then + throw <| IO.userError + s!"arena keep: expected {k} frees, got {after.freeCount - before.freeCount}" + -- `keeper` was promoted out of the scope, so its data is intact: reducing it still sees `2.0`. + let keptSum := (Buffer.toFloatArray (Buffer.reduceSum keeper)).get! 0 + let expectedSum := 2.0 * Float.ofNat n.toNat + Utils.assertApprox "arena kept-buffer survives scope" keptSum expectedSum (tol := 1e-1) + IO.println " promoted buffer survived its arena" + def runGradientAliasingStress : IO Unit := do IO.println "== CUDA tape gradient aliasing regression ==" @@ -237,6 +322,7 @@ def run : IO Unit := do IO.println "=== CUDA runtime stress suite ===" runRngStress runReleaseStress + runArenaStress runGradientAliasingStress runLargeBufferStress runMatmulStress diff --git a/csrc/cuda/common/torchlean_cuda_arena.h b/csrc/cuda/common/torchlean_cuda_arena.h new file mode 100644 index 0000000..df9875c --- /dev/null +++ b/csrc/cuda/common/torchlean_cuda_arena.h @@ -0,0 +1,181 @@ +#pragma once + +#include + +#include +#include +#include +#include + +#include "torchlean_cuda_buffer.h" + +// Scoped device-memory arena ("withCudaArena"). +// +// Why this exists. TorchLean wraps each device allocation in a Lean external object whose finalizer +// returns the memory. In a long *pure* eager loop (a `foldl` of `Buffer -> Buffer` ops with no IO +// sequencing) every intermediate `Buffer` stays GC-*reachable* until the final readback, so the +// finalizer never runs and device memory grows without bound. Explicit `release` cannot help from +// inside a pure carrier: there is no IO point at which to call it. +// +// A scoped arena sidesteps GC reachability entirely. `arena_enter` opens an allocation epoch; every +// buffer allocated while it is open is registered to the epoch. `arena_exit` frees the device data of +// *every* buffer in the epoch -- whether or not it is still reachable -- except a small set of `keep` +// buffers (the step's results), which are promoted to the parent epoch (or untracked at the outermost +// level). This matches the natural phase boundary of a training loop: one LM step / one fold. +// +// Registration model. Each tracked buffer points at a heap-allocated `torchlean_arena_reg` (and the +// reg points back). When a buffer's struct is freed early (its finalizer runs mid-scope), it flips +// the reg's `alive` flag instead of leaving a dangling pointer, so the exit walk skips it. All +// registry mutation is serialized by `g_arena_mutex`; the fast paths (`arena_register` when no epoch +// is open, `arena_unlink` when the buffer is untracked) take no lock. +// +// Threading contract. `arena_enter`/`arena_exit` and the allocations between them run on one driver +// thread; finalizers (which only ever *unlink*) may run on any thread and are mutex-safe. Concurrent +// arenas on multiple threads are not supported in this first cut. + +typedef struct torchlean_arena_reg { + torchlean_cuda_buffer* b; // owning struct; valid only while `alive` + size_t depth; // epoch index this reg belongs to (stack position) + bool alive; // false once the buffer struct has been freed or the reg neutralized +} torchlean_arena_reg; + +typedef struct torchlean_arena_epoch { + torchlean_arena_reg** regs; + size_t count; + size_t cap; +} torchlean_arena_epoch; + +static pthread_mutex_t g_torchlean_arena_mutex = PTHREAD_MUTEX_INITIALIZER; +static torchlean_arena_epoch* g_torchlean_arena_stack = NULL; +static size_t g_torchlean_arena_depth = 0; // number of open epochs +static size_t g_torchlean_arena_cap = 0; + +static inline void torchlean_arena_epoch_push(torchlean_arena_epoch* e, torchlean_arena_reg* r) { + if (e->count == e->cap) { + size_t ncap = e->cap == 0 ? 16 : e->cap * 2; + void* p = realloc(e->regs, ncap * sizeof(torchlean_arena_reg*)); + if (!p) { + lean_internal_panic_out_of_memory(); + } + e->regs = (torchlean_arena_reg**)p; + e->cap = ncap; + } + e->regs[e->count++] = r; +} + +// Open a new allocation epoch. Subsequent allocations on this thread are registered to it until the +// matching `torchlean_arena_exit`. +static inline void torchlean_arena_enter(void) { + pthread_mutex_lock(&g_torchlean_arena_mutex); + if (g_torchlean_arena_depth == g_torchlean_arena_cap) { + size_t ncap = g_torchlean_arena_cap == 0 ? 4 : g_torchlean_arena_cap * 2; + void* p = realloc(g_torchlean_arena_stack, ncap * sizeof(torchlean_arena_epoch)); + if (!p) { + pthread_mutex_unlock(&g_torchlean_arena_mutex); + lean_internal_panic_out_of_memory(); + } + g_torchlean_arena_stack = (torchlean_arena_epoch*)p; + g_torchlean_arena_cap = ncap; + } + torchlean_arena_epoch* e = &g_torchlean_arena_stack[g_torchlean_arena_depth++]; + e->regs = NULL; + e->count = 0; + e->cap = 0; + pthread_mutex_unlock(&g_torchlean_arena_mutex); +} + +// Register a freshly allocated buffer to the current epoch (no-op when no arena is open or the buffer +// holds no device memory). Sets `b->arena_reg`. +static inline void torchlean_arena_register(torchlean_cuda_buffer* b) { + if (g_torchlean_arena_depth == 0 || !b || !b->data) { + return; // fast path: nothing to track + } + pthread_mutex_lock(&g_torchlean_arena_mutex); + if (g_torchlean_arena_depth == 0) { + pthread_mutex_unlock(&g_torchlean_arena_mutex); + return; + } + torchlean_arena_reg* r = (torchlean_arena_reg*)malloc(sizeof(torchlean_arena_reg)); + if (!r) { + pthread_mutex_unlock(&g_torchlean_arena_mutex); + lean_internal_panic_out_of_memory(); + } + r->b = b; + r->depth = g_torchlean_arena_depth - 1; + r->alive = true; + b->arena_reg = r; + torchlean_arena_epoch_push(&g_torchlean_arena_stack[g_torchlean_arena_depth - 1], r); + pthread_mutex_unlock(&g_torchlean_arena_mutex); +} + +// Detach a buffer from the registry because its struct is about to be freed. Flips the reg `alive` +// flag (so a concurrent exit walk skips it) and clears the back-pointer. No-op for untracked buffers. +static inline void torchlean_arena_unlink(torchlean_cuda_buffer* b) { + if (!b || !b->arena_reg) { + return; // fast path: not arena-tracked + } + pthread_mutex_lock(&g_torchlean_arena_mutex); + torchlean_arena_reg* r = (torchlean_arena_reg*)b->arena_reg; + if (r) { + r->alive = false; + } + b->arena_reg = NULL; + pthread_mutex_unlock(&g_torchlean_arena_mutex); +} + +// Close the current epoch. Releases the device data of every still-live buffer allocated in it, +// except the `nkeep` buffers in `keep`, which are promoted to the parent epoch (or untracked at the +// outermost level). `release_data` is the backend's per-buffer reclaim (cache-return on CUDA, free on +// the CPU stub). An unbalanced exit (no open epoch) is ignored. +static inline void torchlean_arena_exit(torchlean_cuda_buffer** keep, size_t nkeep, + bool (*release_data)(torchlean_cuda_buffer*)) { + pthread_mutex_lock(&g_torchlean_arena_mutex); + if (g_torchlean_arena_depth == 0) { + pthread_mutex_unlock(&g_torchlean_arena_mutex); + return; + } + size_t exited = --g_torchlean_arena_depth; + torchlean_arena_epoch e = g_torchlean_arena_stack[exited]; + bool has_parent = g_torchlean_arena_depth > 0; + + // Promote kept buffers out of the exiting epoch so the walk below leaves their data alone. A kept + // buffer is recognized by its reg living at the exiting depth; one tracked elsewhere (or not at all) + // is ignored. + for (size_t k = 0; k < nkeep; ++k) { + torchlean_cuda_buffer* kb = keep[k]; + if (!kb || !kb->arena_reg) { + continue; + } + torchlean_arena_reg* old = (torchlean_arena_reg*)kb->arena_reg; + if (old->depth != exited) { + continue; // belongs to an ancestor epoch; not ours to promote + } + old->alive = false; // neutralize: the walk frees this reg without releasing data + if (has_parent) { + torchlean_arena_reg* nr = (torchlean_arena_reg*)malloc(sizeof(torchlean_arena_reg)); + if (!nr) { + pthread_mutex_unlock(&g_torchlean_arena_mutex); + lean_internal_panic_out_of_memory(); + } + nr->b = kb; + nr->depth = g_torchlean_arena_depth - 1; + nr->alive = true; + kb->arena_reg = nr; + torchlean_arena_epoch_push(&g_torchlean_arena_stack[g_torchlean_arena_depth - 1], nr); + } else { + kb->arena_reg = NULL; // untrack: managed by RC / finalizer from here on + } + } + + for (size_t i = 0; i < e.count; ++i) { + torchlean_arena_reg* r = e.regs[i]; + if (r->alive) { + torchlean_cuda_buffer* b = r->b; // struct still valid while alive + b->arena_reg = NULL; // detach before freeing the reg (finalizer-safe) + (void)release_data(b); + } + free(r); + } + free(e.regs); + pthread_mutex_unlock(&g_torchlean_arena_mutex); +} diff --git a/csrc/cuda/common/torchlean_cuda_buffer.h b/csrc/cuda/common/torchlean_cuda_buffer.h index f8c0768..89fcb26 100644 --- a/csrc/cuda/common/torchlean_cuda_buffer.h +++ b/csrc/cuda/common/torchlean_cuda_buffer.h @@ -46,8 +46,9 @@ static inline uint32_t nat_to_u32_or_panic(b_lean_obj_arg o, const char* msg) { } typedef struct { - size_t size; // number of float32 elements - float* data; // device/host pointer (depending on build) + size_t size; // number of float32 elements + float* data; // device/host pointer (depending on build) + void* arena_reg; // non-NULL while tracked by an open `withCudaArena` scope (see torchlean_cuda_arena.h) } torchlean_cuda_buffer; // Helpers implemented by `torchlean_cuda_tensor.cu` / `torchlean_cuda_tensor_stub.c`. @@ -126,6 +127,12 @@ LEAN_EXPORT uint64_t torchlean_cuda_allocator_free_count(uint32_t u); LEAN_EXPORT uint64_t torchlean_cuda_allocator_device_free_bytes(uint32_t u); LEAN_EXPORT uint64_t torchlean_cuda_allocator_device_total_bytes(uint32_t u); +// Scoped device-memory arena (`withCudaArena`). `enter` opens an allocation epoch; `exit` reclaims +// every device buffer allocated in it except the `keep` array, which is promoted to the parent scope. +// See `torchlean_cuda_arena.h` for the model. Both are IO actions on the Lean side. +LEAN_EXPORT lean_obj_res torchlean_cuda_arena_enter(lean_obj_arg world); +LEAN_EXPORT lean_obj_res torchlean_cuda_arena_exit(b_lean_obj_arg keep, lean_obj_arg world); + #ifdef __cplusplus } // extern "C" #endif diff --git a/csrc/cuda/tensor/torchlean_cuda_tensor.cu b/csrc/cuda/tensor/torchlean_cuda_tensor.cu index b71f6b3..91c4399 100644 --- a/csrc/cuda/tensor/torchlean_cuda_tensor.cu +++ b/csrc/cuda/tensor/torchlean_cuda_tensor.cu @@ -1,6 +1,7 @@ #include #include +#include "torchlean_cuda_arena.h" #include "torchlean_cuda_buffer.h" #include "torchlean_cuda_common.h" #include "torchlean_cuda_deterministic_reductions_env.h" @@ -239,6 +240,7 @@ static void torchlean_cuda_buffer_finalize(void* ptr) { if (!b) { return; } + torchlean_arena_unlink(b); (void)torchlean_cuda_buffer_release_data(b); free(b); } @@ -275,6 +277,7 @@ extern "C" void torchlean_cuda_buffer_drop_unboxed(torchlean_cuda_buffer* b) { if (!b) { return; } + torchlean_arena_unlink(b); (void)torchlean_cuda_buffer_release_data(b); free(b); } @@ -286,6 +289,7 @@ extern "C" torchlean_cuda_buffer* torchlean_cuda_buffer_alloc(size_t n) { } b->size = n; b->data = NULL; + b->arena_reg = NULL; if (n > 0) { b->data = torchlean_cuda_take_cached_block(n); } @@ -302,6 +306,7 @@ extern "C" torchlean_cuda_buffer* torchlean_cuda_buffer_alloc(size_t n) { } if (n > 0) { torchlean_cuda_note_alloc(n); + torchlean_arena_register(b); } return b; } @@ -742,6 +747,31 @@ extern "C" LEAN_EXPORT uint32_t torchlean_runtime_collect_allocator(uint32_t for return 1; } +extern "C" LEAN_EXPORT lean_obj_res torchlean_cuda_arena_enter(lean_obj_arg world) { + (void)world; + torchlean_arena_enter(); + return lean_io_result_mk_ok(lean_box(0)); +} + +extern "C" LEAN_EXPORT lean_obj_res torchlean_cuda_arena_exit(b_lean_obj_arg keep, + lean_obj_arg world) { + (void)world; + size_t nkeep = lean_array_size((lean_object*)keep); + torchlean_cuda_buffer** ptrs = NULL; + if (nkeep > 0) { + ptrs = (torchlean_cuda_buffer**)malloc(nkeep * sizeof(torchlean_cuda_buffer*)); + if (!ptrs) { + lean_internal_panic_out_of_memory(); + } + for (size_t i = 0; i < nkeep; ++i) { + ptrs[i] = torchlean_cuda_buffer_unbox(lean_array_get_core((lean_object*)keep, i)); + } + } + torchlean_arena_exit(ptrs, nkeep, torchlean_cuda_buffer_release_data); + free(ptrs); + return lean_io_result_mk_ok(lean_box(0)); +} + extern "C" LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_zeros(uint32_t n) { torchlean_cuda_buffer* out = torchlean_cuda_buffer_alloc((size_t)n); if (n == 0) { diff --git a/csrc/cuda/tensor/torchlean_cuda_tensor_stub.c b/csrc/cuda/tensor/torchlean_cuda_tensor_stub.c index 9e29aa0..5bfc269 100644 --- a/csrc/cuda/tensor/torchlean_cuda_tensor_stub.c +++ b/csrc/cuda/tensor/torchlean_cuda_tensor_stub.c @@ -1,6 +1,7 @@ #include #include +#include "torchlean_cuda_arena.h" #include "torchlean_cuda_buffer.h" #include "torchlean_cuda_deterministic_reductions_env.h" #include "torchlean_cuda_rng_common.h" @@ -83,6 +84,7 @@ static void torchlean_cuda_buffer_finalize(void* ptr) { if (!b) { return; } + torchlean_arena_unlink(b); (void)torchlean_cuda_buffer_release_data(b); free(b); } @@ -119,6 +121,7 @@ void torchlean_cuda_buffer_drop_unboxed(torchlean_cuda_buffer* b) { if (!b) { return; } + torchlean_arena_unlink(b); (void)torchlean_cuda_buffer_release_data(b); free(b); } @@ -130,6 +133,7 @@ torchlean_cuda_buffer* torchlean_cuda_buffer_alloc(size_t n) { } b->size = n; b->data = NULL; + b->arena_reg = NULL; if (n > 0) { b->data = (float*)malloc(n * sizeof(float)); if (!b->data) { @@ -137,6 +141,7 @@ torchlean_cuda_buffer* torchlean_cuda_buffer_alloc(size_t n) { lean_internal_panic_out_of_memory(); } torchlean_cuda_note_alloc(n); + torchlean_arena_register(b); } return b; } @@ -201,6 +206,30 @@ LEAN_EXPORT uint32_t torchlean_runtime_collect_allocator(uint32_t force) { return 1; } +LEAN_EXPORT lean_obj_res torchlean_cuda_arena_enter(lean_obj_arg world) { + (void)world; + torchlean_arena_enter(); + return lean_io_result_mk_ok(lean_box(0)); +} + +LEAN_EXPORT lean_obj_res torchlean_cuda_arena_exit(b_lean_obj_arg keep, lean_obj_arg world) { + (void)world; + size_t nkeep = lean_array_size((lean_object*)keep); + torchlean_cuda_buffer** ptrs = NULL; + if (nkeep > 0) { + ptrs = (torchlean_cuda_buffer**)malloc(nkeep * sizeof(torchlean_cuda_buffer*)); + if (!ptrs) { + lean_internal_panic_out_of_memory(); + } + for (size_t i = 0; i < nkeep; ++i) { + ptrs[i] = torchlean_cuda_buffer_unbox(lean_array_get_core((lean_object*)keep, i)); + } + } + torchlean_arena_exit(ptrs, nkeep, torchlean_cuda_buffer_release_data); + free(ptrs); + return lean_io_result_mk_ok(lean_box(0)); +} + LEAN_EXPORT lean_obj_res torchlean_cuda_buffer_zeros(uint32_t n) { torchlean_cuda_buffer* out = torchlean_cuda_buffer_alloc((size_t)n); for (size_t i = 0; i < (size_t)n; ++i) { From eed66d607f97f855ecdc1f78e86b49023ea8af3a Mon Sep 17 00:00:00 2001 From: Nicolas Rouquette Date: Tue, 23 Jun 2026 11:34:25 -0700 Subject: [PATCH 2/2] feat(cuda): use-after-arena-free detector + in-suite fork death test A buffer reclaimed by `arena_exit` keeps its Lean external object but has size 0 and freed/cached data; using it afterwards is a use-after-free. Before, that surfaced only as a cryptic "size mismatch (N vs 0)" deep in a kernel -- and slipped through silently when both operands were freed, since the bare size check passes (0 == 0). Add an opt-in detector (TORCHLEAN_ARENA_DEBUG=1): `arena_exit` records the reclaiming epoch in a new `arena_freed_depth` buffer field, and the `require_same_size{2,3}` choke point (which every binary/ternary op already calls) asserts liveness before comparing sizes, so a stale operand panics naming the op and epoch: use-after-arena-free: torchlean_cuda_buffer_add lhs was reclaimed by arena_exit at depth 0 When off (the default) the detector is one predicted branch on a cached env read plus an 8-byte field; the suite runs at the same wall-clock on or off. Regression coverage: `runArenaDetectorDeathTest` in nn_tests_suite (runs on both the CUDA build and the CPU stub). A detected UAF must panic, which cannot be caught in-process, so it forks the suite binary per configuration and inspects the outcome: * positive -- detector on + planted UAF => child aborts with the message * negative -- detector on + a valid promotion reused in a binary op (through the detector's own choke point) => clean exit * control -- detector off + planted UAF => clean exit (silent slip) check.sh --ci-all --cuda: build + test + lint all green. --- NN/Tests/Runtime/Cuda/Stress.lean | 91 +++++++++++++++++++ NN/Tests/Suite.lean | 18 +++- csrc/cuda/common/torchlean_cuda_arena.h | 3 + csrc/cuda/common/torchlean_cuda_buffer.h | 40 ++++++++ csrc/cuda/tensor/torchlean_cuda_tensor.cu | 1 + csrc/cuda/tensor/torchlean_cuda_tensor_stub.c | 1 + 6 files changed, 149 insertions(+), 5 deletions(-) diff --git a/NN/Tests/Runtime/Cuda/Stress.lean b/NN/Tests/Runtime/Cuda/Stress.lean index b02bd88..09f1657 100644 --- a/NN/Tests/Runtime/Cuda/Stress.lean +++ b/NN/Tests/Runtime/Cuda/Stress.lean @@ -318,11 +318,102 @@ def runMatmulStress : IO Unit := do Utils.assertTensorApprox (s := sY2) "matmul stress case2 fp32" yFp322 yRef2 (tol := 7e-3) Utils.assertTensorApprox (s := sY2) "matmul stress case2 fp64" yFp642 yRef2 (tol := 1e-9) +/-- +Planted use-after-free, the subject of `runArenaDetectorDeathTest`. Allocates two buffers inside an +arena, reclaims **both** at `arenaExit`, then uses them in an op — the same hazard as touching a +`release`d buffer. Because reclaimed buffers have `size == 0`, the bare size check (`0 == 0`) lets this +slip through to a launch on freed memory; only the detector (`TORCHLEAN_ARENA_DEBUG=1`) catches it, +naming the epoch. Run only in a forked child (selected by `TORCHLEAN_ARENA_UAF_PROBE=uaf`): under the +detector it `panic`s (which is why it must be forked, not asserted in-process); with the detector off it +returns a silently-wrong size-0 result — exactly the silent corruption the detector closes. -/ +def runArenaUseAfterFreeProbe : IO Unit := do + IO.println "== cuda arena use-after-free probe ==" + let n : UInt32 := 16 + Buffer.arenaEnter + let a := Buffer.full n 3.0 + let b := Buffer.full n 5.0 + let forced := Buffer.size a + Buffer.size b -- force both allocations inside the epoch + if forced != 2 * n then + throw <| IO.userError "uaf probe: operands not allocated" + Buffer.arenaExit #[] -- reclaim BOTH (keep nothing) + -- `a` and `b` are now reclaimed (size 0). The detector asserts liveness before the size check and + -- panics here; with the detector off, `add` slips past `0 == 0` and yields a silently-empty buffer. + let bad := Buffer.add a b + IO.println s!" detector OFF: use-after-free slipped through, result size = {Buffer.size bad} (expected 16)" + +/-- +Valid arena promotion, the negative-case subject of `runArenaDetectorDeathTest`. Promotes one buffer +past the scope and reclaims another, then uses the *promoted* buffer in a **binary** op — which flows +through the same `require_same_size2` choke point the detector guards. A promoted buffer is live +(`arena_freed_depth == 0`), so the detector must not fire and the result must be correct. Selected in a +forked child by `TORCHLEAN_ARENA_UAF_PROBE=valid`. -/ +def runArenaValidPromotionProbe : IO Unit := do + IO.println "== cuda arena valid-promotion probe ==" + let n : UInt32 := 16 + Buffer.arenaEnter + let keep := Buffer.full n 2.0 + let scratch := Buffer.full n 7.0 + let forced := Buffer.size keep + Buffer.size scratch -- force both allocations inside the epoch + if forced != 2 * n then + throw <| IO.userError "valid-promotion probe: operands not allocated" + Buffer.arenaExit #[keep] -- promote `keep`; reclaim `scratch` + -- `keep` is promoted (live). A binary op on it exercises the detector's choke point and must pass. + let ok := Buffer.add keep keep + let s := (Buffer.toFloatArray (Buffer.reduceSum ok)).get! 0 + if s != 4.0 * Float.ofNat n.toNat then + throw <| IO.userError s!"valid-promotion probe: wrong result {s} (expected {4.0 * Float.ofNat n.toNat})" + IO.println s!" promoted buffer reused in a binary op, result sum = {s}" + +/-- +Positive + negative regression test for the arena use-after-free detector. A detected UAF must `panic` +(it cannot be caught in-process), so the suite binary is forked in each configuration via +`/proc/self/exe` and its outcome inspected: + +* **positive** — `TORCHLEAN_ARENA_DEBUG=1` + the planted UAF ⇒ the child aborts with a + `use-after-arena-free` panic; +* **negative (no false positive)** — detector on + a *valid* arena promotion + (`runArenaValidPromotionProbe`, which reuses a promoted buffer in a binary op through the detector's + own choke point) ⇒ the child exits cleanly, so the detector never fires on a kept buffer; +* **control** — detector off + the planted UAF ⇒ the child exits cleanly (the UAF slips through, the + silent corruption the detector closes). + +The forked children re-enter the suite with `TORCHLEAN_ARENA_UAF_PROBE` set (see `NN.Tests.run`) and so +run only the relevant fragment. Linux-only (uses `/proc/self/exe`); skipped with a note elsewhere. -/ +def runArenaDetectorDeathTest : IO Unit := do + IO.println "== cuda arena use-after-free detector (fork death test) ==" + let self : System.FilePath := "/proc/self/exe" + if !(← self.pathExists) then + IO.println " skipped: no /proc/self/exe (fork death test is Linux-only)" + return + let contains (hay needle : String) : Bool := (hay.splitOn needle).length ≥ 2 + let fork (debug : Bool) (mode : String) : IO IO.Process.Output := do + let env := #[("TORCHLEAN_ARENA_UAF_PROBE", some mode)] + let env := if debug then env.push ("TORCHLEAN_ARENA_DEBUG", some "1") else env + IO.Process.output { cmd := self.toString, args := #[], env := env } + -- positive: the detector aborts the planted use-after-free, naming the hazard. + let pos ← fork true "uaf" + if pos.exitCode == 0 then + throw <| IO.userError "arena UAF detector: detector ON did NOT abort the planted use-after-free" + if !(contains (pos.stderr ++ pos.stdout) "use-after-arena-free") then + throw <| IO.userError s!"arena UAF detector: detector ON aborted without the expected message; stderr:\n{pos.stderr}" + IO.println " positive: detector ON aborts the planted use-after-free ✓" + -- negative: a valid promotion under the detector is left untouched (no false positive). + let neg ← fork true "valid" + if neg.exitCode != 0 then + throw <| IO.userError s!"arena UAF detector: false positive on a valid promotion (exit {neg.exitCode}); stderr:\n{neg.stderr}" + IO.println " negative: detector ON leaves a valid arena promotion untouched ✓" + -- control: with the detector off the same use-after-free slips through and the child exits cleanly. + let off ← fork false "uaf" + if off.exitCode != 0 then + throw <| IO.userError s!"arena UAF detector: with the detector off the probe should exit cleanly (exit {off.exitCode})" + IO.println " control: detector OFF leaves the use-after-free undetected, as designed ✓" + def run : IO Unit := do IO.println "=== CUDA runtime stress suite ===" runRngStress runReleaseStress runArenaStress + runArenaDetectorDeathTest runGradientAliasingStress runLargeBufferStress runMatmulStress diff --git a/NN/Tests/Suite.lean b/NN/Tests/Suite.lean index 9c35785..fd84442 100644 --- a/NN/Tests/Suite.lean +++ b/NN/Tests/Suite.lean @@ -42,11 +42,19 @@ def usage : String := ] def run : IO Unit := do - IO.println "== TorchLean: curated tests ==" - Tests.Floats.run - Tests.Rationals.Suite.run - Tests.Cuda.run - IO.println "== TorchLean: all curated tests passed ==" + -- Death-test child modes for the arena use-after-free detector. A forked child (see + -- `Tests.Cuda.Stress.runArenaDetectorDeathTest`) re-enters here with `TORCHLEAN_ARENA_UAF_PROBE` set + -- and runs only the planted UAF (`uaf` — expected to panic under `TORCHLEAN_ARENA_DEBUG=1`) or a + -- valid promotion (`valid` — expected to be left alone), then exits, so the parent can inspect it. + match ← IO.getEnv "TORCHLEAN_ARENA_UAF_PROBE" with + | some "uaf" => Tests.Cuda.Stress.runArenaUseAfterFreeProbe + | some "valid" => Tests.Cuda.Stress.runArenaValidPromotionProbe + | _ => + IO.println "== TorchLean: curated tests ==" + Tests.Floats.run + Tests.Rationals.Suite.run + Tests.Cuda.run + IO.println "== TorchLean: all curated tests passed ==" def main (args : List String) : IO Unit := do match args with diff --git a/csrc/cuda/common/torchlean_cuda_arena.h b/csrc/cuda/common/torchlean_cuda_arena.h index df9875c..c76aa14 100644 --- a/csrc/cuda/common/torchlean_cuda_arena.h +++ b/csrc/cuda/common/torchlean_cuda_arena.h @@ -173,6 +173,9 @@ static inline void torchlean_arena_exit(torchlean_cuda_buffer** keep, size_t nke torchlean_cuda_buffer* b = r->b; // struct still valid while alive b->arena_reg = NULL; // detach before freeing the reg (finalizer-safe) (void)release_data(b); + if (torchlean_arena_debug_enabled()) { + b->arena_freed_depth = exited + 1; // record the reclaiming epoch for the UAF detector + } } free(r); } diff --git a/csrc/cuda/common/torchlean_cuda_buffer.h b/csrc/cuda/common/torchlean_cuda_buffer.h index 89fcb26..3d463f0 100644 --- a/csrc/cuda/common/torchlean_cuda_buffer.h +++ b/csrc/cuda/common/torchlean_cuda_buffer.h @@ -5,6 +5,7 @@ #include #include #include +#include #ifdef __cplusplus extern "C" { @@ -49,6 +50,7 @@ typedef struct { size_t size; // number of float32 elements float* data; // device/host pointer (depending on build) void* arena_reg; // non-NULL while tracked by an open `withCudaArena` scope (see torchlean_cuda_arena.h) + size_t arena_freed_depth; // debug only: 0 = live; else (reclaiming arena depth + 1). See the UAF detector below. } torchlean_cuda_buffer; // Helpers implemented by `torchlean_cuda_tensor.cu` / `torchlean_cuda_tensor_stub.c`. @@ -83,10 +85,43 @@ static inline lean_object* torchlean_cuda_box_four_buffers( return out; } +// Arena use-after-free detector (debug only; opt in with `TORCHLEAN_ARENA_DEBUG=1`). +// +// A device buffer reclaimed by `arena_exit` keeps its struct (Lean still owns the external object) but +// has `size == 0`/`data == NULL`, and in debug mode `arena_freed_depth` records the epoch that freed +// it (depth + 1). Touching it afterwards is the same hazard as using a `release`d buffer. The size +// helpers below are the common choke point for every binary/ternary op, so asserting liveness here +// turns a stale operand into a panic naming the epoch that freed it, instead of a silent corrupt +// kernel launch — and it catches the case the bare size check cannot (both operands freed ⇒ `0 == 0`). +// When the flag is off, the whole detector is one predicted branch on a cached int. +static inline int torchlean_arena_debug_enabled(void) { + static int cached = -1; // -1 = env not yet read + if (cached < 0) { + const char* v = getenv("TORCHLEAN_ARENA_DEBUG"); + cached = (v && v[0] && !(v[0] == '0' && v[1] == '\0')) ? 1 : 0; + } + return cached; +} + +static inline void torchlean_arena_assert_live( + const torchlean_cuda_buffer* b, const char* role, const char* fn) { + if (b->arena_freed_depth != 0) { + char msg[224]; + snprintf(msg, sizeof(msg), + "use-after-arena-free: %s %s was reclaimed by arena_exit at depth %zu", + fn, role, b->arena_freed_depth - 1); + lean_internal_panic(msg); + } +} + static inline void torchlean_cuda_require_same_size2( const torchlean_cuda_buffer* a, const torchlean_cuda_buffer* b, const char* fn) { + if (torchlean_arena_debug_enabled()) { + torchlean_arena_assert_live(a, "lhs", fn); + torchlean_arena_assert_live(b, "rhs", fn); + } if (a->size != b->size) { char msg[192]; snprintf(msg, sizeof(msg), "%s: size mismatch (%zu vs %zu)", fn, a->size, b->size); @@ -99,6 +134,11 @@ static inline void torchlean_cuda_require_same_size3( const torchlean_cuda_buffer* b, const torchlean_cuda_buffer* c, const char* fn) { + if (torchlean_arena_debug_enabled()) { + torchlean_arena_assert_live(a, "arg0", fn); + torchlean_arena_assert_live(b, "arg1", fn); + torchlean_arena_assert_live(c, "arg2", fn); + } if (a->size != b->size || a->size != c->size) { char msg[224]; snprintf(msg, sizeof(msg), "%s: size mismatch (%zu vs %zu vs %zu)", fn, a->size, b->size, diff --git a/csrc/cuda/tensor/torchlean_cuda_tensor.cu b/csrc/cuda/tensor/torchlean_cuda_tensor.cu index 91c4399..cb41799 100644 --- a/csrc/cuda/tensor/torchlean_cuda_tensor.cu +++ b/csrc/cuda/tensor/torchlean_cuda_tensor.cu @@ -290,6 +290,7 @@ extern "C" torchlean_cuda_buffer* torchlean_cuda_buffer_alloc(size_t n) { b->size = n; b->data = NULL; b->arena_reg = NULL; + b->arena_freed_depth = 0; if (n > 0) { b->data = torchlean_cuda_take_cached_block(n); } diff --git a/csrc/cuda/tensor/torchlean_cuda_tensor_stub.c b/csrc/cuda/tensor/torchlean_cuda_tensor_stub.c index 5bfc269..9548fc9 100644 --- a/csrc/cuda/tensor/torchlean_cuda_tensor_stub.c +++ b/csrc/cuda/tensor/torchlean_cuda_tensor_stub.c @@ -134,6 +134,7 @@ torchlean_cuda_buffer* torchlean_cuda_buffer_alloc(size_t n) { b->size = n; b->data = NULL; b->arena_reg = NULL; + b->arena_freed_depth = 0; if (n > 0) { b->data = (float*)malloc(n * sizeof(float)); if (!b->data) {