Task-local state follow-ups: `prepare_state` in finalizers, `context!` restore, `device!` global default

[luraess]

Follow-up to #1104. With #1104, #1092 and #1096, finalizers no longer call the task-local API directly (see the audit below). What's left:

1. Finalizers still go through prepare_state

Every generated HIP/library wrapper (hipFreeAsync, hipStreamDestroy, the MIOpen and rocSPARSE destroys, ...) starts with AMDGPU.prepare_state(). From a finalizer, that:

• creates task-local state, on the default device, on whatever task GC interrupted;
• switches the thread back to that task's context, so a finalizer's HIP.context!(obj.ctx) never reaches the C call (the suite passes, so HIP tolerates the wrong context for these calls);
• drains a sticky HIP error that the interrupted task may not have checked yet.

Proposal: skip it in finalizers. GC.in_finalizer exists from Julia 1.10.

@inline function prepare_state(state = nothing)
    GC.in_finalizer() && return # finalizers pick their context with `HIP.context!`
    state === nothing && (state = task_local_state!())
    # ... unchanged
end

On top of the three PRs, this takes the audit to 0 (last column below). The rest of the suite is unchanged, with no finalizer or free errors.

This only covers finalizers run by GC. An explicit finalize or unsafe_free! still goes through prepare_state, which still overrides HIP.context!.

2. context!(f, ctx) doesn't restore on a fresh task

On a task without state, context!(ctx) creates it on ctx and returns nothing, so the finally never switches back. device!(f, dev) takes the same path:

fetch(Threads.@spawn (AMDGPU.device!(() -> nothing, AMDGPU.device(2)); AMDGPU.device_id()))  # 2, not the default

Fix: create the default state first, so there is always something to restore:

function context!(ctx::HIPContext)
    state = task_local_state!()
    old_ctx = state.context
    if old_ctx != ctx
        HIP.context!(ctx)
        state.device = HIP.device()
        state.context = ctx
    end
    return old_ctx
end

function context!(f::Function, ctx::HIPContext)
    old_ctx = context!(ctx)
    try
        f()
    finally
        old_ctx != ctx && context!(old_ctx)
    end
end

This also drops the ctx.valid check. It can't fail since #605 removed the HIPContext finalizer, and its @warn isn't finalizer-safe anyway.

3. device! sets a process-global default

device! also writes HIP.DEFAULT_DEVICE[], which every new task starts on. #487 did this on purpose, but the docstring says device! only affects the calling task, and the write happens from any task. For example, the spawned task in multi_gpu.jl "Correctly switching HIP context" leaves later test files on that worker on device 2. Options: document it, only update the default from the root task, or add an explicit default_device!. Either way, multi_gpu.jl should restore device 1.

4. pool_cleanup trims the wrong device

pool_cleanup runs HIP.device!(dev) do reclaim() end, but reclaim() reads the task-local AMDGPU.device(), so on a multi-GPU node, idle trimming meant for device 2 trims the cleanup task's device instead. This only runs in interactive sessions. Fix: pass the device to reclaim.

Audit

Full suite with the patch below (mi200 / gfx90a, 2 GCDs, Julia 1.12). Reports are capped at 50 per worker, so the last row is a floor.

Task-local access from a finalizer	main	#1104 + #1092 + #1096	+ item 1
AMDGPU.context! in the HIPStream finalizer	31	0	0
AMDGPU.context() in rocFFT release_plan!	26	0	0
prepare_state() in the generated wrappers	1543	1600	0

Audit patch (not for merge)

--- a/src/tls.jl
+++ b/src/tls.jl
@@ -20,11 +20,29 @@
     end
 end
 
-task_local_state()::Union{Nothing, TaskLocalState} =
+# AUDIT ONLY (not for merge): report task-local state access from finalizers.
+# Finalizers run on whichever task triggered GC, so any hit here reads or
+# rewrites an unrelated task's device.
+const _TLS_AUDIT_HITS = Threads.Atomic{Int}(0)
+function _tls_audit(what::Symbol)
+    GC.in_finalizer() || return
+    n = Threads.atomic_add!(_TLS_AUDIT_HITS, 1) + 1
+    n > 50 && return
+    Core.println(Core.stderr, "AMDGPU-TLS-AUDIT #", n, ": ", what, " called from a finalizer")
+    Base.show_backtrace(Core.stderr, backtrace())
+    Core.println(Core.stderr)
+    return
+end
+
+function task_local_state()::Union{Nothing, TaskLocalState}
+    _tls_audit(:task_local_state)
     get(task_local_storage(), :AMDGPU, nothing)
+end
 
-task_local_state!(args...)::TaskLocalState =
+function task_local_state!(args...)::TaskLocalState
+    _tls_audit(:task_local_state!)
     get!(() -> TaskLocalState(args...), task_local_storage(), :AMDGPU)
+end
 
 Base.copy(state::TaskLocalState) = TaskLocalState(
     state.device, state.context, copy(state.streams))

[luraess]

Items 1, 2 and 4 are in #1106 and #1107 . Two points I'd like feedback on:

• Item 3: should device! keep setting the default device for new tasks? The options are to document it, to only update the default from the root task, or to add an explicit default_device!. I'd document it for now. @pxl-th, since Remove default device stuff #487 made it the default on purpose.
• Item 1 caveat: explicit finalize/unsafe_free! calls still have their HIP.context! overridden by prepare_state. Should prepare_state respect an enclosing HIP.context!, or is that fine as is, given HIP tolerates it?