plan handle cache fix

[gregordecristoforo]

Proposed fix #1053

HandleCache got a max_idle budget that bounds the total number of idle handles across all keys (not just per-key via max_entries), with a new idle_dtors dict holding each idle handle's destructor so that _evict_idle! can reclaim handles for any key once the budget is exceeded, running those destructors outside the lock. Also, release_plan! was changed to capture the plan handle in a local variable and close over that (() -> rocfft_plan_destroy(handle)) instead of closing over the plan and reading plan.handle lazily.

With this implementation the MWE in #1053 gives the following output on LUMI:

➜  AMDGPU.jl git:(gd/plan-handle-cache-fix) ✗ julia --project=. mwe.jl
Initial GPU memory: 63.896 GiB free / 63.984 GiB total
Planning 6000 distinct FFT lengths (16, 18, 20, ... in steps of 2)...
[50/6000] free=63.371 GiB cached_plan_shapes=64 idle_handles=64
used=0 bytes cached(pool reserved)=464 bytes
[100/6000] free=63.293 GiB cached_plan_shapes=64 idle_handles=64
used=0 bytes cached(pool reserved)=864 bytes
[150/6000] free=63.293 GiB cached_plan_shapes=64 idle_handles=64
used=0 bytes cached(pool reserved)=1.234 KiB
[200/6000] free=63.262 GiB cached_plan_shapes=64 idle_handles=64
used=0 bytes cached(pool reserved)=1.625 KiB
...

Is this PR going into the right direction?

[gregordecristoforo]

@luraess Would you be able to have a look?

[github-actions]

AMDGPU.jl Benchmarks

Details

