implment interleave mode with customizable size

Signed-off-by: Łukasz Plewa <[email protected]>

Author: lplewa

include/umf/providers/provider_os_memory.h

@@ -70,6 +70,8 @@ typedef struct umf_os_memory_provider_params_t {
 
     /// Describes how node list is interpreted
     umf_numa_mode_t numa_mode;
+    /// part size for interleave mode - 0 means default (system specific)
+    size_t part_size;
 } umf_os_memory_provider_params_t;
 
 /// @brief OS Memory Provider operation results
@@ -94,6 +96,7 @@ umfOsMemoryProviderParamsDefault(void) {
         NULL,                                       /* numa_list */
         0,                                          /* numa_list_len */
         UMF_NUMA_MODE_DEFAULT,                      /* numa_mode */
+        0                                           /* part_size */
     };
 
     return params;

src/provider/provider_os_memory.c

@@ -16,6 +16,7 @@
 
 #include "base_alloc_global.h"
 #include "provider_os_memory_internal.h"
+#include "utils_concurrency.h"
 #include "utils_log.h"
 
 #include <umf.h>
@@ -28,11 +29,14 @@ typedef struct os_memory_provider_t {
     unsigned protection; // combination of OS-specific protection flags
 
     // NUMA config
-    hwloc_bitmap_t nodeset;
+    hwloc_bitmap_t *nodeset;
+    unsigned nodeset_len;
     char *nodeset_str_buf;
     hwloc_membind_policy_t numa_policy;
     int numa_flags; // combination of hwloc flags
 
+    size_t part_size;
+    size_t alloc_sum; // sum of all allocations - used for manual interleaving
     hwloc_topology_t topo;
 } os_memory_provider_t;
 
@@ -81,30 +85,67 @@ static void os_store_last_native_error(int32_t native_error, int errno_value) {
     TLS_last_native_error.errno_value = errno_value;
 }
 
