Merge pull request #355 from JuliaParallel/sb/has_cuda

expose interface to check CUDA support

Author: vchuravy

docs/src/environment.md

@@ -9,6 +9,5 @@ MPI.Initialized
 MPI.Finalize
 MPI.Finalized
 MPI.universe_size
+MPI.has_cuda
 ```
-
-

docs/src/installation.md

@@ -34,6 +34,7 @@ controlled with the optional environment variables:
 - `JULIA_MPICC`: MPI C compiler (default: `mpicc`)
 - `JULIA_MPIEXEC`: MPI launcher command (default: `mpiexec`)
 - `JULIA_MPIEXEC_ARGS`: Additional arguments to be passed to MPI launcher (only used in the build step and tests).
+- `JULIA_MPI_HAS_CUDA`: override the [`MPI.has_cuda`](@ref) function.
 
 If your MPI installation changes (e.g. it is upgraded by the system, or you switch
 libraries), you will need to re-run `build MPI` at the package prompt.

docs/src/usage.md

@@ -33,6 +33,9 @@ If your MPI implementation has been compiled with CUDA support, then `CuArray`s
 send and receive buffers for point-to-point and collective operations (they may also work
 with one-sided operations, but these are not often supported).
 
+If using Open MPI, the status of CUDA support can be checked via the
+[`MPI.has_cuda()`](@ref) function.
+
 ## Finalizers
 
 In order to ensure MPI routines are called in the correct order at finalization time,

src/cuda.jl

@@ -2,6 +2,7 @@ import .CuArrays: CuArray
 import .CuArrays.CUDAdrv: CuPtr, synchronize
 import .CuArrays.CUDAdrv.Mem: DeviceBuffer
 
+
 function Base.cconvert(::Type{MPIPtr}, buf::CuArray{T}) where T
     Base.cconvert(CuPtr{T}, buf) # returns DeviceBuffer
 end

src/environment.jl

@@ -158,3 +158,28 @@ function Wtime()
     end
 end
 
+
+"""
+    MPI.has_cuda()
+
+Check if the MPI implementation is known to have CUDA support. Currently only Open MPI
+provides a mechanism to check, so it will return `false` with other implementations
+(unless overriden).
+
+This can be overriden by setting the `JULIA_MPI_HAS_CUDA` environment variable to `true`
+or `false`.
+"""
+function has_cuda()
+    flag = get(ENV, "JULIA_MPI_HAS_CUDA", nothing)
+    if flag === nothing
+        # Only Open MPI provides a function to check CUDA support
+        @static if startswith(MPI_LIBRARY_VERSION, "Open MPI")
+            # int MPIX_Query_cuda_support(void)
+            return 0 != ccall((:MPIX_Query_cuda_support, libmpi), Cint, ())
+        else
+            return false
+        end
+    else
+        return parse(Bool, flag)
+    end
+end

test/test_basic.jl

@@ -1,12 +1,25 @@
 using Test
 using MPI
 
+if get(ENV,"JULIA_MPI_TEST_ARRAYTYPE","") == "CuArray"
+    using CuArrays
+    ArrayType = CuArray
+else
+    ArrayType = Array
+end
+
 @test !MPI.Initialized()
 MPI.Init()
 @test MPI.Initialized()
 
 @test 0 <= MPI.Comm_rank(MPI.COMM_WORLD) < MPI.Comm_size(MPI.COMM_WORLD)
 
+@test MPI.has_cuda() isa Bool
+
+if ArrayType != Array
+    @test MPI.has_cuda()
+end
+
 @test !MPI.Finalized()
 MPI.Finalize()
 @test MPI.Finalized()