Benchmark suite	Current: 32dfa70	Previous: 4ff61e3	Ratio
amdgpu/synchronization/context/device	550 ns	567.5 ns	0.97
amdgpu/synchronization/stream/blocking	230 ns	232.5 ns	0.99
amdgpu/synchronization/stream/nonblocking	307.5 ns	312.5 ns	0.98
applications/bitonic_sort	1654649 ns	1650734.75 ns	1.00
applications/convolution	114364 ns	106826.75 ns	1.07
applications/floyd_warshall	11720482 ns	11787759.25 ns	0.99
applications/histogram	786033.75 ns	789913 ns	1.00
applications/prefix_sum	253781 ns	251799.25 ns	1.01
array/accumulate/Float32/1d	78586.25 ns	72711.25 ns	1.08
array/accumulate/Float32/dims=1	267648.75 ns	284589.75 ns	0.94
array/accumulate/Float32/dims=1L	88828.75 ns	81133.75 ns	1.09
array/accumulate/Float32/dims=2	85178.75 ns	70688.75 ns	1.20
array/accumulate/Float32/dims=2L	2752692.5 ns	2754571.75 ns	1.00
array/accumulate/Int64/1d	80371.25 ns	77051.25 ns	1.04
array/accumulate/Int64/dims=1	243926 ns	242136.5 ns	1.01
array/accumulate/Int64/dims=1L	84058.75 ns	83854 ns	1.00
array/accumulate/Int64/dims=2	87288.75 ns	70858.75 ns	1.23
array/accumulate/Int64/dims=2L	2893489.5 ns	2892754.5 ns	1.00
array/broadcast	72758.5 ns	72816.5 ns	1.00
array/construct	2235 ns	2200 ns	1.02
array/copy	37165.5 ns	36975.5 ns	1.01
array/copyto!/cpu_to_gpu	111399 ns	110439.5 ns	1.01
array/copyto!/gpu_to_cpu	112054 ns	110576.75 ns	1.01
array/copyto!/gpu_to_gpu	58845.75 ns	58938.5 ns	1.00
array/iteration/findall/bool	138292 ns	133184.25 ns	1.04
array/iteration/findall/int	148644.5 ns	146924.25 ns	1.01
array/iteration/findfirst/bool	183707.75 ns	183770 ns	1.00
array/iteration/findfirst/int	162694.75 ns	145197 ns	1.12
array/iteration/findmin/1d	110244 ns	109714 ns	1.00
array/iteration/findmin/2d	112094.25 ns	106499 ns	1.05
array/iteration/logical	241226 ns	240093.25 ns	1.00
array/iteration/scalar	287576.5 ns	292796.75 ns	0.98
array/permutedims/2d	71396 ns	70961.25 ns	1.01
array/permutedims/3d	70928.5 ns	70623.75 ns	1.00
array/permutedims/4d	73811 ns	73141.25 ns	1.01
array/random/rand/Float32	38838 ns	44920.75 ns	0.86
array/random/rand/Int64	46345.5 ns	53576 ns	0.87
array/random/rand!/Float32	58715.75 ns	64641 ns	0.91
array/random/rand!/Int64	72351.25 ns	71331.25 ns	1.01
array/random/randn/Float32	80618.75 ns	78961.25 ns	1.02
array/random/randn!/Float32	80113.5 ns	72186.25 ns	1.11
array/reductions/mapreduce/Float32/1d	101131.5 ns	94511.5 ns	1.07
array/reductions/mapreduce/Float32/dims=1	89928.75 ns	75061.25 ns	1.20
array/reductions/mapreduce/Float32/dims=1L	831944.5 ns	832329.5 ns	1.00
array/reductions/mapreduce/Float32/dims=2	87511.25 ns	83001.5 ns	1.05
array/reductions/mapreduce/Float32/dims=2L	143689.75 ns	143060 ns	1.00
array/reductions/mapreduce/Int64/1d	100709 ns	93981.75 ns	1.07
array/reductions/mapreduce/Int64/dims=1	90044 ns	82693.75 ns	1.09
array/reductions/mapreduce/Int64/dims=1L	828552 ns	833094.75 ns	0.99
array/reductions/mapreduce/Int64/dims=2	90451.5 ns	83931.5 ns	1.08
array/reductions/mapreduce/Int64/dims=2L	144284.75 ns	143907.5 ns	1.00
array/reductions/reduce/Float32/1d	100948.75 ns	94144.25 ns	1.07
array/reductions/reduce/Float32/dims=1	90528.75 ns	75216.5 ns	1.20
array/reductions/reduce/Float32/dims=1L	829094.75 ns	832537.25 ns	1.00
array/reductions/reduce/Float32/dims=2	87216.25 ns	83219 ns	1.05
array/reductions/reduce/Float32/dims=2L	143742 ns	143200 ns	1.00
array/reductions/reduce/Int64/1d	98919 ns	91536.25 ns	1.08
array/reductions/reduce/Int64/dims=1	90206.25 ns	83593.5 ns	1.08
array/reductions/reduce/Int64/dims=1L	831364.25 ns	829128.75 ns	1.00
array/reductions/reduce/Int64/dims=2	87133.75 ns	83766 ns	1.04
array/reductions/reduce/Int64/dims=2L	144259.5 ns	143587 ns	1.00
array/reverse/1d	45030.75 ns	45153 ns	1.00
array/reverse/1dL	75326.25 ns	75516 ns	1.00
array/reverse/1dL_inplace	79653.75 ns	79963.75 ns	1.00
array/reverse/1d_inplace	54255.75 ns	60771 ns	0.89
array/reverse/2d	49568.25 ns	43250.75 ns	1.15
array/reverse/2dL	60491 ns	83061 ns	0.73
array/reverse/2dL_inplace	91433.75 ns	89986.25 ns	1.02
array/reverse/2d_inplace	41990.5 ns	55665.75 ns	0.75
array/sorting/1d	330699.75 ns	333063 ns	0.99
gemm/tiled	1886265 ns	1899696.5 ns	0.99
gemm/tiled_unbounded	1899457.5 ns	1922439.5 ns	0.99
integration/byval/reference	40091 ns	39001 ns	1.03
integration/byval/slices=1	40650 ns	39781 ns	1.02
integration/byval/slices=2	146803 ns	147233 ns	1.00
integration/byval/slices=3	245973 ns	245404 ns	1.00
integration/volumerhs	4892521 ns	4962111 ns	0.99
kernel/indexing	57100.75 ns	57025.75 ns	1.00
kernel/indexing_checked	58495.75 ns	56630.75 ns	1.03
kernel/launch	1377.5 ns	1365.25 ns	1.01
kernel/rand	102719 ns	88876 ns	1.16
latency/import	1715513936 ns	1752022784 ns	0.98
latency/precompile	40284399173 ns	39860711432 ns	1.01
latency/ttfp	2332473179 ns	2327930658 ns	1.00
stencil/diffusion3d	1617712.75 ns	1624179.5 ns	1.00
stencil/diffusion3d_checked	1656408.25 ns	1655412.5 ns	1.00

