MPI shared memory socket scope

docs/sphinx/cookbook/shared_memory_allocators.rst

@@ -79,7 +79,11 @@ 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")};
-   auto node_allocator{rm.makeResource("SHARED::mpi3_alloc", traits)};
+   traits.scope = umpire::MemoryResourceTraits::shared_scope::socket; // or node
+   auto node_allocator{rm.makeResource("SHARED::MPI3::mpi3_alloc", traits)};
+
+.. note::
+   Socket scope requires Linux and MPI ranks bound such that each rank's CPU affinity mask maps to a single socket.
 
 See the bottom of this page for a full example of how to use MPI3 Shared Memory Allocators with Umpire.
 
@@ -97,9 +101,9 @@ To create these Shared Memory allocators, you can do the following:
    auto ipc_traits{umpire::get_default_resource_traits("SHARED::POSIX")};
 
    // then create an allocator:
-   auto mpi3_node_allocator{rm.makeResource("SHARED::mpi3_alloc", traits)};
+   auto mpi3_node_allocator{rm.makeResource("SHARED::mpi3_alloc", mpi3_traits)};
    // or
-   auto ipc_node_allocator{rm.makeResource("SHARED::ipc_alloc", traits)};
+   auto ipc_node_allocator{rm.makeResource("SHARED::ipc_alloc", ipc_traits)};
 
    // and allocate with
    mpi3_node_allocator.allocate(1024 * sizeof(double));
@@ -131,4 +135,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,11 +21,11 @@ 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
-  auto mpi3_shm_allocator = rm.makeResource("SHARED::mpi3_alloc", traits);
+  auto mpi3_shm_allocator = rm.makeResource("SHARED::MPI3::mpi3_alloc", traits);
 
   // Get communicator for the allocator
   MPI_Comm shm_comm = umpire::get_communicator_for_allocator(mpi3_shm_allocator, MPI_COMM_WORLD);

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)
@@ -318,6 +321,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();
 

src/umpire/resource/HostMpi3SharedMemoryResource.cpp

@@ -7,26 +7,254 @@
 
 #include "umpire/resource/HostMpi3SharedMemoryResource.hpp"
 
+#include <cerrno>
+#include <cstring>
+#include <fstream>
+#include <set>
+#include <sstream>
+#include <string>
+
+#if defined(__linux__)
+#include <sched.h>
+#include <unistd.h>
+#endif
+
 #include "umpire/resource/MemoryResource.hpp"
 #include "umpire/util/MPI.hpp"
 #include "umpire/util/Macros.hpp"
