Always inspect the task-local context when verifying before freeing. (#1462)

For some reason, the thread-bound context can become desynchronized
from the task-local one. Generally we don't notice this, because we
prefix every API call with a call that synchronizes both. Here, however,
we explicitly didn't to avoid initializing the state as that was thought
to cause the kind of initialization that may have to yield (which is
unsupported when done so from a finalizer).

However, just creating the task local state shouldn't result in yield,
only creating a stream does, like querying `active_state` as was done
before #1383.

Author: maleadt

lib/cufft/fft.jl

@@ -25,6 +25,11 @@ abstract type CuFFTPlan{T<:cufftNumber, K, inplace} <: Plan{T} end
 Base.convert(::Type{cufftHandle}, p::CuFFTPlan) = p.handle
 
 function CUDA.unsafe_free!(plan::CuFFTPlan, stream::CuStream=stream())
+    # verify that the caller has switched contexts
+    if plan.ctx != context()
+      error("Trying to free $plan from an unrelated context")
+    end
+
     cufftDestroy(plan)
     unsafe_free!(plan.workarea, stream)
 end

src/pool.jl

@@ -351,17 +351,9 @@ Releases a buffer `buf` to the memory pool.
   return
 end
 @inline function _free(buf::Mem.DeviceBuffer; stream::Union{Nothing,CuStream})
-    # NOTE: this function is often called from finalizers, from which we can't switch tasks,
-    #       so we need to take care not to call managed functions (i.e. functions that may
-    #       initialize the CUDA context) because querying the active context using
-    #       `current_context()` takes a lock
-
-    # verify that the caller has called `context!` already, which eagerly activates the
-    # context (i.e. doesn't only set it in the state, but configures the CUDA APIs).
-    handle_ref = Ref{CUcontext}()
-    cuCtxGetCurrent(handle_ref)
-    if buf.ctx.handle != handle_ref[]
-      error("Trying to free $buf from a different context than the one it was allocated from ($(handle_ref[]))")
+    # verify that the caller has switched contexts
+    if buf.ctx != context()
+      error("Trying to free $buf from an unrelated context")
     end
 
     dev = current_device()