Don't touch destroyed contexts when releasing their resources

CUDACore/lib/cudadrv/state.jl

@@ -153,22 +153,51 @@ function context!(ctx::CuContext)
         old_ctx = state.context
         if old_ctx != ctx
             activate(ctx)
-            dev = current_device()
-            state.device = dev
-            state.context = ctx
-            disposal_stream(ctx)    # see `prepare_cuda_state`
+            switch_context!(state, current_device(), ctx)
         end
     end
 
     return old_ctx
 end
 
+function switch_context!(state::TaskLocalState, dev::CuDevice, ctx::CuContext)
+    state.device = dev
+    state.context = ctx
+
+    # the task's stream is cached per device, but belongs to a context. after switching
+    # back, the task claims its pooled stream of that context again (see `claim_stream!`).
+    devidx = deviceid(dev)+1
+    current = state.streams[devidx]
+    if current !== nothing && current.ctx !== nothing && current.ctx != ctx
+        state.streams[devidx] = nothing
+    end
+
+    disposal_stream(ctx)    # see `prepare_cuda_state`
+end
+
 @inline function context!(f::F, ctx::CuContext) where {F<:Function}
-    old_ctx = context!(ctx)::Union{CuContext,Nothing}
+    state = task_local_state()
+    if state === nothing
+        context!(ctx)
+        try
+            return f()
+        finally
+            # the task didn't use CUDA before, so let it start over
+            delete!(task_local_storage(), :CUDA)
+        end
+    elseif state.context == ctx
+        return f()
+    end
+
+    # restore the task's stream too, which may not be a pooled one (see `stream!`)
+    devidx = deviceid(state.device)+1
+    old_stream = state.streams[devidx]
+    old_ctx = context!(ctx)::CuContext
     try
         f()
     finally
-        old_ctx !== nothing && old_ctx != ctx && context!(old_ctx)
+        context!(old_ctx)
+        state.streams[devidx] = old_stream
     end
 end
 
@@ -251,9 +280,7 @@ function device!(dev::CuDevice, flags=nothing)
     if state === nothing
         task_local_state!(dev)
     else
-        state.device = dev
-        state.context = ctx
-        disposal_stream(ctx)        # see `prepare_cuda_state`
+        switch_context!(state, dev, ctx)
     end
     activate(ctx)
 

CUDACore/src/compiler/exceptions.jl

@@ -28,6 +28,8 @@ end
 # check the exception flags on every API call, similarly to how CUDA handles errors
 function check_exceptions()
     for (ctx,mem) in exception_infos
+        # destroying a context frees its exception flag
+        isvalid(ctx) || continue
         exception_info = convert(ExceptionInfo, mem)
         if exception_info.status != 0
             # restore the structure

CUDACore/src/graph.jl

@@ -209,7 +209,10 @@ end
 Base.unsafe_convert(::Type{CUgraph}, graph::CuGraph) = graph.handle
 
 function release_now(graph::CuGraph; destroy=cuGraphDestroy)
-    context!(graph.ctx) do
+    # graphs outlive the context they were created in
+    if isvalid(graph.ctx)
+        context!(() -> destroy(graph), graph.ctx)
+    else
         destroy(graph)
     end
     # (not reached when destroying failed: the resource coordinator then retains the object,
@@ -809,7 +812,12 @@ end
 function free_allocations(graph_allocations::GraphAllocations)
     (; ctx, allocations, stream, generation) = graph_allocations
     @lock stream_disposal_lock begin
-        stream, ctx = release_stream(stream, ctx, generation)
+        released = release_stream(stream, ctx, generation)
+        # XXX: graph allocations outlive the context they were made in, like allocations
+        #      from a memory pool (see `free_context`), but the device to free them on
+        #      isn't known here, so they leak when that context has been destroyed
+        released === nothing && return
+        stream, ctx = released
         context!(ctx) do
             for ptr in allocations
                 cuMemFreeAsync(ptr, stream)
@@ -820,10 +828,13 @@ function free_allocations(graph_allocations::GraphAllocations)
 end
 
 function release_now(exec::CuGraphExec; destroy=cuGraphExecDestroy)
