Add has_rocm for OpenMPI

docs/src/reference/library.md

@@ -14,5 +14,6 @@ MPI.MPI_LIBRARY_VERSION_STRING
 ```@docs
 MPI.versioninfo
 MPI.has_cuda
+MPI.has_rocm
 MPI.identify_implementation
 ```

src/environment.jl

@@ -320,7 +320,7 @@ 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`.
@@ -344,3 +344,28 @@ 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`.
+"""
+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

test/test_basic.jl

@@ -12,6 +12,12 @@ 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.Finalized()
 MPI.Finalize()
 @test MPI.Finalized()