-static umf_result_t nodemask_to_hwloc_nodeset(const unsigned *nodelist,
-                                              unsigned long listsize,
-                                              hwloc_bitmap_t *out_nodeset) {
-    if (out_nodeset == NULL) {
-        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
-    }
+static umf_result_t initialize_nodeset(os_memory_provider_t *os_provider,
+                                       const unsigned *nodelist,
+                                       unsigned long listsize,
+                                       int separate_nodes) {
 
-    *out_nodeset = hwloc_bitmap_alloc();
-    if (!*out_nodeset) {
+    unsigned long array_size = (listsize && separate_nodes) ? listsize : 1;
+    os_provider->nodeset =
+        umf_ba_global_alloc(sizeof(*os_provider->nodeset) * array_size);
+
+    if (!os_provider->nodeset) {
         return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
     }
 
+    hwloc_bitmap_t *out_nodeset = os_provider->nodeset;
+    os_provider->nodeset_len = array_size;
     if (listsize == 0) {
+        // Hwloc_set_area_membind fails if empty nodeset is passed so
+        // if no node is specified, just pass all available nodes.
+        // For modes where no node is needed, they will be ignored anyway.
+        out_nodeset[0] = hwloc_bitmap_dup(
+            hwloc_topology_get_complete_nodeset(os_provider->topo));
+        if (!out_nodeset[0]) {
+            goto err_free_list;
+        }
         return UMF_RESULT_SUCCESS;
     }
 
-    for (unsigned long i = 0; i < listsize; i++) {
-        if (hwloc_bitmap_set(*out_nodeset, nodelist[i])) {
-            hwloc_bitmap_free(*out_nodeset);
-            return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
+    for (unsigned long i = 0; i < array_size; i++) {
+        out_nodeset[i] = hwloc_bitmap_alloc();
+        if (!out_nodeset[i]) {
+            for (unsigned long j = 0; j < i; j++) {
+                hwloc_bitmap_free(out_nodeset[j]);
+            }
+            goto err_free_list;
+        }
+    }
+
+    if (separate_nodes) {
+        for (unsigned long i = 0; i < listsize; i++) {
+            if (hwloc_bitmap_set(out_nodeset[i], nodelist[i])) {
+                goto err_free_bitmaps;
+            }
+        }
+    } else {
+        for (unsigned long i = 0; i < listsize; i++) {
+            if (hwloc_bitmap_set(out_nodeset[0], nodelist[i])) {
+                goto err_free_bitmaps;
+            }
         }
     }
 
     return UMF_RESULT_SUCCESS;
+
+err_free_bitmaps:
+    for (unsigned long i = 0; i < array_size; i++) {
+        hwloc_bitmap_free(out_nodeset[i]);
+    }
+err_free_list:
+    umf_ba_global_free(*out_nodeset);
+    os_provider->nodeset_len = 0;
+    return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
 }
 
 umf_result_t os_translate_flags(unsigned in_flags, unsigned max,
@@ -132,42 +173,61 @@ umf_result_t os_translate_flags(unsigned in_flags, unsigned max,
     return UMF_RESULT_SUCCESS;
 }
 
-static umf_result_t translate_numa_mode(umf_numa_mode_t mode, int nodemaskEmpty,
-                                        hwloc_membind_policy_t *numa_policy) {
+static umf_result_t validate_numa_mode(umf_numa_mode_t mode,
+                                       int nodemaskEmpty) {
     switch (mode) {
     case UMF_NUMA_MODE_DEFAULT:
+    case UMF_NUMA_MODE_LOCAL:
         if (!nodemaskEmpty) {
             // nodeset must be empty
             return UMF_RESULT_ERROR_INVALID_ARGUMENT;
         }
-        *numa_policy = HWLOC_MEMBIND_DEFAULT;
         return UMF_RESULT_SUCCESS;
     case UMF_NUMA_MODE_BIND:
-        if (nodemaskEmpty) {
-            // nodeset must not be empty
-            return UMF_RESULT_ERROR_INVALID_ARGUMENT;
-        }
-        *numa_policy = HWLOC_MEMBIND_BIND;
-        return UMF_RESULT_SUCCESS;
     case UMF_NUMA_MODE_INTERLEAVE:
         if (nodemaskEmpty) {
             // nodeset must not be empty
             return UMF_RESULT_ERROR_INVALID_ARGUMENT;
         }
-        *numa_policy = HWLOC_MEMBIND_INTERLEAVE;
         return UMF_RESULT_SUCCESS;
     case UMF_NUMA_MODE_PREFERRED:
-        *numa_policy = HWLOC_MEMBIND_BIND;
         return UMF_RESULT_SUCCESS;
-    case UMF_NUMA_MODE_LOCAL:
-        if (!nodemaskEmpty) {
-            // nodeset must be empty
-            return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    default:
+        assert(0);
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+}
+
+static hwloc_membind_policy_t translate_numa_mode(umf_numa_mode_t mode,
+                                                  int manual) {
+    switch (mode) {
+    case UMF_NUMA_MODE_DEFAULT:
+        return HWLOC_MEMBIND_DEFAULT;
+    case UMF_NUMA_MODE_BIND:
+        return HWLOC_MEMBIND_BIND;
+    case UMF_NUMA_MODE_INTERLEAVE:
+        // In manual mode, we manually implement interleaving,
+        // by binding memory to specific NUMA nodes.
+        if (manual) {
+            return HWLOC_MEMBIND_BIND;
         }
-        *numa_policy = HWLOC_MEMBIND_BIND;
-        return UMF_RESULT_SUCCESS;
+        return HWLOC_MEMBIND_INTERLEAVE;
+    case UMF_NUMA_MODE_PREFERRED:
+        return HWLOC_MEMBIND_BIND;
+    case UMF_NUMA_MODE_LOCAL:
+        return HWLOC_MEMBIND_BIND;
     }
-    return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    assert(0);
+    return -1;
+}
+
+//return 1 if umf will binds memory directly to single NUMA node, based on internal algorithm
+//return 0 if umf will just set numa memory policy, and kernel will decide where to allocate memory
+static int is_dedicated_node_bind(umf_os_memory_provider_params_t *in_params) {
+    if (in_params->numa_mode == UMF_NUMA_MODE_INTERLEAVE) {
+        return in_params->part_size > 0;
+    }
+    return 0;
 }
 
 static int getHwlocMembindFlags(umf_numa_mode_t mode) {
@@ -193,17 +253,19 @@ static umf_result_t translate_params(umf_os_memory_provider_params_t *in_params,
 
     // NUMA config
     int emptyNodeset = in_params->numa_list_len == 0;
-    result = translate_numa_mode(in_params->numa_mode, emptyNodeset,
-                                 &provider->numa_policy);
+    result = validate_numa_mode(in_params->numa_mode, emptyNodeset);
     if (result != UMF_RESULT_SUCCESS) {
         LOG_ERR("incorrect NUMA mode: %u", in_params->numa_mode);
         return result;
     }
 
+    int dedicated_node_bind = is_dedicated_node_bind(in_params);
+    provider->numa_policy =
+        translate_numa_mode(in_params->numa_mode, dedicated_node_bind);
     provider->numa_flags = getHwlocMembindFlags(in_params->numa_mode);
-
-    return nodemask_to_hwloc_nodeset(
-        in_params->numa_list, in_params->numa_list_len, &provider->nodeset);
+    provider->part_size = in_params->part_size;
+    return initialize_nodeset(provider, in_params->numa_list,
+                              in_params->numa_list_len, dedicated_node_bind);
 }
 
 static umf_result_t os_initialize(void *params, void **provider) {
@@ -250,11 +312,13 @@ static umf_result_t os_initialize(void *params, void **provider) {
     if (!os_provider->nodeset_str_buf) {
         LOG_INFO("Allocating memory for printing NUMA nodes failed");
     } else {
-        if (hwloc_bitmap_list_snprintf(os_provider->nodeset_str_buf,
-                                       NODESET_STR_BUF_LEN,
-                                       os_provider->nodeset)) {
-            LOG_INFO("OS provider initialized with NUMA nodes: %s",
-                     os_provider->nodeset_str_buf);
+        LOG_INFO("OS provider initialized with NUMA nodes:");
+        for (unsigned i = 0; i < os_provider->nodeset_len; i++) {
+            if (hwloc_bitmap_list_snprintf(os_provider->nodeset_str_buf,
+                                           NODESET_STR_BUF_LEN,
+                                           os_provider->nodeset[i])) {
+                LOG_INFO("%s", os_provider->nodeset_str_buf);
+            }
         }
     }
 
@@ -280,7 +344,10 @@ static void os_finalize(void *provider) {
         umf_ba_global_free(os_provider->nodeset_str_buf);
     }
 
-    hwloc_bitmap_free(os_provider->nodeset);
+    for (unsigned i = 0; i < os_provider->nodeset_len; i++) {
+        hwloc_bitmap_free(os_provider->nodeset[i]);
+    }
+    umf_ba_global_free(os_provider->nodeset);
     hwloc_topology_destroy(os_provider->topo);
     umf_ba_global_free(os_provider);
 }
@@ -387,6 +454,17 @@ static int os_mmap_aligned(void *hint_addr, size_t length, size_t alignment,
     return 0;
 }
 
+static int get_membind(os_memory_provider_t *provider, size_t size) {
+    if (provider->nodeset_len == 1) {
+        return 0;
+    }
+
+    assert(provider->part_size != 0);
+    size_t s = util_fetch_and_add64(&provider->alloc_sum, size);
+
+    return (s / provider->part_size) % provider->nodeset_len;
+}
+
 static umf_result_t os_alloc(void *provider, size_t size, size_t alignment,
                              void **resultPtr) {
     int ret;
@@ -434,32 +512,31 @@ static umf_result_t os_alloc(void *provider, size_t size, size_t alignment,
     }
 
     errno = 0;
-    if (hwloc_bitmap_iszero(os_provider->nodeset)) {
-        // Hwloc_set_area_membind fails if empty nodeset is passed so if no node is specified,
-        // just pass all available nodes. For modes where no node is needed, they will be
-        // ignored anyway.
-        hwloc_const_nodeset_t complete_nodeset =
-            hwloc_topology_get_complete_nodeset(os_provider->topo);
-        ret = hwloc_set_area_membind(os_provider->topo, addr, size,
-                                     complete_nodeset, os_provider->numa_policy,
-                                     os_provider->numa_flags);
-    } else {
+    unsigned membind = get_membind(os_provider, size);
+    size_t bind_size =
+        os_provider->nodeset_len == 1 ? size : os_provider->part_size;
+
+    do {
+        size_t s = bind_size < size ? bind_size : size;
         ret = hwloc_set_area_membind(
-            os_provider->topo, addr, size, os_provider->nodeset,
+            os_provider->topo, addr, s, os_provider->nodeset[membind++],
             os_provider->numa_policy, os_provider->numa_flags);
-    }
 
-    if (ret) {
-        os_store_last_native_error(UMF_OS_RESULT_ERROR_BIND_FAILED, errno);
-        LOG_PERR("binding memory to NUMA node failed");
-        // TODO: (errno == 0) when hwloc_set_area_membind() fails on Windows - ignore this temporarily
-        if (errno != ENOSYS &&
-            errno != 0) { // ENOSYS - Function not implemented
-            // Do not error out if memory binding is not implemented at all (like in case of WSL on Windows).
-            goto err_unmap;
+        size -= s;
+        membind %= os_provider->nodeset_len;
+        if (ret) {
+            os_store_last_native_error(UMF_OS_RESULT_ERROR_BIND_FAILED, errno);
+            LOG_PERR("binding memory to NUMA node failed");
+            // TODO: (errno == 0) when hwloc_set_area_membind() fails on Windows,
+            // ignore this temporarily
+            if (errno != ENOSYS &&
+                errno != 0) { // ENOSYS - Function not implemented
+                // Do not error out if memory binding is not implemented at all
+                // (like in case of WSL on Windows).
+                goto err_unmap;
+            }
         }
-    }
-
+    } while (size > 0);
     *resultPtr = addr;
 
     return UMF_RESULT_SUCCESS;

src/utils/utils_concurrency.h

@@ -70,6 +70,8 @@ static __inline unsigned char util_mssb_index(long long value) {
     InterlockedExchange64((LONG64 volatile *)object, (LONG64)desired)
 #define util_atomic_increment(object)                                          \
     InterlockedIncrement64((LONG64 volatile *)object)
+#define util_fetch_and_add64(ptr, value)                                       \
+    InterlockedExchangeAdd64((LONG64 *)(ptr), value)
 #else
 #define util_lssb_index(x) ((unsigned char)__builtin_ctzll(x))
 #define util_mssb_index(x) ((unsigned char)(63 - __builtin_clzll(x)))
@@ -87,6 +89,7 @@ static __inline unsigned char util_mssb_index(long long value) {
 
 #define util_atomic_increment(object)                                          \
     __atomic_add_fetch(object, 1, __ATOMIC_ACQ_REL)
+#define util_fetch_and_add64 __sync_fetch_and_add
 #endif
 
 #ifdef __cplusplus

test/provider_os_memory.cpp

@@ -45,7 +45,7 @@ using providerCreateExtParams = std::tuple<umf_memory_provider_ops_t *, void *>;
 
 umf::provider_unique_handle_t
 providerCreateExt(providerCreateExtParams params) {
-    umf_memory_provider_handle_t hProvider;
+    umf_memory_provider_handle_t hProvider = nullptr;
     auto [provider_ops, provider_params] = params;
 
     auto ret =