-    context!(exec.ctx) do
-        allocations = launched_allocations(exec)
-        allocations === nothing || release_now(allocations)
-        # (destroying an executable graph that is still executing is allowed)
+    allocations = launched_allocations(exec)
+    allocations === nothing || release_now(allocations)
+    # (destroying an executable graph that is still executing is allowed)
+    if isvalid(exec.ctx)
+        context!(() -> destroy(exec), exec.ctx)
+    else
+        # executable graphs outlive the context they were instantiated in
         destroy(exec)
     end
     # (not reached when destroying failed: the resource coordinator then retains the object,

CUDACore/src/resources/backends.jl

@@ -139,8 +139,10 @@ end
       # free in stream order. `cuMemFree` would wait for all work on the device to finish,
       # blocking kernel launches from other threads in the meantime. that also works for
       # memory that wasn't allocated from a pool.
-      @lock stream_disposal_lock begin
-        handle, ctx = release_stream(stream, stream_ctx, generation)
+      free_ctx = free_context(mem)
+      free_ctx === nothing || @lock stream_disposal_lock begin
+        handle, ctx = @something(release_stream(stream, stream_ctx, generation),
+                                 (disposal_stream(free_ctx), free_ctx))
         context!(ctx) do
           free_stream = if !mem.async
             # when memory that wasn't allocated from a pool is freed in stream order,
@@ -156,20 +158,18 @@ end
     else
       # without support for freeing memory in stream order, `cuMemFree` is the only option,
       # which waits for all running kernels. defer that until memory is reclaimed.
-      destroy_later(mem) do mem
-        context!(mem.ctx) do
-          free(mem)
-        end
-      end
+      defer_release(free, mem; ctx=mem.ctx, blocking=true)
     end
     owned_allocation && account!(memory_stats(mem.dev), -sizeof(mem))
 end
 @inline function _pool_free(mem::Union{UnifiedMemory,HostMemory}, stream::CuStream,
                             stream_ctx::CuContext, generation::Int, owned_allocation::Bool,
                             attachment::Attachment)
   if mem.pooled
-    @lock stream_disposal_lock begin
-      stream, ctx = release_stream(stream, stream_ctx, generation)
+    free_ctx = free_context(mem)
+    free_ctx === nothing || @lock stream_disposal_lock begin
+      stream, ctx = @something(release_stream(stream, stream_ctx, generation),
+                               (disposal_stream(free_ctx), free_ctx))
       context!(ctx) do
         cuMemFreeAsync(convert(CuPtr{Cvoid}, mem), stream)
       end
@@ -187,6 +187,21 @@ end
   owned_allocation && account!(_host_stats, -sizeof(mem))
 end
 
+# the context to free memory in, or `nothing` if it was freed along with the context it was
+# allocated in. allocations from a memory pool outlive that context, so are freed in another.
+function free_context(mem::DeviceMemory)
+  isvalid(mem.ctx) && return mem.ctx
+  mem.async || return nothing
+  ctx = context(mem.dev)
+  return isvalid(ctx) ? ctx : nothing
+end
+function free_context(mem::Union{UnifiedMemory,HostMemory})
+  isvalid(mem.ctx) && return mem.ctx
+  mem.pooled || return nothing
+  ctx = context(device())
+  return isvalid(ctx) ? ctx : nothing
+end
+
 # freed memory may only be released when reclaiming memory, so do so when running out
 function alloc_or_reclaim(::Type{M}, sz, args...) where {M}
   mem = retry_reclaim(isnothing) do
@@ -291,8 +306,11 @@ cached_size(sz) = sz <= 1<<20 ? nextpow(2, max(sz, 512)) : cld(sz, 1<<20) << 20
 
 function cache_put!(cache::BlockCache{M}, mem::M, stream::CuStream,
                     stream_ctx::CuContext, generation::Int; attach_host::Bool=false) where {M}
+  # (freed along with its context)
+  isvalid(mem.ctx) || return
   idle = @lock stream_disposal_lock begin
-    stream, ctx = release_stream(stream, stream_ctx, generation)
+    stream, ctx = @something(release_stream(stream, stream_ctx, generation),
+                             (disposal_stream(mem.ctx), mem.ctx))
     # (the event needs to be created in the context of the stream it is recorded on)
     context!(ctx) do
       if attach_host
