diff --git a/c++/nda/concepts.hpp b/c++/nda/concepts.hpp index ce0282d1..377a6200 100644 --- a/c++/nda/concepts.hpp +++ b/c++/nda/concepts.hpp @@ -24,6 +24,8 @@ #include "./stdutil/concepts.hpp" #include "./traits.hpp" +#include + #include #include #include @@ -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; }; + template + concept MPISharedMemoryAllocator = requires(A &a) { + { a.allocate(size_t{}, mpi::shared_communicator{}) } noexcept -> std::same_as; + { a.allocate_zero(size_t{}, mpi::shared_communicator{}) } noexcept -> std::same_as; + { a.deallocate(std::declval()) } noexcept; + { A::address_space } -> std::same_as; + }; + /** * @brief Check if a given type satisfies the memory handle concept. * diff --git a/c++/nda/mem/allocators.hpp b/c++/nda/mem/allocators.hpp index f63d58db..c4c2f6df 100644 --- a/c++/nda/mem/allocators.hpp +++ b/c++/nda/mem/allocators.hpp @@ -28,6 +28,8 @@ #include "./memset.hpp" #include "../macros.hpp" +#include + #include #include #include @@ -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 + class leak_check; template - class leak_check : A { + class leak_check : 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 + class leak_check : 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{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{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*>(b.userdata); + } + }; + /** @} */ } // namespace nda::mem diff --git a/c++/nda/mem/handle.hpp b/c++/nda/mem/handle.hpp index 7bbeb9b0..0dd4bebc 100644 --- a/c++/nda/mem/handle.hpp +++ b/c++/nda/mem/handle.hpp @@ -940,6 +940,268 @@ namespace nda::mem { [[nodiscard]] T *data() const noexcept { return _data; } }; + template + struct handle_mpi_shm { + static_assert(std::is_nothrow_destructible_v, "nda::mem::handle_mpi_shm requires the value_type to have a non-throwing destructor"); + + private: + // Pointer to the start of the actual data. + T *_data = nullptr; + + // Size of the data (number of T elements). Invariant: size > 0 iif data != nullptr. + size_t _size = 0; + + // Special userdata used by the allocator. + void *_userdata = nullptr; + + // Allocator to use. +#ifndef NDA_DEBUG_LEAK_CHECK + static inline A allocator; // NOLINT (allocator is not specific to a single instance) +#else + static inline leak_check allocator; // NOLINT (allocator is not specific to a single instance) +#endif + + // For shared ownership (points to a blk_T_t). + mutable std::shared_ptr sptr; + + // Type of the memory block, i.e. a pointer to the data and its size. + using blk_T_t = std::tuple; + + // Release the handled memory (data pointer and size are not set to null here). + static void destruct(blk_T_t b) noexcept { + auto [data, size, win] = b; + + // do nothing if the data is null + if (data == nullptr) return; + + // if needed, call the destructors of the objects stored + if constexpr (A::address_space == Host and !(std::is_trivial_v or nda::is_complex_v)) { + for (size_t i = 0; i < size; ++i) data[i].~T(); + } + + // deallocate the memory block + allocator.deallocate({(char *)data, size * sizeof(T), win}); + } + + // Deleter for the shared pointer. + static void deleter(void *p) noexcept { destruct(*((blk_T_t *)p)); } + + public: + /// Value type of the data. + using value_type = T; + + /// nda::mem::Allocator type. + using allocator_type = A; + + /// nda::mem::AddressSpace in which the memory is allocated. + static constexpr auto address_space = allocator_type::address_space; + + /** + * @brief Get a shared pointer to the memory block. + * @return A copy of the shared pointer stored in the current handle. + */ + std::shared_ptr get_sptr() const { + if (not sptr) sptr.reset(new blk_T_t{_data, _size, _userdata}, deleter); + return sptr; + } + + /** + * @brief Destructor for the handle. + * @details If the shared pointer is set, it does nothing. Otherwise, it explicitly calls the destructor of + * non-trivial objects and deallocates the memory. + */ + ~handle_mpi_shm() noexcept { + if (not sptr and not(is_null())) destruct({_data, _size, _userdata}); + } + + /// Default constructor leaves the handle in a null state (`nullptr` and size 0). + handle_mpi_shm() = default; + + /** + * @brief Move constructor simply copies the pointers and size and resets the source handle to a null state. + * @param h Source handle. + */ + handle_mpi_shm(handle_mpi_shm &&h) noexcept : _data(h._data), _size(h._size), _userdata(h._userdata), sptr(std::move(h.sptr)) { + h._data = nullptr; + h._size = 0; + h._userdata = nullptr; + } + + /** + * @brief Move assignment operator first releases the resources held by the current handle and then moves the + * resources from the source to the current handle. + * + * @param h Source handle. + */ + handle_mpi_shm &operator=(handle_mpi_shm &&h) noexcept { + // release current resources if they are not shared and not null + if (not sptr and not(is_null())) destruct({_data, _size, _userdata}); + + // move the resources from the source handle + _data = h._data; + _size = h._size; + _userdata = h._userdata; + sptr = std::move(h.sptr); + + // reset the source handle to a null state + h._data = nullptr; + h._size = 0; + h._userdata = nullptr; + return *this; + } + + /** + * @brief Copy constructor makes a deep copy of the data from another handle. + * @param h Source handle. + */ + explicit handle_mpi_shm(handle_mpi_shm const &h) : handle_mpi_shm(h.size(), do_not_initialize) { + if (is_null()) return; + if constexpr (std::is_trivially_copyable_v) { + memcpy(_data, h.data(), h.size() * sizeof(T)); + } else { + for (size_t i = 0; i < _size; ++i) new (_data + i) T(h[i]); + } + } + + /** + * @brief Copy assignment operator utilizes the copy constructor and move assignment operator to make a deep copy of + * the data and size from the source handle. + * + * @param h Source handle. + */ + handle_mpi_shm &operator=(handle_mpi_shm const &h) { + *this = handle_mpi_shm{h}; + return *this; + } + + /** + * @brief Construct a handle by making a deep copy of the data from another handle. + * + * @tparam H nda::mem::OwningHandle type. + * @param h Source handle. + template H> + explicit handle_mpi_shm(H const &h) : handle_mpi_shm(h.size(), do_not_initialize) { + if (is_null()) return; + if constexpr (std::is_trivially_copyable_v) { + memcpy((void *)_data, (void *)h.data(), _size * sizeof(T)); + } else { + static_assert(address_space == H::address_space, + "Constructing an nda::mem::handle_mpi_shm from a handle of a different address space requires a trivially copyable value_type"); + for (size_t i = 0; i < _size; ++i) new (_data + i) T(h[i]); + } + } + */ + + + /** + * @brief Assignment operator utilizes another constructor and move assignment to make a deep copy of the data and + * size from the source handle. + * + * @tparam AS Allocator type of the source handle. + * @param h Source handle with a different allocator. + template + handle_mpi_shm &operator=(handle_mpi_shm const &h) { + *this = handle_mpi_shm{h}; + return *this; + } + */ + + /** + * @brief Construct a handle by allocating memory for the data of a given size but without initializing it. + * @param size Size of the data (number of elements). + */ + handle_mpi_shm(long size, do_not_initialize_t) { + if (size == 0) return; + auto b = allocator.allocate(size * sizeof(T)); + if (not b.ptr) throw std::bad_alloc{}; + _data = (T *)b.ptr; + _size = size; + _userdata = b.userdata; + } + + /** + * @brief Construct a handle by allocating memory for the data of a given size and initializing it to zero. + * @param size Size of the data (number of elements). + */ + handle_mpi_shm(long size, init_zero_t) { + if (size == 0) return; + auto b = allocator.allocate_zero(size * sizeof(T)); + if (not b.ptr) throw std::bad_alloc{}; + _data = (T *)b.ptr; + _size = size; + _userdata = b.userdata; + } + + /** + * @brief Construct a handle by allocating memory for the data of a given size and initializing it depending on the + * value type. + * + * @details The data is initialized as follows: + * - If `T` is std::complex and nda::mem::init_dcmplx is true, the data is initialized to zero. + * - If `T` is not trivial and not complex, the data is default constructed by placement new operator calls. + * - Otherwise, the data is not initialized. + * + * @param size Size of the data (number of elements). + */ + handle_mpi_shm(long size) { + if (size == 0) return; + blk_shm_t b; + if constexpr (is_complex_v && init_dcmplx) + b = allocator.allocate_zero(size * sizeof(T)); + else + b = allocator.allocate(size * sizeof(T)); + if (not b.ptr) throw std::bad_alloc{}; + _data = (T *)b.ptr; + _size = size; + _userdata = b.userdata; + + // call placement new for non trivial and non complex types + if constexpr (!std::is_trivial_v and !is_complex_v) { + for (size_t i = 0; i < size; ++i) new (_data + i) T(); + } + } + + /** + * @brief Subscript operator to access the data. + * + * @param i Index of the element to access. + * @return Reference to the element at the given index. + */ + [[nodiscard]] T &operator[](long i) noexcept { return _data[i]; } + + /** + * @brief Subscript operator to access the data. + * + * @param i Index of the element to access. + * @return Const reference to the element at the given index. + */ + [[nodiscard]] T const &operator[](long i) const noexcept { return _data[i]; } + + /** + * @brief Check if the handle is in a null state. + * @return True if the data is a `nullptr` (and the size is 0). + */ + [[nodiscard]] bool is_null() const noexcept { +#ifdef NDA_DEBUG + // check the invariants in debug mode + EXPECTS((_data == nullptr) == (_size == 0)); +#endif + return _data == nullptr; + } + + /** + * @brief Get a pointer to the stored data. + * @return Pointer to the start of the handled memory. + */ + [[nodiscard]] T *data() const noexcept { return _data; } + + /** + * @brief Get the size of the handle. + * @return Number of elements of type `T` in the handled memory. + */ + [[nodiscard]] long size() const noexcept { return _size; } + }; + /** @} */ } // namespace nda::mem diff --git a/c++/nda/mem/policies.hpp b/c++/nda/mem/policies.hpp index 99aa0796..89b6f5f2 100644 --- a/c++/nda/mem/policies.hpp +++ b/c++/nda/mem/policies.hpp @@ -114,6 +114,20 @@ namespace nda { using handle = mem::handle_borrowed; }; + /** + * @brief Memory policy using an nda::mem::handle_heap. + * @tparam Allocator Allocator type to be used. + */ + template + struct mpi_shared_memory { + /** + * @brief Handle type for the policy. + * @tparam T Value type of the data. + */ + template + using handle = mem::handle_mpi_shm; + }; + /** @} */ } // namespace nda diff --git a/c++/nda/shared_array.hpp b/c++/nda/shared_array.hpp new file mode 100644 index 00000000..375a0f0e --- /dev/null +++ b/c++/nda/shared_array.hpp @@ -0,0 +1,41 @@ +// Copyright (c) 2018 Commissariat à l'énergie atomique et aux énergies alternatives (CEA) +// Copyright (c) 2018 Centre national de la recherche scientifique (CNRS) +// Copyright (c) 2018-2024 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: Thomas Hahn, Olivier Parcollet, Nils Wentzell + +/** + * @file + * @brief Provides the class for arrays in MPI shared memory. + */ + +#pragma once + +#include "./basic_array.hpp" + +namespace nda { + + template + class shared_array : public basic_array> { + private: + using Base = basic_array>; + mpi::shared_communicator _c{mpi::communicator{}.split_shared()}; + public: + using Base::Base; + shared_array(mpi::shared_communicator c) : _c(c) { + + }; + }; +} // namespace nda diff --git a/deps/CMakeLists.txt b/deps/CMakeLists.txt index 8101214c..e0e94d74 100644 --- a/deps/CMakeLists.txt +++ b/deps/CMakeLists.txt @@ -92,9 +92,9 @@ external_dependency(h5 # -- MPI -- external_dependency(mpi - GIT_REPO https://github.com/TRIQS/mpi + GIT_REPO https://github.com/Mobellaaj/mpi VERSION 1.3 - GIT_TAG unstable + GIT_TAG bog-shm ) ## Pybind 11 diff --git a/test/c++/nda_mpi_shared.cpp b/test/c++/nda_mpi_shared.cpp new file mode 100644 index 00000000..d4035585 --- /dev/null +++ b/test/c++/nda_mpi_shared.cpp @@ -0,0 +1,38 @@ +// 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 +#include +#include + +// ============================================================== + +TEST(SHM, Allocator) { //NOLINT + nda::basic_array> A(3, 3); + EXPECT_EQ(A.shape(), (shape_t<2>{3, 3})); +} + +TEST(SHM, Constructor) { //NOLINT + mpi::communicator world; + mpi::shared_communicator shm = world.split_shared(); + nda::shared_array A(shm); +} + +MPI_TEST_MAIN;