Test support for MPI_Reduce and MPI_Allreduce for AMDGPU

Author: Keluaa

test/mpi_support_test.jl

@@ -6,17 +6,30 @@ MPI.Init()
 # (or a similar signal) when called, which cannot be handled in Julia in a portable way.
 
 op = ARGS[1]
-if op == "Iallreduce"
-    # Iallreduce is unsupported for CUDA with OpenMPI + UCX
-    # See https://docs.open-mpi.org/en/main/tuning-apps/networking/cuda.html#which-mpi-apis-do-not-work-with-cuda-aware-ucx
+if op == "Allreduce"
+    # Allreduce is unsupported for AMDGPU with UCX
+    send_arr = ArrayType(zeros(Int, 1))
+    recv_arr = ArrayType{Int}(undef, 1)
+    synchronize()
+    MPI.Allreduce!(send_arr, recv_arr, +, MPI.COMM_WORLD)
+
+elseif op == "Iallreduce"
+    # Iallreduce is unsupported for CUDA with OpenMPI 5 + UCX
     send_arr = ArrayType(zeros(Int, 1))
     recv_arr = ArrayType{Int}(undef, 1)
     synchronize()
     req = MPI.Iallreduce!(send_arr, recv_arr, +, MPI.COMM_WORLD)
     MPI.Wait(req)
 
+elseif op == "Reduce"
+    # Reduce is unsupported for AMDGPU with UCX
+    send_arr = ArrayType(zeros(Int, 1))
+    recv_arr = ArrayType{Int}(undef, 1)
+    synchronize()
+    MPI.Reduce!(send_arr, recv_arr, +, MPI.COMM_WORLD; root=0)
+
 elseif op == "Ireduce"
-    # Iallreduce is unsupported for CUDA with OpenMPI + UCX
+    # Ireduce is unsupported for CUDA with OpenMPI 5 + UCX
     send_arr = ArrayType(zeros(Int, 1))
     recv_arr = ArrayType{Int}(undef, 1)
     synchronize()

test/runtests.jl

@@ -74,12 +74,17 @@ end
 function is_mpi_operation_supported(mpi_op, n=nprocs)
     test_file = joinpath(@__DIR__, "mpi_support_test.jl")
     cmd = `$(mpiexec()) -n $n $(Base.julia_cmd()) --startup-file=no $test_file $mpi_op`
-    supported = success(run(ignorestatus(cmd)))
+    cmd = ignorestatus(pipeline(cmd; stderr=devnull))
+    process = run(cmd)
+    supported = success(process)
     !supported && @warn "$mpi_op is unsupported with $backend_name"
+    !supported && !Base.process_signaled(process) && @warn "support check for $mpi_op may be broken as process did not signal"
     return supported
 end
 
 if ArrayType != Array  # we expect that only GPU backends can have unsupported features
+    ENV["JULIA_MPI_TEST_ALLREDUCE"] = is_mpi_operation_supported("Allreduce")
+    ENV["JULIA_MPI_TEST_REDUCE"] = is_mpi_operation_supported("Reduce")
     ENV["JULIA_MPI_TEST_IALLREDUCE"] = is_mpi_operation_supported("Iallreduce")
     ENV["JULIA_MPI_TEST_IREDUCE"] = is_mpi_operation_supported("Ireduce")
 end

test/test_allreduce.jl

@@ -1,5 +1,12 @@
 include("common.jl")
 
+allreduce_supported = get(ENV, "JULIA_MPI_TEST_ALLREDUCE", "true") == "true"
+iallreduce_supported = get(ENV, "JULIA_MPI_TEST_IALLREDUCE", "true") == "true"
+if !allreduce_supported
+    @warn "Skipping all tests in 'test_allreduce.jl' as reductions are unsupported"
+    return
+end
+
 MPI.Init()
 
 comm_size = MPI.Comm_size(MPI.COMM_WORLD)
@@ -13,8 +20,6 @@ else
     operators = [MPI.SUM, +, (x,y) -> 2x+y-x]
 end
 
-iallreduce_supported = get(ENV, "JULIA_MPI_TEST_IALLREDUCE", "true") == "true"
-
 
 for T = [Int]
     for dims = [1, 2, 3]

test/test_reduce.jl

@@ -9,7 +9,12 @@ const can_do_closures =
     Sys.ARCH !== :aarch64 &&
     !startswith(string(Sys.ARCH), "arm")
 
+reduce_supported = get(ENV, "JULIA_MPI_TEST_REDUCE", "true") == "true"
 ireduce_supported = get(ENV, "JULIA_MPI_TEST_IREDUCE", "true") == "true"
+if !reduce_supported
+    @warn "Skipping all tests in 'test_reduce.jl' as reductions are unsupported"
+    return
+end
 
 using DoubleFloats