Multi-Node rocshmem_finalize() bug (#138)

src/CMakeLists.txt

@@ -32,7 +32,7 @@ target_sources(
     backend_bc.cpp
     context_host.cpp
     context_device.cpp
-    mpi_init_singleton.cpp
+    mpi_instance.cpp
     rocshmem_gpu.cpp
     rocshmem.cpp
     team.cpp

src/backend_bc.cpp

@@ -86,18 +86,7 @@ Backend::Backend(MPI_Comm comm) : heap{comm} {
 }
 
 void Backend::init_mpi_once(MPI_Comm comm) {
-  int init_done{};
-  NET_CHECK(MPI_Initialized(&init_done));
-
-  int provided{};
-  if (!init_done) {
-    NET_CHECK(MPI_Init_thread(0, 0, MPI_THREAD_MULTIPLE, &provided));
-    if (provided != MPI_THREAD_MULTIPLE) {
-      fprintf(stderr, "MPI_THREAD_MULTIPLE support disabled.\n");
-    }
-  }
   if (comm == MPI_COMM_NULL) comm = MPI_COMM_WORLD;
-
   NET_CHECK(MPI_Comm_dup(comm, &backend_comm));
   NET_CHECK(MPI_Comm_size(backend_comm, &num_pes));
   NET_CHECK(MPI_Comm_rank(backend_comm, &my_pe));

src/memory/remote_heap_info.hpp

@@ -55,15 +55,8 @@ class CommunicatorMPI {
   CommunicatorMPI(char* heap_base, size_t heap_size,
                   MPI_Comm comm = MPI_COMM_WORLD)
     : comm_{comm} {
-    int initialized;
-    MPI_Initialized(&initialized);
-    if (!initialized) {
-      int provided;
-      MPI_Init_thread(nullptr, nullptr, MPI_THREAD_MULTIPLE, &provided);
-    }
     MPI_Comm_rank(comm_, &my_pe_);
     MPI_Comm_size(comm_, &num_pes_);
-
     heap_window_info_ = WindowInfo(comm_, heap_base, heap_size);
   }
 

src/mpi_instance.cpp

@@ -22,13 +22,11 @@
  * IN THE SOFTWARE.
  *****************************************************************************/
 
-#include "mpi_init_singleton.hpp"
+#include "mpi_instance.hpp"
 
 namespace rocshmem {
 
-MPIInitSingleton* MPIInitSingleton::instance{nullptr};
-
-MPIInitSingleton::MPIInitSingleton() {
+MPIInstance::MPIInstance(MPI_Comm comm) {
   MPI_Initialized(&pre_init_done);
 
   if (!pre_init_done) {
@@ -36,11 +34,15 @@ MPIInitSingleton::MPIInitSingleton() {
     MPI_Init_thread(nullptr, nullptr, MPI_THREAD_MULTIPLE, &provided);
   }
 
-  MPI_Comm_size(MPI_COMM_WORLD, &nprocs_);
-  MPI_Comm_rank(MPI_COMM_WORLD, &my_rank_);
+  if (comm == MPI_COMM_NULL) {
+    comm = MPI_COMM_WORLD;
+  }
+
+  MPI_Comm_size(comm, &nprocs_);
+  MPI_Comm_rank(comm, &my_rank_);
 }
 
-MPIInitSingleton::~MPIInitSingleton() {
+MPIInstance::~MPIInstance() {
   int finalized{0};
   MPI_Finalized(&finalized);
   if (!finalized && !pre_init_done) {
@@ -48,16 +50,8 @@ MPIInitSingleton::~MPIInitSingleton() {
   }
 }
 
-MPIInitSingleton* MPIInitSingleton::GetInstance() {
-  if (!instance) {
-    instance = new MPIInitSingleton();
-    return instance;
-  }
-  return instance;
-}
+int MPIInstance::get_rank() { return my_rank_; }
 
-int MPIInitSingleton::get_rank() { return my_rank_; }
-
-int MPIInitSingleton::get_nprocs() { return nprocs_; }
+int MPIInstance::get_nprocs() { return nprocs_; }
 
 }  // namespace rocshmem

src/mpi_instance.hpp

@@ -22,77 +22,64 @@
  * IN THE SOFTWARE.
  *****************************************************************************/
 
-#ifndef LIBRARY_SRC_MPI_INIT_SINGLETON_HPP_
-#define LIBRARY_SRC_MPI_INIT_SINGLETON_HPP_
+#ifndef LIBRARY_SRC_MPI_INSTANCE_HPP_
+#define LIBRARY_SRC_MPI_INSTANCE_HPP_
 
 #include <mpi.h>
 
 #include <memory>
 
 /**
- * @file mpi_init_singleton.hpp
+ * @file mpi_instance.hpp
  *
  * @brief Contains MPI library initialization code
  */
 
 namespace rocshmem {
 
-class MPIInitSingleton {
- private:
-  /**
-   * @brief Primary constructor
-   */
-  MPIInitSingleton();
+class MPIInstance {
+  public:
+    /**
+     * @brief Primary constructor
+     */
+    MPIInstance(MPI_Comm comm);
 
- public:
-  /**
-   * @brief Destructor
-   */
-  ~MPIInitSingleton();
+    /**
+     * @brief Destructor
+     */
+    ~MPIInstance();
 
-  /**
-   * @brief Invoke singleton construction or return handle
-   *
-   * @return Initialized handle to singleton
-   */
-  static MPIInitSingleton* GetInstance();
+    /**
+     * @brief Accessor for my COMM_WORLD rank identifier
+     *
+     * @return My COMM_WORLD rank identifier
+     */
+    int get_rank();
 
-  /**
-   * @brief Accessor for my COMM_WORLD rank identifier
-   *
-   * @return My COMM_WORLD rank identifier
-   */
-  int get_rank();
+    /**
+     * @brief Accessor for number or processes in COMM_WORLD
+     *
+     * @return Number of processes in COMM_WORLD
+     */
+    int get_nprocs();
 
-  /**
-   * @brief Accessor for number or processes in COMM_WORLD
-   *
-   * @return Number of processes in COMM_WORLD
-   */
-  int get_nprocs();
+  private:
+    /**
+     * @brief My MPI rank identifier
+     */
+    int my_rank_{-1};
 
- private:
-  /**
-   * @brief My MPI rank identifier
-   */
-  int my_rank_{-1};
+    /**
+     * @brief Number of MPI processes
+     */
+    int nprocs_{-1};
 
-  /**
-   * @brief Number of MPI processes
-   */
-  int nprocs_{-1};
-
-  /**
-   * @brief Was MPI initialized before rocshmem_init call
-   */
-  int pre_init_done{0};
-
-  /**
-   * @brief Refers to global variable
-   */
-  static MPIInitSingleton* instance;
+    /**
+     * @brief Was MPI initialized before rocshmem_init call
+     */
+    int pre_init_done{0};
 };
 
 }  // namespace rocshmem
 
-#endif  // LIBRARY_SRC_MPI_INIT_SINGLETON_HPP_
+#endif  // LIBRARY_SRC_MPI_INSTANCE_HPP_

src/reverse_offload/mpi_transport.cpp

@@ -46,16 +46,7 @@ namespace rocshmem {
 
 MPITransport::MPITransport(MPI_Comm comm, Queue* q)
   : queue{q}, Transport{} {
-  int init_done{};
-  NET_CHECK(MPI_Initialized(&init_done));
 
-  int provided{};
-  if (!init_done) {
-    NET_CHECK(MPI_Init_thread(0, 0, MPI_THREAD_MULTIPLE, &provided));
-    if (provided != MPI_THREAD_MULTIPLE) {
-      fprintf(stderr, "MPI_THREAD_MULTIPLE support disabled.\n");
-    }
-  }
   assert(comm != MPI_COMM_NULL);
 
   NET_CHECK(MPI_Comm_dup(comm, &ro_net_comm_world));

src/rocshmem.cpp

@@ -47,7 +47,7 @@
 #include "ipc/backend_ipc.hpp"
 #include "ipc/context_ipc_tmpl_host.hpp"
 #endif
-#include "mpi_init_singleton.hpp"
+#include "mpi_instance.hpp"
 #include "team.hpp"
 #include "templates_host.hpp"
 #include "util.hpp"
@@ -67,6 +67,7 @@ namespace rocshmem {
   }
 
 Backend *backend = nullptr;
+MPIInstance *mpi_instance = nullptr;
 
 rocshmem_ctx_t ROCSHMEM_HOST_CTX_DEFAULT;
 
@@ -86,6 +87,8 @@ rocshmem_ctx_t ROCSHMEM_HOST_CTX_DEFAULT;
 
   rocm_init();
 
+  mpi_instance = new MPIInstance(comm);
+
 #ifdef USE_RO
   CHECK_HIP(hipHostMalloc(&backend, sizeof(ROBackend)));
   backend = new (backend) ROBackend(comm);
@@ -103,7 +106,7 @@ rocshmem_ctx_t ROCSHMEM_HOST_CTX_DEFAULT;
 						 rocshmem_init_attr_t *attr) {
   MPI_Comm comm = MPI_COMM_NULL;
 
-  if ((attr == nullptr) || 
+  if ((attr == nullptr) ||
       ((flags != ROCSHMEM_INIT_WITH_UNIQUEID) &&
        (flags != ROCSHMEM_INIT_WITH_MPI_COMM)) ) {
     fprintf(stderr, "ROCSHMEM_ERROR: %s in file '%s' in line %d\n",
@@ -224,24 +227,21 @@ rocshmem_ctx_t ROCSHMEM_HOST_CTX_DEFAULT;
 }
 
 [[maybe_unused]] __host__ int rocshmem_my_pe() {
-  if(backend == nullptr) {
-    MPIInitSingleton *s = s->GetInstance();
-    return s->get_rank();
-  }
-  else
-  {
-    return backend->getMyPE();
+  if (mpi_instance != nullptr) {
+    return mpi_instance->get_rank();
   }
+
+  fprintf(stderr, "[WARNING] rocshmem_init() has not been called\n");
+  return -1;
 }
 
 [[maybe_unused]] __host__ int rocshmem_n_pes() {
-  if(backend == nullptr) {
-    MPIInitSingleton *s = s->GetInstance();
-    return s->get_nprocs();
-  }
-  else {
-    return backend->getNumPEs();
+  if (mpi_instance != nullptr) {
+    return mpi_instance->get_nprocs();
   }
+
+  fprintf(stderr, "[WARNING] rocshmem_init() has not been called\n");
+  return -1;
 }
 
 [[maybe_unused]] __host__ void *rocshmem_malloc(size_t size) {
@@ -294,7 +294,7 @@ rocshmem_ctx_t ROCSHMEM_HOST_CTX_DEFAULT;
   backend->~Backend();
   CHECK_HIP(hipHostFree(backend));
 
-  delete MPIInitSingleton::GetInstance();
+  delete mpi_instance;
 }
 
 __host__ void rocshmem_query_thread(int *provided) {