MPI support

.github/workflows/CI.yml

@@ -26,6 +26,7 @@ jobs:
           - '1.10'
           - '1.11'
           - '1.12'
+          - '1.13'
         os:
           - ubuntu-latest
           - windows-latest

Project.toml

@@ -1,6 +1,6 @@
 name = "FiniteDiffWENO5"
 uuid = "4b1b4241-afed-47c3-a843-c6e03821b375"
-version = "0.2.0"
+version = "0.2.1"
 authors = ["Hugo Dominguez  <hdomingu@uni-mainz.de>"]
 
 [deps]
@@ -9,13 +9,16 @@ MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
 [weakdeps]
 Chmy = "33a72cf0-4690-46d7-b987-06506c2248b9"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 
 [extensions]
 ChmyExt = ["KernelAbstractions", "Chmy"]
 KAExt = ["KernelAbstractions"]
+MPIExt = "MPI"
 
 [compat]
 Chmy = "0.1"
 KernelAbstractions = "0.9"
+MPI = "0.20"
 MuladdMacro = "0.2"
 julia = "1.10"

ext/ChmyExt.jl

@@ -99,12 +99,16 @@ function WENOScheme(
 
     TFlux = typeof(fl)
     TArray = typeof(du)
+    extent = FiniteDiffWENO5.default_extent(sizes, FiniteDiffWENO5.default_global_periodic(boundary, N))
 
-    return WENOScheme{T, TArray, TFlux, typeof(vcenter), typeof(vperiodic), typeof(form_tag), typeof(boundary)}(
+    topology = FiniteDiffWENO5.NoTopology() # Chmy schemes use unpadded arrays
+    halo_buffers = FiniteDiffWENO5.EmptyHaloBuffers()
+
+    return WENOScheme{T, TArray, TFlux, typeof(vcenter), typeof(vperiodic), typeof(form_tag), typeof(boundary), typeof(extent), typeof(topology), typeof(halo_buffers)}(
         stag = stag, form = form_tag, boundary = boundary, multithreading = multithreading,
         lim_ZS = lim_ZS, fl = fl, fr = fr, du = du, ut = ut, vcenter = vcenter,
-        vperiodic = vperiodic,
-        upwind_mode = upwind_mode,
+        vperiodic = vperiodic, extent = extent, topology = topology,
+        upwind_mode = upwind_mode, halo_buffers = halo_buffers,
     )
 end
 
@@ -162,12 +166,17 @@ function MultiphaseWENOScheme(
         NamedTuple{labels}(ntuple(_ -> Field(backend, grid, Center(), T), direction_count)) :
         nothing
     vperiodic = stag ? FiniteDiffWENO5.velocity_periodicity(boundary, labels) : nothing
+    extent = FiniteDiffWENO5.default_extent(sizes, FiniteDiffWENO5.default_global_periodic(boundary, N))
+    topology = FiniteDiffWENO5.NoTopology() # Chmy schemes use unpadded arrays
+    halo_buffers = FiniteDiffWENO5.EmptyHaloBuffers()
 
     return MultiphaseWENOScheme{
         T, NP, typeof(du), typeof(fl), typeof(vcenter), typeof(vperiodic), typeof(boundary),
+        typeof(extent), typeof(topology), typeof(halo_buffers),
     }(
         stag = stag, boundary = boundary, multithreading = multithreading,
         fl = fl, fr = fr, du = du, ut = ut, vcenter = vcenter, vperiodic = vperiodic,
+        extent = extent, topology = topology, halo_buffers = halo_buffers,
     )
 end
 
@@ -321,6 +330,7 @@ function WENO_step!(
     end
     @assert get_backend(v.x) == backend
 
+    FiniteDiffWENO5.require_no_topology(scheme, "Chmy multiphase WENO_step!")
     launch = Launcher(arch, grid)
     (; fl, fr, ut, du, boundary, χ, γ, ζ, ϵ) = scheme
     nx = grid.axes[1].length
@@ -366,6 +376,7 @@ function WENO_step!(
     @assert get_backend(v.x) == backend
     @assert get_backend(v.y) == backend
 
+    FiniteDiffWENO5.require_no_topology(scheme, "Chmy multiphase WENO_step!")
     launch = Launcher(arch, grid)
     (; fl, fr, ut, du, boundary, χ, γ, ζ, ϵ) = scheme
     nx, ny = map(axis -> axis.length, grid.axes)
@@ -416,6 +427,7 @@ function WENO_step!(
     @assert get_backend(v.y) == backend
     @assert get_backend(v.z) == backend
 
+    FiniteDiffWENO5.require_no_topology(scheme, "Chmy multiphase WENO_step!")
     launch = Launcher(arch, grid)
     (; fl, fr, ut, du, boundary, χ, γ, ζ, ϵ) = scheme
     nx, ny, nz = map(axis -> axis.length, grid.axes)
@@ -480,6 +492,8 @@ function WENO_step!(u::T_field, v::Velocity1D, weno::FiniteDiffWENO5.WENOScheme,
 
     launch = Launcher(arch, grid)
 
+    FiniteDiffWENO5.require_no_topology(weno, "Chmy WENO_step!")
+
     #! do things here for halos and such for clusters for boundaries probably
 
     nx = grid.axes[1].length
@@ -513,9 +527,11 @@ function WENO_step!(u::T_field, v::Velocity1D, weno::FiniteDiffWENO5.WENOScheme,
         interior(u) .= @muladd inv(3.0) .* interior(u) .+ 2.0 / 3.0 .* interior(ut) .- 2.0 / 3.0 .* Δt .* interior(du)
     else
         vupwind = conservative ? v : prepare_velocity_chmy_1D!(weno, v, nx, grid, arch, launch)
-        launch(arch, grid, upwind_update_KA_1D! => (
-            u, vupwind, nx, Δx_, Δt, stag && conservative, boundary, grid,
-        ))
+        launch(
+            arch, grid, upwind_update_KA_1D! => (
+                u, vupwind, nx, Δx_, Δt, stag && conservative, boundary, grid,
+            )
+        )
     end
 
     return nothing
@@ -550,6 +566,8 @@ function WENO_step!(u::T_field, v::Velocity2D, weno::FiniteDiffWENO5.WENOScheme,
 
     launch = Launcher(arch, grid)
 
+    FiniteDiffWENO5.require_no_topology(weno, "Chmy WENO_step!")
+
     #! do things here for halos and such for clusters for boundaries probably
 
     nx = grid.axes[1].length
@@ -588,9 +606,11 @@ function WENO_step!(u::T_field, v::Velocity2D, weno::FiniteDiffWENO5.WENOScheme,
         interior(u) .= @muladd inv(3.0) .* interior(u) .+ 2.0 / 3.0 .* interior(ut) .- 2.0 / 3.0 .* Δt .* interior(du)
     else
         vupwind = conservative ? v : prepare_velocity_chmy_2D!(weno, v, nx, ny, grid, arch, launch)
-        launch(arch, grid, upwind_update_KA_2D! => (
-            u, vupwind, nx, ny, Δx_, Δy_, Δt, stag && conservative, boundary, grid,
-        ))
+        launch(
+            arch, grid, upwind_update_KA_2D! => (
+                u, vupwind, nx, ny, Δx_, Δy_, Δt, stag && conservative, boundary, grid,
+            )
+        )
     end
 
     return nothing
@@ -628,6 +648,8 @@ function WENO_step!(u::T_field, v::Velocity3D, weno::FiniteDiffWENO5.WENOScheme,
 
     launch = Launcher(arch, grid)
 
+    FiniteDiffWENO5.require_no_topology(weno, "Chmy WENO_step!")
+
     #! do things here for halos and such for clusters for boundaries probably
 
     nx = grid.axes[1].length
@@ -671,9 +693,11 @@ function WENO_step!(u::T_field, v::Velocity3D, weno::FiniteDiffWENO5.WENOScheme,
         interior(u) .= @muladd inv(3.0) .* interior(u) .+ 2.0 / 3.0 .* interior(ut) .- 2.0 / 3.0 .* Δt .* interior(du)
     else
         vupwind = conservative ? v : prepare_velocity_chmy_3D!(weno, v, nx, ny, nz, grid, arch, launch)
-        launch(arch, grid, upwind_update_KA_3D! => (
-            u, vupwind, nx, ny, nz, Δx_, Δy_, Δz_, Δt, stag && conservative, boundary, grid,
-        ))
+        launch(
+            arch, grid, upwind_update_KA_3D! => (
+                u, vupwind, nx, ny, nz, Δx_, Δy_, Δz_, Δt, stag && conservative, boundary, grid,
+            )
+        )
     end
 
     return nothing

ext/KAExt.jl

@@ -78,12 +78,16 @@ function WENOScheme(
 
     TFlux = typeof(fl)
     TArray = typeof(du)
+    extent = FiniteDiffWENO5.default_extent(sizes, FiniteDiffWENO5.default_global_periodic(boundary, N))
 
-    return WENOScheme{T, TArray, TFlux, typeof(vcenter), typeof(vperiodic), typeof(form_tag), typeof(boundary)}(
+    topology = FiniteDiffWENO5.NoTopology() # KA schemes use unpadded arrays
+    halo_buffers = FiniteDiffWENO5.EmptyHaloBuffers()
+
+    return WENOScheme{T, TArray, TFlux, typeof(vcenter), typeof(vperiodic), typeof(form_tag), typeof(boundary), typeof(extent), typeof(topology), typeof(halo_buffers)}(
         stag = stag, form = form_tag, boundary = boundary, multithreading = multithreading,
         lim_ZS = lim_ZS, fl = fl, fr = fr, du = du, ut = ut, vcenter = vcenter,
-        vperiodic = vperiodic,
-        upwind_mode = upwind_mode,
+        vperiodic = vperiodic, extent = extent, topology = topology,
+        upwind_mode = upwind_mode, halo_buffers = halo_buffers,
     )
 end
 
@@ -159,12 +163,17 @@ function MultiphaseWENOScheme(
     ut = ntuple(_ -> backend_zeros(sizes), valNP)
     vcenter = stag ? NamedTuple{labels}(ntuple(_ -> backend_zeros(sizes), min(N, 3))) : nothing
     vperiodic = stag ? FiniteDiffWENO5.velocity_periodicity(boundary, labels) : nothing
+    extent = FiniteDiffWENO5.default_extent(sizes, FiniteDiffWENO5.default_global_periodic(boundary, N))
+    topology = FiniteDiffWENO5.NoTopology() # KA schemes use unpadded arrays
+    halo_buffers = FiniteDiffWENO5.EmptyHaloBuffers()
 
     return MultiphaseWENOScheme{
         T, NP, typeof(du), typeof(fl), typeof(vcenter), typeof(vperiodic), typeof(boundary),
+        typeof(extent), typeof(topology), typeof(halo_buffers),
     }(
         stag = stag, boundary = boundary, multithreading = multithreading,
         fl = fl, fr = fr, du = du, ut = ut, vcenter = vcenter, vperiodic = vperiodic,
+        extent = extent, topology = topology, halo_buffers = halo_buffers,
     )
 end
 
@@ -199,6 +208,8 @@ function WENO_step!(u::T_KA, v::NamedTuple{(:x,), <:Tuple{<:AbstractArray{<:Real
     @assert get_backend(u) == backend
     @assert get_backend(v.x) == backend
 
+    FiniteDiffWENO5.require_no_topology(weno, "KernelAbstractions WENO_step!")
+
     #! do things here for halos and such for clusters for boundaries probably
 
     (; ut, du, fl, fr, stag, lim_ZS, boundary, χ, γ, ζ, ϵ, upwind_mode, form, vcenter) = weno
@@ -283,6 +294,8 @@ function WENO_step!(u::T_KA, v::NamedTuple{(:x, :y), <:Tuple{Vararg{AbstractArra
     @assert get_backend(v.x) == backend
     @assert get_backend(v.y) == backend
 
+    FiniteDiffWENO5.require_no_topology(weno, "KernelAbstractions WENO_step!")
+
     #! do things here for halos and such for clusters for boundaries probably
 
     (; ut, du, fl, fr, stag, lim_ZS, boundary, χ, γ, ζ, ϵ, upwind_mode, form) = weno
@@ -374,6 +387,8 @@ function WENO_step!(u::T_KA, v::NamedTuple{(:x, :y, :z), <:Tuple{Vararg{Abstract
     @assert get_backend(v.y) == backend
     @assert get_backend(v.z) == backend
 
+    FiniteDiffWENO5.require_no_topology(weno, "KernelAbstractions WENO_step!")
+
     #! do things here for halos and such for clusters for boundaries probably
 
     nx = size(u, 1)

ext/KAMultiphase1D.jl

@@ -101,6 +101,12 @@ if nameof(@__MODULE__) == :KAExt
             @assert get_backend(phases[q]) == backend
         end
         @assert get_backend(v.x) == backend
+        scheme.topology isa FiniteDiffWENO5.NoTopology || throw(
+            ArgumentError(
+                "KernelAbstractions multiphase WENO_step! does not support a distributed " *
+                    "topology (GPU+MPI is unsupported)"
+            )
+        )
 
         (; fl, fr, ut, du, boundary, χ, γ, ζ, ϵ) = scheme
         nx = size(phases[1], 1)

ext/KAMultiphase2D.jl

@@ -169,6 +169,12 @@ if nameof(@__MODULE__) == :KAExt
         end
         @assert get_backend(v.x) == backend
         @assert get_backend(v.y) == backend
+        scheme.topology isa FiniteDiffWENO5.NoTopology || throw(
+            ArgumentError(
+                "KernelAbstractions multiphase WENO_step! does not support a distributed " *
+                    "topology (GPU+MPI is unsupported)"
+            )
+        )
 
         (; fl, fr, ut, du, boundary, χ, γ, ζ, ϵ) = scheme
         nx, ny = size(phases[1])

ext/KAMultiphase3D.jl

@@ -219,6 +219,12 @@ if nameof(@__MODULE__) == :KAExt
         @assert get_backend(v.x) == backend
         @assert get_backend(v.y) == backend
         @assert get_backend(v.z) == backend
+        scheme.topology isa FiniteDiffWENO5.NoTopology || throw(
+            ArgumentError(
+                "KernelAbstractions multiphase WENO_step! does not support a distributed " *
+                    "topology (GPU+MPI is unsupported)"
+            )
+        )
 
         (; fl, fr, ut, du, boundary, χ, γ, ζ, ϵ) = scheme
         nx, ny, nz = size(phases[1])

ext/MPIExt.jl

@@ -0,0 +1,16 @@
+module MPIExt
+
+using FiniteDiffWENO5
+using FiniteDiffWENO5: AbstractWENOTopology
+
+import FiniteDiffWENO5: weno_ndims, weno_halo, weno_owned_size, weno_global_size,
+    weno_global_offset, weno_periodic, weno_physical_low, weno_physical_high,
+    weno_exchange_halo!, weno_allreduce_max, weno_allreduce_min, weno_cartesian_topology,
+    halo_buffers_for, halo_buffers_for_multiphase
+
+using MPI
+
+include("mpi/topology.jl")
+include("mpi/halo.jl")
+
+end # module MPIExt