Repository: JuliaParallel/MPI.jl
Pull request #821 (2024-02-25 20:47:34)

Add `has_rocm` for OpenMPI
See https://docs.open-mpi.org/en/main/man-openmpi/man3/MPIX_Query_rocm_support.3.html#mpix-query-rocm-support

Seems to be available from v5 of OpenMPI

giordano: Can also document `has_rocm` in https://juliaparallel.org/MPI.jl/stable/knownissues/#Hints-to-ensure-ROCm-aware-MPI-to-be-functional, like `has_cuda` is documented in the CUDA section.
giordano: @avik-pal are you interested in moving this forward?
avik-pal: Completely fell off my radar, I will finish this over the weekend 
JBlaschke: Once this is merged, I can implement `has_gpu`; `has_cuda` and `has_rocm` for Cray MPICH
luraess: Maybe one could add a sentence similar to CUDA (https://github.com/JuliaParallel/MPI.jl/blob/master/docs/src/usage.md) in the ROCm-aware section of the docs? 
avik-pal: > Maybe one could add a sentence similar to CUDA ([`master`/docs/src/usage.md](https://github.com/JuliaParallel/MPI.jl/blob/master/docs/src/usage.md?rgh-link-date=2024-06-21T21%3A20%3A49Z))
> 
> ```
> If using OpenMPI, the status of CUDA support can be checked via the [MPI.has_cuda()](https://github.com/JuliaParallel/MPI.jl/blob/master/docs/src/@ref) function.
> ```
> 
> in the ROCm-aware section of the docs?

done
giordano: Thanks!
giordano: > Can also document `has_rocm` in https://juliaparallel.org/MPI.jl/stable/knownissues/#Hints-to-ensure-ROCm-aware-MPI-to-be-functional, like `has_cuda` is documented in the CUDA section.

Last documentation request and the we should be good to go ☝️ 
avik-pal: Updated the docs
vchuravy: Can you erase on master?
avik-pal: > Can you erase on master?

What do you want me to erase?

giordano: I presume that was meant to be "rebase".

Patch:
diff --git a/docs/src/knownissues.md b/docs/src/knownissues.md
index d3e6a2a..6fc71ed 100644
--- a/docs/src/knownissues.md
+++ b/docs/src/knownissues.md
@@ -180,7 +180,7 @@ Make sure to:
   ```
 - Then in Julia, upon loading MPI and CUDA modules, you can check
   - CUDA version: `CUDA.versioninfo()`
-  - If MPI has CUDA: `MPI.has_cuda()`
+  - If MPI has CUDA: [`MPI.has_cuda()`](@ref)
   - If you are using correct MPI library: `MPI.libmpi`
 
 After that, it may be preferred to run the Julia MPI script (as suggested [here](https://discourse.julialang.org/t/cuda-aware-mpi-works-on-system-but-not-for-julia/75060/11)) launching it from a shell script (as suggested [here](https://discourse.julialang.org/t/cuda-aware-mpi-works-on-system-but-not-for-julia/75060/4)).
@@ -197,6 +197,7 @@ Make sure to:
   ```
 - Then in Julia, upon loading MPI and CUDA modules, you can check
   - AMDGPU version: `AMDGPU.versioninfo()`
+  - If MPI has ROCm: [`MPI.has_rocm()`](@ref)
   - If you are using correct MPI implementation: `MPI.identify_implementation()`
 
 After that, [this script](https://gist.github.com/luraess/c228ec08629737888a18c6a1e397643c) can be used to verify if ROCm-aware MPI is functional (modified after the CUDA-aware version from [here](https://discourse.julialang.org/t/cuda-aware-mpi-works-on-system-but-not-for-julia/75060/11)). It may be preferred to run the Julia ROCm-aware MPI script launching it from a shell script (as suggested [here](https://discourse.julialang.org/t/cuda-aware-mpi-works-on-system-but-not-for-julia/75060/4)).
diff --git a/docs/src/reference/library.md b/docs/src/reference/library.md
index 94e3942..302015f 100644
--- a/docs/src/reference/library.md
+++ b/docs/src/reference/library.md
@@ -14,5 +14,7 @@ MPI.MPI_LIBRARY_VERSION_STRING
 ```@docs
 MPI.versioninfo
 MPI.has_cuda
+MPI.has_rocm
+MPI.has_gpu
 MPI.identify_implementation
 ```
diff --git a/docs/src/usage.md b/docs/src/usage.md
index 2b8ebe4..c57eae1 100644
--- a/docs/src/usage.md
+++ b/docs/src/usage.md
@@ -98,7 +98,8 @@ should confirm your MPI implementation to have the ROCm support (AMDGPU) enabled
 [alltoall\_test\_rocm\_multigpu.jl](https://gist.github.com/luraess/a47931d7fb668bd4348a2c730d5489f4) should confirm 
 your ROCm-aware MPI implementation to use multiple AMD GPUs (one GPU per rank).
 
-The status of ROCm (AMDGPU) support cannot currently be queried.
+If using OpenMPI, the status of ROCm support can be checked via the
+[`MPI.has_rocm()`](@ref) function.
 
 ## Writing MPI tests
 
diff --git a/src/environment.jl b/src/environment.jl
index 52597c4..5408bdd 100644
--- a/src/environment.jl
+++ b/src/environment.jl
@@ -320,13 +320,15 @@ Wtime() = API.MPI_Wtime()
 
 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).
+(unless overriden). For "IBMSpectrumMPI" it will return `true`.
 
 This can be overriden by setting the `JULIA_MPI_HAS_CUDA` environment variable to `true`
 or `false`.
 
 !!! note
     For OpenMPI or OpenMPI-based implementations you first need to call [Init()](@ref).
+
+See also [`MPI.has_rocm`](@ref) for ROCm support.
 """
 function has_cuda()
     flag = get(ENV, "JULIA_MPI_HAS_CUDA", nothing)
@@ -334,7 +336,7 @@ function has_cuda()
         # Only Open MPI provides a function to check CUDA support
         @static if MPI_LIBRARY == "OpenMPI"
             # int MPIX_Query_cuda_support(void)
-            return 0 != ccall((:MPIX_Query_cuda_support, libmpi), Cint, ())
+            return @ccall libmpi.MPIX_Query_cuda_support()::Bool
         elseif MPI_LIBRARY == "IBMSpectrumMPI"
             return true
         else
@@ -344,3 +346,44 @@ function has_cuda()
         return parse(Bool, flag)
     end
 end
+
+"""
+    MPI.has_rocm()
+
+Check if the MPI implementation is known to have ROCm 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_ROCM` environment variable to `true`
+or `false`.
+
+See also [`MPI.has_cuda`](@ref) for CUDA support.
+"""
+function has_rocm()
+    flag = get(ENV, "JULIA_MPI_HAS_ROCM", nothing)
+    if flag === nothing
+        # Only Open MPI provides a function to check ROCm support
+        @static if MPI_LIBRARY == "OpenMPI" && MPI_LIBRARY_VERSION ≥ v"5"
+            # int MPIX_Query_rocm_support(void)
+            return @ccall libmpi.MPIX_Query_rocm_support()::Bool
+        else
+            return false
+        end
+    else
+        return parse(Bool, flag)
+    end
+end
+
+"""
+    MPI.has_gpu()
+
+Checks if the MPI implementation is known to have GPU support. Currently this checks for the
+following GPUs:
+
+1. CUDA: via [`MPI.has_cuda`](@ref)
+2. ROCm: via [`MPI.has_rocm`](@ref)
+
+See also [`MPI.has_cuda`](@ref) and [`MPI.has_rocm`](@ref) for more fine-grained
+checks.
+"""
+has_gpu() = has_cuda() || has_rocm()


Test patch:
diff --git a/test/test_basic.jl b/test/test_basic.jl
index 495c9f8..3da0139 100644
--- a/test/test_basic.jl
+++ b/test/test_basic.jl
@@ -8,10 +8,22 @@ MPI.Init()
 
 @test MPI.has_cuda() isa Bool
 
-if get(ENV,"JULIA_MPI_TEST_ARRAYTYPE","") == "CuArray"
+if get(ENV, "JULIA_MPI_TEST_ARRAYTYPE", "") == "CuArray"
     @test MPI.has_cuda()
 end
 
+@test MPI.has_rocm() isa Bool
+
+if get(ENV, "JULIA_MPI_TEST_ARRAYTYPE", "") == "ROCArray"
+    @test MPI.has_rocm()
+end
+
+@test MPI.has_gpu() isa Bool
+
+if get(ENV, "JULIA_MPI_TEST_ARRAYTYPE", "") == "CuArray" || get(ENV, "JULIA_MPI_TEST_ARRAYTYPE", "") == "ROCArray"
+    @test MPI.has_gpu()
+end
+
 @test !MPI.Finalized()
 MPI.Finalize()
 @test MPI.Finalized()