feat(core): Use NUMA-aware allocator for tile data

Author: timniederhausen

core/CMakeLists.txt

@@ -13,6 +13,7 @@ set(SOURCE_FILES
     src/utils.cpp
     src/performance_counters.cpp
     src/target.cpp
+    src/tile_data.cpp
     src/kernels.cpp
     src/hyperparameters.cpp
     src/cpu/gp_functions.cpp

core/include/gprat/tile_data.hpp

@@ -4,13 +4,59 @@
 #pragma once
 
 #include "gprat/detail/config.hpp"
-#include "gprat/performance_counters.hpp"
 
 #include <hpx/serialization/serialize_buffer.hpp>
 #include <span>
 
 GPRAT_NS_BEGIN
 
+namespace detail
+{
+void *allocate_tile_data(std::size_t num_bytes);
+void deallocate_tile_data(void *p, std::size_t num_bytes);
+
+template <class T>
+struct tile_data_allocator
+{
+    typedef T value_type;
+
+    tile_data_allocator() = default;
+
+    template <class U>
+    constexpr tile_data_allocator(const tile_data_allocator<U> &) noexcept
+    { }
+
+    [[nodiscard]] T *allocate(std::size_t n)
+    {
+        if (n > (std::numeric_limits<std::size_t>::max)() / sizeof(T))
+        {
+            throw std::bad_array_new_length();
+        }
+
+        if (auto p = static_cast<T *>(allocate_tile_data(n * sizeof(T))))
+        {
+            return p;
+        }
+
+        throw std::bad_alloc();
+    }
+
+    void deallocate(T *p, std::size_t n) noexcept { deallocate_tile_data(p, n * sizeof(T)); }
+};
+
+template <class T, class U>
+bool operator==(const tile_data_allocator<T> &, const tile_data_allocator<U> &)
+{
+    return true;
+}
+
+template <class T, class U>
+bool operator!=(const tile_data_allocator<T> &, const tile_data_allocator<U> &)
+{
+    return false;
+}
+}  // namespace detail
+
 /**
  * @brief Non-mutable reference-counted dynamic array of a given type T.
  * This class represents a simple reference-counted non-resizeable buffer with elements of type T.
@@ -25,7 +71,7 @@ template <typename T>
 class const_tile_data
 {
   protected:
-    typedef hpx::serialization::serialize_buffer<T> cpu_buffer_type;
+    typedef hpx::serialization::serialize_buffer<T, detail::tile_data_allocator<T>> cpu_buffer_type;
 
     struct hold_reference
     {
@@ -38,25 +84,12 @@ class const_tile_data
         cpu_buffer_type data_;
     };
 
-    // In case we want pooling down the road...
-    static T *allocate(std::size_t n)
-    {
-        track_tile_data_allocation(n);
-        return new T[n];
-    }
-
-    static void deallocate(T *p) noexcept
-    {
-        track_tile_data_deallocation(0);  // we don't know here
-        delete[] p;
-    }
-
   public:
     const_tile_data() = default;
 
     // Create a new (uninitialized) tile_data of the given size.
     explicit const_tile_data(std::size_t size) :
-        cpu_data_(allocate(size), size, cpu_buffer_type::take, &const_tile_data::deallocate)
+        cpu_data_(size)
     { }
 
     // Create a tile_data which acts as a proxy to a part of the embedded array.
@@ -85,10 +118,12 @@ class const_tile_data
         return { cpu_data_.data(), cpu_data_.size() };
     }
 
+    friend bool operator==(const const_tile_data &a, const const_tile_data &b) noexcept
+    {
+        return a.cpu_data_ == b.cpu_data_;
+    }
+
   protected:
-    // Serialization support: even if all of the code below runs on one
-    // locality only, we need to provide an (empty) implementation for the
-    // serialization as all arguments passed to actions have to support this.
     friend class hpx::serialization::access;
 
     template <typename Archive>

core/src/tile_data.cpp

@@ -0,0 +1,34 @@
+#include "gprat/tile_data.hpp"
+
+#include "gprat/performance_counters.hpp"
+
+#include <hpx/runtime_local/runtime_local.hpp>
+
+GPRAT_NS_BEGIN
+
+namespace detail
+{
+
+void *allocate_tile_data(std::size_t num_bytes)
+{
+    auto &topology = hpx::get_runtime().get_topology();
+    const auto bitmap = topology.cpuset_to_nodeset(topology.get_machine_affinity_mask());
+
+    track_tile_data_allocation(num_bytes);
+    return topology.allocate_membind(num_bytes, bitmap, hpx::threads::hpx_hwloc_membind_policy::membind_firsttouch, 0);
+}
+
+void deallocate_tile_data(void *p, std::size_t num_bytes)
+{
+    track_tile_data_deallocation(num_bytes);
+
+    if (hpx::is_running())
+    {
+        auto &topology = hpx::get_runtime().get_topology();
+        topology.deallocate(p, num_bytes);
+    }
+}
+
+}  // namespace detail
+
+GPRAT_NS_END