This comment was automatically generated by workflow using github-action-benchmark.

[luraess]

Thanks. I checked the proposed changes and iterated on it with the robot. The direction looks right to me (diagnosis matches what I tested, the fix works).

I ran it on MI300 (gfx942, ROCm 7.2.3). The #1053 MWE now holds at 64 idle handles. Free memory goes flat from iteration 100 through 1200, with 195 MiB of drift that plateaus. Your tests pass as well: core 104/104, fft 7/7, hip_rocarray/fft at 268/268. The leak in your test measures 1.02 MiB against the 400 MiB budget.

A few things I noticed.

On framing. Capturing the handle by value in release_plan! looks like a prerequisite of this PR rather than a bug fix. Before it, the destructor always ran synchronously inside push!. That happened ahead of unsafe_free! setting plan.handle = C_NULL, so nothing shipped broken. Might be worth rewording the comment and the test's "used to queue…". There is a second argument for the change though: closing over plan and parking it in a process-global dict would resurrect a finalized object.

1. Eviction order in _evict_idle!. This is the one I'd look at before merging.

Your docstring already notes it is "a count budget, not an LRU". It may be a little sharper than arbitrary, though. Eviction restarts from the front each time and drains until the budget is met, so it tends to hit the same keys repeatedly. I measured this against the real IDLE_HANDLES, over 200 iterations of "N shapes reused every iteration, plus M shapes used once":

max_idle = 64, max_entries = 32
hot=1   cold/iter=4  ->   77 rebuilds of a hot plan (ideal 1)
hot=4   cold/iter=4  ->  800 rebuilds (ideal 4)    <- 100% miss
hot=8   cold/iter=8  -> 1547 rebuilds (ideal 8)
hot=16  cold/iter=4  -> 3200 rebuilds (ideal 16)   <- 100% miss

On main all four rows sit at ideal, since nothing is evicted at all. A rebuild costs about 4.8 ms when rocFFT's RTC cache is warm for that shape, and 14.4 ms when it is cold. At four reused shapes, every lookup misses.

This might be one change together with idle_dtors rather than two. A single Dict{K, Vector{IdleHandle{V}}}, with entries carrying (handle, dtor, seq), would collapse the two dicts into one and make the parallel invariant unrepresentable. It also gives eviction a global order to work from. Evicting the lowest seq brings all four rows back to ideal. Each key's vector is already in ascending seq order, so only its front is a candidate. That keeps the scan at O(number of keys), and the keys are bounded by max_idle + 1.

Diff, if useful (src/cache.jl only, and it takes max_idle = 64 as given since that is the separate question in 2)

diff --git a/src/cache.jl b/src/cache.jl
index bfc41a0d..4884ba02 100644
--- a/src/cache.jl
+++ b/src/cache.jl
@@ -7,13 +7,17 @@
 
 export HandleCache
 
+# One idle handle, together with the destructor that reclaims it and the global
+# insertion order used to pick an eviction victim.
+struct IdleHandle{V}
+    handle::V
+    dtor::Any
+    seq::Int
+end
+
 struct HandleCache{K, V}
     active_handles::Set{Pair{K, V}}
-    idle_handles::Dict{K, Vector{V}}
-    # Destructor closure for every idle handle, kept parallel to `idle_handles`
-    # so that idle handles for *any* key can be reclaimed (not just the one
-    # currently being pushed).
-    idle_dtors::Dict{K, Vector{Any}}
+    idle_handles::Dict{K, Vector{IdleHandle{V}}}
     lock::Base.ThreadSynchronizer
     # TODO when finalizers are run on their own tasks use reentrant lock
 
