make CUDA extension for CUDACore and CUDATools instead of complete CUDA

Author: JoshuaLampert

Project.toml

@@ -56,7 +56,8 @@ UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+CUDACore = "bd0ed864-bdfe-4181-a5ed-ce625a5fdea2"
+CUDATools = "9ec180c6-1c07-47c7-9e6e-ebefa4d1f6d0"
 Convex = "f65535da-76fb-5f13-bab9-19810c17039a"
 ECOS = "e2685f51-7e38-5353-a97d-a921fd2c8199"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
@@ -67,7 +68,7 @@ SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 
 [extensions]
 TrixiAMDGPUExt = "AMDGPU"
-TrixiCUDAExt = "CUDA"
+TrixiCUDAExt = ["CUDACore", "CUDATools"]
 TrixiConvexECOSExt = ["Convex", "ECOS"]
 TrixiMakieExt = "Makie"
 TrixiNLsolveExt = "NLsolve"
@@ -80,7 +81,8 @@ AMDGPU = "2.2.1"
 AcceleratedKernels = "0.4.3"
 Accessors = "0.1.42"
 Adapt = "4.4"
-CUDA = "5.9.1, 6"
+CUDACore = "6"
+CUDATools = "6"
 CodeTracking = "1.0.5, 2, 3"
 ConstructionBase = "1.5.8"
 Convex = "0.16"

ext/TrixiCUDAExt.jl

@@ -1,7 +1,8 @@
 # Package extension for adding CUDA-based features to Trixi.jl
 module TrixiCUDAExt
 
-using CUDA: CUDA, CuArray, CuDeviceArray, CUDABackend, KernelAdaptor, @device_override
+using CUDACore: CUDACore, CuArray, CuDeviceArray, CUDABackend, KernelAdaptor, @device_override
+using CUDATools: CUDATools, NVML, has_nvml
 import Trixi
 
 function Trixi.storage_type(::Type{<:CuArray})
@@ -28,68 +29,68 @@ function Trixi.trixi_backend_info!(setup, ::CUDABackend)
     push!(setup, "Backend" => "KernelAbstractions CUDA")
     # Reimplementation of CUDA.versioninfo() to fit with Trixi's summary box format
     # see https://github.com/JuliaGPU/CUDA.jl/blob/9f56ee20afef4a770a028066d2fa9e7825d258da/src/utilities.jl#L41
-    if !CUDA.functional()
+    if !CUDACore.functional()
         push!(setup, "CUDA" => "CUDA not functional")
         return nothing
     end
 
     push!(setup, "CUDA toolchain:" => "")
-    push!(setup, "  - runtime" => string(CUDA.runtime_version()))
+    push!(setup, "  - runtime" => string(CUDACore.runtime_version()))
     push!(setup,
-          "  - toolkit" => CUDA.local_toolkit ? "local installation" :
+          "  - toolkit" => CUDACore.local_toolkit ? "local installation" :
                            "artifact installation")
 
-    if CUDA.has_nvml()
-        driver_str = string(CUDA.NVML.driver_version())
+    if has_nvml()
+        driver_str = string(NVML.driver_version())
     else
         driver_str = "unknown"
     end
-    driver_str *= " for $(CUDA.driver_version())"
+    driver_str *= " for $(CUDACore.driver_version())"
     push!(setup, "  - driver" => driver_str)
-    push!(setup, "  - compiler" => string(CUDA.compiler_version()))
+    push!(setup, "  - compiler" => string(CUDACore.compiler_version()))
 
     # Skip CUDA libraries
     # Skip Julia packages
     # Skip Toolchain
     # Skip Environment
     # Skip Preferences
 
-    devs = CUDA.devices()
+    devs = CUDACore.devices()
     if isempty(devs)
         push!(setup, "CUDA devices" => "none")
     else
         push!(setup, "CUDA devices:" => "")
     end
     for (i, dev) in enumerate(devs)
         function query_nvml()
-            mig = CUDA.uuid(dev) != CUDA.parent_uuid(dev)
-            nvml_gpu = CUDA.NVML.Device(CUDA.parent_uuid(dev))
-            nvml_dev = CUDA.NVML.Device(CUDA.uuid(dev); mig)
+            mig = CUDACore.uuid(dev) != CUDACore.parent_uuid(dev)
+            nvml_gpu = NVML.Device(CUDACore.parent_uuid(dev))
+            nvml_dev = NVML.Device(CUDACore.uuid(dev); mig)
 
-            str = CUDA.NVML.name(nvml_dev)
-            cap = CUDA.NVML.compute_capability(nvml_gpu)
-            mem = CUDA.NVML.memory_info(nvml_dev)
+            str = NVML.name(nvml_dev)
+            cap = NVML.compute_capability(nvml_gpu)
+            mem = NVML.memory_info(nvml_dev)
 
             (; str, cap, mem)
         end
 
         function query_cuda()
-            str = CUDA.name(dev)
-            cap = CUDA.capability(dev)
-            mem = CUDA.device!(dev) do
+            str = CUDACore.name(dev)
+            cap = CUDACore.capability(dev)
+            mem = CUDACore.device!(dev) do
                 # this requires a device context, so we prefer NVML
-                (free = CUDA.free_memory(), total = CUDA.total_memory())
+                (free = CUDACore.free_memory(), total = CUDACore.total_memory())
             end
             (; str, cap, mem)
         end
 
-        str, cap, mem = if CUDA.has_nvml()
+        str, cap, mem = if has_nvml()
             try
                 query_nvml()
             catch err
-                if !isa(err, CUDA.NVML.NVMLError) ||
+                if !isa(err, NVML.NVMLError) ||
                    !in(err.code,
-                       [CUDA.NVML.ERROR_NOT_SUPPORTED, CUDA.NVML.ERROR_NO_PERMISSION])
+                       [NVML.ERROR_NOT_SUPPORTED, NVML.ERROR_NO_PERMISSION])
                     rethrow()
                 end
                 query_cuda()
@@ -103,7 +104,7 @@ function Trixi.trixi_backend_info!(setup, ::CUDABackend)
 end
 
 function Trixi.trixi_device_memory_use(::CUDABackend)
-    info = CUDA.MemoryInfo()
+    info = CUDACore.MemoryInfo()
     used_bytes = info.total_bytes - info.free_bytes
     return used_bytes
 end