initial changes for mpi shared memory socket scope

Author: kab163

docs/sphinx/cookbook/shared_memory_allocators.rst

@@ -79,6 +79,7 @@ To create an allocator with the MPI3 Shared Memory resource, you can do the foll
 .. code-block:: cpp
 
    auto traits{umpire::get_default_resource_traits("SHARED::MPI3")};
+   traits.scope = umpire::MemoryResourceTraits::shared_scope::socket; // or node
    auto node_allocator{rm.makeResource("SHARED::mpi3_alloc", traits)};
 
 See the bottom of this page for a full example of how to use MPI3 Shared Memory Allocators with Umpire.
@@ -131,4 +132,3 @@ when creating the MPI3 Shared Memory allocator, a name is not needed when alloca
 
 .. literalinclude:: ../../../examples/mpi3_shared_memory.cpp
    :language: cpp
-

examples/mpi3_shared_memory.cpp

@@ -21,7 +21,7 @@ int main(int argc, char** argv)
   auto traits = umpire::get_default_resource_traits("SHARED::MPI3");
   traits.size = 1 * 1024 * 1024; // 1 MB
 
-  // Node scope is required for mpi3 shared memory
+  // Node scope is the default for MPI3 shared memory; socket scope is also supported.
   traits.scope = umpire::MemoryResourceTraits::shared_scope::node;
 
   // Create allocator using MPI3 shared memory

src/umpire/Umpire.cpp

@@ -17,6 +17,9 @@
 
 #include "umpire/ResourceManager.hpp"
 #include "umpire/config.hpp"
+#if defined(UMPIRE_ENABLE_MPI3_SHARED_MEMORY)
+#include "umpire/resource/HostMpi3SharedMemoryResource.hpp"
+#endif
 #include "umpire/resource/HostSharedMemoryResource.hpp"
 #include "umpire/resource/MemoryResource.hpp"
 #if defined(UMPIRE_ENABLE_MPI) && defined(UMPIRE_ENABLE_IPC_SHARED_MEMORY)
@@ -27,6 +30,7 @@
 #include "umpire/strategy/DynamicPoolList.hpp"
 #include "umpire/strategy/QuickPool.hpp"
 #include "umpire/strategy/ResourceAwarePool.hpp"
+#include "umpire/util/mpi_shared.hpp"
 #include "umpire/util/wrap_allocator.hpp"
 
 #if !defined(_MSC_VER)
@@ -318,6 +322,10 @@ MPI_Comm get_communicator_for_allocator(Allocator a, MPI_Comm comm)
   if (auto alloc = dynamic_cast<strategy::DeviceIpcAllocator*>(a.getAllocationStrategy()))
     return alloc->get_scope_communicator();
 #endif
+#if defined(UMPIRE_ENABLE_MPI3_SHARED_MEMORY)
+  if (auto resource = dynamic_cast<resource::HostMpi3SharedMemoryResource*>(a.getAllocationStrategy()))
+    return resource->getSharedCommunicator();
+#endif
 
   std::map<int, MPI_Comm>& cached_communicators = get_cached_communicators();
 