@@ -23,56 +27,59 @@ struct HandleCache{K, V}
     # cache for workloads that use a large number of distinct keys (e.g. rocFFT
     # plans for many different shapes), which would otherwise leak one handle
     # per key forever since no single key ever reaches `max_entries`. See #1053.
-    # Inert for caches keyed on something with few values (e.g. `HIPContext`),
-    # where the per-key `max_entries` budget is the one that matters.
     max_idle::Int
+    # Stamps each idle handle so `_evict_idle!` can evict in least-recently-cached
+    # order across all keys, instead of whichever key `Dict` iteration visits first.
+    seq::Base.RefValue{Int}
 
     function HandleCache{K, V}(max_entries::Int = 32, max_idle::Int = 64) where {K, V}
         new{K,V}(
             Set{Pair{K, V}}(),
-            Dict{K, Vector{V}}(),
-            Dict{K, Vector{Any}}(),
+            Dict{K, Vector{IdleHandle{V}}}(),
             Base.ThreadSynchronizer(),
-            max_entries, max_idle)
+            max_entries, max_idle, Ref(0))
     end
 end
 
+# Total number of idle handles across all keys.
+# Must be called while holding `cache.lock`.
+total_idle(cache::HandleCache) = sum(length, values(cache.idle_handles); init = 0)
+
 # Take an idle handle for `key` out of the cache, or `nothing` if there is none.
+# Its destructor is dropped: the caller owns the handle again, and hands us a
+# fresh destructor when it comes back through `push!`.
 # Must be called while holding `cache.lock`.
 function _take_idle!(cache::HandleCache{K, V}, key) where {K, V}
-    handles = get(cache.idle_handles, key, nothing)
-    (handles ≡ nothing || isempty(handles)) && return nothing
-    handle = pop!(handles)
-    pop!(cache.idle_dtors[key])
-    if isempty(handles)
-        delete!(cache.idle_handles, key)
-        delete!(cache.idle_dtors, key)
-    end
-    return handle
+    entries = get(cache.idle_handles, key, nothing)
+    entries ≡ nothing && return nothing
+    entry = pop!(entries)
+    isempty(entries) && delete!(cache.idle_handles, key)
+    return entry.handle
 end
 
 # Evict idle handles until at most `max_idle` remain across all keys, returning
-# the destructor closures of the evicted handles to be run by the caller outside
-# the lock. This is a count budget, not an LRU: keys are visited in `Dict`
-# iteration order (arbitrary after deletions) and each key's oldest handles go
-# first. Must be called while holding `cache.lock`.
+# the destructors of the evicted handles to be run by the caller outside the
+# lock. Victims are chosen in least-recently-cached order across all keys, so a
+# key that keeps being reused survives churn from keys that are used once.
+# Each key's vector is in ascending `seq` order, so only its front is a
+# candidate, making the scan O(number of keys) per eviction.
+# Must be called while holding `cache.lock`.
 function _evict_idle!(cache::HandleCache{K, V}) where {K, V}
     evicted = Any[]
-    total = sum(length, values(cache.idle_handles); init = 0)
-    total ≤ cache.max_idle && return evicted
-    for key in collect(keys(cache.idle_handles))
-        handles = cache.idle_handles[key]
-        dtors = cache.idle_dtors[key]
-        while !isempty(handles) && total > cache.max_idle
-            popfirst!(handles)
-            push!(evicted, popfirst!(dtors))
-            total -= 1
-        end
-        if isempty(handles)
-            delete!(cache.idle_handles, key)
-            delete!(cache.idle_dtors, key)
+    total = total_idle(cache)
+    while total > cache.max_idle
+        victim, oldest = nothing, typemax(Int)
+        for (key, entries) in cache.idle_handles
+            if entries[1].seq < oldest
+                victim, oldest = key, entries[1].seq
+            end
         end
-        total ≤ cache.max_idle && break
+        victim ≡ nothing && break
+
+        entries = cache.idle_handles[victim]
+        push!(evicted, popfirst!(entries).dtor)
+        isempty(entries) && delete!(cache.idle_handles, victim)
+        total -= 1
     end
     return evicted
 end
