Numa aware alloc

Author: lucjaulmes

CMakeLists.txt

@@ -323,6 +323,21 @@ if (ENABLE_MULTITHREADING)
     OpenMP::OpenMP_C
   )
 
+  # Find NUMA - location of NUMA headers
+  if (WIN32)
+    compile_option(NUMA_AWARE 0)
+    message(WARNING "Building on Windows, QuEST will not be aware of numa locality")
+  else()
+    include(FindPkgConfig)
+    pkg_search_module(NUMA numa IMPORTED_TARGET GLOBAL)
+    compile_option(NUMA_AWARE ${NUMA_FOUND})
+    if (${NUMA_FOUND})
+      target_link_libraries(QuEST PRIVATE PkgConfig::NUMA)
+    else()
+      message(WARNING "libnuma not found, QuEST will not be aware of numa locality")
+    endif()
+  endif()
+
   if (VERBOSE_LIB_NAME)
     string(CONCAT LIB_NAME ${LIB_NAME} "+mt")
   endif()

quest/src/api/qureg.cpp

@@ -151,7 +151,7 @@ Qureg validateAndCreateCustomQureg(int numQubits, int isDensMatr, int useDistrib
     Qureg qureg = qureg_populateNonHeapFields(numQubits, isDensMatr, useDistrib, useGpuAccel, useMultithread);
 
     // always allocate CPU memory
-    qureg.cpuAmps = cpu_allocArray(qureg.numAmpsPerNode); // nullptr if failed
+    qureg.cpuAmps = cpu_allocNumaArray(qureg.numAmpsPerNode); // nullptr if failed
 
     // conditionally allocate GPU memory and communication buffers (even if numNodes == 1).
     // note that in distributed settings but where useDistrib=false, each node will have a
@@ -331,7 +331,7 @@ void destroyQureg(Qureg qureg) {
     validate_quregFields(qureg, __func__);
 
     // free CPU memory
-    cpu_deallocArray(qureg.cpuAmps);
+    cpu_deallocNumaArray(qureg.cpuAmps, qureg.numAmpsPerNode);
 
     // free CPU communication buffer
     if (qureg.isDistributed)

quest/src/cpu/cpu_config.cpp

@@ -14,6 +14,7 @@
 #include <vector>
 #include <cstring>
 #include <cstdlib>
+#include <cstdint>
 
 using std::vector;
 
@@ -34,6 +35,12 @@ using std::vector;
     #include <omp.h>
 #endif
 
+#if NUMA_AWARE
+    #include <sys/mman.h>
+    #include <unistd.h>
+    #include <numaif.h>
+    #include <numa.h>
+#endif // NUMA_AWARE
 
 
 /*
@@ -105,23 +112,69 @@ int cpu_getCurrentNumThreads() {
  * MEMORY ALLOCATION
  */
 
+#if NUMA_AWARE
+unsigned long get_page_size() {
+    static unsigned long page_size = 0;
+    if (!page_size) {
+        page_size = sysconf(_SC_PAGESIZE);
+        if (page_size == ~0UL) {
+            perror("Failed to get the page size");
+        }
+    }
+    return page_size;
+}
+
+unsigned long get_numa_nodes() {
+    static int n_nodes = 0;
+    if (!n_nodes) {
+        n_nodes = numa_num_configured_nodes();
+        if (n_nodes < 1) {
+            perror("Failed to get the numa node count");
+        }
+    }
+    return n_nodes;
+}
+#endif
 
 qcomp* cpu_allocArray(qindex length) {
+    return (qcomp*) calloc(length, sizeof(qcomp));
+}
 
-    /// @todo
-    /// here, we calloc the entire array in a serial setting, rather than one malloc 
-    /// followed by threads subsequently memset'ing their own partitions. The latter
-    /// approach would distribute the array pages across NUMA nodes, accelerating 
-    /// their subsequent access by the same threads (via NUMA's first-touch policy).
-    /// We have so far foregone this optimisation since a thread's memory-access pattern
-    /// in many of the QuEST functions is non-trivial, and likely to be inconsistent 
-    /// with the memset pattern. As such, I expect the benefit is totally occluded
-    /// and only introduces potential new bugs - but this should be tested and confirmed!
-
-    // we call calloc over malloc in order to fail immediately if mem isn't available;
-    // caller must handle nullptr result
 
-    return (qcomp*) calloc(length, sizeof(qcomp));
+qcomp* cpu_allocNumaArray(qindex length) {
+#if !NUMA_AWARE
+    return cpu_allocArray(length);
+#else
+    unsigned long page_size = get_page_size();
+    int n_nodes = get_numa_nodes();
+
+    qindex size = length * sizeof(qcomp);
+    int pages = (size + page_size - 1) / page_size;
+    void *addr = mmap(NULL, pages * page_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+    if (n_nodes == 1) {
+        return reinterpret_cast<qcomp*>(addr);
+    }
+
+    // distribution strategy: floor_pages per node, distribute remain_pages as spread out as possible
+    int floor_pages = pages / n_nodes;
+    int spread_pages = pages % n_nodes;
+
+    uintptr_t pos = (uintptr_t)addr;
+    for (int node = 0, shift = n_nodes; node < n_nodes; ++node) {
+        shift -= spread_pages;
+        int node_pages = floor_pages + (shift <= 0);
+
+        unsigned long node_mask = 1UL << node;
+        mbind((void*)pos, node_pages * page_size, MPOL_BIND, &node_mask, sizeof(node_mask) * 8, 0);
+
+        pos += node_pages * page_size;
+        if (shift <= 0) {
+            shift += n_nodes;
+        }
+    }
+
+    return reinterpret_cast<qcomp*>(addr);
+#endif // NUMA_AWARE
 }
 
 
@@ -132,6 +185,23 @@ void cpu_deallocArray(qcomp* arr) {
 }
 
 
+void cpu_deallocNumaArray(qcomp* arr, qindex length) {
+    if (arr == nullptr) {
+        return;
+    }
+
+#if !NUMA_AWARE
+    return cpu_deallocArray(arr);
+#else
+    unsigned long page_size = get_page_size();
+    qindex size = length * sizeof(qcomp);
+    int pages = (size + page_size - 1) / page_size;
+
+    munmap(arr, pages * page_size);
+#endif // NUMA_AWARE
+}
+
+
 qcomp** cpu_allocAndInitMatrixWrapper(qcomp* arr, qindex dim) {
 
     // do not allocate if arr alloc failed (caller will handle)

quest/src/cpu/cpu_config.hpp

@@ -46,6 +46,9 @@ int cpu_getCurrentNumThreads();
 qcomp* cpu_allocArray(qindex length);
 void cpu_deallocArray(qcomp* arr);
 
+qcomp* cpu_allocNumaArray(qindex length);
+void cpu_deallocNumaArray(qcomp* arr, qindex length);
+
 qcomp** cpu_allocAndInitMatrixWrapper(qcomp* arr, qindex dim);
 void cpu_deallocMatrixWrapper(qcomp** wrapper);