From aeb15de4a26c5a8a21dbaf7804ef0e2e10778120 Mon Sep 17 00:00:00 2001 From: Ludovic Raess Date: Wed, 23 Sep 2026 22:44:22 +0200 Subject: [PATCH 1/2] Keep stream and rocFFT plan destructors out of task-local state The HIPStream finalizer and the rocFFT plan destructor switched context with the task-local `AMDGPU.context!`. Finalizers run on whichever task triggers GC, so a stream finalizer running on a fresh task pinned that task to the stream's device. Under ParallelTestRunner this carried from one test file to the next and made `multi_gpu.jl` start on device 2. Use `HIP.context!`, which only saves and restores the thread's HIP context. Co-Authored-By: Claude Opus 5.5 --- src/fft/wrappers.jl | 3 ++- src/hip/stream.jl | 3 ++- test/core/tls.jl | 10 ++++++++++ 3 files changed, 14 insertions(+), 2 deletions(-) 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..8f9fdf89e 100644 --- a/test/core/tls.jl +++ b/test/core/tls.jl @@ -54,4 +54,14 @@ end # Must return true without segfaulting on an already-finalized stream. @test AMDGPU.HIP.isdone(s) == true end + + @testset "Finalizer leaves task-local state alone" begin + # Finalizers run on whichever task triggers GC. If the stream finalizer + # initializes that task's state, the task is stuck on the stream's device. + s = HIPStream() + @test fetch(@async begin + finalize(s) + AMDGPU.task_local_state() ≡ nothing + end) + end end From de6ecfe4397ac8d7942218e47f1a88028f470c36 Mon Sep 17 00:00:00 2001 From: Ludovic Raess Date: Wed, 23 Sep 2026 23:31:13 +0200 Subject: [PATCH 2/2] Test that the stream finalizer keeps the running task's device --- test/core/tls.jl | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/test/core/tls.jl b/test/core/tls.jl index 8f9fdf89e..e530e7b24 100644 --- a/test/core/tls.jl +++ b/test/core/tls.jl @@ -55,13 +55,15 @@ end @test AMDGPU.HIP.isdone(s) == true end - @testset "Finalizer leaves task-local state alone" begin - # Finalizers run on whichever task triggers GC. If the stream finalizer - # initializes that task's state, the task is stuck on the stream's device. - s = HIPStream() - @test fetch(@async begin - finalize(s) - AMDGPU.task_local_state() ≡ nothing - 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