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 specyfic)
+    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_size;
     hwloc_topology_t topo;
 } os_memory_provider_t;
 
@@ -83,28 +87,64 @@ static void os_store_last_native_error(int32_t native_error, int 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;
-    }
+                                              os_memory_provider_t *os_provider,
+                                              int is_separate) {
 
-    *out_nodeset = hwloc_bitmap_alloc();
-    if (!*out_nodeset) {
+    unsigned long array_size = listsize && is_separate ? 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 (is_separate) {
+        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 +172,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;
+    }
+    assert(0);
+    return -1;
+}
+
+//return 1 if membinding logic is performed at umf side
+//return 0 if membinding logic is performed by kernel
+static int advancedMembind(umf_os_memory_provider_params_t *in_params) {
+    if (in_params->numa_mode == UMF_NUMA_MODE_INTERLEAVE) {
+        return in_params->part_size > 0;
     }
-    return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    return 0;
 }
 
 static int getHwlocMembindFlags(umf_numa_mode_t mode) {
@@ -193,17 +252,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 manual_mode = advancedMembind(in_params);
+    provider->numa_policy =
+        translate_numa_mode(in_params->numa_mode, manual_mode);
     provider->numa_flags = getHwlocMembindFlags(in_params->numa_mode);
-
+    provider->part_size = in_params->part_size;
     return nodemask_to_hwloc_nodeset(
-        in_params->numa_list, in_params->numa_list_len, &provider->nodeset);
+        in_params->numa_list, in_params->numa_list_len, provider, manual_mode);
 }
 
 static umf_result_t os_initialize(void *params, void **provider) {
@@ -252,7 +313,7 @@ static umf_result_t os_initialize(void *params, void **provider) {
     } else {
         if (hwloc_bitmap_list_snprintf(os_provider->nodeset_str_buf,
                                        NODESET_STR_BUF_LEN,
-                                       os_provider->nodeset)) {
+                                       *os_provider->nodeset)) {
             LOG_INFO("OS provider initialized with NUMA nodes: %s",
                      os_provider->nodeset_str_buf);
         }
@@ -280,7 +341,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 +451,15 @@ 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;
+    }
+
+    size_t s = util_fetch_and_add64(&provider->alloc_size, 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 +507,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 implecmented at all
+                // (like in case of WSL on Windows).
+                goto err_unmap;
+            }
         }
-    }
-
+    } while (size > 0);
     *resultPtr = addr;
 
     return UMF_RESULT_SUCCESS;

src/utils/utils_posix_concurrency.c

@@ -12,6 +12,15 @@
 
 #include "utils_concurrency.h"
 
+// TODO: implement all atomics
+#ifndef _MSC_VER
+#define util_fetch_and_add64 __sync_fetch_and_add
+#else
+#include <windows.h>
+#define util_fetch_and_add64(ptr, value)                                       \
+    InterlockedExchangeAdd64((LONG64 *)(ptr), value)
+#endif
+
 size_t util_mutex_get_size(void) { return sizeof(pthread_mutex_t); }
 
 os_mutex_t *util_mutex_init(void *ptr) {

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 =