+#if defined(UMPIRE_ENABLE_NUMA)
+#include "umpire/util/numa.hpp"
+#endif
 #include "umpire/util/error.hpp"
 
 namespace umpire {
 namespace resource {
 
+namespace {
+
+constexpr int IGNORE_KEY{0};
+
+std::string get_mpi_error_message(int error_code)
+{
+  char buffer[MPI_MAX_ERROR_STRING];
+  int length{0};
+  const int status = MPI_Error_string(error_code, buffer, &length);
+
+  if (status != MPI_SUCCESS) {
+    return fmt::format("MPI error code {} (MPI_Error_string failed with code {})", error_code, status);
+  }
+
+  return std::string{buffer, static_cast<std::size_t>(length)};
+}
+
+void check_mpi_call(int error_code, const char* call_name)
+{
+  if (error_code != MPI_SUCCESS) {
+    UMPIRE_ERROR(runtime_error, fmt::format("{} failed: {}", call_name, get_mpi_error_message(error_code)));
+  }
+}
+
+#if defined(__linux__)
+/*!
+ * \brief Map an operating-system CPU index to the socket color used for
+ * communicator splitting.
+ *
+ * The helper first reads Linux topology information and then falls back to the
+ * NUMA helper when available.
+ *
+ * \param cpu CPU index from the current rank's affinity mask.
+ * \param color Output socket identifier used as the MPI_Comm_split color when
+ *        the lookup succeeds.
+ *
+ * \return true when a socket identifier was found for \p cpu, false otherwise.
+ */
+bool try_get_socket_color_for_cpu(int cpu, int& color)
+{
+  std::ifstream package_id_file{"/sys/devices/system/cpu/cpu" + std::to_string(cpu) + "/topology/physical_package_id"};
+  if (package_id_file) {
+    package_id_file >> color;
+    if (package_id_file) {
+      UMPIRE_LOG(Debug, "Using Linux physical_package_id " << color << " for cpu " << cpu);
+      return true;
+    }
+
+    UMPIRE_LOG(Debug, "Could not read physical_package_id for cpu " << cpu);
+  } else {
+    UMPIRE_LOG(Debug, "Could not open physical_package_id for cpu " << cpu);
+  }
+
+#if defined(UMPIRE_ENABLE_NUMA)
+  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());
+  }
+#endif
+
+  UMPIRE_USE_VAR(cpu);
+  UMPIRE_USE_VAR(color);
+  return false;
+}
+
+/*!
+ * \brief Format the set of socket colors found in a rank's affinity mask.
+ *
+ * \param colors Socket identifiers collected while inspecting the affinity
+ *        mask.
+ *
+ * \return Comma-separated list used in diagnostic messages.
+ */
+std::string format_socket_colors(const std::set<int>& colors)
+{
+  std::ostringstream ss;
+
+  bool first = true;
+  for (int color : colors) {
+    if (!first) {
+      ss << ", ";
+    }
+    ss << color;
+    first = false;
+  }
+
+  return ss.str();
+}
+#endif
+
+/*!
+ * \brief Determine the single socket color implied by the current rank's CPU
+ * affinity.
+ *
+ * Socket-scoped shared memory requires each rank to be pinned to exactly one
+ * socket so that all ranks with the same socket color can be grouped into one
+ * shared-memory communicator.
+ *
+ * \return Socket color used to split the node-local communicator.
+ */
+int get_socket_color_from_affinity()
+{
+#if defined(__linux__)
+  const long cpu_count = sysconf(_SC_NPROCESSORS_CONF);
+  if (cpu_count <= 0) {
+    UMPIRE_ERROR(runtime_error, "Could not determine the number of configured CPUs for socket scope");
+  }
+
+  const auto set_size = CPU_ALLOC_SIZE(cpu_count);
+  cpu_set_t* cpu_set = CPU_ALLOC(cpu_count);
+  if (!cpu_set) {
+    UMPIRE_ERROR(runtime_error, "Failed to allocate a CPU affinity mask for socket scope");
+  }
+
+  CPU_ZERO_S(set_size, cpu_set);
+
+  if (sched_getaffinity(0, set_size, cpu_set) != 0) {
+    const std::string message =
+        fmt::format("sched_getaffinity failed while determining socket scope: {}", std::strerror(errno));
+    CPU_FREE(cpu_set);
+    UMPIRE_ERROR(runtime_error, message);
+  }
+
+  std::set<int> socket_colors;
+
+  for (int cpu = 0; cpu < cpu_count; ++cpu) {
+    if (!CPU_ISSET_S(cpu, set_size, cpu_set)) {
+      continue;
+    }
+
+    int color{-1};
+    if (!try_get_socket_color_for_cpu(cpu, color)) {
+      const std::string message =
+          fmt::format("Unable to determine a socket color for cpu {} in the rank affinity mask", cpu);
+      CPU_FREE(cpu_set);
+      UMPIRE_ERROR(runtime_error, message);
+    }
+
+    socket_colors.insert(color);
+  }
+
+  CPU_FREE(cpu_set);
+
+  if (socket_colors.empty()) {
+    UMPIRE_ERROR(runtime_error, "The rank CPU affinity mask is empty; cannot determine socket scope");
+  }
+
+  if (socket_colors.size() != 1) {
+    UMPIRE_ERROR(runtime_error,
+                 fmt::format("Socket-scoped MPI3 shared memory requires each MPI rank to be bound to exactly one "
+                             "socket. This rank's affinity mask spans socket colors: {}",
+                             format_socket_colors(socket_colors)));
+  }
+
+  return *socket_colors.begin();
+#else
+  UMPIRE_ERROR(runtime_error,
+               "shared_scope::socket for MPI3 shared memory requires Linux CPU affinity information");
+#endif
+}
+
+/*!
+ * \brief Build the communicator that defines which ranks share allocations.
+ *
+ * \param comm Parent communicator used to discover ranks that may share host
+ *        memory.
+ * \param scope Requested sharing scope from MemoryResourceTraits. Node scope
+ *        keeps one communicator per node, while socket scope further splits the
+ *        node-local communicator using socket affinity.
+ *
+ * \return Communicator containing exactly the ranks that participate in each
+ * shared-memory allocation for the resource.
+ */
+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) {
+    check_mpi_call(MPI_Comm_split_type(comm, MPI_COMM_TYPE_SHARED, IGNORE_KEY, MPI_INFO_NULL, &shared_comm),
+                   "MPI_Comm_split_type");
+    return shared_comm;
+  }
+
+  if (scope == MemoryResourceTraits::shared_scope::socket) {
+    const int color = get_socket_color_from_affinity();
+
+    MPI_Comm node_comm{MPI_COMM_NULL};
+    check_mpi_call(MPI_Comm_split_type(comm, MPI_COMM_TYPE_SHARED, IGNORE_KEY, MPI_INFO_NULL, &node_comm),
+                   "MPI_Comm_split_type");
+    UMPIRE_LOG(Debug, "Creating socket-scoped shared communicator with color " << color);
+
+    const int split_status = MPI_Comm_split(node_comm, color, IGNORE_KEY, &shared_comm);
+    if (split_status != MPI_SUCCESS) {
+      if (node_comm != MPI_COMM_NULL) {
+        MPI_Comm_free(&node_comm);
+      }
+      check_mpi_call(split_status, "MPI_Comm_split");
+    }
+
+    check_mpi_call(MPI_Comm_free(&node_comm), "MPI_Comm_free");
+
+    return shared_comm;
+  }
+
+  UMPIRE_ERROR(runtime_error,
+               fmt::format("Unsupported shared communicator scope: {}", to_string(scope)));
+}
+
+} // end anonymous namespace
+
 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);
