add mempolicy for split mode

include/umf/mempolicy.h

@@ -25,9 +25,21 @@ typedef enum umf_mempolicy_membind_t {
     /// Bind memory to namespace
     UMF_MEMPOLICY_BIND,
     /// Prefer memory from namespace but fallback to other memory if not available
-    UMF_MEMPOLICY_PREFERRED
+    UMF_MEMPOLICY_PREFERRED,
+    /// Allocation will be split evenly across nodes specified in nodemask.
+    /// umf_mempolicy_split_partition_t can be used to specify different distribution.
+    UMF_MEMPOLICY_SPLIT
 } umf_mempolicy_membind_t;
 
+/// user defined partition for UMF_MEMPOLICY_SPLIT mode
+typedef struct umf_mempolicy_split_partition_t {
+    /// The weight of the partition, representing the proportion of
+    /// the allocation that should be assigned to this NUMA node.
+    unsigned weight;
+    /// The NUMA node where the pages assigned to this partition will be bound.
+    unsigned target;
+} umf_mempolicy_split_partition_t;
+
 ///
 /// @brief Creates a new memory policy
 /// @param bind memory binding policy
@@ -52,6 +64,18 @@ umf_result_t umfMempolicyDestroy(umf_mempolicy_handle_t hPolicy);
 ///
 umf_result_t umfMempolicySetInterleavePartSize(umf_mempolicy_handle_t hPolicy,
                                                size_t partSize);
+
+///
+/// @brief Sets custom split partitions
+/// @param hPolicy handle to memory policy
+/// @param partList ordered array of partitions
+/// @param partListLen length of the partList array
+/// @return UMF_RESULT_SUCCESS on success or appropriate error code on failure.
+///
+umf_result_t
+umfMempolicySetCustomSplitPartitions(umf_mempolicy_handle_t hPolicy,
+                                     umf_mempolicy_split_partition_t *partList,
+                                     size_t partListLen);
 #ifdef __cplusplus
 }
 #endif

src/libumf.def.in

@@ -36,9 +36,10 @@ EXPORTS
     umfMemoryProviderPurgeForce
     umfMemoryProviderPurgeLazy
     umfMemoryProviderPutIPCHandle
-    umfMempolicyCreate;
-    umfMempolicyDestroy;
-    umfMempolicySetInterleavePartSize;
+    umfMempolicyCreate
+    umfMempolicyDestroy
+    umfMempolicySetCustomSplitPartitions
+    umfMempolicySetInterleavePartSize
     umfMemspaceDestroy
     umfOpenIPCHandle    
     umfOsMemoryProviderOps

src/libumf.map

