WIP: MPI shared memory

Author: hmenke

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>
@@ -61,6 +63,9 @@ namespace nda::mem {
 
     /// Size of the memory block in bytes.
     size_t s = 0;
+
+    /// Pointer to special data required by the allocator.
+    void *userdata = nullptr;
   };
 
   /**
@@ -653,6 +658,91 @@ 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 shared_allocator {
+    public:
+    /// Default constructor.
+    shared_allocator()                                    = default;
+
+    /// Deleted copy constructor.
+    shared_allocator(shared_allocator const &)            = delete;
+
+    /// Default move constructor.
+    shared_allocator(shared_allocator &&)                 = default;
+
+    /// Deleted copy assignment operator.
+    shared_allocator &operator=(shared_allocator const &) = delete;
+
+    /// Default move assignment operator.
+    shared_allocator &operator=(shared_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.
+     * @return nda::mem::blk_t memory block.
+     */
+    static blk_t allocate(size_t s) noexcept {
+      return allocate(s, mpi::communicator{}.split_shared());
+    }
+
+    /**
+     * @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_t allocate(MPI_Aint s, mpi::shared_communicator shm) 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.
+     * @return nda::mem::blk_t memory block.
+     */
+    static blk_t allocate_zero(size_t s) noexcept {
+      return allocate_zero(s, mpi::communicator{}.split_shared());
+    }
+
+    /**
+     * @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_t allocate_zero(MPI_Aint s, mpi::shared_communicator shm) 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_t b) noexcept {
+      delete static_cast<mpi::shared_window<char>*>(b.userdata);
+    }
+  };
+
   /** @} */
 
 } // namespace nda::mem

deps/CMakeLists.txt

@@ -92,9 +92,9 @@ external_dependency(h5
 
 # -- MPI --
 external_dependency(mpi
-  GIT_REPO https://github.com/TRIQS/mpi
-  VERSION 1.3
-  GIT_TAG unstable
+  GIT_REPO https://github.com/hmenke/mpi
+  VERSION 1.2
+  GIT_TAG shm
 )
 
 ## Pybind 11

test/c++/nda_mpi_shared.cpp

@@ -0,0 +1,31 @@
+// Copyright (c) 2020-2023 Simons Foundation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0.txt
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// Authors: Olivier Parcollet, Nils Wentzell
+
+#define NDA_DEBUG_LEAK_CHECK
+
+#include "./test_common.hpp"
+
+#include <nda/basic_array.hpp>
+#include <nda/mem.hpp>
+
+// ==============================================================
+
+TEST(SHM, Allocator) { //NOLINT
+  nda::basic_array<long, 2, nda::C_layout, 'A', nda::heap_basic<nda::mem::shared_allocator>> A(3, 3);
+  EXPECT_EQ(A.shape(), (shape_t<2>{3, 3}));
+}
+
+MPI_TEST_MAIN;