Provide KokkosComm::initialize/finalize

docs/api/core.rst

@@ -1,12 +1,15 @@
 Core
-====
+****
+
+MPI API Support
+===============
 
 .. list-table:: MPI API Support
     :widths: 40 30 15
     :header-rows: 1
 
     * - MPI
-      - ``KokkosComm::``
+      - KokkosComm
       - ``Kokkos::View``
     * - ``MPI_Send``
       - ``send`` or ``send<CommMode::Standard>``
@@ -33,6 +36,43 @@ Core
       - ``reduce``
       - ✓
 
+
+Initialization and finalization
+-------------------------------
+
+KokkosComm provides a unified interface for initializing and finalizing both Kokkos and MPI.
+
+.. Attention:: It is mandatory to use KokkosComm's initialization and finalization functions instead of their respective Kokkos and MPI counterparts.
+
+.. cpp:function:: void KokkosComm::initialize(int &argc, char *argv[], int mpi_required_thread_lvl)
+.. cpp:function:: void KokkosComm::initialize(int &argc, char *argv[])
+
+    Initializes the MPI execution environment with the required MPI thread level support (``MPI_THREAD_MULTIPLE`` if unspecified), and then initializes the Kokkos execution environment. This function also strips ``--kokkos-*`` flags to prevent Kokkos from printing them on all MPI ranks.
+
+    :param argc: Non-negative value representing the number of command-line arguments passed to the program.
+    :param argv: Pointer to the first element of an array of ``argc + 1`` pointers, of which the last one is null and the previous, if any, point to null-terminated multi-byte strings that represent the arguments passed to the program.
+    :param mpi_required_thread_lvl: Level of desired MPI thread support.
+
+    **Requirements:**
+
+    * ``KokkosComm::initialize`` has the same combined requirements as ``MPI_Init`` and ``Kokkos::initialize``.
+    * ``KokkosComm::initialize`` must be called in place of ``MPI_Init`` and ``Kokkos::initialize``.
+    * User-initiated MPI objects cannot be constructed, and MPI functions cannot be called until after ``KokkosComm::initialize`` is called.
+    * User-initiated Kokkos objects cannot be constructed until after ``KokkosComm::initialize`` is called.
+
+.. cpp:function:: void KokkosComm::finalize()
+
+    Terminates the Kokkos and MPI execution environments.
+
+    Programs are ill-formed if they do not call this function *after* calling ``KokkosComm::initialize``.
+
+    **Requirements:**
+
+    * ``KokkosComm::finalize`` has the same combined requirements as ``MPI_Finalize`` and ``Kokkos::finalize``.
+    * ``KokkosComm::finalize`` must be called in place of ``MPI_Finalize`` and ``Kokkos::finalize``.
+    * ``KokkosComm::finalize`` must be called after user-initialized Kokkos objects are out of scope.
+
+
 Point-to-point
 --------------
 

perf_tests/test_main.cpp

@@ -15,6 +15,7 @@
 //@HEADER
 
 #include "KokkosComm_include_mpi.hpp"
+#include "KokkosComm.hpp"
 
 #include <Kokkos_Core.hpp>
 #include <benchmark/benchmark.h>
@@ -32,15 +33,12 @@ class NullReporter : public ::benchmark::BenchmarkReporter {
 // The main is rewritten to allow for MPI initializing and for selecting a
 // reporter according to the process rank
 int main(int argc, char **argv) {
-  MPI_Init(&argc, &argv);
-
-  int rank;
-  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
-
-  Kokkos::initialize();
+  KokkosComm::initialize(argc, argv);
 
   ::benchmark::Initialize(&argc, argv);
 
+  int rank;
+  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
   if (rank == 0)
     // root process will use a reporter from the usual set provided by
     // ::benchmark
@@ -51,7 +49,6 @@ int main(int argc, char **argv) {
     ::benchmark::RunSpecifiedBenchmarks(&null);
   }
 
-  Kokkos::finalize();
-  MPI_Finalize();
+  KokkosComm::finalize();
   return 0;
-}
+}

src/KokkosComm.hpp

@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include "KokkosComm_include_mpi.hpp"
 #include "KokkosComm_collective.hpp"
 #include "KokkosComm_version.hpp"
 #include "KokkosComm_isend.hpp"
@@ -28,8 +29,49 @@
 
 #include <Kokkos_Core.hpp>
 
+#include <algorithm>
+#include <cstdio>
+#include <string_view>
+
 namespace KokkosComm {
 
+inline void initialize(int &argc, char *argv[], int mpi_required_thread_lvl) {
+  int flag;
+  MPI_Initialized(&flag);
+  // Eagerly abort if MPI has already been initialized
+  if (0 != flag) {
+    MPI_Abort(MPI_COMM_WORLD, -1);
+  }
+
+  int provided;
+  MPI_Init_thread(&argc, &argv, mpi_required_thread_lvl, &provided);
+  // Abort if MPI failed to provide the required thread level
+  if (mpi_required_thread_lvl != provided) {
+    MPI_Abort(MPI_COMM_WORLD, -1);
+  }
+
+  int rank;
+  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
+  // Strip "--help" and "--kokkos-help" from the flags passed to Kokkos if we are not on rank 0 to prevent Kokkos
+  // from printing the help message multiple times.
+  if (0 != rank) {
+    auto *help_it = std::find_if(
+        argv, argv + argc, [](std::string_view const &x) { return x.find("--kokkos-") != std::string_view::npos; });
+    if (help_it != argv + argc) {
+      std::swap(*help_it, *(argv + argc - 1));
+      --argc;
+    }
+  }
+  Kokkos::initialize(argc, argv);
+}
+
+inline void initialize(int &argc, char *argv[]) { initialize(argc, argv, MPI_THREAD_MULTIPLE); }
+
+inline void finalize() {
+  Kokkos::finalize();
+  MPI_Finalize();
+}
+
 template <CommMode SendMode = CommMode::Default, KokkosExecutionSpace ExecSpace, KokkosView SendView>
 Req isend(const ExecSpace &space, const SendView &sv, int dest, int tag, MPI_Comm comm) {
   return Impl::isend<SendMode>(space, sv, dest, tag, comm);

unit_tests/test_main.cpp

@@ -81,7 +81,8 @@ int main(int argc, char *argv[]) {
   // Intialize google test
   ::testing::InitGoogleTest(&argc, argv);
 
-  MPI_Init(&argc, &argv);
+  KokkosComm::initialize(argc, argv);
+
   int rank, size;
   MPI_Comm_rank(MPI_COMM_WORLD, &rank);
   MPI_Comm_size(MPI_COMM_WORLD, &size);
@@ -90,8 +91,6 @@ int main(int argc, char *argv[]) {
               << KOKKOSCOMM_VERSION_PATCH << ")\n";
   }
 
-  Kokkos::initialize();
-
   // Intialize google test
   ::testing::InitGoogleTest(&argc, argv);
 
@@ -105,9 +104,8 @@ int main(int argc, char *argv[]) {
   // run tests
   auto exit_code = RUN_ALL_TESTS();
 
-  // Finalize MPI before exiting
-  Kokkos::finalize();
-  MPI_Finalize();
+  // Finalize KokkosComm before exiting
+  KokkosComm::finalize();
 
   return exit_code;
 }