Allow precompiling CUDA code. (#1865)

Author: maleadt

lib/cudnn/src/cuDNN.jl

@@ -151,22 +151,21 @@ end
 
 function __init__()
     precompiling = ccall(:jl_generating_output, Cint, ()) != 0
-    precompiling && return
 
     CUDA.functional() || return
 
     global libcudnn
     if CUDA_Runtime == CUDA_Runtime_jll
         if !CUDNN_jll.is_available()
-            @error "cuDNN is not available for your platform ($(Base.BinaryPlatforms.triplet(CUDNN_jll.host_platform)))"
+            precompiling || @error "cuDNN is not available for your platform ($(Base.BinaryPlatforms.triplet(CUDNN_jll.host_platform)))"
             return
         end
         libcudnn = CUDNN_jll.libcudnn
     else
         dirs = CUDA_Runtime.find_toolkit()
         path = CUDA_Runtime.get_library(dirs, "cudnn"; optional=true)
         if path === nothing
-            @error "cuDNN is not available on your system (looked in $(join(dirs, ", ")))"
+            precompiling || @error "cuDNN is not available on your system (looked in $(join(dirs, ", ")))"
             return
         end
         libcudnn = path

lib/custatevec/src/cuStateVec.jl

@@ -103,22 +103,21 @@ end
 
 function __init__()
     precompiling = ccall(:jl_generating_output, Cint, ()) != 0
-    precompiling && return
 
     CUDA.functional() || return
 
     global libcustatevec
     if CUDA_Runtime == CUDA_Runtime_jll
         if !cuQuantum_jll.is_available()
-            @error "cuQuantum is not available for your platform ($(Base.BinaryPlatforms.triplet(cuQuantum_jll.host_platform)))"
+            precompiling || @error "cuQuantum is not available for your platform ($(Base.BinaryPlatforms.triplet(cuQuantum_jll.host_platform)))"
             return
         end
         libcustatevec = cuQuantum_jll.libcustatevec
     else
         dirs = CUDA_Runtime.find_toolkit()
         path = CUDA_Runtime.get_library(dirs, "custatevec"; optional=true)
         if path === nothing
-            @error "cuQuantum is not available on your system (looked for custatevec in $(join(dirs, ", ")))"
+            precompiling || @error "cuQuantum is not available on your system (looked for custatevec in $(join(dirs, ", ")))"
             return
         end
         libcustatevec = path

lib/cutensor/src/cuTENSOR.jl

@@ -87,22 +87,21 @@ end
 
 function __init__()
     precompiling = ccall(:jl_generating_output, Cint, ()) != 0
-    precompiling && return
 
     CUDA.functional() || return
 
     global libcutensor
     if CUDA_Runtime == CUDA_Runtime_jll
         if !CUTENSOR_jll.is_available()
-            @error "cuTENSOR is not available for your platform ($(Base.BinaryPlatforms.triplet(CUTENSOR_jll.host_platform)))"
+            precompiling || @error "cuTENSOR is not available for your platform ($(Base.BinaryPlatforms.triplet(CUTENSOR_jll.host_platform)))"
             return
         end
         libcutensor = CUTENSOR_jll.libcutensor
     else
         dirs = CUDA_Runtime.find_toolkit()
         path = CUDA_Runtime.get_library(dirs, "cutensor"; optional=true)
         if path === nothing
-            @error "cuTENSOR is not available on your system (looked in $(join(dirs, ", ")))"
+            precompiling || @error "cuTENSOR is not available on your system (looked in $(join(dirs, ", ")))"
             return
         end
         libcutensor = path

lib/cutensornet/src/cuTensorNet.jl

@@ -104,22 +104,21 @@ end
 
 function __init__()
     precompiling = ccall(:jl_generating_output, Cint, ()) != 0
-    precompiling && return
 
     CUDA.functional() || return
 
     global libcutensornet
     if CUDA_Runtime == CUDA_Runtime_jll
         if !cuQuantum_jll.is_available()
-            @error "cuQuantum is not available for your platform ($(Base.BinaryPlatforms.triplet(cuQuantum_jll.host_platform)))"
+            precompiling || @error "cuQuantum is not available for your platform ($(Base.BinaryPlatforms.triplet(cuQuantum_jll.host_platform)))"
             return
         end
         libcutensornet = cuQuantum_jll.libcutensornet
     else
         dirs = CUDA_Runtime.find_toolkit()
         path = CUDA_Runtime.get_library(dirs, "cutensornet"; optional=true)
         if path === nothing
-            @error "cuQuantum is not available on your system (looked for cutensornet in $(join(dirs, ", ")))"
+            precompiling || @error "cuQuantum is not available on your system (looked for cutensornet in $(join(dirs, ", ")))"
             return
         end
         libcutensornet = path

src/initialization.jl

@@ -30,7 +30,6 @@ end
 
 function __init__()
     precompiling = ccall(:jl_generating_output, Cint, ()) != 0
-    precompiling && return
 
     # TODO: make errors here (and in submodules/subpackages like cuBLAS and cuDNN) fatal,
     #       and remove functional(), once people sufficiently use weak dependencies.
@@ -124,12 +123,15 @@ function __init__()
     end
 
     # register device overrides
-    eval(Expr(:block, overrides...))
-    empty!(overrides)
-    @require SpecialFunctions="276daf66-3868-5448-9aa4-cd146d93841b" begin
-        include("device/intrinsics/special_math.jl")
+    if !precompiling
         eval(Expr(:block, overrides...))
         empty!(overrides)
+
+        @require SpecialFunctions="276daf66-3868-5448-9aa4-cd146d93841b" begin
+            include("device/intrinsics/special_math.jl")
+            eval(Expr(:block, overrides...))
+            empty!(overrides)
+        end
     end
 
     # ensure that operations executed by the REPL back-end finish before returning,