Skip to content

Instantly share code, notes, and snippets.

@luraess
Last active July 19, 2023 08:24
Show Gist options
  • Save luraess/ed93cc09ba04fe16f63b4219c1811566 to your computer and use it in GitHub Desktop.
Save luraess/ed93cc09ba04fe16f63b4219c1811566 to your computer and use it in GitHub Desktop.
CUDA-aware MPI multi-GPU test
using MPI
using CUDA
MPI.Init()
comm = MPI.COMM_WORLD
rank = MPI.Comm_rank(comm)
# select device
comm_l = MPI.Comm_split_type(comm, MPI.COMM_TYPE_SHARED, rank)
rank_l = MPI.Comm_rank(comm_l)
gpu_id = CUDA.device!(rank_l)
# select device
size = MPI.Comm_size(comm)
dst = mod(rank+1, size)
src = mod(rank-1, size)
println("rank=$rank rank_loc=$rank_l (gpu_id=$gpu_id), size=$size, dst=$dst, src=$src")
N = 4
send_mesg = CuArray{Float64}(undef, N)
recv_mesg = CuArray{Float64}(undef, N)
fill!(send_mesg, Float64(rank))
CUDA.synchronize()
rank==0 && println("start sending...")
MPI.Sendrecv!(send_mesg, dst, 0, recv_mesg, src, 0, comm)
println("recv_mesg on proc $rank_l: $recv_mesg")
rank==0 && println("done.")
@luraess
Copy link
Author

luraess commented Nov 29, 2022

Hi, it could be you're MPI has no CUDA-aware support. For CUDA-aware MPI to work, you should use system MPI that was linked against local CUDA install during compilation. If you are using system MPI that is supposed to be CUDA-aware, you can check it's functionality by typing MPI.has_cuda() (see here).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment