Keep stream and rocFFT plan destructors out of task-local state

src/fft/wrappers.jl

@@ -33,7 +33,8 @@ function release_plan!(plan)
     function destroy()
         handle != C_NULL && Threads.atomic_add!(N_PLANS_DESTROYED, 1)
         # Pin to `ctx`, since eviction may run this under a different context.
-        AMDGPU.context!(() -> rocfft_plan_destroy(handle), ctx)
+        # HIP-level, since eviction can run in a finalizer.
+        HIP.context!(() -> rocfft_plan_destroy(handle), ctx)
     end
     push!(destroy, IDLE_HANDLES, key, value)
 end

src/hip/stream.jl

@@ -29,7 +29,8 @@ function HIPStream(priority::Symbol = :normal)
     stream = HIPStream(stream_ref[], priority, d, HIPContext(d), true)
     return finalizer(stream) do s
         Base.@atomic s.valid = false
-        AMDGPU.context!(s.ctx) do
+        # Not `AMDGPU.context!`: finalizers run on arbitrary tasks.
+        HIP.context!(s.ctx) do
             hipStreamDestroy(s.stream)
         end
     end

test/core/tls.jl

@@ -54,4 +54,16 @@ end
         # Must return true without segfaulting on an already-finalized stream.
         @test AMDGPU.HIP.isdone(s) == true
     end
+
+    if length(AMDGPU.devices()) > 1
+        @testset "Stream finalizer keeps the running task's device" begin
+            # Finalizers run on whichever task triggers GC. The stream finalizer
+            # must not move that task onto the stream's device.
+            default = fetch(@async AMDGPU.device())
+            other = first(d for d in AMDGPU.devices() if d != default)
+            s = AMDGPU.device!(() -> HIPStream(), other)
+            @test s.device == other
+            @test fetch(@async (finalize(s); AMDGPU.device())) == default
+        end
+    end
 end