Transition GPU support to weak-deps and extensions (#714)

Co-authored-by: Simon Byrne <[email protected]>

Author: vchuravy

Project.toml

@@ -16,7 +16,21 @@ Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 Sockets = "6462fe0b-24de-5631-8697-dd941f90decc"
 
+[extras]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+
+[weakdeps]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+
+[extensions]
+AMDGPUExt = "AMDGPU"
+CUDAExt = "CUDA"
+
 [compat]
+AMDGPU = "0.3, 0.4"
+CUDA = "3, 4"
 DocStringExtensions = "0.8, 0.9"
 MPIPreferences = "0.1.6"
 Requires = "~0.5, 1.0"

src/rocm.jl renamed to ext/AMDGPUExt.jl

@@ -1,4 +1,9 @@
-import .AMDGPU
+module AMDGPUExt
+
+import MPI
+isdefined(Base, :get_extension) ? (import AMDGPU) : (import ..AMDGPU)
+import MPI: MPIPtr, Buffer, Datatype
+
 
 function Base.cconvert(::Type{MPIPtr}, A::AMDGPU.ROCArray{T}) where T
     A
@@ -19,3 +24,5 @@ end
 function Buffer(arr::AMDGPU.ROCArray)
     Buffer(arr, Cint(length(arr)), Datatype(eltype(arr)))
 end
+
+end # AMDGPUExt

src/cuda.jl renamed to ext/CUDAExt.jl

@@ -1,4 +1,8 @@
-import .CUDA
+module CUDAExt
+
+import MPI 
+isdefined(Base, :get_extension) ? (import CUDA) : (import ..CUDA)
+import MPI: MPIPtr, Buffer, Datatype
 
 function Base.cconvert(::Type{MPIPtr}, buf::CUDA.CuArray{T}) where T
     Base.cconvert(CUDA.CuPtr{T}, buf) # returns DeviceBuffer
@@ -19,3 +23,5 @@ end
 function Buffer(arr::CUDA.CuArray)
     Buffer(arr, Cint(length(arr)), Datatype(eltype(arr)))
 end
+
+end #CUDAExt

src/MPI.jl

@@ -1,7 +1,6 @@
 module MPI
 
 using Libdl, Serialization
-using Requires
 using DocStringExtensions
 import MPIPreferences
 
@@ -80,6 +79,10 @@ include("misc.jl")
 
 include("deprecated.jl")
 
+if !isdefined(Base, :get_extension)
+    using Requires
+end
+
 function __init__()
     MPIPreferences.check_unchanged()
 
@@ -136,8 +139,10 @@ function __init__()
 
     run_load_time_hooks()
 
-    @require AMDGPU="21141c5a-9bdb-4563-92ae-f87d6854732e" include("rocm.jl")
-    @require CUDA="052768ef-5323-5732-b1bb-66c8b64840ba" include("cuda.jl")
+    @static if !isdefined(Base, :get_extension)
+        @require AMDGPU="21141c5a-9bdb-4563-92ae-f87d6854732e" include("../ext/AMDGPUExt.jl")
+        @require CUDA="052768ef-5323-5732-b1bb-66c8b64840ba" include("../ext/CUDAExt.jl")
+    end
 end
 
 end

test/runtests.jl

@@ -6,7 +6,11 @@ using MPIPreferences
 using DoubleFloats
 if get(ENV, "JULIA_MPI_TEST_ARRAYTYPE", "") == "CuArray"
     import CUDA
-    CUDA.version()
+    if isdefined(CUDA, :versioninfo)
+        CUDA.versioninfo()
+    else
+        CUDA.version()
+    end
     CUDA.precompile_runtime()
     ArrayType = CUDA.CuArray
 elseif get(ENV,"JULIA_MPI_TEST_ARRAYTYPE","") == "ROCArray"