@@ -109,13 +116,16 @@ function Base.push!(f::Function, cache::HandleCache{K, V}, key::K, handle::V) wh
             """)
         delete!(cache.active_handles, key => handle)
 
-        handles = get!(() -> V[], cache.idle_handles, key)
+        entries = get(cache.idle_handles, key, nothing)
         # Keep the original semantics: save while the key holds ≤ `max_entries`,
         # so a key can grow to `max_entries + 1`.
-        saved = length(handles) ≤ cache.max_entries
+        saved = (entries ≡ nothing ? 0 : length(entries)) ≤ cache.max_entries
         if saved
-            push!(handles, handle)
-            push!(get!(() -> Any[], cache.idle_dtors, key), f)
+            # Only materialise the vector once we know we are filling it, so a
+            # key is present in `idle_handles` iff it holds at least one handle.
+            entries ≡ nothing &&
+                (entries = cache.idle_handles[key] = IdleHandle{V}[])
+            push!(entries, IdleHandle{V}(handle, f, (cache.seq[] += 1)))
         end
 
         # Enforce the global idle-handle budget.

Checked on MI300: invariants hold over 200 randomised trials x 400 ops (max_entries 0-5, max_idle 0-12, with no lost or double-freed handles, the budget never exceeded, and no empty vectors left behind). The table above becomes 1 / 3 / 4 / 8 rebuilds against the real IDLE_HANDLES, and hip_rocarray/fft stays green.

The destructor's context could then be a two-liner in release_plan!. An evicted plan may now be destroyed under whichever context happens to be current, where before it ran inline under the one that built the key. library_state wraps its destructor in context!(tls.context), presumably for this reason:

ctx, handle = AMDGPU.context(), plan.handle
push!(() -> AMDGPU.context!(() -> rocfft_plan_destroy(handle), ctx), IDLE_HANDLES, key, value)

There is a fuller version: one dtor(key, handle) on the cache, with no per-handle closures, which is roughly the TODO you removed. It touches all six construction sites though, so probably not this PR.

2. max_idle and the HIPContext caches. The comment says it is inert for those, but I don't think it is. max_entries allows 33 idle handles per key, so two contexts already exceed 64. I hit that on a 2-GCD allocation. rocBLAS, rocSPARSE, MIOpen, hipTENSOR and rocRAND all share this default, and a rocBLAS handle is more expensive to rebuild than an FFT plan. One option would be max_idle = typemax(Int) as the default, with an explicit budget only on rocFFT's IDLE_HANDLES.

Related: a count may be a weak proxy for what is actually being bounded. #1053's numbers work out to roughly 10 MB per plan. The small 1-D plans here are nearer 3 MB, and 64 large 3-D plans would be different again. Making the number settable could help until the memory-pressure TODO lands.

3. Tests.

• The closure introspection asserts a local variable name. Renaming handle to h in release_plan! makes fieldnames return (:h,), so the test fails with nothing else changed. The leak subtest below it seems to cover the same ground, so I'd probably drop it.
• churn(200:2:598) costs about 66 s. Given the margin is 1.02 MiB against 400 MiB, churn(200:2:400) should give the same signal in half the time.
• Once eviction has a defined order, a test pinning it would be useful: a hot key surviving cold churn. That is the part that can regress quietly.
• sum(length, values(...); init = 0) appears three times. A total_idle(cache) helper would DRY it, and give the tests an accessor instead of reaching into fields.

Small one: the try/catch now swallows a failing rocfft_plan_destroy that used to throw. That seems reasonable for the deferred case. Worth noting that @error allocates and takes the logging lock, and push! is reached from a finalizer, so @debug or a rate limit might be safer.

Separately, and pre-existing: release_plan! builds its key from AMDGPU.context() at finalization time, not the context the plan was created in. On multi-GPU that could already file a plan under the wrong device. Probably its own issue rather than something for this PR.

[gregordecristoforo]

Thanks @luraess for the detailed feedback! I'll probably have all day on Tuesday next week to implement it.

Separately, and pre-existing: release_plan! builds its key from AMDGPU.context() at finalization time, not the context the plan was created in. On multi-GPU that could already file a plan under the wrong device. Probably its own issue rather than something for this PR.

I am happy to look at this in a separate PR. I can use LUMI for multi-GPU testing.