Add mpi_isroot() (#46)

Co-authored-by: Joshua Lampert <51029046+JoshuaLampert@users.noreply.github.com>
Co-authored-by: Benedict <135045760+benegee@users.noreply.github.com>
Co-authored-by: Hendrik Ranocha <ranocha@users.noreply.github.com>
Co-authored-by: Valentin Churavy <v.churavy@gmail.com>

Author: 4 people

ext/TrixiBaseMPIExt.jl

@@ -1,20 +1,28 @@
 # Package extension for adding MPI-based features to TrixiBase.jl
 module TrixiBaseMPIExt
 
-# Load package extension code on Julia v1.9 and newer
-if isdefined(Base, :get_extension)
-    using MPI: MPI
-end
+using MPI
 import TrixiBase
 
-# This is a really working version - assuming the same
+function __init__()
+    TrixiBase.__MPI__AVAILABLE__[] = true
+end
+
+# These are really working functions - assuming the same
 # communication pattern etc. used in Trixi.jl.
-function TrixiBase.mpi_isparallel(::Val{:MPIExt})
+function TrixiBase.mpi_isparallel_internal()
     if MPI.Initialized()
         return MPI.Comm_size(MPI.COMM_WORLD) > 1
     else
         return false
     end
 end
 
+function TrixiBase.mpi_isroot_internal()
+    if MPI.Initialized()
+        return MPI.Comm_rank(MPI.COMM_WORLD) == 0
+    else
+        return true
+    end
+end
 end

src/TrixiBase.jl

@@ -3,6 +3,7 @@ module TrixiBase
 using ChangePrecision: ChangePrecision
 using TimerOutputs: TimerOutput, TimerOutputs
 
+include("mpi.jl")
 include("trixi_include.jl")
 include("trixi_timeit.jl")
 

src/mpi.jl

@@ -0,0 +1,26 @@
+# These are just dummy functions. We only implement real
+# versions if MPI.jl is loaded to avoid letting TrixiBase.jl
+# depend explicitly on MPI.jl.
+
+# This will be true if TrixiBaseMPIExt is loaded
+const __MPI__AVAILABLE__ = Ref{Bool}(false)
+
+# These functions are defined in the TrixiBaseMPIExt extension
+function mpi_isparallel_internal end
+function mpi_isroot_internal end
+
+function mpi_isparallel()
+    if __MPI__AVAILABLE__[]
+        return mpi_isparallel_internal()
+    else
+        return false
+    end
+end
+
+function mpi_isroot()
+    if __MPI__AVAILABLE__[]
+        return mpi_isroot_internal()
+    else
+        return true
+    end
+end

src/trixi_include.jl

@@ -46,7 +46,7 @@ function trixi_include(mapexpr::Function, mod::Module, elixir::AbstractString; k
     end
 
     # Print information on potential wait time only in non-parallel case
-    if !mpi_isparallel(Val{:MPIExt}())
+    if !mpi_isparallel()
         @info "You just called `trixi_include`. Julia may now compile the code, please be patient."
     end
     Base.include(ex -> mapexpr(replace_assignments(insert_maxiters(ex); kwargs...)),
@@ -203,8 +203,3 @@ function find_assignment(expr, destination)
 
     result
 end
-
-# This is just a dummy function. We only implement a real
-# version if MPI.jl is loaded to avoid letting TrixiBase.jl
-# depend explicitly on MPI.jl.
-mpi_isparallel(x) = false

test/Project.toml

@@ -1,11 +1,13 @@
 [deps]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
+MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TrixiTest = "0a316866-cbd0-4425-8bcb-08103b2c1f26"
 
 [compat]
 Aqua = "0.7, 0.8"
 ExplicitImports = "1.0.1"
+MPI = "0.20"
 Test = "1"
 TrixiTest = "0.1"

test/runtests.jl

@@ -2,6 +2,8 @@ include("test_util.jl")
 
 @testset verbose=true "TrixiBase.jl Tests" begin
     include("test_aqua.jl")
+    include("test_mpi.jl")
+    run(`$(mpiexec()) -n 2 $(Base.julia_cmd()) --threads=1 $(abspath("test_mpi.jl"))`)
     include("trixi_include.jl")
     include("test_timers.jl")
 end;

test/test_mpi.jl

@@ -0,0 +1,14 @@
+
+include("test_util.jl")
+
+@testset verbose=true "MPI helper functions" begin
+    @test TrixiBase.mpi_isparallel() == false
+    @test TrixiBase.mpi_isroot() == true
+
+    using MPI
+    @test TrixiBase.mpi_isparallel() == false
+    @test TrixiBase.mpi_isroot() == true
+    MPI.Init()
+    @test TrixiBase.mpi_isparallel() == (MPI.Comm_size(MPI.COMM_WORLD) > 1)
+    @test TrixiBase.mpi_isroot() == (MPI.Comm_rank(MPI.COMM_WORLD) == 0)
+end