-  MPI_Comm_rank(m_shared_comm, &m_local_rank);
+  m_shared_comm = create_shared_communicator(util::MPI::getCommunicator(), traits.scope);
+  check_mpi_call(MPI_Comm_rank(m_shared_comm, &m_local_rank), "MPI_Comm_rank");
 
   // Free the comm at exit during cleanup in MPI_Finalize. We pass the m_shared_comm
   // by turning it into an int (as for Fortran) and then decoding that in the callback.
   int keyval = 0;
-  MPI_Comm_create_keyval(MPI_COMM_NULL_COPY_FN, free_comm, &keyval, nullptr);
-  MPI_Comm_set_attr(MPI_COMM_SELF, keyval, (void*)(intptr_t)MPI_Comm_c2f(m_shared_comm));
+  check_mpi_call(MPI_Comm_create_keyval(MPI_COMM_NULL_COPY_FN, free_comm, &keyval, nullptr), "MPI_Comm_create_keyval");
+  check_mpi_call(MPI_Comm_set_attr(MPI_COMM_SELF, keyval, (void*)(intptr_t)MPI_Comm_c2f(m_shared_comm)),
+                 "MPI_Comm_set_attr");
 }
 
 HostMpi3SharedMemoryResource::~HostMpi3SharedMemoryResource()
@@ -36,14 +264,15 @@ HostMpi3SharedMemoryResource::~HostMpi3SharedMemoryResource()
 
 void* HostMpi3SharedMemoryResource::allocate(std::size_t bytes)
 {
-  void* ptr;
-  MPI_Win win;
+  void* ptr{nullptr};
+  MPI_Win win{MPI_WIN_NULL};
   MPI_Aint local_size = (m_local_rank != 0) ? 0 : bytes;
   MPI_Aint size = bytes;
   int disp{sizeof(unsigned char)};
 
-  MPI_Win_allocate_shared(local_size, disp, MPI_INFO_NULL, m_shared_comm, &ptr, &win);
-  MPI_Win_shared_query(win, 0, &size, &disp, &ptr);
+  check_mpi_call(MPI_Win_allocate_shared(local_size, disp, MPI_INFO_NULL, m_shared_comm, &ptr, &win),
+                 "MPI_Win_allocate_shared");
+  check_mpi_call(MPI_Win_shared_query(win, 0, &size, &disp, &ptr), "MPI_Win_shared_query");
   m_shared_windows[ptr] = win;
 
   return ptr;
@@ -53,7 +282,7 @@ void HostMpi3SharedMemoryResource::deallocate(void* ptr, std::size_t)
 {
   auto window = m_shared_windows.find(ptr);
   if (window != m_shared_windows.end()) {
-    MPI_Win_free(&(window->second));
+    check_mpi_call(MPI_Win_free(&(window->second)), "MPI_Win_free");
     m_shared_windows.erase(window);
   } else {
     UMPIRE_ERROR(umpire::unknown_pointer_error, "");
@@ -73,14 +302,18 @@ 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))
 {
   // Interpret attribute_val as a MPI_Fint comm number.
   const auto comm_number = (MPI_Fint)(intptr_t)(attribute_val);
   MPI_Comm comm_to_free = MPI_Comm_f2c(comm_number);
-  MPI_Comm_free(&comm_to_free);
-  return MPI_SUCCESS;
+  return MPI_Comm_free(&comm_to_free);
 }
 
 } // end of namespace resource