Perform low-level clean-up in the correct context.

lib/cudadrv/context.jl

@@ -117,6 +117,24 @@ function Base.pop!(::Type{CuContext})
     CuContext(handle_ref[])
 end
 
+# perform some finalizer actions in a context
+macro finalize_in_ctx(ctx, body)
+    # XXX: should this not integrate with the high-level context management from state.jl?
+    #      it might be good that the driver API wrappers don't need that runtime-esque
+    #      state management, but it might be confusing that `context()` doesn't work here.
+    quote
+        ctx = $(esc(ctx))
+        if isvalid(ctx)
+            push!(CuContext, ctx)
+            try
+                $(esc(body))
+            finally
+                pop!(CuContext)
+            end
+        end
+    end
+end
+
 """
     activate(ctx::CuContext)
 

lib/cudadrv/events.jl

@@ -26,9 +26,7 @@ mutable struct CuEvent
 end
 
 function unsafe_destroy!(e::CuEvent)
-    if isvalid(e.ctx)
-        cuEventDestroy_v2(e)
-    end
+    @finalize_in_ctx e.ctx cuEventDestroy_v2(e)
 end
 
 Base.unsafe_convert(::Type{CUevent}, e::CuEvent) = e.handle

lib/cudadrv/memory.jl

@@ -5,7 +5,7 @@ export Mem, attribute, attribute!, memory_type, is_managed
 module Mem
 
 using ..CUDA
-using ..CUDA: @enum_without_prefix, CUstream, CUdevice, CuDim3, CUarray, CUarray_format
+using ..CUDA: @enum_without_prefix, CUstream, CUdevice, CuDim3, CUarray, CUarray_format, @finalize_in_ctx
 using ..CUDA.APIUtils
 
 using Base: @deprecate_binding
@@ -656,9 +656,9 @@ function __unpin(ptr::Ptr{Nothing})
 
     @spinlock __pin_lock begin
         pin_count = @inbounds __pin_count[key] -= 1
-        if pin_count == 0 && CUDA.isvalid(ctx)
+        if pin_count == 0
             buf = @inbounds __pins[key]
-            Mem.unregister(buf)
+            @finalize_in_ctx ctx Mem.unregister(buf)
             delete!(__pins, key)
             delete!(__pin_count, key)
         end

lib/cudadrv/module.jl

@@ -65,9 +65,7 @@ mutable struct CuModule
 end
 
 function unsafe_unload!(mod::CuModule)
-    if isvalid(mod.ctx)
-        cuModuleUnload(mod)
-    end
+    @finalize_in_ctx mod.ctx cuModuleUnload(mod)
 end
 
 Base.unsafe_convert(::Type{CUmodule}, mod::CuModule) = mod.handle

lib/cudadrv/module/linker.jl

@@ -44,9 +44,7 @@ mutable struct CuLink
 end
 
 function unsafe_destroy!(link::CuLink)
-    if isvalid(link.ctx)
-        cuLinkDestroy(link)
-    end
+    @finalize_in_ctx link.ctx cuLinkDestroy(link)
 end
 
 Base.unsafe_convert(::Type{CUlinkState}, link::CuLink) = link.handle

lib/cudadrv/pool.jl

@@ -28,9 +28,7 @@ function CuMemoryPool(dev::CuDevice)
 end
 
 function unsafe_destroy!(pool::CuMemoryPool)
-    if isvalid(pool.ctx)
-        cuMemPoolDestroy(pool)
-    end
+    @finalize_in_ctx pool.ctx cuMemPoolDestroy(pool)
 end
 
 Base.unsafe_convert(::Type{CUmemoryPool}, pool::CuMemoryPool) = pool.handle

lib/cudadrv/stream.jl

@@ -39,9 +39,7 @@ function CuStream(; flags::CUstream_flags=STREAM_DEFAULT,
 end
 
 function unsafe_destroy!(s::CuStream)
-    if isvalid(s.ctx)
-        cuStreamDestroy_v2(s)
-    end
+    @finalize_in_ctx s.ctx cuStreamDestroy_v2(s)
 end
 
 function Base.show(io::IO, stream::CuStream)