@@ -318,7 +336,10 @@ function cache_take!(cache::BlockCache{M}, ctx::CuContext, sz::Int) where {M}
   block = @lock cache.lock begin
     blocks = get(cache.blocks, (ctx, sz), nothing)
     blocks === nothing && return nothing
-    i = findlast(block -> isdone(block.idle), blocks)
+    # (memory may have been last used in another context, whose destruction ended the work)
+    i = findlast(blocks) do block
+      block.idle.ctx == ctx || isvalid(block.idle.ctx) ? isdone(block.idle) : true
+    end
     i === nothing && return nothing
     Base.@atomic cache.bytes -= sz
     popat!(blocks, i)
@@ -343,7 +364,9 @@ function cache_release!(cache::BlockCache)
     end
     block === nothing && break
     attempt_release(ReleaseAction(block, nothing, true) do block
-      synchronize(block.idle)
+      # (freed along with its context, whose work has ended too)
+      isvalid(block.mem.ctx) || return
+      isvalid(block.idle.ctx) && synchronize(block.idle)
       context!(block.mem.ctx) do
         free(block.mem)
       end

CUDACore/src/resources/lifecycle.jl

@@ -113,7 +113,8 @@ without calling `f` from the garbage collector. Many CUDA API calls that release
 wait for all running kernels to finish, which can stall or deadlock the thread the garbage
 collector happens to run on. Instead, the finalizer only queues `obj`, and a regular task
 calls `f(obj)` later on, in the context `ctx` that was active when registering. Pass
-`ctx=nothing` if `f` activates the right context itself.
+`ctx=nothing` if `f` activates the right context itself. If `ctx` has been destroyed by
+then, `f` is not called, as destroying a context releases its resources.
 
 With `blocking=true` (the default), `f` is only called when CUDA.jl reclaims memory, i.e.,
 when an allocation runs out of memory, or when calling `CUDA.reclaim()`. Only use
@@ -148,10 +149,12 @@ end
 
 destroy_later(f, obj) = defer_release(f, obj; blocking=true)
 
+# destroying a context releases its resources, and their handles may be reused afterwards,
+# so resources of a destroyed context are dropped instead of released
 function run_release(action::ReleaseAction)
   if action.ctx === nothing
     Base.invokelatest(action.f, action.obj)
-  else
+  elseif isvalid(action.ctx)
     context!(action.ctx) do
       Base.invokelatest(action.f, action.obj)
     end
@@ -238,7 +241,8 @@ function release_held!(; reclaim=false, select=Returns(true))
       continue
     end
     try
-      if gate !== nothing && !isdone(gate)
+      # (the work of a destroyed context has ended)
+      if gate !== nothing && isvalid(gate.ctx) && !isdone(gate)
         hold!(held)
         continue
       end

CUDACore/src/resources/managed.jl

@@ -220,11 +220,17 @@ function prepare_host_access!(managed::Managed)
           stream = disposal_stream(ctx)
         else
           fence = nothing
-          if !recycled(managed) && isvalid(managed.stream)
+          # this also handles owners that have been recycled or destroyed
+          released = release_stream(managed.stream, managed.stream_ctx, managed.generation)
+          if released === nothing
+            # the stream's context was destroyed, so attach in the one the memory lives in
+            # (memory that was freed along with it can't be accessed anymore)
+            ctx = @something(free_context(managed.mem), throw(UndefRefError()))
+            released = (disposal_stream(ctx), ctx)
+          elseif !recycled(managed) && isvalid(managed.stream)
             check_capture(managed.stream)
           end
-          # this also handles owners that have been recycled or destroyed
-          stream, ctx = release_stream(managed.stream, managed.stream_ctx, managed.generation)
+          stream, ctx = released
         end
         context!(ctx) do
           fence === nothing || cuStreamWaitEvent(stream, fence::CuEvent, 0)

CUDACore/src/resources/pins.jl

@@ -17,15 +17,15 @@ const pending_registration_bytes = Threads.Atomic{Int}(0)
 
 function release_now(reg::RetiredRegistration)
   Threads.atomic_add!(pending_registration_bytes, sizeof(reg.mem))
-  hold!(ReleaseAction(unregister_retired!, reg, reg.mem.ctx, true))
+  hold!(ReleaseAction(unregister_retired!, reg, nothing, true))
 end
 release_registration(reg::RetiredRegistration) = discard(reg)
 
 function unregister_retired!(reg::RetiredRegistration)
   if reg.counted
     __unpin(pointer(reg.mem), reg.mem.ctx)
-  else
-    unregister(reg.mem)
+  elseif isvalid(reg.mem.ctx)   # (a destroyed context has no registrations left)
+    context!(() -> unregister(reg.mem), reg.mem.ctx)
   end
   Threads.atomic_sub!(pending_registration_bytes, sizeof(reg.mem))
   return
@@ -193,7 +193,8 @@ function __unpin(ptr::Ptr, ctx::CuContext)
         pop!(__pinned_memory, key)
     end
     try
-        context!(ctx) do
+        # (a destroyed context has no registrations left)
+        isvalid(ctx) && context!(ctx) do
             unregister(mem)
         end
     catch

CUDACore/src/resources/streams.jl

@@ -1,4 +1,8 @@
-release_now(event::CuEvent) = unsafe_destroy!(event)
+function release_now(event::CuEvent)
+  # (resources of a destroyed context are dropped, see `run_release`)
+  isvalid(event.ctx) && unsafe_destroy!(event)
+  return
+end
 
 # memory that was last used on a stream may be retired after that stream, and released
 # after the stream has been destroyed. it then needs to wait for the work that was submitted
@@ -7,6 +11,11 @@ const stream_disposal_lock = ReentrantLock()
 function release_now(stream::CuStream)
   @lock stream_disposal_lock begin
     isvalid(stream) || return
+    if !isvalid(stream.ctx)
+      # destroyed along with its context, which also ended its work
+      Base.@atomic stream.valid = false
+      return
+    end
     context!(stream.ctx) do
       event = CuEvent(EVENT_DISABLE_TIMING)
       cuEventRecord(event, stream)
@@ -32,6 +41,8 @@ function disposal_stream(ctx::CuContext)
   @lock disposal_streams_lock get!(disposal_streams, ctx) do
     # this is also when CUDA is first used, so start releasing retired resources periodically
     start_retired_drainer()
+    # forget the streams of destroyed contexts, which were destroyed along with them
+    filter!(((other, _),) -> isvalid(other), disposal_streams)
     context!(ctx) do
       handle = Ref{CUstream}()
       cuStreamCreate(handle, STREAM_NON_BLOCKING)
@@ -41,9 +52,11 @@ function disposal_stream(ctx::CuContext)
 end
 
 # a stream to release memory on that was last used on `stream` (during `generation`), along
-# with its context. needs to be called with `stream_disposal_lock` held, which also keeps
-# the stream from being handed to another task (see `claim_stream!`).
+# with its context, or `nothing` if that context has been destroyed, which ended the work.
+# needs to be called with `stream_disposal_lock` held, which also keeps the stream from
+# being handed to another task (see `claim_stream!`).
 function release_stream(stream::CuStream, ctx::CuContext, generation::Int)
+  isvalid(something(stream.ctx, ctx)) || return nothing
   if generation != CUDACore.generation(stream)
     # the stream has been handed to another task, so the work has finished. don't use the
     # stream, which would wait for the new owner's work, or end up in its capture.
@@ -81,7 +94,8 @@ end
 on_per_thread_stream(stream::CuStream) = stream.handle == CU_STREAM_PER_THREAD
 function synchronize_and(f, ctx::CuContext)
   return x -> begin
-    context!(device_synchronize, ctx)
+    # (the work of a destroyed context has ended)
+    isvalid(ctx) && context!(device_synchronize, ctx)
     f(x)
   end
 end
@@ -106,7 +120,10 @@ function release_now(retired::RetiredOwner)
   end
   event = try
     @lock stream_disposal_lock begin
-      handle, ctx = release_stream(stream, ctx, retired.managed.generation)
+      released = release_stream(stream, ctx, retired.managed.generation)
+      # (the work of a destroyed context has ended)
+      released === nothing && return
+      handle, ctx = released
       context!(ctx) do
         event = CuEvent(EVENT_DISABLE_TIMING)
         cuEventRecord(event, handle)