diff --git a/src/fft/wrappers.jl b/src/fft/wrappers.jl index d7ac0257c..90337ec6a 100644 --- a/src/fft/wrappers.jl +++ b/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 diff --git a/src/hip/stream.jl b/src/hip/stream.jl index b2c1dbc9b..3f89f5bdb 100644 --- a/src/hip/stream.jl +++ b/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 diff --git a/test/core/tls.jl b/test/core/tls.jl index b0adc452e..e530e7b24 100644 --- a/test/core/tls.jl +++ b/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