Add option to partition on a sub-communicator (one per rank) (#150)

* Subdivide comm for partitioning

* Add command line switch

* Free sub and shm comms

src/main.cpp

@@ -53,15 +53,18 @@ void solve(int argc, char* argv[])
 {
   po::options_description desc("Allowed options");
   bool mem_profile;
+  bool use_subcomm;
   desc.add_options()("help,h", "print usage message")(
       "problem_type", po::value<std::string>()->default_value("poisson"),
       "problem (poisson, cgpoisson, or elasticity)")(
       "mesh_type", po::value<std::string>()->default_value("cube"),
       "mesh (cube or unstructured)")(
       "memory_profiling", po::bool_switch(&mem_profile)->default_value(false),
-      "turn on memory logging")("scaling_type",
-                                po::value<std::string>()->default_value("weak"),
-                                "scaling (weak or strong)")(
+      "turn on memory logging")(
+      "subcomm_partition", po::bool_switch(&use_subcomm)->default_value(false),
+      "Use sub-communicator for partitioning")(
+      "scaling_type", po::value<std::string>()->default_value("weak"),
+      "scaling (weak or strong)")(
       "output", po::value<std::string>()->default_value(""),
       "output directory (no output unless this is set)")(
       "ndofs", po::value<std::size_t>()->default_value(50000),
@@ -127,8 +130,9 @@ void solve(int argc, char* argv[])
   dolfinx::common::Timer t0("ZZZ Create Mesh");
   if (mesh_type == "cube")
   {
-    mesh = std::make_shared<dolfinx::mesh::Mesh<double>>(create_cube_mesh(
-        MPI_COMM_WORLD, ndofs, strong_scaling, ndofs_per_node, order));
+    mesh = std::make_shared<dolfinx::mesh::Mesh<double>>(
+        create_cube_mesh(MPI_COMM_WORLD, ndofs, strong_scaling, ndofs_per_node,
+                         order, use_subcomm));
   }
   else
   {

src/mesh.cpp

@@ -73,7 +73,7 @@ std::int64_t num_pdofs(std::int64_t i, std::int64_t j, std::int64_t k,
 
 dolfinx::mesh::Mesh<double>
 create_cube_mesh(MPI_Comm comm, std::size_t target_dofs, bool target_dofs_total,
-                 std::size_t dofs_per_node, int order)
+                 std::size_t dofs_per_node, int order, bool use_subcomm)
 {
   // Get number of processes
   const std::size_t num_processes = dolfinx::MPI::size(comm);
@@ -157,12 +157,32 @@ create_cube_mesh(MPI_Comm comm, std::size_t target_dofs, bool target_dofs_total,
 #error "No mesh partitioner has been selected"
 #endif
 
+  MPI_Comm sub_comm;
+
+  if (use_subcomm)
+  {
+    // Create a sub-communicator for mesh partitioning
+    MPI_Comm shm_comm;
+    // Get a local comm on each node
+    MPI_Comm_split_type(comm, MPI_COMM_TYPE_SHARED, 0, MPI_INFO_NULL,
+                        &shm_comm);
+    int shm_comm_rank = dolfinx::MPI::rank(shm_comm);
+    MPI_Comm_free(&shm_comm);
+    // Create a comm across nodes, using rank 0 of the local comm on each node
+    int color = (shm_comm_rank == 0) ? 0 : MPI_UNDEFINED;
+    MPI_Comm_split(comm, color, 0, &sub_comm);
+  }
+  else
+    MPI_Comm_dup(comm, &sub_comm);
+
   auto cell_part = dolfinx::mesh::create_cell_partitioner(
       dolfinx::mesh::GhostMode::none, graph_part);
   auto mesh = dolfinx::mesh::create_box(
-      comm, {{{0.0, 0.0, 0.0}, {1.0, 1.0, 1.0}}}, {Nx, Ny, Nz},
+      comm, sub_comm, {{{0.0, 0.0, 0.0}, {1.0, 1.0, 1.0}}}, {Nx, Ny, Nz},
       dolfinx::mesh::CellType::tetrahedron, cell_part);
 
+  MPI_Comm_free(&sub_comm);
+
   if (dolfinx::MPI::rank(mesh.comm()) == 0)
   {
     std::cout << "UnitCube (" << Nx << "x" << Ny << "x" << Nz

src/mesh.h

@@ -21,7 +21,7 @@ class Mesh;
 
 dolfinx::mesh::Mesh<double>
 create_cube_mesh(MPI_Comm comm, std::size_t target_dofs, bool target_dofs_total,
-                 std::size_t dofs_per_node, int order);
+                 std::size_t dofs_per_node, int order, bool use_subcomm);
 
 std::shared_ptr<dolfinx::mesh::Mesh<double>>
 create_spoke_mesh(MPI_Comm comm, std::size_t target_dofs,