From caf3f9654aaa6f8f02e190edfdc4396a6d6da9fa Mon Sep 17 00:00:00 2001
From: Yao Lu <luyaocns@gmail.com>
Date: Wed, 3 Jun 2020 16:08:53 +0800
Subject: [PATCH] improve mpi

---
 Project.toml  |  3 ++-
 src/GCMAES.jl |  2 +-
 src/util.jl   |  4 +++-
 test/mpi.jl   | 12 +++---------
 4 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/Project.toml b/Project.toml
index 7108538..dd69b29 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "GCMAES"
 uuid = "4aa9d100-eb0f-11e8-15f1-25748831eb3b"
 authors = ["Yao Lu <luyaocns@gmail.com>"]
-version = "0.1.9"
+version = "0.1.10"
 
 [deps]
 BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
@@ -9,6 +9,7 @@ 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"
diff --git a/src/GCMAES.jl b/src/GCMAES.jl
index 8e45e66..48db628 100644
--- a/src/GCMAES.jl
+++ b/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")
diff --git a/src/util.jl b/src/util.jl
index 80691e0..42f3119 100644
--- a/src/util.jl
+++ b/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
 
diff --git a/test/mpi.jl b/test/mpi.jl
index e0e2dfe..f702f4e 100644
--- a/test/mpi.jl
+++ b/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()
\ No newline at end of file
+rm("CMAES.bson", force = true)
\ No newline at end of file