Use a function decorator macro to transition back into GC unsafe domain.

Author: maleadt

lib/cublas/CUBLAS.jl

@@ -210,12 +210,7 @@ function log_message(ptr)
     return
 end
 
-function _log_message(blob)
-    # see @gcsafe_ccall documentation
-    @static if VERSION < v"1.9"
-        GC.safepoint()
-    end
-
+@gcunsafe_callback function _log_message(blob)
     # the message format isn't documented, but it looks like a message starts with a capital
     # and the severity (e.g. `I!`), and subsequent lines start with a lowercase mark (`!i`)
     #

lib/cudadrv/occupancy.jl

@@ -36,11 +36,7 @@ end
 # HACK: callback function for `launch_configuration` on platforms without support for
 #       trampolines as used by `@cfunction` (JuliaLang/julia#27174, JuliaLang/julia#32154)
 _shmem_cb = nothing
-function _shmem_cint_cb(x::Cint)
-    # see @gcsafe_ccall documentation
-    @static if VERSION < v"1.9"
-        GC.safepoint()
-    end
+@gcunsafe_callback function _shmem_cint_cb(x::Cint)
     Cint(something(_shmem_cb)(x))
 end
 _shmem_cb_lock = Threads.ReentrantLock()
@@ -64,11 +60,7 @@ function launch_configuration(fun::CuFunction; shmem::Union{Integer,Base.Callabl
     if isa(shmem, Integer)
         cuOccupancyMaxPotentialBlockSize(blocks_ref, threads_ref, fun, C_NULL, shmem, max_threads)
     elseif Sys.ARCH == :x86 || Sys.ARCH == :x86_64
-        function shmem_cint(threads)
-            # see @gcsafe_ccall documentation
-            @static if VERSION < v"1.9"
-                GC.safepoint()
-            end
+        @gcunsafe_callback function shmem_cint(threads)
             Cint(shmem(threads))
         end
         cb = @cfunction($shmem_cint, Cint, (Cint,))

lib/cudnn/src/cuDNN.jl

@@ -137,12 +137,7 @@ function log_message(sev, udata, dbg_ptr, ptr)
     return
 end
 
-function _log_message(sev, dbg, str)
-    # see @gcsafe_ccall documentation
-    @static if VERSION < v"1.9"
-        GC.safepoint()
-    end
-
+@gcunsafe_callback function _log_message(sev, dbg, str)
     lines = split(str, '\0')
     msg = join(lines, '\n')
     if sev == CUDNN_SEV_INFO

lib/cupti/wrappers.jl

@@ -12,13 +12,8 @@ end
 # multiple subscribers aren't supported, so make sure we only call CUPTI once
 const callback_lock = ReentrantLock()
 
-function callback(userdata::Ptr{Cvoid}, domain::CUpti_CallbackDomain,
-                  id::CUpti_CallbackId, data_ptr::Ptr{Cvoid})
-    # see @gcsafe_ccall documentation
-    @static if VERSION < v"1.9"
-        GC.safepoint()
-    end
-
+@gcunsafe_callback function callback(userdata::Ptr{Cvoid}, domain::CUpti_CallbackDomain,
+                                     id::CUpti_CallbackId, data_ptr::Ptr{Cvoid})
     cfg = Base.unsafe_pointer_to_objref(userdata)::CallbackConfig
 
     # decode the callback data
@@ -131,15 +126,10 @@ end
 const activity_lock = ReentrantLock()
 const activity_config = Ref{Union{Nothing,ActivityConfig}}(nothing)
 
-function request_buffer(dest_ptr, sz_ptr, max_num_records_ptr)
+@gcunsafe_callback function request_buffer(dest_ptr, sz_ptr, max_num_records_ptr)
     # this function is called by CUPTI, but directly from the application, so it should be
     # fine to perform I/O or allocate memory here.
 
-    # see @gcsafe_ccall documentation
-    @static if VERSION < v"1.9"
-        GC.safepoint()
-    end
-
     dest = Base.unsafe_wrap(Array, dest_ptr, 1)
     sz = Base.unsafe_wrap(Array, sz_ptr, 1)
     max_num_records = Base.unsafe_wrap(Array, max_num_records_ptr, 1)
@@ -167,7 +157,7 @@ function request_buffer(dest_ptr, sz_ptr, max_num_records_ptr)
     return
 end
 
-function complete_buffer(ctx_handle, stream_id, buf_ptr, sz, valid_sz)
+@gcunsafe_callback function complete_buffer(ctx_handle, stream_id, buf_ptr, sz, valid_sz)
     # this function is called by a CUPTI worker thread while our application may be waiting
     # for `cuptiActivityFlushAll` to complete. that means we cannot do I/O here, or we could
     # yield while the application cannot make any progress.
@@ -176,11 +166,6 @@ function complete_buffer(ctx_handle, stream_id, buf_ptr, sz, valid_sz)
     # to prevent this, we call `sizehint!` in `request_buffer`.
     # XXX: `sizehint!` isn't a guarantee; use `resize!` and a cursor?
 
-    # see @gcsafe_ccall documentation
-    @static if VERSION < v"1.9"
-        GC.safepoint()
-    end
-
     cfg = activity_config[]
     if cfg !== nothing
         push!(cfg.results, (ctx_handle, stream_id, buf_ptr, sz, valid_sz))

lib/custatevec/src/cuStateVec.jl

@@ -87,12 +87,7 @@ end
 
 ## logging
 
-function log_message(log_level, function_name, message)
-    # see @gcsafe_ccall documentation
-    @static if VERSION < v"1.9"
-        GC.safepoint()
-    end
-
+@gcunsafe_callback function log_message(log_level, function_name, message)
     function_name = unsafe_string(function_name)
     message = unsafe_string(message)
     output = if isempty(message)

lib/cutensor/src/cuTENSOR.jl

@@ -70,12 +70,7 @@ end
 
 ## logging
 
-function log_message(log_level, function_name, message)
-    # see @gcsafe_ccall documentation
-    @static if VERSION < v"1.9"
-        GC.safepoint()
-    end
-
+@gcunsafe_callback function log_message(log_level, function_name, message)
     function_name = unsafe_string(function_name)
     message = unsafe_string(message)
     output = if isempty(message)

lib/cutensornet/src/cuTensorNet.jl

@@ -88,12 +88,7 @@ end
 
 ## logging
 
-function log_message(log_level, function_name, message)
-    # see @gcsafe_ccall documentation
-    @static if VERSION < v"1.9"
-        GC.safepoint()
-    end
-
+@gcunsafe_callback function log_message(log_level, function_name, message)
     function_name = unsafe_string(function_name)
     message = unsafe_string(message)
     output = if isempty(message)

lib/utils/call.jl

@@ -1,6 +1,7 @@
 # utilities for calling foreign functionality more conveniently
 
-export @checked, with_workspace, with_workspaces, @debug_ccall, @gcsafe_ccall
+export @checked, with_workspace, with_workspaces,
+       @debug_ccall, @gcsafe_ccall, @gcunsafe_callback
 
 
 ## function wrapper for checking the return value of a function
@@ -166,10 +167,8 @@ render_arg(io, arg::Union{<:Base.RefValue, AbstractArray}) = summary(io, arg)
 # TODO: replace with JuliaLang/julia#49933 once merged
 
 # note that this is generally only safe with functions that do not call back into Julia.
-# when callbacks occur, the code should ensure the GC is not running:
-# - on 1.10 and later, everything is fine because of safepoint_on_entry
-# - on 1.9, @cfunction-based callbacks are fine because they transition to gc_unsafe
-# - on 1.8 and earlier, the code should explicitly call GC.safepoint()!
+# when callbacks occur, the code should ensure the GC is not running by wrapping the code
+# in the `@gcunsafe` macro
 
 function ccall_macro_lower(func, rettype, types, args, nreq)
     # instead of re-using ccall or Expr(:foreigncall) to perform argument conversion,
@@ -213,7 +212,50 @@ function ccall_macro_lower(func, rettype, types, args, nreq)
     end
 end
 
+"""
+    @gcsafe_ccall ...
+
+Call a foreign function just like `@ccall`, but marking it safe for the GC to run. This is
+useful for functions that may block, so that the GC isn't blocked from running, but may also
+be required to prevent deadlocks (see JuliaGPU/CUDA.jl#2261).
+
+Note that this is generally only safe with non-Julia C functions that do not call back into
+Julia. When using callbacks, the code should make sure to transition back into GC-unsafe
+mode using the `@gcunsafe` macro.
+"""
 macro gcsafe_ccall(expr)
     ccall_macro_lower(Base.ccall_macro_parse(expr)...)
 end
 
+"""
+    @gcunsafe_callback function callback(...)
+        ...
+    end
+
+Mark a callback function as unsafe for the GC to run. This is normally the default for
+Julia code, and is meant to be used in combination with `@gcsafe_ccall`.
+"""
+macro gcunsafe_callback(ex)
+    if VERSION >= v"1.9"
+        # on 1.9+, `@cfunction` already transitions to GC-unsafe mode
+        return esc(ex)
+    end
+
+    # parse the function definition
+    @assert Meta.isexpr(ex, :function)
+    sig = ex.args[1]
+    @assert Meta.isexpr(sig, :call)
+    body = ex.args[2]
+    @assert Meta.isexpr(body, :block)
+
+    gcunsafe_body = quote
+        gc_state = @ccall(jl_gc_unsafe_enter()::Int8)
+        try
+            $(ex)
+        finally
+            @ccall(jl_gc_unsafe_leave(gc_state::Int8)::Cvoid)
+        end
+    end
+
+    return esc(Expr(:function, sig, gcunsafe_body))
+end