diff --git a/lib/cufft/fft.jl b/lib/cufft/fft.jl index e66e576480..320fecb58f 100644 --- a/lib/cufft/fft.jl +++ b/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 diff --git a/src/pool.jl b/src/pool.jl index b627937de3..911a5682aa 100644 --- a/src/pool.jl +++ b/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()