MPI shared memory socket scope

.github/hubcast.yml

@@ -5,7 +5,7 @@ Repo:
 
   check_name: gitlab-ci
 
-  check_types: [child-pipelines, jobs]
+  check_types: [pipelines, child-pipelines]
 
   delete_closed: true
 

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::mpi3_alloc", mpi3_traits)};
    // or
-   auto ipc_node_allocator{rm.makeResource("SHARED::ipc_alloc", traits)};
+   auto ipc_node_allocator{rm.makeResource("SHARED::POSIX::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

@@ -1,12 +1,12 @@
 #include <mpi.h>
 
 #include <iostream>
+#include <string>
 
 #include "umpire/Allocator.hpp"
 #include "umpire/ResourceManager.hpp"
 #include "umpire/Umpire.hpp"
 #include "umpire/config.hpp"
-#include "umpire/resource/HostMpi3SharedMemoryResource.hpp"
 #include "umpire/strategy/NamedAllocationStrategy.hpp"
 #include "umpire/util/MemoryResourceTraits.hpp"
 
@@ -16,16 +16,37 @@ int main(int argc, char** argv)
 
   auto& rm = umpire::ResourceManager::getInstance();
 
+  int world_rank = 0;
+  MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);
+
   // Use MPI3 shared memory resource
   // Note: Could also use "SHARED"
   auto traits = umpire::get_default_resource_traits("SHARED::MPI3");
   traits.size = 1 * 1024 * 1024; // 1 MB
 
-  // Node scope is required for mpi3 shared memory
-  traits.scope = umpire::MemoryResourceTraits::shared_scope::node;
+  // Node scope is the default for MPI3 shared memory; socket scope is also supported.
+  // Pass `--socket` to request socket scope (falls back to node scope when the rank is not pinned to one socket).
+  const bool request_socket_scope = (argc > 1) && (std::string{argv[1]} == "--socket");
+  traits.scope = request_socket_scope ? umpire::MemoryResourceTraits::shared_scope::socket
+                                      : umpire::MemoryResourceTraits::shared_scope::node;
+
+  if (traits.scope == umpire::MemoryResourceTraits::shared_scope::socket) {
+    std::string reason;
+	  if (!umpire::affinity_maps_to_single_socket(reason)) {
+      if (world_rank == 0) {
+        std::cerr << "Requested socket-scoped MPI3 shared memory, but CPU affinity does not map to a single socket: "
+                  << reason << "\nFalling back to node-scoped shared memory.\n";
+      }
+      traits.scope = umpire::MemoryResourceTraits::shared_scope::node;
+    } else {
+      if (world_rank == 0) {
+        std::cout << "Running with socket-scope!" << std::endl;
+      }
+    }
+  }
 
   // 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);
@@ -47,6 +68,7 @@ int main(int argc, char** argv)
 
   mpi3_shm_allocator.deallocate(data);
 
+  // Since we called get_communicator_for_allocator...
   umpire::cleanup_cached_communicators();
   MPI_Finalize();
 

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();
 
@@ -352,6 +359,16 @@ void cleanup_cached_communicators()
 }
 #endif
 
+bool affinity_maps_to_single_socket(std::string& reason)
+{
+#if defined(UMPIRE_ENABLE_MPI3_SHARED_MEMORY)
+  return resource::affinity_maps_to_single_socket(reason);
+#else
+  reason = "MPI3 shared memory support is disabled";
+  return false;
+#endif
+}
+
 void register_external_allocation(void* ptr, util::AllocationRecord record)
 {
   umpire::event::record([&](auto& event) {

src/umpire/Umpire.hpp

@@ -193,19 +193,36 @@ umpire::MemoryResourceTraits get_default_resource_traits(const std::string& name
  */
 void* find_pointer_from_name(Allocator allocator, const std::string& name);
 
-#if defined(UMPIRE_ENABLE_MPI)
-/*!
- * \brief Return the MPI communicator for a shared memory allocator.
- *
+	#if defined(UMPIRE_ENABLE_MPI)
+	/*!
+	 * \brief Return the MPI communicator for a shared memory allocator.
+	 *
  * NOTE: Using this function will REQUIRE users to call the
  * cleanup_cached_communicators() function to avoid memory leaks.
  */
-MPI_Comm get_communicator_for_allocator(Allocator a, MPI_Comm comm);
-void cleanup_cached_communicators();
-#endif
-
-void register_external_allocation(void* ptr, util::AllocationRecord record);
-util::AllocationRecord deregister_external_allocation(void* ptr);
+	MPI_Comm get_communicator_for_allocator(Allocator a, MPI_Comm comm);
+	void cleanup_cached_communicators();
+	#endif
+
+	/*!
+	 * \brief Determine whether the calling rank's CPU affinity mask spans exactly
+	 * one socket.
+	 *
+	 * Socket-scoped MPI3 shared memory requires each rank pinned to exactly one
+	 * socket so that ranks can be grouped into socket-local communicators.
+	 *
+	 * When MPI3 shared memory support is disabled (or the platform does not
+	 * provide the required affinity information), this function returns false and
+	 * populates \p reason with a diagnostic message.
+	 *
+	 * \param reason Diagnostic message describing why the check failed.
+	 *
+	 * \return true when the affinity mask is non-empty and maps to a single socket.
+	 */
+	bool affinity_maps_to_single_socket(std::string& reason);
+
+	void register_external_allocation(void* ptr, util::AllocationRecord record);
+	util::AllocationRecord deregister_external_allocation(void* ptr);
 
 /*!
  * \brief Returns the Camp resource associated with a particular allocation