Take care not to spawn tasks during precompilation. (#2226)

JuliaGPU · Jan 8, 2024 · 27c3b37 · 27c3b37
1 parent 6e15ab6
commit 27c3b37
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 21 deletions.
diff --git a/lib/cublas/CUBLAS.jl b/lib/cublas/CUBLAS.jl
@@ -237,28 +237,28 @@ end
 
 function __init__()
     precompiling = ccall(:jl_generating_output, Cint, ()) != 0
-    precompiling && return
 
     if !CUDA_Runtime.is_available()
-        #@error "cuBLAS is not available"
+        #precompiling || @error "cuBLAS is not available"
         return
     end
 
     # register a log callback
-    log_cond[] = Base.AsyncCondition() do async_cond
-        blob = ""
-        while true
-            message_length = log_cursor[]
-            blob = unsafe_string(pointer(log_buffer), message_length)
-            if Threads.atomic_cas!(log_cursor, message_length, UInt(0)) == message_length
-                break
+    if !Sys.iswindows() && # NVIDIA bug #3321130 &&
+       !precompiling && (isdebug(:init, CUBLAS) || Base.JLOptions().debug_level >= 2)
+        log_cond[] = Base.AsyncCondition() do async_cond
+            blob = ""
+            while true
+                message_length = log_cursor[]
+                blob = unsafe_string(pointer(log_buffer), message_length)
+                if Threads.atomic_cas!(log_cursor, message_length, UInt(0)) == message_length
+                    break
+                end
             end
+            _log_message(blob)
+            return
         end
-        _log_message(blob)
-        return
-    end
-    if (isdebug(:init, CUBLAS) || Base.JLOptions().debug_level >= 2) &&
-       !Sys.iswindows() # NVIDIA bug #3321130
+
         callback = @cfunction(log_message, Nothing, (Cstring,))
         cublasSetLoggerCallback(callback)
     end

diff --git a/lib/cudnn/src/cuDNN.jl b/lib/cudnn/src/cuDNN.jl
@@ -176,7 +176,7 @@ function __init__()
     end
 
     # register a log callback
-    if isdebug(:init, cuDNN) || Base.JLOptions().debug_level >= 2
+    if !precompiling && (isdebug(:init, cuDNN) || Base.JLOptions().debug_level >= 2)
         log_cond[] = Base.AsyncCondition() do async_cond
             message = Base.@lock log_lock popfirst!(log_messages)
             _log_message(message...)

diff --git a/lib/custatevec/src/cuStateVec.jl b/lib/custatevec/src/cuStateVec.jl
@@ -128,7 +128,7 @@ function __init__()
     end
 
     # register a log callback
-    if isdebug(:init, cuStateVec) || Base.JLOptions().debug_level >= 2
+    if !precompiling && (isdebug(:init, cuStateVec) || Base.JLOptions().debug_level >= 2)
         callback = @cfunction(log_message, Nothing, (Int32, Cstring, Cstring))
         custatevecLoggerSetCallback(callback)
         custatevecLoggerOpenFile(Sys.iswindows() ? "NUL" : "/dev/null")

diff --git a/lib/cutensor/src/cuTENSOR.jl b/lib/cutensor/src/cuTENSOR.jl
@@ -110,7 +110,7 @@ function __init__()
     end
 
     # register a log callback
-    if isdebug(:init, cuTENSOR) || Base.JLOptions().debug_level >= 2
+    if !precompiling && (isdebug(:init, cuTENSOR) || Base.JLOptions().debug_level >= 2)
         callback = @cfunction(log_message, Nothing, (Int32, Cstring, Cstring))
         cutensorLoggerSetCallback(callback)
         cutensorLoggerOpenFile(Sys.iswindows() ? "NUL" : "/dev/null")

diff --git a/src/initialization.jl b/src/initialization.jl
@@ -168,14 +168,12 @@ function __init__()
 
     # warn about old, deprecated environment variables
     if haskey(ENV, "JULIA_CUDA_USE_BINARYBUILDER") && !local_toolkit
-        @error """JULIA_CUDA_USE_BINARYBUILDER is deprecated, and CUDA.jl always uses artifacts now.
-                  To use a local installation, use overrides or preferences to customize the artifact.
-                  Please check the CUDA.jl or Pkg.jl documentation for more details."""
+        @error """JULIA_CUDA_USE_BINARYBUILDER is deprecated. Call `CUDA.jl.set_runtime_version!` to use a local toolkit."""
         # we do not warn about this when we're already using the new preference,
         # because during the transition clusters will be deploying both mechanisms.
     end
     if haskey(ENV, "JULIA_CUDA_VERSION")
-        @error """JULIA_CUDA_VERSION is deprecated. Call `CUDA.jl.set_runtime_version!` to use a different version instead."""
+        @error """JULIA_CUDA_VERSION is deprecated. Call `CUDA.jl.set_runtime_version!` to use a different version."""
     end
 
     if !local_toolkit