@@ -329,8 +337,8 @@ MPI_Comm get_communicator_for_allocator(Allocator a, MPI_Comm comm)
   if (cached_comm != cached_communicators.end()) {
     c = cached_comm->second;
   } else {
-    if (scope == MemoryResourceTraits::shared_scope::node) {
-      MPI_Comm_split_type(comm, MPI_COMM_TYPE_SHARED, 0, MPI_INFO_NULL, &c);
+    if (scope == MemoryResourceTraits::shared_scope::node || scope == MemoryResourceTraits::shared_scope::socket) {
+      c = util::create_shared_communicator(comm, scope);
     } else {
       c = MPI_COMM_NULL;
     }

src/umpire/resource/HostMpi3SharedMemoryResource.cpp

@@ -10,6 +10,7 @@
 #include "umpire/resource/MemoryResource.hpp"
 #include "umpire/util/MPI.hpp"
 #include "umpire/util/Macros.hpp"
+#include "umpire/util/mpi_shared.hpp"
 #include "umpire/util/error.hpp"
 
 namespace umpire {
@@ -18,8 +19,7 @@ namespace resource {
 HostMpi3SharedMemoryResource::HostMpi3SharedMemoryResource(const std::string& name, int id, MemoryResourceTraits traits)
     : MemoryResource{name, id, traits}
 {
-  constexpr int IGNORE_KEY{0};
-  MPI_Comm_split_type(util::MPI::getCommunicator(), MPI_COMM_TYPE_SHARED, IGNORE_KEY, MPI_INFO_NULL, &m_shared_comm);
+  m_shared_comm = util::create_shared_communicator(util::MPI::getCommunicator(), traits.scope);
   MPI_Comm_rank(m_shared_comm, &m_local_rank);
 
   // Free the comm at exit during cleanup in MPI_Finalize. We pass the m_shared_comm
@@ -73,6 +73,11 @@ Platform HostMpi3SharedMemoryResource::getPlatform() noexcept
   return Platform::host;
 }
 
+MPI_Comm HostMpi3SharedMemoryResource::getSharedCommunicator() const noexcept
+{
+  return m_shared_comm;
+}
+
 int HostMpi3SharedMemoryResource::free_comm(MPI_Comm UMPIRE_UNUSED_ARG(comm), int UMPIRE_UNUSED_ARG(keyval),
                                             void* attribute_val, void* UMPIRE_UNUSED_ARG(extra_state))
 {

src/umpire/resource/HostMpi3SharedMemoryResource.hpp

@@ -4,8 +4,8 @@
 //
 // SPDX-License-Identifier: (MIT)
 //////////////////////////////////////////////////////////////////////////////
-#ifndef __Host_Shared_Memory_Resource_HPP
-#define __Host_Shared_Memory_Resource_HPP
+#ifndef UMPIRE_Host_Mpi3_Shared_Memory_Resource_HPP
+#define UMPIRE_Host_Mpi3_Shared_Memory_Resource_HPP
 
 #include <map>
 #include <memory>
@@ -32,6 +32,8 @@ class HostMpi3SharedMemoryResource : public MemoryResource {
 
   Platform getPlatform() noexcept override;
 
+  MPI_Comm getSharedCommunicator() const noexcept;
+
  private:
   static int free_comm(MPI_Comm comm, int keyval, void* attribute_val, void* extra_state);
 
@@ -42,4 +44,4 @@ class HostMpi3SharedMemoryResource : public MemoryResource {
 
 } // end of namespace resource
 } // end of namespace umpire
-#endif // __Host_Shared_Memory_Resource_HPP
+#endif // UMPIRE_Host_Mpi3_Shared_Memory_Resource_HPP

src/umpire/resource/HostMpi3SharedMemoryResourceFactory.cpp

@@ -28,8 +28,9 @@ std::unique_ptr<resource::MemoryResource> HostMpi3SharedMemoryResourceFactory::c
 std::unique_ptr<resource::MemoryResource> HostMpi3SharedMemoryResourceFactory::create(const std::string& name, int id,
                                                                                       MemoryResourceTraits traits)
 {
-  if (traits.scope != MemoryResourceTraits::shared_scope::node) {
-    UMPIRE_ERROR(runtime_error, "HostMpi3SharedMemoryResource only supports shared_scope::node");
+  if (traits.scope != MemoryResourceTraits::shared_scope::node &&
+      traits.scope != MemoryResourceTraits::shared_scope::socket) {
+    UMPIRE_ERROR(runtime_error, "HostMpi3SharedMemoryResource only supports shared_scope::node or ::socket");
   }
   return util::make_unique<HostMpi3SharedMemoryResource>(name, id, traits);
 }

src/umpire/util/CMakeLists.txt

@@ -22,6 +22,7 @@ set (umpire_util_headers
   MemoryResourceTraits.hpp
   MemoryMap.hpp
   MemoryMap.inl
+  mpi_shared.hpp
   OutputBuffer.hpp
   Platform.hpp
   allocation_statistics.hpp
@@ -43,6 +44,7 @@ set (umpire_util_sources
   io.cpp
   Logger.cpp
   MPI.cpp
+  mpi_shared.cpp
   OutputBuffer.cpp
   allocation_statistics.cpp
   detect_vendor.cpp)

src/umpire/util/mpi_shared.cpp

@@ -0,0 +1,121 @@
+//////////////////////////////////////////////////////////////////////////////
+// Copyright (c) 2016-26, Lawrence Livermore National Security, LLC and Umpire
+// project contributors. See the COPYRIGHT file for details.
+//
+// SPDX-License-Identifier: (MIT)
+//////////////////////////////////////////////////////////////////////////////
+
+#include "umpire/util/mpi_shared.hpp"
+
+#if defined(UMPIRE_ENABLE_MPI)
+
+#include <fstream>
+#include <string>
+
+#if defined(__linux__)
+#include <sched.h>
+#endif
+
+#include "umpire/util/Macros.hpp"
+#if defined(UMPIRE_ENABLE_NUMA)
+#include "umpire/util/numa.hpp"
+#endif
+#include "umpire/util/error.hpp"
+
+namespace umpire {
+namespace util {
+
+namespace {
+
+constexpr int IGNORE_KEY{0};
+
+bool get_socket_color_from_package_id(int& color)
+{
+#if defined(__linux__)
+  const int cpu = sched_getcpu();
+  if (cpu < 0) {
+    UMPIRE_LOG(Debug, "sched_getcpu failed while determining socket color");
+    return false;
+  }
+
+  std::ifstream package_id_file{"/sys/devices/system/cpu/cpu" + std::to_string(cpu) +
+                                "/topology/physical_package_id"};
+  if (!package_id_file) {
+    UMPIRE_LOG(Debug, "Could not open physical_package_id for cpu " << cpu);
+    return false;
+  }
+
+  package_id_file >> color;
+  if (!package_id_file) {
+    UMPIRE_LOG(Debug, "Could not read physical_package_id for cpu " << cpu);
+    return false;
+  }
+
+  UMPIRE_LOG(Debug, "Using Linux physical_package_id " << color << " for cpu " << cpu);
+  return true;
+#else
+  UMPIRE_USE_VAR(color);
+  return false;
+#endif
+}
+
+bool get_socket_color_from_numa(int& color)
+{
+#if defined(UMPIRE_ENABLE_NUMA) && defined(__linux__)
+  const int cpu = sched_getcpu();
+  if (cpu < 0) {
+    UMPIRE_LOG(Debug, "sched_getcpu failed while determining NUMA fallback color");
+    return false;
+  }
+
+  try {
+    color = numa::node_of_cpu(cpu);
+    UMPIRE_LOG(Debug, "Using NUMA node " << color << " as socket color fallback for cpu " << cpu);
+    return true;
+  } catch (const std::exception& e) {
+    UMPIRE_LOG(Debug, "NUMA fallback failed for cpu " << cpu << ": " << e.what());
+    return false;
+  }
+#else
+  UMPIRE_USE_VAR(color);
+  return false;
+#endif
+}
+
+} // end anonymous namespace
+
+MPI_Comm create_shared_communicator(MPI_Comm comm, MemoryResourceTraits::shared_scope scope)
+{
+  MPI_Comm shared_comm{MPI_COMM_NULL};
+
+  if (scope == MemoryResourceTraits::shared_scope::node) {
+    MPI_Comm_split_type(comm, MPI_COMM_TYPE_SHARED, IGNORE_KEY, MPI_INFO_NULL, &shared_comm);
+    return shared_comm;
+  }
+
+  if (scope == MemoryResourceTraits::shared_scope::socket) {
+    MPI_Comm node_comm{MPI_COMM_NULL};
+    MPI_Comm_split_type(comm, MPI_COMM_TYPE_SHARED, IGNORE_KEY, MPI_INFO_NULL, &node_comm);
+
+    int color{-1};
+    if (!get_socket_color_from_package_id(color) && !get_socket_color_from_numa(color)) {
+      MPI_Comm_free(&node_comm);
+      UMPIRE_ERROR(runtime_error,
+                   "Unable to determine a socket color for shared_scope::socket. "
+                   "Expected Linux CPU topology data or a NUMA fallback.");
+    }
+
+    UMPIRE_LOG(Debug, "Creating socket-scoped shared communicator with color " << color);
+    MPI_Comm_split(node_comm, color, IGNORE_KEY, &shared_comm);
+    MPI_Comm_free(&node_comm);
+    return shared_comm;
+  }
+
+  UMPIRE_ERROR(runtime_error,
+               fmt::format("Unsupported shared communicator scope: {}", to_string(scope)));
+}
+
+} // end of namespace util
+} // end of namespace umpire
+
+#endif

src/umpire/util/mpi_shared.hpp

@@ -0,0 +1,27 @@
+//////////////////////////////////////////////////////////////////////////////
+// Copyright (c) 2016-26, Lawrence Livermore National Security, LLC and Umpire
+// project contributors. See the COPYRIGHT file for details.
+//
+// SPDX-License-Identifier: (MIT)
+//////////////////////////////////////////////////////////////////////////////
+#ifndef UMPIRE_mpi_shared_HPP
+#define UMPIRE_mpi_shared_HPP
+
+#include "umpire/config.hpp"
+#include "umpire/util/MemoryResourceTraits.hpp"
+
+#if defined(UMPIRE_ENABLE_MPI)
+#include "mpi.h"
+#endif
+
+namespace umpire {
+namespace util {
+
+#if defined(UMPIRE_ENABLE_MPI)
+MPI_Comm create_shared_communicator(MPI_Comm comm, MemoryResourceTraits::shared_scope scope);
+#endif
+
+} // end of namespace util
+} // end of namespace umpire
+
+#endif // UMPIRE_mpi_shared_HPP

src/umpire/util/numa.cpp

@@ -58,6 +58,19 @@ int get_location(void* ptr)
   return numa_node;
 }
 
+int node_of_cpu(int cpu)
+{
+  if (numa_available() < 0)
+    UMPIRE_ERROR(runtime_error, "libnuma is unusable.");
+
+  const int numa_node = numa_node_of_cpu(cpu);
+  if (numa_node < 0) {
+    UMPIRE_ERROR(runtime_error, fmt::format("numa::node_of_cpu error for cpu {}", cpu));
+  }
+
+  return numa_node;
+}
+
 std::vector<int> get_host_nodes()
 {
   if (numa_available() < 0)