@@ -33,6 +33,7 @@ UMF_1.0 {
         umfMemoryProviderPutIPCHandle;
         umfMempolicyCreate;
         umfMempolicyDestroy;
+        umfMempolicySetCustomSplitPartitions;
         umfMempolicySetInterleavePartSize;
         umfMemspaceCreateFromNumaArray;
         umfMemspaceDestroy;

src/memory_targets/memory_target_numa.c

@@ -20,6 +20,7 @@
 #include "memory_target_numa.h"
 #include "mempolicy_internal.h"
 #include "topology.h"
+#include "utils_assert.h"
 #include "utils_log.h"
 
 struct numa_memory_target_t {
@@ -75,14 +76,33 @@ static umf_result_t numa_memory_provider_create_from_memspace(
         switch (policy->type) {
         case UMF_MEMPOLICY_INTERLEAVE:
             params.numa_mode = UMF_NUMA_MODE_INTERLEAVE;
-            params.part_size = policy->ops.part_size;
+            params.part_size = policy->ops.interleave.part_size;
             break;
         case UMF_MEMPOLICY_BIND:
             params.numa_mode = UMF_NUMA_MODE_BIND;
             break;
         case UMF_MEMPOLICY_PREFERRED:
             params.numa_mode = UMF_NUMA_MODE_PREFERRED;
             break;
+        case UMF_MEMPOLICY_SPLIT:
+            params.numa_mode = UMF_NUMA_MODE_SPLIT;
+
+            // compile time check to ensure we can just cast
+            // umf_mempolicy_split_partition_t to
+            // umf_numa_split_partition_t
+            COMPILE_ERROR_ON(sizeof(umf_mempolicy_split_partition_t) !=
+                             sizeof(umf_numa_split_partition_t));
+            COMPILE_ERROR_ON(
+                offsetof(umf_mempolicy_split_partition_t, weight) !=
+                offsetof(umf_numa_split_partition_t, weight));
+            COMPILE_ERROR_ON(
+                offsetof(umf_mempolicy_split_partition_t, target) !=
+                offsetof(umf_numa_split_partition_t, target));
+
+            params.partitions =
+                (umf_numa_split_partition_t *)policy->ops.split.part;
+            params.partitions_len = policy->ops.split.part_len;
+            break;
         default:
             return UMF_RESULT_ERROR_INVALID_ARGUMENT;
         }

src/mempolicy.c

@@ -25,8 +25,11 @@ umf_result_t umfMempolicyCreate(umf_mempolicy_membind_t bind,
 
     (*policy)->type = bind;
     if (bind == UMF_MEMPOLICY_INTERLEAVE) {
-        (*policy)->ops.part_size = 0;
+        (*policy)->ops.interleave.part_size = 0;
+    } else if (bind == UMF_MEMPOLICY_SPLIT) {
+        (*policy)->ops.split.part_len = 0;
     }
+
     return UMF_RESULT_SUCCESS;
 }
 
@@ -45,6 +48,22 @@ umf_result_t umfMempolicySetInterleavePartSize(umf_mempolicy_handle_t policy,
         return UMF_RESULT_ERROR_INVALID_ARGUMENT;
     }
 
-    policy->ops.part_size = partSize;
+    policy->ops.interleave.part_size = partSize;
+    return UMF_RESULT_SUCCESS;
+}
+
+umf_result_t
+umfMempolicySetCustomSplitPartitions(umf_mempolicy_handle_t policy,
+                                     umf_mempolicy_split_partition_t *partList,
+                                     size_t partListLen) {
+    if (policy == NULL) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+    if (policy->type != UMF_MEMPOLICY_SPLIT) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    policy->ops.split.part = partList;
+    policy->ops.split.part_len = partListLen;
     return UMF_RESULT_SUCCESS;
 }

src/mempolicy_internal.h

@@ -19,7 +19,13 @@ extern "C" {
 typedef struct umf_mempolicy_t {
     umf_mempolicy_membind_t type;
     union {
-        size_t part_size;
+        struct {
+            size_t part_size;
+        } interleave;
+        struct {
+            umf_mempolicy_split_partition_t *part;
+            size_t part_len;
+        } split;
     } ops;
 } umf_mempolicy_t;
 

src/utils/utils_assert.h

@@ -56,6 +56,8 @@ extern "C" {
         }                                                                      \
     } while (0)
 
+#define COMPILE_ERROR_ON(cond) ((void)sizeof(char[(cond) ? -1 : 1]))
+
 #ifdef __cplusplus
 }
 #endif

test/memspaces/mempolicy.cpp

@@ -36,6 +36,7 @@ TEST_F(test, mempolicyDefaultPreferred) {
     ASSERT_NE(ProviderInternal, nullptr);
     EXPECT_EQ(ProviderInternal->numa_policy, HWLOC_MEMBIND_BIND);
     EXPECT_EQ(ProviderInternal->numa_flags, HWLOC_MEMBIND_BYNODESET);
+    EXPECT_EQ(ProviderInternal->mode, UMF_NUMA_MODE_PREFERRED);
     umfMemoryProviderDestroy(hProvider);
 }
 
@@ -59,6 +60,7 @@ TEST_F(test, mempolicyDefaultBind) {
     EXPECT_EQ(ProviderInternal->numa_policy, HWLOC_MEMBIND_BIND);
     EXPECT_EQ(ProviderInternal->numa_flags,
               HWLOC_MEMBIND_BYNODESET | HWLOC_MEMBIND_STRICT);
+    EXPECT_EQ(ProviderInternal->mode, UMF_NUMA_MODE_BIND);
     umfMemoryProviderDestroy(hProvider);
 }
 
@@ -82,6 +84,7 @@ TEST_F(test, mempolicyDefaultInterleave) {
     EXPECT_EQ(ProviderInternal->numa_policy, HWLOC_MEMBIND_INTERLEAVE);
     EXPECT_EQ(ProviderInternal->numa_flags, HWLOC_MEMBIND_BYNODESET);
     EXPECT_EQ(ProviderInternal->part_size, 0);
+    EXPECT_EQ(ProviderInternal->mode, UMF_NUMA_MODE_INTERLEAVE);
     umfMemoryProviderDestroy(hProvider);
 }
 
@@ -109,5 +112,79 @@ TEST_F(test, mempolicyInterleavePartSize) {
     EXPECT_EQ(ProviderInternal->numa_flags,
               HWLOC_MEMBIND_BYNODESET | HWLOC_MEMBIND_STRICT);
     EXPECT_EQ(ProviderInternal->part_size, part_size);
+    EXPECT_EQ(ProviderInternal->mode, UMF_NUMA_MODE_INTERLEAVE);
     umfMemoryProviderDestroy(hProvider);
 }
+
+TEST_F(test, mempolicyDefaultSplit) {
+    umf_memory_provider_handle_t hProvider = nullptr;
+    umf_mempolicy_handle_t hPolicy = nullptr;
+
+    umf_result_t ret = umfMempolicyCreate(UMF_MEMPOLICY_SPLIT, &hPolicy);
+    ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
+
+    ret = umfMemoryProviderCreateFromMemspace(umfMemspaceHostAllGet(), hPolicy,
+                                              &hProvider);
+    ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
+    ASSERT_NE(hProvider, nullptr);
+    ret = umfMempolicyDestroy(hPolicy);
+    ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
+
+    os_memory_provider_t *ProviderInternal =
+        (os_memory_provider_t *)providerGetPriv(hProvider);
+    ASSERT_NE(ProviderInternal, nullptr);
+    EXPECT_EQ(ProviderInternal->numa_policy, HWLOC_MEMBIND_BIND);
+    EXPECT_EQ(ProviderInternal->numa_flags,
+              HWLOC_MEMBIND_BYNODESET | HWLOC_MEMBIND_STRICT);
+    EXPECT_EQ(ProviderInternal->partitions_len, ProviderInternal->nodeset_len);
+    EXPECT_EQ(ProviderInternal->mode, UMF_NUMA_MODE_SPLIT);
+    umfMemoryProviderDestroy(hProvider);
+}
+
+TEST_F(test, mempolicyCustomSplit) {
+    umf_memory_provider_handle_t hProvider = nullptr;
+    umf_mempolicy_handle_t hPolicy = nullptr;
+
+    umf_result_t ret = umfMempolicyCreate(UMF_MEMPOLICY_SPLIT, &hPolicy);
+    ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
+    umf_mempolicy_split_partition_t part[] = {{1, 0}, {1, 0}};
+
+    ret = umfMempolicySetCustomSplitPartitions(hPolicy, part, 2);
+    ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
+
+    ret = umfMemoryProviderCreateFromMemspace(umfMemspaceHostAllGet(), hPolicy,
+                                              &hProvider);
+    ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
+    ASSERT_NE(hProvider, nullptr);
+    ret = umfMempolicyDestroy(hPolicy);
+    ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
+
+    os_memory_provider_t *ProviderInternal =
+        (os_memory_provider_t *)providerGetPriv(hProvider);
+    ASSERT_NE(ProviderInternal, nullptr);
+    EXPECT_EQ(ProviderInternal->numa_policy, HWLOC_MEMBIND_BIND);
+    EXPECT_EQ(ProviderInternal->numa_flags,
+              HWLOC_MEMBIND_BYNODESET | HWLOC_MEMBIND_STRICT);
+    EXPECT_EQ(ProviderInternal->partitions_len, 2);
+    EXPECT_EQ(ProviderInternal->mode, UMF_NUMA_MODE_SPLIT);
+    EXPECT_EQ(ProviderInternal->partitions_weight_sum, 2);
+    EXPECT_EQ(ProviderInternal->partitions[0].target,
+              ProviderInternal->partitions[1].target);
+    umfMemoryProviderDestroy(hProvider);
+}
+
+TEST_F(test, mempolicySplitNegative) {
+    umf_mempolicy_handle_t hPolicy = nullptr;
+
+    umf_result_t ret = umfMempolicyCreate(UMF_MEMPOLICY_BIND, &hPolicy);
+    ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
+    umf_mempolicy_split_partition_t part[] = {{1, 0}, {1, 0}};
+
+    ret = umfMempolicySetCustomSplitPartitions(hPolicy, part, 2);
+    ASSERT_EQ(ret, UMF_RESULT_ERROR_INVALID_ARGUMENT);
+
+    ret = umfMempolicySetCustomSplitPartitions(NULL, part, 2);
+    ASSERT_EQ(ret, UMF_RESULT_ERROR_INVALID_ARGUMENT);
+    ret = umfMempolicyDestroy(hPolicy);
+    ASSERT_EQ(ret, UMF_RESULT_SUCCESS);
+}