From 6c8e02d4a3990a7729fa3b08f4c4f62178edb411 Mon Sep 17 00:00:00 2001
From: Henri Menke <henri@henrimenke.de>
Date: Thu, 23 Nov 2023 15:18:25 +0100
Subject: [PATCH] WIP: MPI shared memory

---
 c++/nda/mem/allocators.hpp  | 46 +++++++++++++++++++++++++++++++++++++
 deps/CMakeLists.txt         |  4 ++--
 test/c++/nda_mpi_shared.cpp | 27 ++++++++++++++++++++++
 3 files changed, 75 insertions(+), 2 deletions(-)
 create mode 100644 test/c++/nda_mpi_shared.cpp

diff --git a/c++/nda/mem/allocators.hpp b/c++/nda/mem/allocators.hpp
index 80f3c61cb..c6a9f95d7 100644
--- a/c++/nda/mem/allocators.hpp
+++ b/c++/nda/mem/allocators.hpp
@@ -25,6 +25,8 @@
 #include <numeric>
 #include <concepts>
 
+#include <mpi/mpi.hpp>
+
 #include "address_space.hpp"
 #include "../macros.hpp"
 
@@ -47,6 +49,7 @@ namespace nda::mem {
   struct blk_t {
     char *ptr = nullptr;
     size_t s  = 0;
+    void *userdata = nullptr;
   };
 
   // -------------------------  Malloc allocator ----------------------------
@@ -335,4 +338,47 @@ namespace nda::mem {
     auto const &histogram() const noexcept { return hist; }
   };
 
+  // -------------------------  MPI shared memory allocator ----------------------------
+  //
+  // Allocates the same amount of memory on each shared memory island
+  //
+  class shared_allocator {
+    public:
+    shared_allocator()                                    = default;
+    shared_allocator(shared_allocator const &)            = delete;
+    shared_allocator(shared_allocator &&)                 = default;
+    shared_allocator &operator=(shared_allocator const &) = delete;
+    shared_allocator &operator=(shared_allocator &&)      = default;
+
+    static constexpr auto address_space = Host;
+
+    static blk_t allocate(size_t s) noexcept {
+      return allocate(s, mpi::communicator{}.split_shared());
+    }
+
+    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
+    }
+
+    static blk_t allocate_zero(size_t s) noexcept {
+      return allocate_zero(s, mpi::communicator{}.split_shared());
+    }
+
+    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
+    }
+
+    static void deallocate(blk_t b) noexcept {
+      delete static_cast<mpi::shared_window<char>*>(b.userdata);
+    }
+  };
+
 } // namespace nda::mem
diff --git a/deps/CMakeLists.txt b/deps/CMakeLists.txt
index 507e7a33d..3ab65d6ef 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/hmenke/mpi
   VERSION 1.2
-  GIT_TAG unstable
+  GIT_TAG shm
 )
 
 ## Pybind 11
diff --git a/test/c++/nda_mpi_shared.cpp b/test/c++/nda_mpi_shared.cpp
new file mode 100644
index 000000000..1ba555142
--- /dev/null
+++ b/test/c++/nda_mpi_shared.cpp
@@ -0,0 +1,27 @@
+// 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
+
+#include "./test_common.hpp"
+
+// ==============================================================
+
+TEST(SHM, Allocator) { //NOLINT
+  if (!mpi::has_env) GTEST_SKIP() << "MPI is unavailable";
+  nda::basic_array<long, 2, C_layout, 'A', nda::heap_basic<nda::mem::shared_allocator>> A(3, 3);
+  EXPECT_EQ(A.shape(), (nda::shape_t<2>{3, 3}));
+}
+
+MAKE_MAIN_MPI