Skip to content

Commit

Permalink
Always inspect the task-local context when verifying before freeing.
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
maleadt committed May 9, 2022
1 parent 554dcc4 commit 9928855
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 11 deletions.
5 changes: 5 additions & 0 deletions lib/cufft/fft.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
14 changes: 3 additions & 11 deletions src/pool.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down

0 comments on commit 9928855

Please sign in to comment.