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,67 @@ Core
       - ``reduce``
       - ✓
 
+
+Initialization and finalization
+-------------------------------
+
+KokkosComm provides a unified interface for initializing and finalizing both Kokkos and MPI.
+
+.. Attention:: It is strongly recommended to use KokkosComm's initialization and finalization functions instead of their respective Kokkos and MPI counterparts. However, users have two options for using KokkosComm:
+
+    1. Initialize/finalize KokkosComm using ``KokkosComm::{initialize,finalize}``. In this case, it is **forbidden** to call ``MPI_{Init,Init_thread,Finalize}`` and ``Kokkos::{initialize,finalize}`` (otherwise, the application will abort).
+    2. Initialize/finalize MPI and Kokkos manually through their respective interfaces. In this case, it is **forbidden** to call ``KokkosComm::{initialize,finalize}`` (otherwise, the application will abort).
+
+.. cpp:enum-class:: KokkosComm::ThreadSupportLevel
+
+    A scoped enum to wrap the MPI thread support levels.
+
+    .. cpp:enumerator:: KokkosComm::ThreadSupportLevel::Single
+
+        Only one thread will execute.
+
+    .. cpp:enumerator:: KokkosComm::ThreadSupportLevel::Funneled
+
+        The process may be multi-threaded, but only the main thread will make MPI calls (all MPI calls are funneled to the main thread).
+
+    .. cpp:enumerator:: KokkosComm::ThreadSupportLevel::Serialized
+
+        The process may be multi-threaded, and multiple threads may make MPI calls, but only one at a time: MPI calls are not made concurrently from two distinct threads (all MPI calls are serialized).
+
+    .. cpp:enumerator:: KokkosComm::ThreadSupportLevel::Multiple
+
+        Multiple threads may call MPI, with no restrictions.
+
+
+.. cpp:function:: void KokkosComm::initialize(int &argc, char *argv[], KokkosComm::ThreadSupportLevel required)
+.. 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-help`` 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 required: Level of desired MPI thread support.
+
+    **Requirements:**
+
+    * ``KokkosComm::initialize`` has the same combined requirements as ``MPI_{Init,Init_thread}`` 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,71 @@
 
 #include <Kokkos_Core.hpp>
 
+#include <algorithm>
+#include <cstdio>
+#include <string_view>
+
 namespace KokkosComm {
 
+enum class ThreadSupportLevel {
+  Single     = MPI_THREAD_SINGLE,
+  Funneled   = MPI_THREAD_FUNNELED,
+  Serialized = MPI_THREAD_SERIALIZED,
+  Multiple   = MPI_THREAD_MULTIPLE,
+};
+
+inline void initialize(int &argc, char *argv[], ThreadSupportLevel required) {
+  int flag;
+  MPI_Initialized(&flag);
+  // Forbid calling this function if MPI has already been initialized
+  if (0 != flag) {
+    MPI_Abort(MPI_COMM_WORLD, -1);
+  }
+
+  // Forbid calling this function if Kokkos has already been initialized
+  if (Kokkos::is_initialized()) {
+    std::abort();
+  }
+
+  int provided;
+  MPI_Init_thread(&argc, &argv, static_cast<int>(required), &provided);
+  // Abort if MPI failed to provide the required thread support level
+  if (static_cast<int>(required) < 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) {
+    if (auto *help_it = std::find_if(argv, argv + argc, [](std::string_view const &x) { return x == "--kokkos-help"; });
+        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, ThreadSupportLevel::Multiple); }
+
+inline void finalize() {
+  // Forbid calling this function if Kokkos has already been finalized or isn't yet initialized
+  if (Kokkos::is_finalized() || !Kokkos::is_initialized()) {
+    MPI_Abort(MPI_COMM_WORLD, -1);
+  }
+  Kokkos::finalize();
+
+  int flag;
+  MPI_Finalized(&flag);
+  // Forbid calling this function if MPI has already been finalized
+  if (0 != flag) {
+    std::abort();
+  }
+  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;
 }