Add DistributedArrays

Project.toml

@@ -1,6 +1,6 @@
 name = "MUMPS"
 uuid = "55d2b088-9f4e-11e9-26c0-150b02ea6a46"
-authors = ["Dominique Orban <dominique.orban@gmail.com>","William R Sweeney <wrsweeney2@gmail.com>", "Alexis Montoison <alexis.montoison@polymtl.ca>"]
+authors = ["Dominique Orban <dominique.orban@gmail.com>", "William R Sweeney <wrsweeney2@gmail.com>", "Alexis Montoison <alexis.montoison@polymtl.ca>"]
 version = "1.5.5"
 
 [deps]
@@ -17,6 +17,7 @@ julia = "1.6"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+DistributedArrays = "aaf54ef3-cdf8-58ed-94cc-d582ad619b94"
 
 [targets]
-test = ["Test"]
+test = ["Test", "DistributedArrays"]

src/MUMPS.jl

@@ -50,7 +50,7 @@ include("interface.jl")
 include("convenience.jl")
 include("icntl_alibis.jl")
 include("printing.jl")
-
+include("distributed.jl")
 include("exported_methods.jl")
 
 end

src/distributed.jl

@@ -0,0 +1,37 @@
+using MPI
+
+supports_distributed() = Base.find_package("DistributedArrays") !== nothing
+
+function is_darray(A)
+  if !supports_distributed()
+    return false
+  end
+  try
+    @eval import DistributedArrays
+  catch
+    return false
+  end
+  return A isa DistributedArrays.DArray
+end
+
+function gather_on_root(A; root::Integer = 0, comm = MPI.COMM_WORLD)
+  if !supports_distributed()
+    return A
+  end
+  try
+    @eval import DistributedArrays
+  catch
+    return A
+  end
+
+  if A isa DistributedArrays.DArray
+    rank = MPI.Comm_rank(comm)
+    if rank == root
+      return Array(A)
+    else
+      return nothing
+    end
+  else
+    return A
+  end
+end

test/mumps_test_distributed_helpers.jl

@@ -0,0 +1,28 @@
+"""Tests for the distributed-array helpers.
+
+These tests require `DistributedArrays` to be available in the test environment.
+"""
+
+using Test
+using MPI
+using MUMPS
+import DistributedArrays
+
+comm = MPI.COMM_WORLD
+rank = MPI.Comm_rank(comm)
+
+@testset "distributed helpers" begin
+  # create a small local array and distribute it
+  localA = reshape(collect(1:12), 3, 4)
+  darr = DistributedArrays.distribute(localA)
+
+  g = MUMPS.gather_on_root(darr; root = 0, comm = comm)
+
+  if rank == 0
+    @test isa(g, Array)
+    @test size(g) == size(localA)
+    @test all(g .== localA)
+  else
+    @test g === nothing
+  end
+end

test/runtests.jl

@@ -73,5 +73,9 @@ end
   include("mumps_test_scalapack.jl")
 end
 
+@testset "distributed helpers: " begin
+  include("mumps_test_distributed_helpers.jl")
+end
+
 MPI.Barrier(comm)
 MPI.Finalize()