DBM: Call MPI_Dims_create in dbm_miniapp

cp2k · Jun 26, 2022 · 80ba986 · 80ba986
1 parent b8ba7d1
commit 80ba986
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 6 deletions.
diff --git a/src/dbm/dbm_miniapp.c b/src/dbm/dbm_miniapp.c
@@ -184,21 +184,22 @@ int main(int argc, char *argv[]) {
   const int nranks = dbm_mpi_comm_size(world_comm);
   const int my_rank = dbm_mpi_comm_rank(world_comm);
 
-  if (my_rank == 0) {
-    printf("MPI-ranks: %i  OpenMP-threads: %i\n\n", nranks,
-           omp_get_max_threads());
-  }
-
   if (offload_get_device_count() > 0) {
     offload_set_chosen_device(my_rank % offload_get_device_count());
   }
 
   // Create 2D cart.
-  const int dims[2] = {nranks, 1};
+  int dims[2] = {0, 0};
+  dbm_mpi_dims_create(nranks, 2, dims);
   const int periods[2] = {true, true};
   dbm_mpi_comm_t comm =
       dbm_mpi_cart_create(world_comm, 2, dims, periods, false);
 
+  if (my_rank == 0) {
+    printf("MPI-ranks: %i  MPI-cart: %i x %i  OpenMP-threads: %i\n\n", nranks,
+           dims[0], dims[1], omp_get_max_threads());
+  }
+
   bechmark_multiply(4, 4, 4, comm);
 
   bechmark_multiply(128, 4, 4, comm);

diff --git a/src/dbm/dbm_mpi.c b/src/dbm/dbm_mpi.c
@@ -115,6 +115,21 @@ int dbm_mpi_comm_size(const dbm_mpi_comm_t comm) {
 #endif
 }
 
+/*******************************************************************************
+ * \brief Wrapper around MPI_Dims_create.
+ * \author Ole Schuett
+ ******************************************************************************/
+void dbm_mpi_dims_create(const int nnodes, const int ndims, int dims[]) {
+#if defined(__parallel)
+  CHECK(MPI_Dims_create(nnodes, ndims, dims));
+#else
+  dims[0] = nnodes;
+  for (int i = 1; i < ndims; i++) {
+    dims[i] = 1;
+  }
+#endif
+}
+
 /*******************************************************************************
  * \brief Wrapper around MPI_Cart_create.
  * \author Ole Schuett

diff --git a/src/dbm/dbm_mpi.h b/src/dbm/dbm_mpi.h
@@ -60,6 +60,12 @@ int dbm_mpi_comm_rank(const dbm_mpi_comm_t comm);
  ******************************************************************************/
 int dbm_mpi_comm_size(const dbm_mpi_comm_t comm);
 
+/*******************************************************************************
+ * \brief Wrapper around MPI_Dims_create.
+ * \author Ole Schuett
+ ******************************************************************************/
+void dbm_mpi_dims_create(const int nnodes, const int ndims, int dims[]);
+
 /*******************************************************************************
  * \brief Wrapper around MPI_Cart_create.
  * \author Ole Schuett