improve mpi

Project.toml

@@ -1,14 +1,15 @@
 name = "GCMAES"
 uuid = "4aa9d100-eb0f-11e8-15f1-25748831eb3b"
 authors = ["Yao Lu <[email protected]>"]
-version = "0.1.9"
+version = "0.1.10"
 
 [deps]
 BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"

src/GCMAES.jl

@@ -3,7 +3,7 @@ __precompile__(true)
 module GCMAES
 
 using Printf, Distributed, LinearAlgebra
-using Dates, Random, Statistics
+using Dates, Random, Statistics, Pkg
 using Requires, BSON, FastClosures
 
 include("util.jl")

src/util.jl

@@ -76,11 +76,13 @@ macro mpirun(ex)
     !inmpi() && return esc(ex)
     quote
         @eval using MPI
+        if "Elemental" in keys(Pkg.installed())
+            @eval using Elemental
+        end
         !MPI.Initialized() && MPI.Init()
         MPI.Barrier(MPI.COMM_WORLD)
         $(esc(ex))
         MPI.Barrier(MPI.COMM_WORLD)
-        atexit(MPI.Finalize)
     end
 end
 

test/mpi.jl

@@ -1,11 +1,7 @@
 using GCMAES
 using Zygote
-using MPI
-using Elemental
 using Test
 
-MPI.Initialized() || MPI.Init()
-
 rastrigin(x) = 10length(x) + sum(x.^2 .- 10 .* cos.(2π .* x))
 
 D = 5000
@@ -14,10 +10,8 @@ x0 = fill(0.3f0, D)
 lo = fill(-5.12f0, D)
 hi = fill(5.12f0, D)
 
-GCMAES.minimize(rastrigin, x0, σ0, lo, hi, λ = 200, maxiter = 200)
-
-GCMAES.minimize(rastrigin, x0, σ0, lo, hi, λ = 200, maxiter = 200, autodiff = true)
+@mpirun GCMAES.minimize(rastrigin, x0, σ0, lo, hi, λ = 200, maxiter = 200)
 
-rm("CMAES.bson", force = true)
+@mpirun GCMAES.minimize(rastrigin, x0, σ0, lo, hi, λ = 200, maxiter = 200, autodiff = true)
 
-MPI.Finalize()
+rm("CMAES.bson", force = true)