Switch back after device!(f, dev) on a fresh task

src/memory.jl

@@ -312,7 +312,8 @@ function pool_cleanup()
             pool_mark!(dev, false)
 
             if get(idle_counters, did, 0) >= 5
-                HIP.device!(dev) do
+                # `reclaim` trims the pool of the task-local device.
+                AMDGPU.device!(dev) do
                     reclaim()
                 end
             end

src/tls.jl

@@ -118,32 +118,24 @@ end
 context() = task_local_state!().context
 
 function context!(ctx::HIPContext)
-    state = task_local_state()
-    if state ≡ nothing
-        old_ctx = nothing
+    # Start from the default state, so that `context!(f, ctx)` has a context
+    # to switch back to even on a task that had no state yet.
+    state = task_local_state!()
+    old_ctx = state.context
+    if old_ctx != ctx
         HIP.context!(ctx)
-        task_local_state!(HIP.device(), ctx)
-    else
-        old_ctx = state.context
-        if old_ctx != ctx
-            HIP.context!(ctx)
-            state.device = HIP.device()
-            state.context = ctx
-        end
+        state.device = HIP.device()
+        state.context = ctx
     end
     return old_ctx
 end
 
 function context!(f::Function, ctx::HIPContext)
-    if ctx.valid
-        old_ctx = context!(ctx)
-        return try
-            f()
-        finally
-            old_ctx ≢ nothing && old_ctx != ctx && context!(old_ctx)
-        end
-    else
-        @warn "CTX not valid"
+    old_ctx = context!(ctx)
+    return try
+        f()
+    finally
+        old_ctx != ctx && context!(old_ctx)
     end
 end
 

test/hip_rocarray/multi_gpu.jl

@@ -30,6 +30,15 @@ else
         @test d1 == AMDGPU.device()
     end
 
+    @testset "Scoped switch on a fresh task" begin
+        # `device!(f, dev)` must switch back even on a task that had no state yet.
+        default = fetch(@async AMDGPU.device())
+        other = first(d for d in AMDGPU.devices() if d != default)
+        inside, after = fetch(@async (AMDGPU.device!(AMDGPU.device, other), AMDGPU.device()))
+        @test inside == other
+        @test after == default
+    end
+
     @testset "Arrays" begin
         d1 = AMDGPU.device()