Parallel simulations using MPI.jl

Project.toml

@@ -1,7 +1,7 @@
 name = "WaterLily"
 uuid = "ed894a53-35f9-47f1-b17f-85db9237eebd"
 authors = ["Gabriel Weymouth <gabriel.weymouth@gmail.com>"]
-version = "1.2"
+version = "1.2.0"
 
 [deps]
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
@@ -18,12 +18,14 @@ TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 ReadVTK = "dc215faf-f008-4882-a9f7-a79a826fadc3"
 WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 
 [extensions]
 WaterLilyAMDGPUExt = "AMDGPU"
 WaterLilyCUDAExt = "CUDA"
+WaterLilyMPIExt = "MPI"
 WaterLilyReadVTKExt = "ReadVTK"
 WaterLilyWriteVTKExt = "WriteVTK"
 
@@ -41,9 +43,12 @@ julia = "1.6"
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
 GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
+MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 Makie = "537997a7-5e4e-5d89-9595-2241ea00577e"
 PerformanceTestTools = "dc46b164-d16f-48ec-a853-60448fc869fe"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
@@ -53,4 +58,5 @@ UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
 WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 
 [targets]
-test = ["Test", "BenchmarkTools", "CUDA", "AMDGPU", "GPUArrays", "WriteVTK", "ReadVTK"]
+test = ["Test", "BenchmarkTools", "CUDA", "AMDGPU", "GPUArrays", "WriteVTK", "ReadVTK", "MPI",
+        "FileIO", "JLD2"]

examples/Project.toml

@@ -0,0 +1,19 @@
+[deps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
+GR = "28b8d3ca-fb5f-59d9-8090-bfdbd6d07a71"
+GeometryBasics = "5c1252a2-5f33-56bf-86c9-59e7332b4326"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
+Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
+Meshing = "e6723b4c-ebff-59f1-b4b7-d97aa5274f73"
+Optim = "429524aa-4258-5aef-a3af-852621145aeb"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+ReadVTK = "dc215faf-f008-4882-a9f7-a79a826fadc3"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228"
+WaterLily = "ed894a53-35f9-47f1-b17f-85db9237eebd"
+WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"

examples/TwoD_CircleMPI.jl

@@ -0,0 +1,54 @@
+#mpiexecjl --project=. -n 2 julia TwoD_CircleMPI.jl
+
+using WaterLily
+using WriteVTK
+using MPI
+using StaticArrays
+
+# make a writer with some attributes, now we need to apply the BCs when writting
+velocity(a::Simulation) = a.flow.u |> Array;
+pressure(a::Simulation) = a.flow.p |> Array;
+_body(a::Simulation) = (measure_sdf!(a.flow.σ, a.body);
+                        a.flow.σ |> Array;)
+vorticity(a::Simulation) = (@inside a.flow.σ[I] = 
+                            WaterLily.curl(3,I,a.flow.u)*a.L/a.U;
+                            WaterLily.perBC!(a.flow.σ,());
+                            a.flow.σ |> Array;)
+_vbody(a::Simulation) = a.flow.V |> Array;
+mu0(a::Simulation) = a.flow.μ₀ |> Array;
+ranks(a::Simulation) = (a.flow.σ.=0; 
+                        @inside a.flow.σ[I] = me()+1;
+                        WaterLily.perBC!(a.flow.σ,());
+                        a.flow.σ |> Array;)
+
+custom_attrib = Dict(
+    "u" => velocity,
+    "p" => pressure,
+    "d" => _body,
+    "ω" => vorticity,
+    "v" => _vbody,
+    "μ₀" => mu0,
+    "rank" => ranks
+)# this maps what to write to the name in the file
+
+"""Flow around a circle"""
+function circle(dims,center,radius;Re=250,U=1,psolver=MultiLevelPoisson,mem=Array)
+    body = AutoBody((x,t)->√sum(abs2, x .- center) - radius)
+    Simulation(dims, (U,0), radius; ν=U*radius/Re, body, exitBC=false, mem=mem, psolver=psolver)
+end
+
+# local grid size
+L = 2^6
+
+# init the MPI grid and the simulation
+r = init_mpi((L,2L))
+sim = circle((L,2L),SA[L/2,L+2],L/8;mem=MPIArray) #use MPIArray to use extension
+
+wr = vtkWriter("WaterLily-MPI-circle";attrib=custom_attrib,dir="vtk_data",
+               extents=get_extents(sim.flow.p))
+for _ in 1:50
+    sim_step!(sim,sim_time(sim)+1.0,verbose=true)
+    write!(wr,sim)
+end
+close(wr)
+finalize_mpi()

examples/TwoD_CircleMPIArray.jl

@@ -0,0 +1,13 @@
+using WaterLily
+
+# circle simulation
+function circle(n,m;Re=250,U=1)
+    radius, center = m/8, m/2
+    body = AutoBody((x,t)->√sum(abs2, x .- center) - radius)
+    Simulation((n,m), (U,0), radius; ν=U*radius/Re, body, mem=MPIArray)
+end
+
+# test on sim
+include("/home/marin/Workspace/Tutorials-WaterLily/src/TwoD_plots.jl")
+sim = circle(3*2^6,2^7)
+sim_gif!(sim,duration=10,clims=(-5,5),plotbody=true)