Support using a local compiler through CUDA_Runtime_Discovery (#3080)

Author: apozharski

CUDACore/Project.toml

@@ -45,7 +45,7 @@ BFloat16s = "0.5, 0.6"
 CEnum = "0.2, 0.3, 0.4, 0.5"
 CUDA_Compiler_jll = "0.3, 0.4"
 CUDA_Driver_jll = "13"
-CUDA_Runtime_Discovery = "1"
+CUDA_Runtime_Discovery = "2"
 CUDA_Runtime_jll = "0.21"
 ChainRulesCore = "1"
 EnzymeCore = "0.8.2"

CUDACore/lib/cudadrv/version.jl

@@ -100,7 +100,7 @@ other tools. This is version separately from the CUDA Runtime, in order to ensur
 compatibility with the driver, and make sure we use the latest compatible version regardless
 of the selected runtime.
 """
-compiler_version() = CUDA_Compiler_jll.cuda_version
+compiler_version() = CUDA_Compiler.cuda_version
 
 
 ## helpers

CUDACore/src/CUDACore.jl

@@ -31,8 +31,6 @@ using LLVMLoopInfo
 
 using CUDA_Driver_jll
 
-using CUDA_Compiler_jll
-
 import CUDA_Runtime_jll
 const local_toolkit = CUDA_Runtime_jll.host_platform["cuda_local"] == "true"
 if local_toolkit
@@ -44,6 +42,15 @@ else
 end
 
 import Preferences
+const local_compiler = Preferences.@load_preference("local_compiler", "false") == "true"
+
+if local_compiler
+    using CUDA_Runtime_Discovery
+    const CUDA_Compiler = CUDA_Runtime_Discovery
+else
+    import CUDA_Compiler_jll
+    const CUDA_Compiler = CUDA_Compiler_jll
+end
 
 using Libdl
 

CUDACore/src/compiler/compilation.jl

@@ -32,7 +32,7 @@ function GPUCompiler.link_libraries!(@nospecialize(job::CUDACompilerJob), mod::L
         return
     end
 
-    lib = parse(LLVM.Module, read(libdevice))
+    lib = parse(LLVM.Module, read(CUDA_Compiler.libdevice))
 
     # override libdevice's triple and datalayout to avoid warnings
     triple!(lib, triple(mod))
@@ -339,7 +339,7 @@ function compile(@nospecialize(job::CompilerJob))
         "--output-file", ptxas_output,
         ptx_input
     ])
-    proc, log = run_and_collect(`$(ptxas()) $ptxas_opts`)
+    proc, log = run_and_collect(`$(CUDA_Compiler.ptxas()) $ptxas_opts`)
     log = strip(log)
     if !success(proc)
         reason = proc.termsignal > 0 ? "ptxas received signal $(proc.termsignal)" :
@@ -370,12 +370,12 @@ function compile(@nospecialize(job::CompilerJob))
         append!(nvlink_opts, [
             "--verbose", "--extra-warnings",
             "--arch", arch,
-            "--library-path", dirname(libcudadevrt),
+            "--library-path", dirname(CUDA_Compiler.libcudadevrt),
             "--library", "cudadevrt",
             "--output-file", nvlink_output,
             ptxas_output
         ])
-        proc, log = run_and_collect(`$(nvlink()) $nvlink_opts`)
+        proc, log = run_and_collect(`$(CUDA_Compiler.nvlink()) $nvlink_opts`)
         log = strip(log)
         if !success(proc)
             reason = proc.termsignal > 0 ? "nvlink received signal $(proc.termsignal)" :

CUDATools/src/reflection.jl

@@ -126,7 +126,7 @@ function disassemble_cubin(io::IO, cubin::Vector{Cchar}; raw::Bool)
         write(cubin_io, cubin)
         flush(cubin_io)
 
-        cmd = `$(CUDA_Compiler_jll.nvdisasm()) --print-code --print-line-info $cubin_path`
+        cmd = `$(CUDACore.CUDA_Compiler.nvdisasm()) --print-code --print-line-info $cubin_path`
         for line in readlines(cmd)
             if !raw
                 # nvdisasm output is pretty verbose;

CUDATools/src/utilities.jl

@@ -21,7 +21,12 @@ function versioninfo(io::IO=stdout)
         print(io, "- unknown driver")
     end
     println(io, " for $(driver_version().major).$(driver_version().minor)")
-    println(io, "- compiler $(compiler_version().major).$(compiler_version().minor)")
+    print(io, "- compiler $(compiler_version().major).$(compiler_version().minor), ")
+    if CUDACore.local_compiler
+        println(io, "local installation")
+    else
+        println(io, "artifact installation")
+    end
     println(io)
 
     println(io, "CUDA libraries: ")

lib/cublas/Project.toml

@@ -29,7 +29,7 @@ Adapt = "4.4"
 BFloat16s = "0.5, 0.6"
 CEnum = "0.2, 0.3, 0.4, 0.5"
 CUDACore = "=6.0.0"
-CUDA_Runtime_Discovery = "1"
+CUDA_Runtime_Discovery = "2"
 CUDA_Runtime_jll = "0.21"
 EnzymeCore = "0.8.2"
 GPUArrays = "11.4.1"

lib/cudnn/Project.toml

@@ -15,6 +15,6 @@ CUDNN_jll = "62b44479-cb7b-5706-934f-f13b2eb2e645"
 [compat]
 CEnum = "0.2, 0.3, 0.4, 0.5"
 CUDACore = "=6.0.0"
-CUDA_Runtime_Discovery = "0.2, 0.3, 1"
+CUDA_Runtime_Discovery = "2"
 CUDNN_jll = "9.20"
 julia = "1.10"

lib/cufft/Project.toml

@@ -20,7 +20,7 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 AbstractFFTs = "0.5, 1.0"
 CEnum = "0.2, 0.3, 0.4, 0.5"
 CUDACore = "=6.0.0"
-CUDA_Runtime_Discovery = "1"
+CUDA_Runtime_Discovery = "2"
 CUDA_Runtime_jll = "0.21"
 GPUToolbox = "0.3, 1"
 LinearAlgebra = "1"

lib/cupti/Project.toml

@@ -13,7 +13,7 @@ GPUToolbox = "096a3bc2-3ced-46d0-87f4-dd12716f4bfc"
 [compat]
 CEnum = "0.2, 0.3, 0.4, 0.5"
 CUDACore = "=6.0.0"
-CUDA_Runtime_Discovery = "1"
+CUDA_Runtime_Discovery = "2"
 CUDA_Runtime_jll = "0.21"
 GPUToolbox = "1.1"
 julia = "1.10"