Test for MPI.Irecv/MPI.Isend/MPI.Wait (#518)

Co-authored-by: Valentin Churavy <v.churavy@gmail.com>
Co-authored-by: Lucas Wilcox <lucas@swirlee.com>
Co-authored-by: Mosè Giordano <765740+giordano@users.noreply.github.com>

Author: 4 people

test/integration/MPI/nonblocking_halo.jl

@@ -0,0 +1,75 @@
+using MPI
+using Enzyme
+using Test
+
+
+function halo(x)
+    np = MPI.Comm_size(MPI.COMM_WORLD)
+    rank = MPI.Comm_rank(MPI.COMM_WORLD)
+    requests = Vector{MPI.Request}()
+    if rank != 0
+        buf = @view x[1:1]
+        push!(requests, MPI.Isend(x[2:2], MPI.COMM_WORLD; dest = rank - 1, tag = 0))
+        push!(requests, MPI.Irecv!(buf, MPI.COMM_WORLD; source = rank - 1, tag = 0))
+    end
+    if rank != np - 1
+        buf = @view x[end:end]
+        push!(requests, MPI.Isend(x[(end - 1):(end - 1)], MPI.COMM_WORLD; dest = rank + 1, tag = 0))
+        push!(requests, MPI.Irecv!(buf, MPI.COMM_WORLD; source = rank + 1, tag = 0))
+    end
+    for request in requests
+        MPI.Wait(request)
+    end
+    return nothing
+end
+
+MPI.Init()
+np = MPI.Comm_size(MPI.COMM_WORLD)
+rank = MPI.Comm_rank(MPI.COMM_WORLD)
+nl = rank == 0 ? 0 : 2
+nr = rank == np - 1 ? 0 : 2
+nlocal = nr + nl + 1
+
+x = zeros(nlocal)
+fill!(x, Float64(rank))
+halo(x)
+MPI.Barrier(MPI.COMM_WORLD)
+
+@test x[nl + 1] == Float64(rank)      # Local
+if rank != 0
+    @test x[1] == Float64(rank - 1)   # Recv
+    @test x[2] == Float64(rank)       # Send
+end
+if rank != np - 1
+    @test x[end] == Float64(rank + 1) # Recv
+    @test x[end - 1] == Float64(rank) # Send
+end
+
+dx = zeros(nlocal)
+fill!(dx, Float64(rank))
+autodiff(Reverse, halo, Duplicated(x, dx))
+MPI.Barrier(MPI.COMM_WORLD)
+
+@test dx[nl + 1] == Float64(rank)                  # Local -> no change
+if rank != 0
+    @test dx[1] == 0.0                             # Recv -> Send & zero'd
+    @test dx[2] == Float64(rank + rank - 1)        # Send -> += Recv
+end
+if rank != np - 1
+    @test dx[end] == 0.0                           # Recv -> Send & zero'd
+    @test dx[end - 1] == Float64(rank + rank + 1)  # Send -> += Recv
+end
+
+fill!(dx, Float64(rank))
+autodiff(Forward, halo, Duplicated(x, dx))
+MPI.Barrier(MPI.COMM_WORLD)
+
+@test dx[nl + 1] == Float64(rank)
+if rank != 0
+    @test dx[1] == Float64(rank - 1)
+    @test dx[2] == Float64(rank)
+end
+if rank != np - 1
+    @test dx[end] == Float64(rank + 1)
+    @test dx[end - 1] == Float64(rank)
+end

test/integration/MPI/runtests.jl

@@ -2,6 +2,16 @@ using MPI
 using Enzyme
 using Test
 
+# Current MPI support (needs to be tested from Julia)
+# - MPI_Ssend
+# - MPI_Waitall
+# - MPI_Barrier/MPI_Probe
+# - MPI_Allreduce
+# - MPI_Bcast
+# - MPI_Reduce
+# - MPI_Gather/MPI_Scatter
+# - MPI_Allgather
+
 # Query functions MPI_Comm_size/MPI_Comm_rank
 @testset "queries" for np in (1, 2, 4)
     run(`$(mpiexec()) -n $np $(Base.julia_cmd()) --project=$(@__DIR__) $(joinpath(@__DIR__, "queries.jl"))`)
@@ -11,3 +21,8 @@ end
 @testset "blocking_ring" for np in (1, 2, 4)
     run(`$(mpiexec()) -n $np $(Base.julia_cmd()) --project=$(@__DIR__) $(joinpath(@__DIR__, "blocking_ring.jl"))`)
 end
+
+# Test MPI_Irecv/MPI_Isend/MPI_Wait with a non-blocking halo exchange pattern
+VERSION >= v"1.11.0" && @testset "nonblocking_halo" for np in (1, 2, 4)
+    run(`$(mpiexec()) -n $np $(Base.julia_cmd()) --project=$(@__DIR__) $(joinpath(@__DIR__, "nonblocking_halo.jl"))`)
+end