WIP: MPI shared memory

c++/nda/concepts.hpp

@@ -24,6 +24,8 @@
 #include "./stdutil/concepts.hpp"
 #include "./traits.hpp"
 
+#include <mpi/communicator.hpp>
+
 #include <array>
 #include <concepts>
 #include <type_traits>
@@ -159,6 +161,7 @@ namespace nda {
     /// @cond
     // Forward declarations.
     struct blk_t;
+    struct blk_shm_t;
     enum class AddressSpace;
     /// @endcond
 
@@ -181,6 +184,14 @@ namespace nda {
       { A::address_space } -> std::same_as<AddressSpace const &>;
     };
 
+    template <typename A>
+    concept MPISharedMemoryAllocator = requires(A &a) {
+      { a.allocate(size_t{}, mpi::shared_communicator{}) } noexcept -> std::same_as<blk_shm_t>;
+      { a.allocate_zero(size_t{}, mpi::shared_communicator{}) } noexcept -> std::same_as<blk_shm_t>;
+      { a.deallocate(std::declval<blk_shm_t>()) } noexcept;
+      { A::address_space } -> std::same_as<AddressSpace const &>;
+    };
+
     /**
      * @brief Check if a given type satisfies the memory handle concept.
      *

c++/nda/mem/allocators.hpp

@@ -28,6 +28,8 @@
 #include "./memset.hpp"
 #include "../macros.hpp"
 
+#include <mpi/mpi.hpp>
+
 #include <algorithm>
 #include <cstddef>
 #include <cstdint>
@@ -63,6 +65,18 @@ namespace nda::mem {
     size_t s = 0;
   };
 
+  /// Memory block consisting of a pointer, its size and the MPI shared memory window managing it.
+  struct blk_shm_t {
+    /// Pointer to the memory block.
+    char *ptr = nullptr;
+
+    /// Size of the memory block in bytes.
+    size_t s = 0;
+
+    /// Pointer to the MPI shared memory window.
+    void *userdata = nullptr;
+  };
+
   /**
    * @brief Custom allocator that uses nda::mem::malloc to allocate memory.
    * @tparam AdrSp nda::mem::AddressSpace in which the memory is allocated.
@@ -460,8 +474,10 @@ namespace nda::mem {
    *
    * @tparam A nda::mem::Allocator type to wrap.
    */
+  template <typename Allocator>
+  class leak_check;
   template <Allocator A>
-  class leak_check : A {
+  class leak_check<A> : A {
     // Total memory used by the allocator.
     long memory_used = 0;
 
@@ -558,6 +574,104 @@ namespace nda::mem {
     [[nodiscard]] long get_memory_used() const noexcept { return memory_used; }
   };
 
+  template <MPISharedMemoryAllocator A>
+  class leak_check<A> : A {
+    // Total memory used by the allocator.
+    long memory_used = 0;
+
+    public:
+    /// nda::mem::AddressSpace in which the memory is allocated.
+    static constexpr auto address_space = A::address_space;
+
+    /// Default constructor.
+    leak_check() = default;
+
+    /// Deleted copy constructor.
+    leak_check(leak_check const &) = delete;
+
+    /// Default move constructor.
+    leak_check(leak_check &&) = default;
+
+    /// Deleted copy assignment operator.
+    leak_check &operator=(leak_check const &) = delete;
+
+    /// Default move assignment operator.
+    leak_check &operator=(leak_check &&) = default;
+
+    /**
+     * @brief Destructor that checks for memory leaks.
+     * @details In debug mode, it aborts the program if there is a memory leak.
+     */
+    ~leak_check() {
+      if (!empty()) {
+#ifndef NDEBUG
+        std::cerr << "Memory leak in allocator: " << memory_used << " bytes leaked\n";
+        std::abort();
+#endif
+      }
+    }
+
+    /**
+     * @brief Allocate memory and update the total memory used.
+     *
+     * @param s Size in bytes of the memory to allocate.
+     * @return nda::mem::blk_t memory block.
+     */
+    blk_shm_t allocate(size_t s, mpi::shared_communicator c = mpi::communicator{}.split_shared()) {
+      blk_shm_t b = A::allocate(s, c);
+      memory_used += b.s;
+      return b;
+    }
+
+    /**
+     * @brief Allocate memory, set it to zero and update the total memory used.
+     *
+     * @param s Size in bytes of the memory to allocate.
+     * @return nda::mem::blk_t memory block.
+     */
+    blk_shm_t allocate_zero(size_t s, mpi::shared_communicator c = mpi::communicator{}.split_shared()) {
+      blk_shm_t b = A::allocate_zero(s, c);
+      memory_used += b.s;
+      return b;
+    }
+
+    /**
+     * @brief Deallocate memory and update the total memory used.
+     * @details In debug mode, it aborts the program if the total memory used is smaller than zero.
+     * @param b nda::mem::blk_t memory block to deallocate.
+     */
+    void deallocate(blk_shm_t b) noexcept {
+      memory_used -= b.s;
+      if (memory_used < 0) {
+#ifndef NDEBUG
+        std::cerr << "Memory used by allocator < 0: Memory block to be deleted: b.s = " << b.s << ", b.ptr = " << (void *)b.ptr << "\n";
+        std::abort();
+#endif
+      }
+      A::deallocate(b);
+    }
+
+    /**
+     * @brief Check if the base allocator is empty.
+     * @return True if no memory is currently being used.
+     */
+    [[nodiscard]] bool empty() const { return (memory_used == 0); }
+
+    /**
+     * @brief Check if a given nda::mem::blk_t memory block is owned by the base allocator.
+     *
+     * @param b nda::mem::blk_t memory block.
+     * @return True if the base allocator owns the memory block.
+     */
+    [[nodiscard]] bool owns(blk_t b) const noexcept { return A::owns(b); }
+
+    /**
+     * @brief Get the total memory used by the base allocator.
+     * @return The size of the memory which has been allocated and not yet deallocated.
+     */
+    [[nodiscard]] long get_memory_used() const noexcept { return memory_used; }
+  };
+
   /**
    * @brief Wrap an allocator to gather statistics about memory allocation.
    *
@@ -653,6 +767,71 @@ namespace nda::mem {
     }
   };
 
+  /**
+   * @brief Custom allocator that uses mpi::shared_window to allocate memory.
+   * @tparam AdrSp nda::mem::AddressSpace in which the memory is allocated.
+   *
+   * Allocates the same amount of memory on each shared memory island.
+   */
+  class mpi_shm_allocator {
+    public:
+    /// Default constructor.
+    mpi_shm_allocator()                                    = default;
+
+    /// Deleted copy constructor.
+    mpi_shm_allocator(mpi_shm_allocator const &)            = delete;
+
+    /// Default move constructor.
+    mpi_shm_allocator(mpi_shm_allocator &&)                 = default;
+
+    /// Deleted copy assignment operator.
+    mpi_shm_allocator &operator=(mpi_shm_allocator const &) = delete;
+
+    /// Default move assignment operator.
+    mpi_shm_allocator &operator=(mpi_shm_allocator &&)      = default;
+
+    /// MPI shared memory always lives in the Host address space.
+    static constexpr auto address_space = Host;
+
+    /**
+     * @brief Allocate memory using mpi::shared_window.
+     *
+     * @param s Size in bytes of the memory to allocate.
+     * @param shm MPI shared memory communicator.
+     * @return nda::mem::blk_t memory block.
+     */
+    static blk_shm_t allocate(MPI_Aint s, mpi::shared_communicator shm = mpi::communicator{}.split_shared()) noexcept {
+      auto *win = new mpi::shared_window<char>{shm, shm.rank() == 0 ? s : 0};
+      return {(char *)win->base(0), (std::size_t)s, (void *)win}; // NOLINT
+    }
+
+    /**
+     * @brief Allocate memory and set it to zero.
+     *
+     * @param s Size in bytes of the memory to allocate.
+     * @param shm MPI shared memory communicator.
+     * @return nda::mem::blk_t memory block.
+     */
+    static blk_shm_t allocate_zero(MPI_Aint s, mpi::shared_communicator shm = mpi::communicator{}.split_shared()) noexcept {
+      auto *win = new mpi::shared_window<char>{shm, shm.rank() == 0 ? s : 0};
+      char *baseptr = win->base(0);
+      win->fence();
+      if (shm.rank() == 0) {
+          std::memset(baseptr, 0, s);
+      }
+      win->fence();
+      return {baseptr, (std::size_t)s, (void *)win}; // NOLINT
+    }
+
+    /**
+     * @brief Deallocate memory using mpi::shared_window.
+     * @param b nda::mem::blk_t memory block to deallocate.
+     */
+    static void deallocate(blk_shm_t b) noexcept {
+        delete static_cast<mpi::shared_window<char>*>(b.userdata);
+    }
+  };
+
   /** @} */
 
 } // namespace nda::mem