From e6e6fa9b7b29bce1f407429513020166eed412f1 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 17 Apr 2026 16:25:43 +0200 Subject: [PATCH 1/9] Tighten per-worker GPU memory budget to 4 GiB. The previous 2 GiB/worker estimate ignored CUDA context and library overhead plus peak allocations inside individual tests, letting the cap float up to 64+ workers on large systems. A 4 GiB budget accounts for ~1 GiB baseline (context + loaded libraries) and leaves room for multi-GiB peak test allocations. Also reset the device before querying free memory so the budget reflects true capacity, and `@info`-log the computed budget (device, free memory, gpu_jobs, cpu_threads, cpu_free) so users can see why the cap landed where it did. Co-Authored-By: Claude Opus 4.7 (1M context) --- test/runtests.jl | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/test/runtests.jl b/test/runtests.jl index 7905c5c006..929f8706c0 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -101,15 +101,19 @@ end ## GPU-memory-based parallelism -# pick the first visible device and use its free memory to cap worker count. +# Cap worker count by how much of the primary device's free memory each worker +# claims. A CUDA worker needs its own context + libraries (~0.5–1 GiB baseline) +# plus room for peak per-test allocations; 4 GiB is the per-worker budget. # (Set `CUDA_VISIBLE_DEVICES` to choose which device is used.) +const gpu_memory_per_worker = 4 * 2^30 first_gpu = first(devices()) gpu_free = device!(first_gpu) do - mem = CUDA.free_memory() device_reset!() - mem + CUDA.free_memory() end -gpu_jobs = max(1, Int(gpu_free) ÷ (2 * 2^30)) +gpu_jobs = max(1, Int(gpu_free) ÷ gpu_memory_per_worker) + +@info "Parallelism budget" device = "$(CUDA.name(first_gpu)) ($(deviceid(first_gpu)))" gpu_free = Base.format_bytes(gpu_free) gpu_jobs cpu_threads = Sys.CPU_THREADS cpu_free = Base.format_bytes(Sys.free_memory()) if args.jobs === nothing default_jobs = min(ParallelTestRunner.default_njobs(), gpu_jobs) From 9436327c91f2f1d783cb2c4881ff721e19acf495 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sat, 18 Apr 2026 09:56:29 +0200 Subject: [PATCH 2/9] Relax memory requirement. --- test/runtests.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/test/runtests.jl b/test/runtests.jl index 929f8706c0..7da266a0aa 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -105,11 +105,12 @@ end # claims. A CUDA worker needs its own context + libraries (~0.5–1 GiB baseline) # plus room for peak per-test allocations; 4 GiB is the per-worker budget. # (Set `CUDA_VISIBLE_DEVICES` to choose which device is used.) -const gpu_memory_per_worker = 4 * 2^30 +const gpu_memory_per_worker = 2 * 2^30 first_gpu = first(devices()) gpu_free = device!(first_gpu) do + mem = CUDA.free_memory() device_reset!() - CUDA.free_memory() + mem end gpu_jobs = max(1, Int(gpu_free) ÷ gpu_memory_per_worker) From 514805e68b6e5ce52470ab749136706d09de5466 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sat, 18 Apr 2026 12:05:50 +0200 Subject: [PATCH 3/9] cuSOLVER: free cached workspace between tests; cap pool cache size. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two related changes to reclaim GPU memory that was effectively leaking: 1. Refactor cuSOLVER state to use an object-bound finalizer The fat cusolverDnHandle (and analogous sparse/mg wrappers) held a `workspace_gpu::CuVector{UInt8}` that can grow to hundreds of MiB. Its finalizer was attached to `current_task()`, which on test workers never dies — so `CUDA.reclaim()` could empty idle handle caches but couldn't touch the pinned workspace. Probing showed ~200 MiB of pool memory per cuSOLVER-using worker sitting there permanently. Wrap the raw handles in `mutable struct` and attach the finalizer to the wrapper instead. Add a `pre_reclaim_hooks` mechanism in CUDACore so libraries can drop their TLS state on reclaim; `reclaim()` then runs these hooks + GC + regular hooks, so finalizers fire and the underlying buffers are released before the pool is trimmed. cuSOLVER registers a hook that clears its dense/sparse/mg TLS entries. 2. Cap the test-worker memory pool release threshold at 64 MiB The default `MEMPOOL_ATTR_RELEASE_THRESHOLD` is `typemax(UInt64)` — freed stream-ordered buffers stay cached indefinitely, inflating NVML readings. Set a 64 MiB cap on test workers so released buffers go back to the driver quickly. Trades a per-call pool-alloc cost for peak GPU RSS tracking the actual live working set. Impact on library tests: big wins where workspace was hoarded (e.g. cutensornet/contractions 1508 → 642 MiB), small/neutral elsewhere. 19813 library tests pass cleanly with no errors or warnings. Co-Authored-By: Claude Opus 4.7 (1M context) --- CUDACore/src/memory.jl | 35 ++++++++++++++ lib/cusolver/src/cuSOLVER.jl | 93 ++++++++++++++++++++++++++---------- test/setup.jl | 10 ++++ 3 files changed, 112 insertions(+), 26 deletions(-) diff --git a/CUDACore/src/memory.jl b/CUDACore/src/memory.jl index 819ea5dc5b..c09ce5819d 100644 --- a/CUDACore/src/memory.jl +++ b/CUDACore/src/memory.jl @@ -458,6 +458,12 @@ end const reclaim_hooks = Any[] +# Hooks that run *before* `reclaim_hooks`, followed by a forced GC. Intended +# for library bindings to drop their task-local state so that finalizers can +# run (freeing e.g. cached workspace buffers) before the HandleCache and +# pool-trim stages release their own resources. +const pre_reclaim_hooks = Any[] + """ retry_reclaim(retry_if) do # code that may fail due to insufficient GPU memory @@ -505,6 +511,17 @@ end for hook in reclaim_hooks hook() end + elseif phase == 7 + # drop library-held TLS state (fat handles, cached workspaces) and + # finalize so the underlying buffers actually get freed + for hook in pre_reclaim_hooks + hook() + end + GC.gc(true) + for hook in reclaim_hooks + hook() + end + trim(pool_create(state.device)) else break end @@ -517,6 +534,14 @@ end for hook in reclaim_hooks hook() end + elseif phase == 4 + for hook in pre_reclaim_hooks + hook() + end + GC.gc(true) + for hook in reclaim_hooks + hook() + end else break end @@ -797,6 +822,16 @@ actually reclaimed. """ function reclaim(sz::Int=typemax(Int)) dev = device() + # Drop library-held state first (TLS-pinned fat handles, workspace caches, + # etc.), then run a full GC so finalizers release the underlying buffers + # and return idle handles to their caches. Only then do we tear down the + # caches themselves and trim the pool. + for hook in pre_reclaim_hooks + hook() + end + if !isempty(pre_reclaim_hooks) + GC.gc(true) + end for hook in reclaim_hooks hook() end diff --git a/lib/cusolver/src/cuSOLVER.jl b/lib/cusolver/src/cuSOLVER.jl index f5a075ae27..b52edae786 100644 --- a/lib/cusolver/src/cuSOLVER.jl +++ b/lib/cusolver/src/cuSOLVER.jl @@ -71,16 +71,27 @@ end const idle_dense_handles = HandleCache{CuContext,cusolverDnHandle_t}(dense_handle_ctor, dense_handle_dtor) -# fat handle, includes a cache -struct cusolverDnHandle - handle::cusolverDnHandle_t - workspace_gpu::CuVector{UInt8} - workspace_cpu::Vector{UInt8} - info::CuVector{Cint} +# fat handle, holds the raw cuSOLVER handle together with reusable workspace +# and info buffers. Mutable so the finalizer can be attached to the object +# itself: once no one references this struct (e.g. after the owning task dies +# or we clear it from task-local storage on reclaim), GC runs the finalizer +# which returns the buffers to the pool and the raw handle to the idle cache. +mutable struct cusolverDnHandle + const handle::cusolverDnHandle_t + const ctx::CuContext + const workspace_gpu::CuVector{UInt8} + const workspace_cpu::Vector{UInt8} + const info::CuVector{Cint} end Base.unsafe_convert(::Type{Ptr{cusolverDnContext}}, handle::cusolverDnHandle) = handle.handle +function dense_handle_finalizer(dh::cusolverDnHandle) + CUDACore.unsafe_free!(dh.workspace_gpu) + CUDACore.unsafe_free!(dh.info) + push!(idle_dense_handles, dh.ctx, dh.handle) +end + function dense_handle() cuda = CUDACore.active_state() @@ -97,13 +108,9 @@ function dense_handle() workspace_gpu = CuVector{UInt8}(undef, 0) workspace_cpu = Vector{UInt8}(undef, 0) info = CuVector{Cint}(undef, 1) - fat_handle = cusolverDnHandle(new_handle, workspace_gpu, workspace_cpu, info) - - finalizer(current_task()) do task - CUDACore.unsafe_free!(workspace_gpu) - CUDACore.unsafe_free!(info) - push!(idle_dense_handles, cuda.context, new_handle) - end + fat_handle = cusolverDnHandle(new_handle, cuda.context, workspace_gpu, + workspace_cpu, info) + finalizer(dense_handle_finalizer, fat_handle) cusolverDnSetStream(new_handle, cuda.stream) @@ -141,11 +148,24 @@ end const idle_sparse_handles = HandleCache{CuContext,cusolverSpHandle_t}(sparse_handle_ctor, sparse_handle_dtor) +# mutable wrapper so the raw sparse handle is released via an object-bound +# finalizer (see `cusolverDnHandle` for rationale). +mutable struct cusolverSpHandle + const handle::cusolverSpHandle_t + const ctx::CuContext +end +Base.unsafe_convert(::Type{cusolverSpHandle_t}, handle::cusolverSpHandle) = + handle.handle + +function sparse_handle_finalizer(sh::cusolverSpHandle) + push!(idle_sparse_handles, sh.ctx, sh.handle) +end + function sparse_handle() cuda = CUDACore.active_state() # every task maintains library state per device - LibraryState = @NamedTuple{handle::cusolverSpHandle_t, stream::CuStream} + LibraryState = @NamedTuple{handle::cusolverSpHandle, stream::CuStream} states = get!(task_local_storage(), :CUSOLVER_sparse) do Dict{CuContext,LibraryState}() end::Dict{CuContext,LibraryState} @@ -153,13 +173,12 @@ function sparse_handle() # get or create handle @noinline function new_state(cuda) new_handle = pop!(idle_sparse_handles, cuda.context) - finalizer(current_task()) do task - push!(idle_sparse_handles, cuda.context, new_handle) - end + wrapped = cusolverSpHandle(new_handle, cuda.context) + finalizer(sparse_handle_finalizer, wrapped) cusolverSpSetStream(new_handle, cuda.stream) - (; handle=new_handle, cuda.stream) + (; handle=wrapped, cuda.stream) end state = get!(states, cuda.context) do new_state(cuda) @@ -194,11 +213,26 @@ end::Vector{CuDevice} ndevices() = length(devices()) +# mutable wrapper so the mg handle is destroyed when its owning state struct +# becomes unreachable, rather than being pinned for the lifetime of the task. +mutable struct cusolverMgHandle + const handle::cusolverMgHandle_t + const ctx::CuContext +end +Base.unsafe_convert(::Type{cusolverMgHandle_t}, handle::cusolverMgHandle) = + handle.handle + +function mg_handle_finalizer(mh::cusolverMgHandle) + context!(mh.ctx; skip_destroyed=true) do + cusolverMgDestroy(mh.handle) + end +end + function mg_handle() cuda = CUDACore.active_state() # every task maintains library state per set of devices - LibraryState = @NamedTuple{handle::cusolverMgHandle_t} + LibraryState = @NamedTuple{handle::cusolverMgHandle} states = get!(task_local_storage(), :CUSOLVERmg) do Dict{UInt,LibraryState}() end::Dict{UInt,LibraryState} @@ -214,17 +248,13 @@ function mg_handle() @noinline function new_state(cuda) # we can't reuse cusolverMg handles because they can only be assigned devices once new_handle = cusolverMgCreate() - - finalizer(current_task()) do task - context!(cuda.context; skip_destroyed=true) do - cusolverMgDestroy(new_handle) - end - end + wrapped = cusolverMgHandle(new_handle, cuda.context) + finalizer(mg_handle_finalizer, wrapped) devs = convert.(Cint, devices()) cusolverMgDeviceSelect(new_handle, length(devs), devs) - (; handle=new_handle) + (; handle=wrapped) end state = get!(states, key) do new_state(cuda) @@ -260,9 +290,20 @@ function __init__() end end + # Drop task-local cuSOLVER state on reclaim so the cached workspace + # CuVectors (sometimes hundreds of MiB) become GC-able. + push!(CUDACore.pre_reclaim_hooks, drop_library_state!) + _initialized[] = true end +function drop_library_state!() + delete!(task_local_storage(), :CUSOLVER_dense) + delete!(task_local_storage(), :CUSOLVER_sparse) + delete!(task_local_storage(), :CUSOLVERmg) + return +end + include("precompile.jl") # deprecated binding for backwards compatibility diff --git a/test/setup.jl b/test/setup.jl index 841e99a371..2b9a49fb6b 100644 --- a/test/setup.jl +++ b/test/setup.jl @@ -46,6 +46,16 @@ end # precompile the runtime library CUDA.precompile_runtime() +# Cap the amount of memory the CUDA pool caches (default: unbounded). Freed +# stream-ordered buffers now go back to the driver almost immediately, which +# keeps each test's GPU RSS close to its peak live allocation rather than the +# running max across the test's lifetime. Trades per-call pool-alloc cost +# for lower NVML-reported memory. +let dev = device(), pool = CUDACore.pool_create(dev) + CUDACore.attribute!(pool, CUDACore.MEMPOOL_ATTR_RELEASE_THRESHOLD, + UInt64(64 * 2^20)) +end + ## custom test record capturing CUDA-specific statistics From 2178e2faf4473810316e5f630805288eca335d3e Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sat, 18 Apr 2026 12:15:26 +0200 Subject: [PATCH 4/9] Apply cuSOLVER-style fat-handle refactor to cuBLAS and cuSPARSE. Same pattern as the prior cuSOLVER commit: wrap the raw handle in a mutable struct with an object-bound finalizer, and register a `pre_reclaim_hooks` entry that clears the library's TLS state. When CUDA.reclaim() runs, the fat-handle wrapper becomes unreferenced, GC collects it, its finalizer returns the raw handle to the idle cache, and the subsequent HandleCache hook destroys the idle handle properly. Memory savings for cuBLAS/cuSPARSE are smaller than cuSOLVER (no cached workspace CuVector; just the ~handle), but this removes the last place library state stays pinned across the worker's lifetime, and makes the lifecycle consistent across libraries. Co-Authored-By: Claude Opus 4.7 (1M context) --- lib/cublas/src/cuBLAS.jl | 52 +++++++++++++++++++++++++++++------- lib/cusparse/src/cuSPARSE.jl | 31 +++++++++++++++++---- 2 files changed, 68 insertions(+), 15 deletions(-) diff --git a/lib/cublas/src/cuBLAS.jl b/lib/cublas/src/cuBLAS.jl index 17de4f7bd8..a045cb6914 100644 --- a/lib/cublas/src/cuBLAS.jl +++ b/lib/cublas/src/cuBLAS.jl @@ -82,11 +82,25 @@ function handle_dtor(ctx, handle) end const idle_handles = HandleCache{CuContext,cublasHandle_t}(handle_ctor, handle_dtor) +# mutable wrapper so the raw handle is released via an object-bound finalizer: +# when TLS state is cleared (e.g. on reclaim) and GC runs, the wrapper is +# collected and its finalizer returns the handle to the idle cache instead +# of the handle being pinned for the entire lifetime of the owning task. +mutable struct cublasHandle + const handle::cublasHandle_t + const ctx::CuContext +end +Base.unsafe_convert(::Type{cublasHandle_t}, handle::cublasHandle) = handle.handle + +function handle_finalizer(h::cublasHandle) + push!(idle_handles, h.ctx, h.handle) +end + function handle() cuda = CUDACore.active_state() # every task maintains library state per device - LibraryState = @NamedTuple{handle::cublasHandle_t, stream::CuStream, math_mode::CUDACore.MathMode} + LibraryState = @NamedTuple{handle::cublasHandle, stream::CuStream, math_mode::CUDACore.MathMode} states = get!(task_local_storage(), :CUBLAS) do Dict{CuContext,LibraryState}() end::Dict{CuContext,LibraryState} @@ -94,15 +108,14 @@ function handle() # get library state @noinline function new_state(cuda) new_handle = pop!(idle_handles, cuda.context) - finalizer(current_task()) do task - push!(idle_handles, cuda.context, new_handle) - end + wrapped = cublasHandle(new_handle, cuda.context) + finalizer(handle_finalizer, wrapped) cublasSetStream_v2(new_handle, cuda.stream) cublasSetPointerMode_v2(new_handle, CUBLAS_POINTER_MODE_DEVICE) math_mode!(new_handle, cuda.math_mode) - (; handle=new_handle, cuda.stream, cuda.math_mode) + (; handle=wrapped, cuda.stream, cuda.math_mode) end state = get!(states, cuda.context) do new_state(cuda) @@ -144,6 +157,17 @@ end const idle_xt_handles = HandleCache{Vector{CuContext},cublasXtHandle_t}(xt_handle_ctor, xt_handle_dtor) +# mutable wrapper for the xt handle, see `cublasHandle` for rationale. +mutable struct cublasXtHandle + const handle::cublasXtHandle_t + const ctxs::Vector{CuContext} +end +Base.unsafe_convert(::Type{cublasXtHandle_t}, h::cublasXtHandle) = h.handle + +function xt_handle_finalizer(h::cublasXtHandle) + push!(idle_xt_handles, h.ctxs, h.handle) +end + function devices!(devs::Vector{CuDevice}) task_local_storage(:CUBLASxt_devices, sort(devs; by=deviceid)) return @@ -160,7 +184,7 @@ function xt_handle() cuda = CUDACore.active_state() # every task maintains library state per set of devices - LibraryState = @NamedTuple{handle::cublasXtHandle_t} + LibraryState = @NamedTuple{handle::cublasXtHandle} states = get!(task_local_storage(), :CUBLASxt) do Dict{UInt,LibraryState}() end::Dict{UInt,LibraryState} @@ -177,9 +201,8 @@ function xt_handle() ctxs = [context(dev) for dev in devices()] new_handle = pop!(idle_xt_handles, ctxs) - finalizer(current_task()) do task - push!(idle_xt_handles, ctxs, new_handle) - end + wrapped = cublasXtHandle(new_handle, ctxs) + finalizer(xt_handle_finalizer, wrapped) # if we're using the stream-ordered allocator, # make sure allocations are visible on all devices @@ -193,7 +216,7 @@ function xt_handle() devs = convert.(Cint, devices()) cublasXtDeviceSelect(new_handle, length(devs), devs) - (; handle=new_handle) + (; handle=wrapped) end state = get!(states, key) do new_state(cuda) @@ -278,9 +301,18 @@ function __init__() atexit(flush_log_messages) end + # drop task-local cuBLAS state on reclaim so handles can be destroyed + push!(CUDACore.pre_reclaim_hooks, drop_library_state!) + _initialized[] = true end +function drop_library_state!() + delete!(task_local_storage(), :CUBLAS) + delete!(task_local_storage(), :CUBLASxt) + return +end + include("precompile.jl") # deprecated binding for backwards compatibility diff --git a/lib/cusparse/src/cuSPARSE.jl b/lib/cusparse/src/cuSPARSE.jl index e5d73d477a..c7f45fa6ce 100644 --- a/lib/cusparse/src/cuSPARSE.jl +++ b/lib/cusparse/src/cuSPARSE.jl @@ -72,11 +72,25 @@ function handle_dtor(ctx, handle) end const idle_handles = HandleCache{CuContext,cusparseHandle_t}(handle_ctor, handle_dtor) +# mutable wrapper so the raw handle is released via an object-bound finalizer: +# when TLS state is cleared (e.g. on reclaim) and GC runs, the wrapper is +# collected and its finalizer returns the handle to the idle cache instead +# of the handle being pinned for the entire lifetime of the owning task. +mutable struct cusparseHandle + const handle::cusparseHandle_t + const ctx::CuContext +end +Base.unsafe_convert(::Type{cusparseHandle_t}, h::cusparseHandle) = h.handle + +function handle_finalizer(h::cusparseHandle) + push!(idle_handles, h.ctx, h.handle) +end + function handle() cuda = CUDACore.active_state() # every task maintains library state per device - LibraryState = @NamedTuple{handle::cusparseHandle_t, stream::CuStream} + LibraryState = @NamedTuple{handle::cusparseHandle, stream::CuStream} states = get!(task_local_storage(), :CUSPARSE) do Dict{CuContext,LibraryState}() end::Dict{CuContext,LibraryState} @@ -84,13 +98,12 @@ function handle() # get library state @noinline function new_state(cuda) new_handle = pop!(idle_handles, cuda.context) - finalizer(current_task()) do task - push!(idle_handles, cuda.context, new_handle) - end + wrapped = cusparseHandle(new_handle, cuda.context) + finalizer(handle_finalizer, wrapped) cusparseSetStream(new_handle, cuda.stream) - (; handle=new_handle, cuda.stream) + (; handle=wrapped, cuda.stream) end state = get!(states, cuda.context) do new_state(cuda) @@ -128,9 +141,17 @@ function __init__() libcusparse = CUDA_Runtime_jll.libcusparse end + # drop task-local cuSPARSE state on reclaim so handles can be destroyed + push!(CUDACore.pre_reclaim_hooks, drop_library_state!) + _initialized[] = true end +function drop_library_state!() + delete!(task_local_storage(), :CUSPARSE) + return +end + # KernelAbstractions integration import KernelAbstractions as KA KA.get_backend(::CuSparseVector) = CUDACore.CUDAKernels.CUDABackend() From 056050b27a76c4a61db7981a28b486c1d635e7e2 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sat, 18 Apr 2026 15:18:57 +0200 Subject: [PATCH 5/9] Lower memory budget. --- test/runtests.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/runtests.jl b/test/runtests.jl index 7da266a0aa..66984d1abb 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -105,7 +105,7 @@ end # claims. A CUDA worker needs its own context + libraries (~0.5–1 GiB baseline) # plus room for peak per-test allocations; 4 GiB is the per-worker budget. # (Set `CUDA_VISIBLE_DEVICES` to choose which device is used.) -const gpu_memory_per_worker = 2 * 2^30 +const gpu_memory_per_worker = 1 * 2^30 first_gpu = first(devices()) gpu_free = device!(first_gpu) do mem = CUDA.free_memory() From 6ac0a50bb1f7c23919aa5b6a180d879b65bc7075 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sat, 18 Apr 2026 21:21:12 +0200 Subject: [PATCH 6/9] Unify reclaim machinery behind a typed registry. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the two untyped hook arrays (pre_reclaim_hooks, reclaim_hooks) and the three copies of the phase ladder with a single Reclaimable abstract type, a ReclaimLevel enum, and one reclaim(level) entry point. Libraries register their HandleCache and TaskLocalCache instances from __init__ and implement drop!/purge! via dispatch; CUDACore invokes them with per-hook error isolation, modelled on Base.atexit_hooks. retry_reclaim walks ReclaimLevel one rung at a time. Convert cuDNN, cuRAND, cuTENSOR, cuTensorNet, and cuStateVec from the task-bound finalizer pattern (finalizer(current_task(), …)) to the object-bound fat-handle pattern used by cuBLAS/cuSPARSE/cuSOLVER. The old task-bound variant only released resources once the owning task was garbage-collected, so those libraries previously leaked memory under pressure; they now participate in RECLAIM_DROP_STATE. Co-Authored-By: Claude Opus 4.7 (1M context) --- CUDACore/src/CUDACore.jl | 1 + CUDACore/src/memory.jl | 200 +++++++++++++---------------- CUDACore/src/utils/cache.jl | 15 +-- CUDACore/src/utils/reclaim.jl | 119 +++++++++++++++++ lib/cublas/src/cuBLAS.jl | 32 ++--- lib/cudnn/src/cuDNN.jl | 33 +++-- lib/curand/src/cuRAND.jl | 40 ++++-- lib/curand/src/cuda_integration.jl | 51 ++------ lib/cusolver/src/cuSOLVER.jl | 42 +++--- lib/cusparse/src/cuSPARSE.jl | 18 +-- lib/custatevec/src/cuStateVec.jl | 37 +++--- lib/cutensor/src/cuTENSOR.jl | 33 +++-- lib/cutensornet/src/cuTensorNet.jl | 33 +++-- test/core/pool.jl | 2 +- 14 files changed, 389 insertions(+), 267 deletions(-) create mode 100644 CUDACore/src/utils/reclaim.jl diff --git a/CUDACore/src/CUDACore.jl b/CUDACore/src/CUDACore.jl index f14c56954c..fa2bfa2b45 100644 --- a/CUDACore/src/CUDACore.jl +++ b/CUDACore/src/CUDACore.jl @@ -73,6 +73,7 @@ include("pointer.jl") # core utilities include("utils/call.jl") +include("utils/reclaim.jl") include("utils/cache.jl") include("utils/struct_size.jl") include("../lib/cudadrv/CUDAdrv.jl") diff --git a/CUDACore/src/memory.jl b/CUDACore/src/memory.jl index c09ce5819d..8919186c55 100644 --- a/CUDACore/src/memory.jl +++ b/CUDACore/src/memory.jl @@ -456,105 +456,97 @@ function Base.showerror(io::IO, err::OutOfGPUMemoryError) end end -const reclaim_hooks = Any[] +## reclaim escalation +# +# `Reclaimable`/`register_reclaimable!`/`TaskLocalCache`/`drop!`/`purge!` +# are defined in utils/reclaim.jl. Here we add the ladder that drives +# them along with the allocator-specific sync/trim steps. + +""" + ReclaimLevel + +Escalation levels shared by `reclaim(level)` and `retry_reclaim`, ordered +from cheapest to most aggressive: + +| Level | Action | +| :--- | :--- | +| `RECLAIM_SYNC` | synchronize the current task's stream | +| `RECLAIM_DEVICE_SYNC` | synchronize the whole device | +| `RECLAIM_GC_MINOR` | minor Julia GC (and device sync) | +| `RECLAIM_GC_FULL` | full Julia GC (and device sync) | +| `RECLAIM_POOL_TRIM` | trim unused memory from the pool | +| `RECLAIM_PURGE_CACHES` | empty `HandleCache`s (destroy idle handles) | +| `RECLAIM_DROP_STATE` | also drop task-local library state, then GC+purge+trim | + +Steps that don't apply to the current allocator (e.g. stream sync on a +non-stream-ordered device) are silently skipped. +""" +@enum ReclaimLevel::Int begin + RECLAIM_SYNC = 1 + RECLAIM_DEVICE_SYNC = 2 + RECLAIM_GC_MINOR = 3 + RECLAIM_GC_FULL = 4 + RECLAIM_POOL_TRIM = 5 + RECLAIM_PURGE_CACHES = 6 + RECLAIM_DROP_STATE = 7 +end -# Hooks that run *before* `reclaim_hooks`, followed by a forced GC. Intended -# for library bindings to drop their task-local state so that finalizers can -# run (freeing e.g. cached workspace buffers) before the HandleCache and -# pool-trim stages release their own resources. -const pre_reclaim_hooks = Any[] """ retry_reclaim(retry_if) do # code that may fail due to insufficient GPU memory end -Run a block of code repeatedly until it successfully allocates the memory it needs. -Retries are only attempted when calling `retry_if` with the current return value is true. -At each try, more and more memory is freed from the CUDA memory pool. When that is not -possible anymore, the latest returned value will be returned. +Run a block of code repeatedly while `retry_if(ret)` holds for its return +value, escalating one `ReclaimLevel` between attempts. Returns the final +(or most recent) return value of the block. -This function is intended for use with CUDA APIs, which sometimes allocate (outside of the -CUDA memory pool) and return a specific error code when failing to. It is similar to -`Base.retry`, but deals with return values instead of exceptions for performance reasons. +This is intended for CUDA APIs that allocate outside the pool and report +failure via a status code. It's like `Base.retry`, but works on return +values instead of exceptions for performance reasons. """ @inline function retry_reclaim(f, retry_if) - ret = f() - if retry_if(ret) + ret = f() + retry_if(ret) || return ret return retry_reclaim_slow(f, retry_if, ret) - else +end + +@noinline function retry_reclaim_slow(f, retry_if, ret) + state = active_state() + sync = stream_ordered(state.device) + for level in instances(ReclaimLevel) + reclaim_step(level, state, sync) + ret = f() + retry_if(ret) || return ret + end return ret - end end -## slow path, incrementally reclaiming more memory until we succeed -@noinline function retry_reclaim_slow(f, retry_if, orig_ret) - state = active_state() - is_stream_ordered = stream_ordered(state.device) - phase = 1 - while true - if is_stream_ordered - if phase == 1 - synchronize(state.stream) - elseif phase == 2 - device_synchronize() - elseif phase == 3 +# single step of the reclaim ladder; no-op on allocator modes where it +# doesn't apply. Split out so `retry_reclaim` can escalate one rung at a +# time while `reclaim()` can run the whole ladder cumulatively. +function reclaim_step(level::ReclaimLevel, state, sync::Bool) + if level == RECLAIM_SYNC + sync && synchronize(state.stream) + elseif level == RECLAIM_DEVICE_SYNC + sync && device_synchronize() + elseif level == RECLAIM_GC_MINOR GC.gc(false) - device_synchronize() - elseif phase == 4 + sync && device_synchronize() + elseif level == RECLAIM_GC_FULL GC.gc(true) - device_synchronize() - elseif phase == 5 - # in case we had a release threshold configured - trim(pool_create(state.device)) - elseif phase == 6 - for hook in reclaim_hooks - hook() - end - elseif phase == 7 - # drop library-held TLS state (fat handles, cached workspaces) and - # finalize so the underlying buffers actually get freed - for hook in pre_reclaim_hooks - hook() - end + sync && device_synchronize() + elseif level == RECLAIM_POOL_TRIM + sync && trim(pool_create(state.device)) + elseif level == RECLAIM_PURGE_CACHES + foreach_reclaimable(purge!) + elseif level == RECLAIM_DROP_STATE + foreach_reclaimable(drop!) GC.gc(true) - for hook in reclaim_hooks - hook() - end - trim(pool_create(state.device)) - else - break - end - else - if phase == 1 - GC.gc(false) - elseif phase == 2 - GC.gc(true) - elseif phase == 3 - for hook in reclaim_hooks - hook() - end - elseif phase == 4 - for hook in pre_reclaim_hooks - hook() - end - GC.gc(true) - for hook in reclaim_hooks - hook() - end - else - break - end - end - phase += 1 - - ret = f() - if !retry_if(ret) - return ret + foreach_reclaimable(purge!) + sync && trim(pool_create(state.device)) end - end - - return orig_ret + return end @@ -814,34 +806,24 @@ end end """ - reclaim([sz=typemax(Int)]) + reclaim([level::ReclaimLevel = RECLAIM_DROP_STATE]) + +Free GPU memory by walking the reclaim ladder up to `level`. Use this before +calling into functionality that does not use the CUDA memory pool. Returns +`nothing`. -Reclaims `sz` bytes of cached memory. Use this to free GPU memory before calling into -functionality that does not use the CUDA memory pool. Returns the number of bytes -actually reclaimed. +The default drops task-local library state, runs a full GC so handle +wrappers finalize and return their raw handles to caches, then destroys +those caches and trims the pool. """ -function reclaim(sz::Int=typemax(Int)) - dev = device() - # Drop library-held state first (TLS-pinned fat handles, workspace caches, - # etc.), then run a full GC so finalizers release the underlying buffers - # and return idle handles to their caches. Only then do we tear down the - # caches themselves and trim the pool. - for hook in pre_reclaim_hooks - hook() - end - if !isempty(pre_reclaim_hooks) - GC.gc(true) - end - for hook in reclaim_hooks - hook() - end - if stream_ordered(dev) - device_synchronize() - synchronize(context()) - trim(pool_create(dev)) - else - 0 - end +function reclaim(level::ReclaimLevel = RECLAIM_DROP_STATE) + state = active_state() + sync = stream_ordered(state.device) + for l in instances(ReclaimLevel) + l <= level || break + reclaim_step(l, state, sync) + end + return end @@ -944,7 +926,9 @@ macro timed(ex) gpu_bytes=gpu_mem_stats.alloc_bytes, gpu_memtime=gpu_mem_stats.total_time, gpu_memstats=gpu_mem_stats) end end -@public @allocated, @time, @timed, used_memory, cached_memory, pool_status, reclaim +@public @allocated, @time, @timed, used_memory, cached_memory, pool_status, reclaim, + ReclaimLevel, RECLAIM_SYNC, RECLAIM_DEVICE_SYNC, RECLAIM_GC_MINOR, RECLAIM_GC_FULL, + RECLAIM_POOL_TRIM, RECLAIM_PURGE_CACHES, RECLAIM_DROP_STATE """ used_memory() diff --git a/CUDACore/src/utils/cache.jl b/CUDACore/src/utils/cache.jl index f87a1607db..bd1eef2c81 100644 --- a/CUDACore/src/utils/cache.jl +++ b/CUDACore/src/utils/cache.jl @@ -2,7 +2,7 @@ export HandleCache -struct HandleCache{K,V} +struct HandleCache{K,V} <: Reclaimable ctor dtor @@ -15,16 +15,15 @@ struct HandleCache{K,V} max_entries::Int function HandleCache{K,V}(ctor, dtor; max_entries::Int=32) where {K,V} - obj = new{K,V}(ctor, dtor, Set{Pair{K,V}}(), Dict{K,Vector{V}}(), - Base.ThreadSynchronizer(), max_entries) - - # register a hook to wipe the current context's cache when under memory pressure - push!(reclaim_hooks, ()->empty!(obj)) - - return obj + return new{K,V}(ctor, dtor, Set{Pair{K,V}}(), Dict{K,Vector{V}}(), + Base.ThreadSynchronizer(), max_entries) end end +# destroying idle handles is the `purge!` step of reclaim; individual caches +# must be registered in the owning library's __init__ (see reclaim.jl). +purge!(cache::HandleCache) = empty!(cache) + # remove a handle from the cache, or create a new one function Base.pop!(cache::HandleCache{K,V}, key::K) where {K,V} # check the cache diff --git a/CUDACore/src/utils/reclaim.jl b/CUDACore/src/utils/reclaim.jl new file mode 100644 index 0000000000..ef369110df --- /dev/null +++ b/CUDACore/src/utils/reclaim.jl @@ -0,0 +1,119 @@ +# reclaim registry +# +# Libraries (cuBLAS, cuDNN, …) hold GPU resources that the memory subsystem +# needs to release under pressure. Two kinds of resources exist: +# +# - live state referenced by a running task (typically via +# `task_local_storage`), e.g. fat handles wrapping library handles plus +# workspace buffers. Released via `drop!` — the library clears its TLS +# references so the wrappers become GC-eligible; their object-bound +# finalizers then return the raw handles to a `HandleCache`; +# +# - idle resources cached for reuse, e.g. the `HandleCache` of previously- +# returned library handles. Released via `purge!` — the cache is +# emptied and each entry's destructor runs. +# +# `Reclaimable` unifies both. Instances are registered via +# `register_reclaimable!` — typically from a library's `__init__`, since +# mutations to this registry performed during the precompilation of a +# downstream package don't carry over to module load (see +# `register_reclaimable!` for details). + +abstract type Reclaimable end + +""" + CUDACore.drop!(r::Reclaimable) + +Release references to live state so the associated GC-managed wrappers +become collectible (and their finalizers can run). Default: do nothing. +""" +drop!(::Reclaimable) = nothing + +""" + CUDACore.purge!(r::Reclaimable) + +Destroy idle cached resources owned by `r`. Default: do nothing. +""" +purge!(::Reclaimable) = nothing + +const reclaimables = Reclaimable[] +const reclaimables_lock = ReentrantLock() + +""" + CUDACore.register_reclaimable!(r::Reclaimable) + +Register `r` so that `reclaim` invokes its `drop!` and `purge!` methods. +Idempotent: re-registering the same instance is a no-op. Returns `r`. + +Call this from a package's `__init__`: top-level constructor calls are +captured by Julia's precompilation cache, but mutations they perform on +dependencies (like pushing into CUDACore's registry) do not persist to +module load time. +""" +function register_reclaimable!(r::Reclaimable) + @lock reclaimables_lock begin + r in reclaimables || push!(reclaimables, r) + end + return r +end + +# invoke f(r) per Reclaimable with per-hook error isolation, modelled +# after Base's `atexit_hooks`: one bad hook mustn't prevent the others. +function foreach_reclaimable(f) + @lock reclaimables_lock begin + for r in reclaimables + try + f(r) + catch ex + @error "reclaim callback failed" type=typeof(r) exception=(ex, catch_backtrace()) + end + end + end +end + + +## task-local state helper + +""" + TaskLocalCache{K,V}(key::Symbol) + +Declarative marker for a library's per-task cache stored under `key` in +`task_local_storage()`. Must be `register_reclaimable!`'d from the owning +package's `__init__` so that `RECLAIM_DROP_STATE` clears the current +task's entry, letting the stored values (typically mutable handle wrappers) +be garbage-collected. Their finalizers then return the underlying +resources to a `HandleCache`. + +Usage: + + const state_cache = CUDACore.TaskLocalCache{CuContext, LibraryState}(:CUBLAS) + + function __init__() + ... + CUDACore.register_reclaimable!(state_cache) + end + + function handle() + states = CUDACore.task_dict(state_cache) + state = get!(() -> new_state(...), states, key) + ... + end + +Only the current task's storage is touched: Julia's `IdDict`-backed TLS +isn't safe for concurrent mutation across threads, so cross-task drops are +intentionally deferred to task GC (values become unreachable when the +task is collected). +""" +struct TaskLocalCache{K,V} <: Reclaimable + key::Symbol + TaskLocalCache{K,V}(key::Symbol) where {K,V} = new{K,V}(key) +end + +@inline function task_dict(s::TaskLocalCache{K,V}) where {K,V} + get!(() -> Dict{K,V}(), task_local_storage(), s.key)::Dict{K,V} +end + +function drop!(s::TaskLocalCache) + delete!(task_local_storage(), s.key) + return +end diff --git a/lib/cublas/src/cuBLAS.jl b/lib/cublas/src/cuBLAS.jl index a045cb6914..c6751023f9 100644 --- a/lib/cublas/src/cuBLAS.jl +++ b/lib/cublas/src/cuBLAS.jl @@ -96,14 +96,13 @@ function handle_finalizer(h::cublasHandle) push!(idle_handles, h.ctx, h.handle) end +const LibraryState = @NamedTuple{handle::cublasHandle, stream::CuStream, math_mode::CUDACore.MathMode} +const state_cache = CUDACore.TaskLocalCache{CuContext, LibraryState}(:CUBLAS) + function handle() cuda = CUDACore.active_state() - # every task maintains library state per device - LibraryState = @NamedTuple{handle::cublasHandle, stream::CuStream, math_mode::CUDACore.MathMode} - states = get!(task_local_storage(), :CUBLAS) do - Dict{CuContext,LibraryState}() - end::Dict{CuContext,LibraryState} + states = CUDACore.task_dict(state_cache) # get library state @noinline function new_state(cuda) @@ -180,14 +179,13 @@ end::Vector{CuDevice} ndevices() = length(devices()) +const XtLibraryState = @NamedTuple{handle::cublasXtHandle} +const xt_state_cache = CUDACore.TaskLocalCache{UInt, XtLibraryState}(:CUBLASxt) + function xt_handle() cuda = CUDACore.active_state() - # every task maintains library state per set of devices - LibraryState = @NamedTuple{handle::cublasXtHandle} - states = get!(task_local_storage(), :CUBLASxt) do - Dict{UInt,LibraryState}() - end::Dict{UInt,LibraryState} + states = CUDACore.task_dict(xt_state_cache) # for performance, don't use a tuple of contexts to index the TLS key = zero(UInt) @@ -301,18 +299,16 @@ function __init__() atexit(flush_log_messages) end - # drop task-local cuBLAS state on reclaim so handles can be destroyed - push!(CUDACore.pre_reclaim_hooks, drop_library_state!) + # wire up reclaim (precompile-captured constructors can't push into + # CUDACore's registry themselves) + CUDACore.register_reclaimable!(idle_handles) + CUDACore.register_reclaimable!(idle_xt_handles) + CUDACore.register_reclaimable!(state_cache) + CUDACore.register_reclaimable!(xt_state_cache) _initialized[] = true end -function drop_library_state!() - delete!(task_local_storage(), :CUBLAS) - delete!(task_local_storage(), :CUBLASxt) - return -end - include("precompile.jl") # deprecated binding for backwards compatibility diff --git a/lib/cudnn/src/cuDNN.jl b/lib/cudnn/src/cuDNN.jl index 93bf92f502..caa3d4405b 100644 --- a/lib/cudnn/src/cuDNN.jl +++ b/lib/cudnn/src/cuDNN.jl @@ -76,25 +76,37 @@ function handle_dtor(ctx, handle) end const idle_handles = HandleCache{CuContext,cudnnHandle_t}(handle_ctor, handle_dtor) +# mutable wrapper so the raw handle is released via an object-bound +# finalizer: when TLS state is cleared on reclaim (or the owning task is +# collected) and GC runs, the wrapper is collected and its finalizer +# returns the handle to the idle cache. +mutable struct cudnnHandle + const handle::cudnnHandle_t + const ctx::CuContext +end +Base.unsafe_convert(::Type{cudnnHandle_t}, h::cudnnHandle) = h.handle + +function handle_finalizer(h::cudnnHandle) + push!(idle_handles, h.ctx, h.handle) +end + +const LibraryState = @NamedTuple{handle::cudnnHandle, stream::CuStream} +const state_cache = CUDACore.TaskLocalCache{CuContext, LibraryState}(:cuDNN) + function handle() cuda = CUDACore.active_state() - # every task maintains library state per device - LibraryState = @NamedTuple{handle::cudnnHandle_t, stream::CuStream} - states = get!(task_local_storage(), :cuDNN) do - Dict{CuContext,LibraryState}() - end::Dict{CuContext,LibraryState} + states = CUDACore.task_dict(state_cache) # get library state @noinline function new_state(cuda) new_handle = pop!(idle_handles, cuda.context) - finalizer(current_task()) do task - push!(idle_handles, cuda.context, new_handle) - end + wrapped = cudnnHandle(new_handle, cuda.context) + finalizer(handle_finalizer, wrapped) cudnnSetStream(new_handle, cuda.stream) - (; handle=new_handle, cuda.stream) + (; handle=wrapped, cuda.stream) end state = get!(states, cuda.context) do new_state(cuda) @@ -174,6 +186,9 @@ function __init__() cudnnSetCallback(typemax(UInt32), C_NULL, callback) end + CUDACore.register_reclaimable!(idle_handles) + CUDACore.register_reclaimable!(state_cache) + _initialized[] = true end diff --git a/lib/curand/src/cuRAND.jl b/lib/curand/src/cuRAND.jl index 55a66002b9..0b4d88e39f 100644 --- a/lib/curand/src/cuRAND.jl +++ b/lib/curand/src/cuRAND.jl @@ -56,31 +56,38 @@ function handle_dtor(ctx, handle) end const idle_library_rngs = HandleCache{CuContext,LibraryRNG}(handle_ctor, handle_dtor) +# wrapper owning a LibraryRNG borrowed from `idle_library_rngs`. Held in +# task-local storage so that, on reclaim or task GC, the wrapper becomes +# unreachable and its finalizer returns the RNG to the idle cache. From +# there, `purge!` on the cache would free the underlying generator. +mutable struct BorrowedLibraryRNG + const rng::LibraryRNG + const ctx::CuContext +end + +function library_rng_finalizer(b::BorrowedLibraryRNG) + push!(idle_library_rngs, b.ctx, b.rng) +end + +const library_state_cache = CUDACore.TaskLocalCache{CuContext, BorrowedLibraryRNG}(:CURAND) + function library_rng() cuda = CUDACore.active_state() - # every task maintains library state per device - LibraryState = @NamedTuple{rng::LibraryRNG} - states = get!(task_local_storage(), :CURAND) do - Dict{CuContext,LibraryState}() - end::Dict{CuContext,LibraryState} + states = CUDACore.task_dict(library_state_cache) - # get library state @noinline function new_state(cuda) new_rng = pop!(idle_library_rngs, cuda.context) - finalizer(current_task()) do task - push!(idle_library_rngs, cuda.context, new_rng) - end - + wrapped = BorrowedLibraryRNG(new_rng, cuda.context) + finalizer(library_rng_finalizer, wrapped) Random.seed!(new_rng) - - (; rng=new_rng) + wrapped end - state = get!(states, cuda.context) do + borrowed = get!(states, cuda.context) do new_state(cuda) end - return state.rng + return borrowed.rng end @@ -103,6 +110,11 @@ function __init__() libcurand = CUDA_Runtime_jll.libcurand end + CUDACore.register_reclaimable!(idle_library_rngs) + CUDACore.register_reclaimable!(library_state_cache) + CUDACore.register_reclaimable!(native_state_cache) + CUDACore.register_reclaimable!(gpuarrays_state_cache) + _initialized[] = true end diff --git a/lib/curand/src/cuda_integration.jl b/lib/curand/src/cuda_integration.jl index cae4bd3e2e..56ffb8f282 100644 --- a/lib/curand/src/cuda_integration.jl +++ b/lib/curand/src/cuda_integration.jl @@ -4,36 +4,17 @@ using CUDACore: AnyCuArray, CuArray, CuContext, active_state using CUDACore: GPUArrays -## native RNG handle cache (kernel-based Philox2x32) +## native RNG (kernel-based Philox2x32) +# +# Holds no GPU memory, just two UInt32s, so no HandleCache is needed: each +# task constructs its own and TLS-drop / task-GC frees it. -function native_rng_ctor(ctx) - context!(ctx) do - NativeRNG() - end -end -function native_rng_dtor(ctx, rng) end -const idle_native_rngs = HandleCache{CuContext,NativeRNG}(native_rng_ctor, native_rng_dtor) +const native_state_cache = CUDACore.TaskLocalCache{CuContext, NativeRNG}(:cuRAND_NativeRNG) function native_rng() cuda = active_state() - - LibraryState = @NamedTuple{rng::NativeRNG} - states = get!(task_local_storage(), :cuRAND_NativeRNG) do - Dict{CuContext,LibraryState}() - end::Dict{CuContext,LibraryState} - - @noinline function new_state(cuda) - new_rng = pop!(idle_native_rngs, cuda.context) - finalizer(current_task()) do task - push!(idle_native_rngs, cuda.context, new_rng) - end - (; rng=new_rng) - end - state = get!(states, cuda.context) do - new_state(cuda) - end - - return state.rng + states = CUDACore.task_dict(native_state_cache) + get!(() -> NativeRNG(), states, cuda.context) end @@ -41,24 +22,16 @@ end ## batched-kernel RNG). Used by Random.rand!/randn!(::AnyCuArray) when no rng ## is supplied, and exposed via CUDA.RNG / CUDA.gpuarrays_rng(). +const gpuarrays_state_cache = CUDACore.TaskLocalCache{CuContext, GPUArrays.RNG{CuArray}}(:cuRAND_DefaultRNG) + function gpuarrays_rng() cuda = active_state() - - LibraryState = @NamedTuple{rng::GPUArrays.RNG{CuArray}} - states = get!(task_local_storage(), :cuRAND_DefaultRNG) do - Dict{CuContext,LibraryState}() - end::Dict{CuContext,LibraryState} - - @noinline function new_state(cuda) + states = CUDACore.task_dict(gpuarrays_state_cache) + get!(states, cuda.context) do new_rng = GPUArrays.RNG{CuArray}() Random.seed!(new_rng) - (; rng=new_rng) + new_rng end - state = get!(states, cuda.context) do - new_state(cuda) - end - - return state.rng end diff --git a/lib/cusolver/src/cuSOLVER.jl b/lib/cusolver/src/cuSOLVER.jl index b52edae786..7838ba0f54 100644 --- a/lib/cusolver/src/cuSOLVER.jl +++ b/lib/cusolver/src/cuSOLVER.jl @@ -92,14 +92,13 @@ function dense_handle_finalizer(dh::cusolverDnHandle) push!(idle_dense_handles, dh.ctx, dh.handle) end +const DenseLibraryState = @NamedTuple{handle::cusolverDnHandle, stream::CuStream} +const dense_state_cache = CUDACore.TaskLocalCache{CuContext, DenseLibraryState}(:CUSOLVER_dense) + function dense_handle() cuda = CUDACore.active_state() - # every task maintains library state per device - LibraryState = @NamedTuple{handle::cusolverDnHandle, stream::CuStream} - states = get!(task_local_storage(), :CUSOLVER_dense) do - Dict{CuContext,LibraryState}() - end::Dict{CuContext,LibraryState} + states = CUDACore.task_dict(dense_state_cache) # get library state @noinline function new_state(cuda) @@ -161,14 +160,13 @@ function sparse_handle_finalizer(sh::cusolverSpHandle) push!(idle_sparse_handles, sh.ctx, sh.handle) end +const SparseLibraryState = @NamedTuple{handle::cusolverSpHandle, stream::CuStream} +const sparse_state_cache = CUDACore.TaskLocalCache{CuContext, SparseLibraryState}(:CUSOLVER_sparse) + function sparse_handle() cuda = CUDACore.active_state() - # every task maintains library state per device - LibraryState = @NamedTuple{handle::cusolverSpHandle, stream::CuStream} - states = get!(task_local_storage(), :CUSOLVER_sparse) do - Dict{CuContext,LibraryState}() - end::Dict{CuContext,LibraryState} + states = CUDACore.task_dict(sparse_state_cache) # get or create handle @noinline function new_state(cuda) @@ -228,14 +226,13 @@ function mg_handle_finalizer(mh::cusolverMgHandle) end end +const MgLibraryState = @NamedTuple{handle::cusolverMgHandle} +const mg_state_cache = CUDACore.TaskLocalCache{UInt, MgLibraryState}(:CUSOLVERmg) + function mg_handle() cuda = CUDACore.active_state() - # every task maintains library state per set of devices - LibraryState = @NamedTuple{handle::cusolverMgHandle} - states = get!(task_local_storage(), :CUSOLVERmg) do - Dict{UInt,LibraryState}() - end::Dict{UInt,LibraryState} + states = CUDACore.task_dict(mg_state_cache) # derive a key from the active and selected devices key = hash(cuda.context) @@ -290,20 +287,15 @@ function __init__() end end - # Drop task-local cuSOLVER state on reclaim so the cached workspace - # CuVectors (sometimes hundreds of MiB) become GC-able. - push!(CUDACore.pre_reclaim_hooks, drop_library_state!) + CUDACore.register_reclaimable!(idle_dense_handles) + CUDACore.register_reclaimable!(idle_sparse_handles) + CUDACore.register_reclaimable!(dense_state_cache) + CUDACore.register_reclaimable!(sparse_state_cache) + CUDACore.register_reclaimable!(mg_state_cache) _initialized[] = true end -function drop_library_state!() - delete!(task_local_storage(), :CUSOLVER_dense) - delete!(task_local_storage(), :CUSOLVER_sparse) - delete!(task_local_storage(), :CUSOLVERmg) - return -end - include("precompile.jl") # deprecated binding for backwards compatibility diff --git a/lib/cusparse/src/cuSPARSE.jl b/lib/cusparse/src/cuSPARSE.jl index c7f45fa6ce..c198120220 100644 --- a/lib/cusparse/src/cuSPARSE.jl +++ b/lib/cusparse/src/cuSPARSE.jl @@ -86,14 +86,13 @@ function handle_finalizer(h::cusparseHandle) push!(idle_handles, h.ctx, h.handle) end +const LibraryState = @NamedTuple{handle::cusparseHandle, stream::CuStream} +const state_cache = CUDACore.TaskLocalCache{CuContext, LibraryState}(:CUSPARSE) + function handle() cuda = CUDACore.active_state() - # every task maintains library state per device - LibraryState = @NamedTuple{handle::cusparseHandle, stream::CuStream} - states = get!(task_local_storage(), :CUSPARSE) do - Dict{CuContext,LibraryState}() - end::Dict{CuContext,LibraryState} + states = CUDACore.task_dict(state_cache) # get library state @noinline function new_state(cuda) @@ -141,17 +140,12 @@ function __init__() libcusparse = CUDA_Runtime_jll.libcusparse end - # drop task-local cuSPARSE state on reclaim so handles can be destroyed - push!(CUDACore.pre_reclaim_hooks, drop_library_state!) + CUDACore.register_reclaimable!(idle_handles) + CUDACore.register_reclaimable!(state_cache) _initialized[] = true end -function drop_library_state!() - delete!(task_local_storage(), :CUSPARSE) - return -end - # KernelAbstractions integration import KernelAbstractions as KA KA.get_backend(::CuSparseVector) = CUDACore.CUDAKernels.CUDABackend() diff --git a/lib/custatevec/src/cuStateVec.jl b/lib/custatevec/src/cuStateVec.jl index 988dbd7a65..1e037956d4 100644 --- a/lib/custatevec/src/cuStateVec.jl +++ b/lib/custatevec/src/cuStateVec.jl @@ -46,34 +46,38 @@ function handle_dtor(ctx, handle) end const idle_handles = HandleCache{CuContext,custatevecHandle_t}(handle_ctor, handle_dtor) -# fat handle, includes a cache -struct cuStateVecHandle - handle::custatevecHandle_t - cache::CuVector{UInt8} +# fat handle: bundles the raw cuStateVec handle with a cache buffer and a +# context reference. Mutable so an object-bound finalizer can release both +# the buffer and the handle when the wrapper becomes unreachable (on +# reclaim or task GC). +mutable struct cuStateVecHandle + const handle::custatevecHandle_t + const ctx::CuContext + const cache::CuVector{UInt8} end Base.unsafe_convert(::Type{Ptr{custatevecContext}}, handle::cuStateVecHandle) = handle.handle +function handle_finalizer(h::cuStateVecHandle) + CUDACore.unsafe_free!(h.cache) + push!(idle_handles, h.ctx, h.handle) +end + +const LibraryState = @NamedTuple{handle::cuStateVecHandle, stream::CuStream} +const state_cache = CUDACore.TaskLocalCache{CuContext, LibraryState}(:CUQUANTUM) + function handle() cuda = CUDACore.active_state() - # every task maintains library state per device - LibraryState = @NamedTuple{handle::cuStateVecHandle, stream::CuStream} - states = get!(task_local_storage(), :CUQUANTUM) do - Dict{CuContext,LibraryState}() - end::Dict{CuContext,LibraryState} + states = CUDACore.task_dict(state_cache) # get library state @noinline function new_state(cuda) new_handle = pop!(idle_handles, cuda.context) cache = CuVector{UInt8}(undef, 0) - fat_handle = cuStateVecHandle(new_handle, cache) - - finalizer(current_task()) do task - CUDACore.unsafe_free!(cache) - push!(idle_handles, cuda.context, new_handle) - end + fat_handle = cuStateVecHandle(new_handle, cuda.context, cache) + finalizer(handle_finalizer, fat_handle) custatevecSetStream(new_handle, cuda.stream) @@ -146,6 +150,9 @@ function __init__() custatevecLoggerSetLevel(5) end + CUDACore.register_reclaimable!(idle_handles) + CUDACore.register_reclaimable!(state_cache) + _initialized[] = true end diff --git a/lib/cutensor/src/cuTENSOR.jl b/lib/cutensor/src/cuTENSOR.jl index a054951b4d..00a7f35e10 100644 --- a/lib/cutensor/src/cuTENSOR.jl +++ b/lib/cutensor/src/cuTENSOR.jl @@ -58,23 +58,35 @@ function handle_dtor(ctx, handle) end const idle_handles = HandleCache{CuContext,cutensorHandle_t}(handle_ctor, handle_dtor) +# mutable wrapper so the raw handle is released via an object-bound +# finalizer: when TLS state is cleared on reclaim (or the owning task is +# collected) and GC runs, the wrapper is collected and its finalizer +# returns the handle to the idle cache. +mutable struct cutensorHandle + const handle::cutensorHandle_t + const ctx::CuContext +end +Base.unsafe_convert(::Type{cutensorHandle_t}, h::cutensorHandle) = h.handle + +function handle_finalizer(h::cutensorHandle) + push!(idle_handles, h.ctx, h.handle) +end + +const LibraryState = @NamedTuple{handle::cutensorHandle} +const state_cache = CUDACore.TaskLocalCache{CuContext, LibraryState}(:cuTENSOR) + function handle() cuda = CUDACore.active_state() - # every task maintains library state per device - LibraryState = @NamedTuple{handle::cutensorHandle_t} - states = get!(task_local_storage(), :cuTENSOR) do - Dict{CuContext,LibraryState}() - end::Dict{CuContext,LibraryState} + states = CUDACore.task_dict(state_cache) # get library state @noinline function new_state(cuda) new_handle = pop!(idle_handles, cuda.context) - finalizer(current_task()) do task - push!(idle_handles, cuda.context, new_handle) - end + wrapped = cutensorHandle(new_handle, cuda.context) + finalizer(handle_finalizer, wrapped) - (; handle=new_handle) + (; handle=wrapped) end state = get!(states, cuda.context) do new_state(cuda) @@ -134,6 +146,9 @@ function __init__() cutensorLoggerSetLevel(5) end + CUDACore.register_reclaimable!(idle_handles) + CUDACore.register_reclaimable!(state_cache) + _initialized[] = true end diff --git a/lib/cutensornet/src/cuTensorNet.jl b/lib/cutensornet/src/cuTensorNet.jl index 2317e249e9..2746a90d95 100644 --- a/lib/cutensornet/src/cuTensorNet.jl +++ b/lib/cutensornet/src/cuTensorNet.jl @@ -50,23 +50,35 @@ function handle_dtor(ctx, handle) end const idle_handles = HandleCache{CuContext,cutensornetHandle_t}(handle_ctor, handle_dtor) +# mutable wrapper so the raw handle is released via an object-bound +# finalizer: when TLS state is cleared on reclaim (or the owning task is +# collected) and GC runs, the wrapper is collected and its finalizer +# returns the handle to the idle cache. +mutable struct cutensornetHandle + const handle::cutensornetHandle_t + const ctx::CuContext +end +Base.unsafe_convert(::Type{cutensornetHandle_t}, h::cutensornetHandle) = h.handle + +function handle_finalizer(h::cutensornetHandle) + push!(idle_handles, h.ctx, h.handle) +end + +const LibraryState = @NamedTuple{handle::cutensornetHandle} +const state_cache = CUDACore.TaskLocalCache{CuContext, LibraryState}(:cuTensorNet) + function handle() cuda = CUDACore.active_state() - # every task maintains library state per device - LibraryState = @NamedTuple{handle::cutensornetHandle_t} - states = get!(task_local_storage(), :cuTensorNet) do - Dict{CuContext,LibraryState}() - end::Dict{CuContext,LibraryState} + states = CUDACore.task_dict(state_cache) # get library state @noinline function new_state(cuda) new_handle = pop!(idle_handles, cuda.context) - finalizer(current_task()) do task - push!(idle_handles, cuda.context, new_handle) - end + wrapped = cutensornetHandle(new_handle, cuda.context) + finalizer(handle_finalizer, wrapped) - (; handle=new_handle) + (; handle=wrapped) end state = get!(states, cuda.context) do new_state(cuda) @@ -126,6 +138,9 @@ function __init__() cutensornetLoggerSetLevel(5) end + CUDACore.register_reclaimable!(idle_handles) + CUDACore.register_reclaimable!(state_cache) + _initialized[] = true end diff --git a/test/core/pool.jl b/test/core/pool.jl index 3a556af7dc..83df4592c5 100644 --- a/test/core/pool.jl +++ b/test/core/pool.jl @@ -31,8 +31,8 @@ end end @testset "reclaim" begin - CUDA.reclaim(1024) CUDA.reclaim() + CUDA.reclaim(CUDACore.RECLAIM_GC_FULL) @test CUDACore.retry_reclaim(isequal(42)) do 42 From bb1b69aae333b9c8d6d6aa7c0bffa5e40c98fb61 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sun, 19 Apr 2026 10:00:44 +0200 Subject: [PATCH 7/9] Add RECLAIM_PURGE_IDLE fast-path rung to the reclaim ladder. Inspired by PyTorch's `FreeMemoryCallback` (see c10/cuda/CUDACachingAllocator.cpp), which invokes registered callbacks before any expensive allocation work. `RECLAIM_PURGE_IDLE` sits at the bottom of the ladder and runs `purge!` on every `Reclaimable` without touching sync, GC, or pool trim. If a library happens to be holding idle cached handles from dead tasks, we get them back immediately; if not, it's a cheap no-op and we escalate. The existing `RECLAIM_PURGE_CACHES` rung stays in place at its former position: after `RECLAIM_GC_{MINOR,FULL}` have run, handle wrappers from dead tasks get finalized and their raw handles end up in the caches. That rung now catches those new entries before we'd have to escalate all the way to `RECLAIM_DROP_STATE` (which also clears the current task's TLS). Co-Authored-By: Claude Opus 4.7 (1M context) --- CUDACore/src/memory.jl | 23 +++++++++++++++++------ 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/CUDACore/src/memory.jl b/CUDACore/src/memory.jl index 8919186c55..fc25b59c51 100644 --- a/CUDACore/src/memory.jl +++ b/CUDACore/src/memory.jl @@ -470,18 +470,28 @@ from cheapest to most aggressive: | Level | Action | | :--- | :--- | +| `RECLAIM_PURGE_IDLE` | empty `HandleCache`s (fast path, no sync/GC) | | `RECLAIM_SYNC` | synchronize the current task's stream | | `RECLAIM_DEVICE_SYNC` | synchronize the whole device | | `RECLAIM_GC_MINOR` | minor Julia GC (and device sync) | | `RECLAIM_GC_FULL` | full Julia GC (and device sync) | | `RECLAIM_POOL_TRIM` | trim unused memory from the pool | -| `RECLAIM_PURGE_CACHES` | empty `HandleCache`s (destroy idle handles) | +| `RECLAIM_PURGE_CACHES` | empty `HandleCache`s again (catches GC-populated entries)| | `RECLAIM_DROP_STATE` | also drop task-local library state, then GC+purge+trim | +`RECLAIM_PURGE_IDLE` and `RECLAIM_PURGE_CACHES` run the same action +(`purge!` on every `Reclaimable`). The first is an opportunistic fast +path before any sync/GC cost: if a library is sitting on cached idle +handles, we can release them immediately. The second runs after GC has +had a chance to populate caches via finalizers, catching those new +entries without escalating to `RECLAIM_DROP_STATE` (which would also +drop the current task's live state). + Steps that don't apply to the current allocator (e.g. stream sync on a non-stream-ordered device) are silently skipped. """ @enum ReclaimLevel::Int begin + RECLAIM_PURGE_IDLE = 0 RECLAIM_SYNC = 1 RECLAIM_DEVICE_SYNC = 2 RECLAIM_GC_MINOR = 3 @@ -526,7 +536,9 @@ end # doesn't apply. Split out so `retry_reclaim` can escalate one rung at a # time while `reclaim()` can run the whole ladder cumulatively. function reclaim_step(level::ReclaimLevel, state, sync::Bool) - if level == RECLAIM_SYNC + if level == RECLAIM_PURGE_IDLE || level == RECLAIM_PURGE_CACHES + foreach_reclaimable(purge!) + elseif level == RECLAIM_SYNC sync && synchronize(state.stream) elseif level == RECLAIM_DEVICE_SYNC sync && device_synchronize() @@ -538,8 +550,6 @@ function reclaim_step(level::ReclaimLevel, state, sync::Bool) sync && device_synchronize() elseif level == RECLAIM_POOL_TRIM sync && trim(pool_create(state.device)) - elseif level == RECLAIM_PURGE_CACHES - foreach_reclaimable(purge!) elseif level == RECLAIM_DROP_STATE foreach_reclaimable(drop!) GC.gc(true) @@ -927,8 +937,9 @@ macro timed(ex) end end @public @allocated, @time, @timed, used_memory, cached_memory, pool_status, reclaim, - ReclaimLevel, RECLAIM_SYNC, RECLAIM_DEVICE_SYNC, RECLAIM_GC_MINOR, RECLAIM_GC_FULL, - RECLAIM_POOL_TRIM, RECLAIM_PURGE_CACHES, RECLAIM_DROP_STATE + ReclaimLevel, RECLAIM_PURGE_IDLE, RECLAIM_SYNC, RECLAIM_DEVICE_SYNC, + RECLAIM_GC_MINOR, RECLAIM_GC_FULL, RECLAIM_POOL_TRIM, RECLAIM_PURGE_CACHES, + RECLAIM_DROP_STATE """ used_memory() From f488c76b8ddda0699feeb876f90c689e026f81cf Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 23 Apr 2026 22:19:04 +0200 Subject: [PATCH 8/9] Rename handle types to avoid name collisions. --- lib/cublas/src/cuBLAS.jl | 22 ++++++++++---------- lib/cudnn/src/cuDNN.jl | 10 +++++----- lib/cusolver/src/cuSOLVER.jl | 32 +++++++++++++++--------------- lib/cusparse/src/cuSPARSE.jl | 10 +++++----- lib/custatevec/src/cuStateVec.jl | 10 +++++----- lib/cutensor/src/cuTENSOR.jl | 10 +++++----- lib/cutensornet/src/cuTensorNet.jl | 10 +++++----- 7 files changed, 52 insertions(+), 52 deletions(-) diff --git a/lib/cublas/src/cuBLAS.jl b/lib/cublas/src/cuBLAS.jl index c6751023f9..ab4c82e2d9 100644 --- a/lib/cublas/src/cuBLAS.jl +++ b/lib/cublas/src/cuBLAS.jl @@ -86,17 +86,17 @@ const idle_handles = HandleCache{CuContext,cublasHandle_t}(handle_ctor, handle_d # when TLS state is cleared (e.g. on reclaim) and GC runs, the wrapper is # collected and its finalizer returns the handle to the idle cache instead # of the handle being pinned for the entire lifetime of the owning task. -mutable struct cublasHandle +mutable struct Handle const handle::cublasHandle_t const ctx::CuContext end -Base.unsafe_convert(::Type{cublasHandle_t}, handle::cublasHandle) = handle.handle +Base.unsafe_convert(::Type{cublasHandle_t}, handle::Handle) = handle.handle -function handle_finalizer(h::cublasHandle) +function handle_finalizer(h::Handle) push!(idle_handles, h.ctx, h.handle) end -const LibraryState = @NamedTuple{handle::cublasHandle, stream::CuStream, math_mode::CUDACore.MathMode} +const LibraryState = @NamedTuple{handle::Handle, stream::CuStream, math_mode::CUDACore.MathMode} const state_cache = CUDACore.TaskLocalCache{CuContext, LibraryState}(:CUBLAS) function handle() @@ -107,7 +107,7 @@ function handle() # get library state @noinline function new_state(cuda) new_handle = pop!(idle_handles, cuda.context) - wrapped = cublasHandle(new_handle, cuda.context) + wrapped = Handle(new_handle, cuda.context) finalizer(handle_finalizer, wrapped) cublasSetStream_v2(new_handle, cuda.stream) @@ -156,14 +156,14 @@ end const idle_xt_handles = HandleCache{Vector{CuContext},cublasXtHandle_t}(xt_handle_ctor, xt_handle_dtor) -# mutable wrapper for the xt handle, see `cublasHandle` for rationale. -mutable struct cublasXtHandle +# mutable wrapper for the xt handle, see `Handle` for rationale. +mutable struct XtHandle const handle::cublasXtHandle_t const ctxs::Vector{CuContext} end -Base.unsafe_convert(::Type{cublasXtHandle_t}, h::cublasXtHandle) = h.handle +Base.unsafe_convert(::Type{cublasXtHandle_t}, h::XtHandle) = h.handle -function xt_handle_finalizer(h::cublasXtHandle) +function xt_handle_finalizer(h::XtHandle) push!(idle_xt_handles, h.ctxs, h.handle) end @@ -179,7 +179,7 @@ end::Vector{CuDevice} ndevices() = length(devices()) -const XtLibraryState = @NamedTuple{handle::cublasXtHandle} +const XtLibraryState = @NamedTuple{handle::XtHandle} const xt_state_cache = CUDACore.TaskLocalCache{UInt, XtLibraryState}(:CUBLASxt) function xt_handle() @@ -199,7 +199,7 @@ function xt_handle() ctxs = [context(dev) for dev in devices()] new_handle = pop!(idle_xt_handles, ctxs) - wrapped = cublasXtHandle(new_handle, ctxs) + wrapped = XtHandle(new_handle, ctxs) finalizer(xt_handle_finalizer, wrapped) # if we're using the stream-ordered allocator, diff --git a/lib/cudnn/src/cuDNN.jl b/lib/cudnn/src/cuDNN.jl index caa3d4405b..0cb891501b 100644 --- a/lib/cudnn/src/cuDNN.jl +++ b/lib/cudnn/src/cuDNN.jl @@ -80,17 +80,17 @@ const idle_handles = HandleCache{CuContext,cudnnHandle_t}(handle_ctor, handle_dt # finalizer: when TLS state is cleared on reclaim (or the owning task is # collected) and GC runs, the wrapper is collected and its finalizer # returns the handle to the idle cache. -mutable struct cudnnHandle +mutable struct Handle const handle::cudnnHandle_t const ctx::CuContext end -Base.unsafe_convert(::Type{cudnnHandle_t}, h::cudnnHandle) = h.handle +Base.unsafe_convert(::Type{cudnnHandle_t}, h::Handle) = h.handle -function handle_finalizer(h::cudnnHandle) +function handle_finalizer(h::Handle) push!(idle_handles, h.ctx, h.handle) end -const LibraryState = @NamedTuple{handle::cudnnHandle, stream::CuStream} +const LibraryState = @NamedTuple{handle::Handle, stream::CuStream} const state_cache = CUDACore.TaskLocalCache{CuContext, LibraryState}(:cuDNN) function handle() @@ -101,7 +101,7 @@ function handle() # get library state @noinline function new_state(cuda) new_handle = pop!(idle_handles, cuda.context) - wrapped = cudnnHandle(new_handle, cuda.context) + wrapped = Handle(new_handle, cuda.context) finalizer(handle_finalizer, wrapped) cudnnSetStream(new_handle, cuda.stream) diff --git a/lib/cusolver/src/cuSOLVER.jl b/lib/cusolver/src/cuSOLVER.jl index 7838ba0f54..8bbbcfa5e6 100644 --- a/lib/cusolver/src/cuSOLVER.jl +++ b/lib/cusolver/src/cuSOLVER.jl @@ -76,23 +76,23 @@ const idle_dense_handles = # itself: once no one references this struct (e.g. after the owning task dies # or we clear it from task-local storage on reclaim), GC runs the finalizer # which returns the buffers to the pool and the raw handle to the idle cache. -mutable struct cusolverDnHandle +mutable struct DnHandle const handle::cusolverDnHandle_t const ctx::CuContext const workspace_gpu::CuVector{UInt8} const workspace_cpu::Vector{UInt8} const info::CuVector{Cint} end -Base.unsafe_convert(::Type{Ptr{cusolverDnContext}}, handle::cusolverDnHandle) = +Base.unsafe_convert(::Type{Ptr{cusolverDnContext}}, handle::DnHandle) = handle.handle -function dense_handle_finalizer(dh::cusolverDnHandle) +function dense_handle_finalizer(dh::DnHandle) CUDACore.unsafe_free!(dh.workspace_gpu) CUDACore.unsafe_free!(dh.info) push!(idle_dense_handles, dh.ctx, dh.handle) end -const DenseLibraryState = @NamedTuple{handle::cusolverDnHandle, stream::CuStream} +const DenseLibraryState = @NamedTuple{handle::DnHandle, stream::CuStream} const dense_state_cache = CUDACore.TaskLocalCache{CuContext, DenseLibraryState}(:CUSOLVER_dense) function dense_handle() @@ -107,7 +107,7 @@ function dense_handle() workspace_gpu = CuVector{UInt8}(undef, 0) workspace_cpu = Vector{UInt8}(undef, 0) info = CuVector{Cint}(undef, 1) - fat_handle = cusolverDnHandle(new_handle, cuda.context, workspace_gpu, + fat_handle = DnHandle(new_handle, cuda.context, workspace_gpu, workspace_cpu, info) finalizer(dense_handle_finalizer, fat_handle) @@ -148,19 +148,19 @@ const idle_sparse_handles = HandleCache{CuContext,cusolverSpHandle_t}(sparse_handle_ctor, sparse_handle_dtor) # mutable wrapper so the raw sparse handle is released via an object-bound -# finalizer (see `cusolverDnHandle` for rationale). -mutable struct cusolverSpHandle +# finalizer (see `DnHandle` for rationale). +mutable struct SpHandle const handle::cusolverSpHandle_t const ctx::CuContext end -Base.unsafe_convert(::Type{cusolverSpHandle_t}, handle::cusolverSpHandle) = +Base.unsafe_convert(::Type{cusolverSpHandle_t}, handle::SpHandle) = handle.handle -function sparse_handle_finalizer(sh::cusolverSpHandle) +function sparse_handle_finalizer(sh::SpHandle) push!(idle_sparse_handles, sh.ctx, sh.handle) end -const SparseLibraryState = @NamedTuple{handle::cusolverSpHandle, stream::CuStream} +const SparseLibraryState = @NamedTuple{handle::SpHandle, stream::CuStream} const sparse_state_cache = CUDACore.TaskLocalCache{CuContext, SparseLibraryState}(:CUSOLVER_sparse) function sparse_handle() @@ -171,7 +171,7 @@ function sparse_handle() # get or create handle @noinline function new_state(cuda) new_handle = pop!(idle_sparse_handles, cuda.context) - wrapped = cusolverSpHandle(new_handle, cuda.context) + wrapped = SpHandle(new_handle, cuda.context) finalizer(sparse_handle_finalizer, wrapped) cusolverSpSetStream(new_handle, cuda.stream) @@ -213,20 +213,20 @@ ndevices() = length(devices()) # mutable wrapper so the mg handle is destroyed when its owning state struct # becomes unreachable, rather than being pinned for the lifetime of the task. -mutable struct cusolverMgHandle +mutable struct MgHandle const handle::cusolverMgHandle_t const ctx::CuContext end -Base.unsafe_convert(::Type{cusolverMgHandle_t}, handle::cusolverMgHandle) = +Base.unsafe_convert(::Type{cusolverMgHandle_t}, handle::MgHandle) = handle.handle -function mg_handle_finalizer(mh::cusolverMgHandle) +function mg_handle_finalizer(mh::MgHandle) context!(mh.ctx; skip_destroyed=true) do cusolverMgDestroy(mh.handle) end end -const MgLibraryState = @NamedTuple{handle::cusolverMgHandle} +const MgLibraryState = @NamedTuple{handle::MgHandle} const mg_state_cache = CUDACore.TaskLocalCache{UInt, MgLibraryState}(:CUSOLVERmg) function mg_handle() @@ -245,7 +245,7 @@ function mg_handle() @noinline function new_state(cuda) # we can't reuse cusolverMg handles because they can only be assigned devices once new_handle = cusolverMgCreate() - wrapped = cusolverMgHandle(new_handle, cuda.context) + wrapped = MgHandle(new_handle, cuda.context) finalizer(mg_handle_finalizer, wrapped) devs = convert.(Cint, devices()) diff --git a/lib/cusparse/src/cuSPARSE.jl b/lib/cusparse/src/cuSPARSE.jl index c198120220..296f95248d 100644 --- a/lib/cusparse/src/cuSPARSE.jl +++ b/lib/cusparse/src/cuSPARSE.jl @@ -76,17 +76,17 @@ const idle_handles = HandleCache{CuContext,cusparseHandle_t}(handle_ctor, handle # when TLS state is cleared (e.g. on reclaim) and GC runs, the wrapper is # collected and its finalizer returns the handle to the idle cache instead # of the handle being pinned for the entire lifetime of the owning task. -mutable struct cusparseHandle +mutable struct Handle const handle::cusparseHandle_t const ctx::CuContext end -Base.unsafe_convert(::Type{cusparseHandle_t}, h::cusparseHandle) = h.handle +Base.unsafe_convert(::Type{cusparseHandle_t}, h::Handle) = h.handle -function handle_finalizer(h::cusparseHandle) +function handle_finalizer(h::Handle) push!(idle_handles, h.ctx, h.handle) end -const LibraryState = @NamedTuple{handle::cusparseHandle, stream::CuStream} +const LibraryState = @NamedTuple{handle::Handle, stream::CuStream} const state_cache = CUDACore.TaskLocalCache{CuContext, LibraryState}(:CUSPARSE) function handle() @@ -97,7 +97,7 @@ function handle() # get library state @noinline function new_state(cuda) new_handle = pop!(idle_handles, cuda.context) - wrapped = cusparseHandle(new_handle, cuda.context) + wrapped = Handle(new_handle, cuda.context) finalizer(handle_finalizer, wrapped) cusparseSetStream(new_handle, cuda.stream) diff --git a/lib/custatevec/src/cuStateVec.jl b/lib/custatevec/src/cuStateVec.jl index 1e037956d4..590315b8f5 100644 --- a/lib/custatevec/src/cuStateVec.jl +++ b/lib/custatevec/src/cuStateVec.jl @@ -50,20 +50,20 @@ const idle_handles = HandleCache{CuContext,custatevecHandle_t}(handle_ctor, hand # context reference. Mutable so an object-bound finalizer can release both # the buffer and the handle when the wrapper becomes unreachable (on # reclaim or task GC). -mutable struct cuStateVecHandle +mutable struct Handle const handle::custatevecHandle_t const ctx::CuContext const cache::CuVector{UInt8} end -Base.unsafe_convert(::Type{Ptr{custatevecContext}}, handle::cuStateVecHandle) = +Base.unsafe_convert(::Type{Ptr{custatevecContext}}, handle::Handle) = handle.handle -function handle_finalizer(h::cuStateVecHandle) +function handle_finalizer(h::Handle) CUDACore.unsafe_free!(h.cache) push!(idle_handles, h.ctx, h.handle) end -const LibraryState = @NamedTuple{handle::cuStateVecHandle, stream::CuStream} +const LibraryState = @NamedTuple{handle::Handle, stream::CuStream} const state_cache = CUDACore.TaskLocalCache{CuContext, LibraryState}(:CUQUANTUM) function handle() @@ -76,7 +76,7 @@ function handle() new_handle = pop!(idle_handles, cuda.context) cache = CuVector{UInt8}(undef, 0) - fat_handle = cuStateVecHandle(new_handle, cuda.context, cache) + fat_handle = Handle(new_handle, cuda.context, cache) finalizer(handle_finalizer, fat_handle) custatevecSetStream(new_handle, cuda.stream) diff --git a/lib/cutensor/src/cuTENSOR.jl b/lib/cutensor/src/cuTENSOR.jl index 00a7f35e10..dbdc9a635e 100644 --- a/lib/cutensor/src/cuTENSOR.jl +++ b/lib/cutensor/src/cuTENSOR.jl @@ -62,17 +62,17 @@ const idle_handles = HandleCache{CuContext,cutensorHandle_t}(handle_ctor, handle # finalizer: when TLS state is cleared on reclaim (or the owning task is # collected) and GC runs, the wrapper is collected and its finalizer # returns the handle to the idle cache. -mutable struct cutensorHandle +mutable struct Handle const handle::cutensorHandle_t const ctx::CuContext end -Base.unsafe_convert(::Type{cutensorHandle_t}, h::cutensorHandle) = h.handle +Base.unsafe_convert(::Type{cutensorHandle_t}, h::Handle) = h.handle -function handle_finalizer(h::cutensorHandle) +function handle_finalizer(h::Handle) push!(idle_handles, h.ctx, h.handle) end -const LibraryState = @NamedTuple{handle::cutensorHandle} +const LibraryState = @NamedTuple{handle::Handle} const state_cache = CUDACore.TaskLocalCache{CuContext, LibraryState}(:cuTENSOR) function handle() @@ -83,7 +83,7 @@ function handle() # get library state @noinline function new_state(cuda) new_handle = pop!(idle_handles, cuda.context) - wrapped = cutensorHandle(new_handle, cuda.context) + wrapped = Handle(new_handle, cuda.context) finalizer(handle_finalizer, wrapped) (; handle=wrapped) diff --git a/lib/cutensornet/src/cuTensorNet.jl b/lib/cutensornet/src/cuTensorNet.jl index 2746a90d95..4adabf17b1 100644 --- a/lib/cutensornet/src/cuTensorNet.jl +++ b/lib/cutensornet/src/cuTensorNet.jl @@ -54,17 +54,17 @@ const idle_handles = HandleCache{CuContext,cutensornetHandle_t}(handle_ctor, han # finalizer: when TLS state is cleared on reclaim (or the owning task is # collected) and GC runs, the wrapper is collected and its finalizer # returns the handle to the idle cache. -mutable struct cutensornetHandle +mutable struct Handle const handle::cutensornetHandle_t const ctx::CuContext end -Base.unsafe_convert(::Type{cutensornetHandle_t}, h::cutensornetHandle) = h.handle +Base.unsafe_convert(::Type{cutensornetHandle_t}, h::Handle) = h.handle -function handle_finalizer(h::cutensornetHandle) +function handle_finalizer(h::Handle) push!(idle_handles, h.ctx, h.handle) end -const LibraryState = @NamedTuple{handle::cutensornetHandle} +const LibraryState = @NamedTuple{handle::Handle} const state_cache = CUDACore.TaskLocalCache{CuContext, LibraryState}(:cuTensorNet) function handle() @@ -75,7 +75,7 @@ function handle() # get library state @noinline function new_state(cuda) new_handle = pop!(idle_handles, cuda.context) - wrapped = cutensornetHandle(new_handle, cuda.context) + wrapped = Handle(new_handle, cuda.context) finalizer(handle_finalizer, wrapped) (; handle=wrapped) From ce6fefd1faf7364d0639b4186a5307a21683eabe Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 14 May 2026 10:16:25 +0200 Subject: [PATCH 9/9] Simplify reclaim ladder and address review feedback. - Collapse 8 ReclaimLevel rungs into 4 (PURGE/SYNC/GC/DROP); each level is a complete reclaim at that aggressiveness, so reclaim(level) runs the matching step directly instead of walking the whole ladder. - foreach_reclaimable now snapshots the registry under the lock and runs callbacks outside it (Base.atexit-style), so a slow or re-entrant callback can't pin other threads. - Drop the closure allocation in task_dict. - Add a note in the ReclaimLevel docstring on why DROP is safe (descriptors/plans are context-bound, user-held wrappers protect their raw handles) and that callbacks must not switch the active device. - Guard the test pool-cap setup with stream_ordered(dev). - Fix the stale "4 GiB" worker-budget comment. - Expand the reclaim testset to cover every level and verify retry_reclaim's escalation behavior. Co-Authored-By: Claude Opus 4.7 (1M context) --- CUDACore/src/memory.jl | 123 ++++++++++++++-------------------- CUDACore/src/utils/reclaim.jl | 51 +++++++------- test/core/pool.jl | 16 ++++- test/runtests.jl | 3 +- test/setup.jl | 21 +++--- 5 files changed, 107 insertions(+), 107 deletions(-) diff --git a/CUDACore/src/memory.jl b/CUDACore/src/memory.jl index fc25b59c51..dd5af3560a 100644 --- a/CUDACore/src/memory.jl +++ b/CUDACore/src/memory.jl @@ -459,46 +459,39 @@ end ## reclaim escalation # # `Reclaimable`/`register_reclaimable!`/`TaskLocalCache`/`drop!`/`purge!` -# are defined in utils/reclaim.jl. Here we add the ladder that drives -# them along with the allocator-specific sync/trim steps. +# are defined in utils/reclaim.jl. Here we add the ladder that drives them +# along with the allocator-specific sync/trim steps. +# +# Registered `drop!`/`purge!` callbacks must not switch the active device: +# the device is captured once per `reclaim` / `retry_reclaim` call. """ ReclaimLevel -Escalation levels shared by `reclaim(level)` and `retry_reclaim`, ordered -from cheapest to most aggressive: - -| Level | Action | -| :--- | :--- | -| `RECLAIM_PURGE_IDLE` | empty `HandleCache`s (fast path, no sync/GC) | -| `RECLAIM_SYNC` | synchronize the current task's stream | -| `RECLAIM_DEVICE_SYNC` | synchronize the whole device | -| `RECLAIM_GC_MINOR` | minor Julia GC (and device sync) | -| `RECLAIM_GC_FULL` | full Julia GC (and device sync) | -| `RECLAIM_POOL_TRIM` | trim unused memory from the pool | -| `RECLAIM_PURGE_CACHES` | empty `HandleCache`s again (catches GC-populated entries)| -| `RECLAIM_DROP_STATE` | also drop task-local library state, then GC+purge+trim | - -`RECLAIM_PURGE_IDLE` and `RECLAIM_PURGE_CACHES` run the same action -(`purge!` on every `Reclaimable`). The first is an opportunistic fast -path before any sync/GC cost: if a library is sitting on cached idle -handles, we can release them immediately. The second runs after GC has -had a chance to populate caches via finalizers, catching those new -entries without escalating to `RECLAIM_DROP_STATE` (which would also -drop the current task's live state). - -Steps that don't apply to the current allocator (e.g. stream sync on a +Escalation levels shared by `reclaim(level)` and `retry_reclaim`, from +cheapest to most aggressive: + +| Level | Action | +| :--- | :--- | +| `RECLAIM_PURGE` | empty `HandleCache`s (no GC, no sync) | +| `RECLAIM_SYNC` | synchronize the device (lets async deallocs finish) | +| `RECLAIM_GC` | run a full Julia GC, then sync + purge + trim | +| `RECLAIM_DROP` | also drop task-local library state, then GC + … | + +`RECLAIM_DROP` clears the calling task's library state — see +[`register_reclaimable!`](@ref). It assumes user-held descriptors/plans +are tied to a context, not to a specific library-handle instance (true +for all libraries CUDA.jl wraps). Live handles the user holds a +reference to are unaffected: the wrapper keeps the raw handle alive. + +Steps that don't apply to the current allocator (e.g. trim on a non-stream-ordered device) are silently skipped. """ @enum ReclaimLevel::Int begin - RECLAIM_PURGE_IDLE = 0 - RECLAIM_SYNC = 1 - RECLAIM_DEVICE_SYNC = 2 - RECLAIM_GC_MINOR = 3 - RECLAIM_GC_FULL = 4 - RECLAIM_POOL_TRIM = 5 - RECLAIM_PURGE_CACHES = 6 - RECLAIM_DROP_STATE = 7 + RECLAIM_PURGE = 0 + RECLAIM_SYNC = 1 + RECLAIM_GC = 2 + RECLAIM_DROP = 3 end @@ -522,39 +515,36 @@ values instead of exceptions for performance reasons. end @noinline function retry_reclaim_slow(f, retry_if, ret) - state = active_state() - sync = stream_ordered(state.device) + dev = active_state().device + so = stream_ordered(dev) for level in instances(ReclaimLevel) - reclaim_step(level, state, sync) + reclaim_step(level, dev, so) ret = f() retry_if(ret) || return ret end return ret end -# single step of the reclaim ladder; no-op on allocator modes where it -# doesn't apply. Split out so `retry_reclaim` can escalate one rung at a -# time while `reclaim()` can run the whole ladder cumulatively. -function reclaim_step(level::ReclaimLevel, state, sync::Bool) - if level == RECLAIM_PURGE_IDLE || level == RECLAIM_PURGE_CACHES +# Each level is a complete reclaim at that aggressiveness — `reclaim(level)` +# just runs the matching step. `retry_reclaim` walks the levels in order to +# bisect on alloc failure. GC.gc(true) drains pending finalizers before +# returning, so the post-GC purge sees caches populated by wrapper finalizers. +function reclaim_step(level::ReclaimLevel, dev::CuDevice, stream_ordered::Bool) + if level == RECLAIM_PURGE foreach_reclaimable(purge!) elseif level == RECLAIM_SYNC - sync && synchronize(state.stream) - elseif level == RECLAIM_DEVICE_SYNC - sync && device_synchronize() - elseif level == RECLAIM_GC_MINOR - GC.gc(false) - sync && device_synchronize() - elseif level == RECLAIM_GC_FULL + stream_ordered && device_synchronize() + elseif level == RECLAIM_GC GC.gc(true) - sync && device_synchronize() - elseif level == RECLAIM_POOL_TRIM - sync && trim(pool_create(state.device)) - elseif level == RECLAIM_DROP_STATE + stream_ordered && device_synchronize() + foreach_reclaimable(purge!) + stream_ordered && trim(pool_create(dev)) + elseif level == RECLAIM_DROP foreach_reclaimable(drop!) GC.gc(true) + stream_ordered && device_synchronize() foreach_reclaimable(purge!) - sync && trim(pool_create(state.device)) + stream_ordered && trim(pool_create(dev)) end return end @@ -816,23 +806,16 @@ end end """ - reclaim([level::ReclaimLevel = RECLAIM_DROP_STATE]) - -Free GPU memory by walking the reclaim ladder up to `level`. Use this before -calling into functionality that does not use the CUDA memory pool. Returns -`nothing`. + reclaim([level::ReclaimLevel = RECLAIM_DROP]) -The default drops task-local library state, runs a full GC so handle -wrappers finalize and return their raw handles to caches, then destroys -those caches and trims the pool. +Free GPU memory at the given [`ReclaimLevel`](@ref). The default drops +task-local library state, runs a full GC so handle wrappers finalize and +return their raw handles to caches, then destroys those caches and trims +the pool. Returns `nothing`. """ -function reclaim(level::ReclaimLevel = RECLAIM_DROP_STATE) - state = active_state() - sync = stream_ordered(state.device) - for l in instances(ReclaimLevel) - l <= level || break - reclaim_step(l, state, sync) - end +function reclaim(level::ReclaimLevel = RECLAIM_DROP) + dev = active_state().device + reclaim_step(level, dev, stream_ordered(dev)) return end @@ -937,9 +920,7 @@ macro timed(ex) end end @public @allocated, @time, @timed, used_memory, cached_memory, pool_status, reclaim, - ReclaimLevel, RECLAIM_PURGE_IDLE, RECLAIM_SYNC, RECLAIM_DEVICE_SYNC, - RECLAIM_GC_MINOR, RECLAIM_GC_FULL, RECLAIM_POOL_TRIM, RECLAIM_PURGE_CACHES, - RECLAIM_DROP_STATE + RECLAIM_PURGE, RECLAIM_SYNC, RECLAIM_GC, RECLAIM_DROP """ used_memory() diff --git a/CUDACore/src/utils/reclaim.jl b/CUDACore/src/utils/reclaim.jl index ef369110df..61eb013a11 100644 --- a/CUDACore/src/utils/reclaim.jl +++ b/CUDACore/src/utils/reclaim.jl @@ -13,11 +13,7 @@ # returned library handles. Released via `purge!` — the cache is # emptied and each entry's destructor runs. # -# `Reclaimable` unifies both. Instances are registered via -# `register_reclaimable!` — typically from a library's `__init__`, since -# mutations to this registry performed during the precompilation of a -# downstream package don't carry over to module load (see -# `register_reclaimable!` for details). +# `Reclaimable` unifies both. Register instances via `register_reclaimable!`. abstract type Reclaimable end @@ -42,13 +38,9 @@ const reclaimables_lock = ReentrantLock() """ CUDACore.register_reclaimable!(r::Reclaimable) -Register `r` so that `reclaim` invokes its `drop!` and `purge!` methods. -Idempotent: re-registering the same instance is a no-op. Returns `r`. - -Call this from a package's `__init__`: top-level constructor calls are -captured by Julia's precompilation cache, but mutations they perform on -dependencies (like pushing into CUDACore's registry) do not persist to -module load time. +Register `r` so `reclaim` invokes its `drop!` and `purge!` methods. +Idempotent; returns `r`. Call from the owning package's `__init__` +(registry mutations performed during precompile don't survive to load). """ function register_reclaimable!(r::Reclaimable) @lock reclaimables_lock begin @@ -57,16 +49,16 @@ function register_reclaimable!(r::Reclaimable) return r end -# invoke f(r) per Reclaimable with per-hook error isolation, modelled -# after Base's `atexit_hooks`: one bad hook mustn't prevent the others. +# Snapshot the registry under the lock, then run callbacks unlocked +# (Base.atexit-style): callbacks may take a while and may transitively +# trigger reclaim, so we don't want to hold the lock across them. function foreach_reclaimable(f) - @lock reclaimables_lock begin - for r in reclaimables - try - f(r) - catch ex - @error "reclaim callback failed" type=typeof(r) exception=(ex, catch_backtrace()) - end + snapshot = @lock reclaimables_lock copy(reclaimables) + for r in snapshot + try + f(r) + catch ex + @error "reclaim callback failed" type=typeof(r) exception=(ex, catch_backtrace()) end end end @@ -79,10 +71,10 @@ end Declarative marker for a library's per-task cache stored under `key` in `task_local_storage()`. Must be `register_reclaimable!`'d from the owning -package's `__init__` so that `RECLAIM_DROP_STATE` clears the current -task's entry, letting the stored values (typically mutable handle wrappers) -be garbage-collected. Their finalizers then return the underlying -resources to a `HandleCache`. +package's `__init__` so that `RECLAIM_DROP` clears the current task's +entry, letting the stored values (typically mutable handle wrappers) be +garbage-collected. Their finalizers then return the underlying resources +to a `HandleCache`. Usage: @@ -110,7 +102,14 @@ struct TaskLocalCache{K,V} <: Reclaimable end @inline function task_dict(s::TaskLocalCache{K,V}) where {K,V} - get!(() -> Dict{K,V}(), task_local_storage(), s.key)::Dict{K,V} + # spelled out instead of get!(()->Dict, ...) to avoid the closure alloc + tls = task_local_storage() + d = get(tls, s.key, nothing) + if d === nothing + d = Dict{K,V}() + tls[s.key] = d + end + return d::Dict{K,V} end function drop!(s::TaskLocalCache) diff --git a/test/core/pool.jl b/test/core/pool.jl index 83df4592c5..4bf55d91d9 100644 --- a/test/core/pool.jl +++ b/test/core/pool.jl @@ -31,15 +31,29 @@ end end @testset "reclaim" begin + # every level should run without erroring, on any allocator + for level in instances(CUDACore.ReclaimLevel) + CUDA.reclaim(level) + end CUDA.reclaim() - CUDA.reclaim(CUDACore.RECLAIM_GC_FULL) + # `retry_reclaim` returns the block's result, retrying while `retry_if` holds @test CUDACore.retry_reclaim(isequal(42)) do 42 end == 42 @test CUDACore.retry_reclaim(isequal(42)) do 41 end == 41 + + # `retry_reclaim` escalates: count how many calls it takes to stop retrying + let n = Ref(0) + ret = CUDACore.retry_reclaim(ret -> ret < 2) do + n[] += 1 + n[] + end + @test ret == 2 + @test n[] == 2 + end end @testset "pool_status" begin diff --git a/test/runtests.jl b/test/runtests.jl index 66984d1abb..66e58ebeaf 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -103,7 +103,8 @@ end # Cap worker count by how much of the primary device's free memory each worker # claims. A CUDA worker needs its own context + libraries (~0.5–1 GiB baseline) -# plus room for peak per-test allocations; 4 GiB is the per-worker budget. +# plus room for peak per-test allocations; with pool caching capped in +# `test/setup.jl`, 1 GiB is enough. # (Set `CUDA_VISIBLE_DEVICES` to choose which device is used.) const gpu_memory_per_worker = 1 * 2^30 first_gpu = first(devices()) diff --git a/test/setup.jl b/test/setup.jl index 2b9a49fb6b..a5587ec128 100644 --- a/test/setup.jl +++ b/test/setup.jl @@ -46,14 +46,19 @@ end # precompile the runtime library CUDA.precompile_runtime() -# Cap the amount of memory the CUDA pool caches (default: unbounded). Freed -# stream-ordered buffers now go back to the driver almost immediately, which -# keeps each test's GPU RSS close to its peak live allocation rather than the -# running max across the test's lifetime. Trades per-call pool-alloc cost -# for lower NVML-reported memory. -let dev = device(), pool = CUDACore.pool_create(dev) - CUDACore.attribute!(pool, CUDACore.MEMPOOL_ATTR_RELEASE_THRESHOLD, - UInt64(64 * 2^20)) +# Cap the amount of memory the CUDA pool keeps cached after frees (default: +# unbounded). Above this watermark, freed stream-ordered buffers go back to +# the driver, which keeps each test's GPU RSS close to its peak live +# allocation rather than the running max across the test's lifetime. The +# threshold trades per-alloc pool-refill cost for lower NVML-reported memory; +# tune up if test wall-time regresses, down if a worker's RSS budget is tight. +const pool_release_threshold = 256 * 2^20 # 256 MiB +let dev = device() + if CUDACore.stream_ordered(dev) + pool = CUDACore.pool_create(dev) + CUDACore.attribute!(pool, CUDACore.MEMPOOL_ATTR_RELEASE_THRESHOLD, + UInt64(pool_release_threshold)) + end end