Add memspace "highest bandwidth"

kswiecicki · kswiecicki · commit 7c35d416e7ed · 2024-04-25T18:33:00.000+02:00
This memspace contains memory targets ordered by the highest bandwidth
value (targets higher in the ordering are prioritised).
Querying the bandwidth value requires HMAT support on the platform,
calling umfMemspaceHighestBandwidthGet() will return NULL if it's not
supported.
Introduce UMF_MEMSPACE_HIGHEST_BANDWIDTH env var. This environment variable
is used to forcefully assign NUMA nodes to the memspace "highest bandwidth".
The ordering of memory targets within this memspace will correspond to the
sequence of NUMA nodes specified in the environment variable.
diff --git a/include/umf/memspace.h b/include/umf/memspace.h
@@ -56,6 +56,11 @@ umf_memspace_handle_t umfMemspaceHostAllGet(void);
 ///
 umf_memspace_handle_t umfMemspaceHighestCapacityGet(void);
 
+/// \brief Retrieves predefined highest bandwidth memspace.
+/// \return highest bandwidth memspace handle on success or NULL on failure.
+///
+umf_memspace_handle_t umfMemspaceHighestBandwidthGet(void);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -80,7 +80,8 @@ set(UMF_SOURCES_LINUX
     memory_targets/memory_target_numa.c
     memspaces/memspace_numa.c
     memspaces/memspace_host_all.c
-    memspaces/memspace_highest_capacity.c)
+    memspaces/memspace_highest_capacity.c
+    memspaces/memspace_highest_bandwidth.c)
 
 set(UMF_SOURCES_WINDOWS ${UMF_SOURCES_WINDOWS} provider/provider_os_memory.c
                         provider/provider_os_memory_windows.c)
diff --git a/src/libumf.map b/src/libumf.map
@@ -29,6 +29,7 @@ UMF_1.0 {
         umfMemoryProviderPutIPCHandle;
         umfMemspaceCreateFromNumaArray;
         umfMemspaceDestroy;
+        umfMemspaceHighestBandwidthGet;
         umfMemspaceHighestCapacityGet;
         umfMemspaceHostAllGet;
         umfOpenIPCHandle;
diff --git a/src/libumf_linux.c b/src/libumf_linux.c
@@ -29,6 +29,7 @@ void __attribute__((destructor)) umfDestroy(void) {
     umfMemoryTrackerDestroy(t);
     umfMemspaceHostAllDestroy();
     umfMemspaceHighestCapacityDestroy();
+    umfMemspaceHighestBandwidthDestroy();
     umfDestroyTopology();
 }
 
diff --git a/src/memory_target.c b/src/memory_target.c
@@ -83,3 +83,14 @@ umf_result_t umfMemoryTargetGetCapacity(umf_memory_target_handle_t memoryTarget,
     assert(capacity);
     return memoryTarget->ops->get_capacity(memoryTarget->priv, capacity);
 }
+
+umf_result_t
+umfMemoryTargetGetBandwidth(umf_memory_target_handle_t srcMemoryTarget,
+                            umf_memory_target_handle_t dstMemoryTarget,
+                            size_t *bandwidth) {
+    assert(srcMemoryTarget);
+    assert(dstMemoryTarget);
+    assert(bandwidth);
+    return srcMemoryTarget->ops->get_bandwidth(
+        srcMemoryTarget->priv, dstMemoryTarget->priv, bandwidth);
+}
diff --git a/src/memory_target.h b/src/memory_target.h
@@ -37,6 +37,10 @@ umf_result_t umfMemoryTargetClone(umf_memory_target_handle_t memoryTarget,
                                   umf_memory_target_handle_t *outHandle);
 umf_result_t umfMemoryTargetGetCapacity(umf_memory_target_handle_t memoryTarget,
                                         size_t *capacity);
+umf_result_t
+umfMemoryTargetGetBandwidth(umf_memory_target_handle_t srcMemoryTarget,
+                            umf_memory_target_handle_t dstMemoryTarget,
+                            size_t *bandwidth);
 
 #ifdef __cplusplus
 }
diff --git a/src/memory_target_ops.h b/src/memory_target_ops.h
@@ -41,6 +41,8 @@ typedef struct umf_memory_target_ops_t {
         umf_memory_provider_handle_t *provider);
 
     umf_result_t (*get_capacity)(void *memoryTarget, size_t *capacity);
+    umf_result_t (*get_bandwidth)(void *srcMemoryTarget, void *dstMemoryTarget,
+                                  size_t *bandwidth);
 } umf_memory_target_ops_t;
 
 #ifdef __cplusplus
diff --git a/src/memory_targets/memory_target_numa.c b/src/memory_targets/memory_target_numa.c
@@ -143,12 +143,49 @@ static umf_result_t numa_get_capacity(void *memTarget, size_t *capacity) {
     return UMF_RESULT_SUCCESS;
 }
 
+static umf_result_t numa_get_bandwidth(void *srcMemoryTarget,
+                                       void *dstMemoryTarget,
+                                       size_t *bandwidth) {
+    hwloc_topology_t topology = umfGetTopology();
+    if (!topology) {
+        return UMF_RESULT_ERROR_NOT_SUPPORTED;
+    }
+
+    hwloc_obj_t srcNumaNode = hwloc_get_obj_by_type(
+        topology, HWLOC_OBJ_NUMANODE,
+        ((struct numa_memory_target_t *)srcMemoryTarget)->physical_id);
+    if (!srcNumaNode) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    hwloc_obj_t dstNumaNode = hwloc_get_obj_by_type(
+        topology, HWLOC_OBJ_NUMANODE,
+        ((struct numa_memory_target_t *)dstMemoryTarget)->physical_id);
+    if (!dstNumaNode) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    struct hwloc_location initiator = {.location.cpuset = srcNumaNode->cpuset,
+                                       .type = HWLOC_LOCATION_TYPE_CPUSET};
+    hwloc_uint64_t value = 0;
+    int ret = hwloc_memattr_get_value(topology, HWLOC_MEMATTR_ID_BANDWIDTH,
+                                      dstNumaNode, &initiator, 0, &value);
+    if (ret) {
+        return (errno == EINVAL) ? UMF_RESULT_ERROR_NOT_SUPPORTED
+                                 : UMF_RESULT_ERROR_UNKNOWN;
+    }
+
+    *bandwidth = value;
+    return UMF_RESULT_SUCCESS;
+}
+
 struct umf_memory_target_ops_t UMF_MEMORY_TARGET_NUMA_OPS = {
     .version = UMF_VERSION_CURRENT,
     .initialize = numa_initialize,
     .finalize = numa_finalize,
     .pool_create_from_memspace = numa_pool_create_from_memspace,
     .clone = numa_clone,
     .get_capacity = numa_get_capacity,
+    .get_bandwidth = numa_get_bandwidth,
     .memory_provider_create_from_memspace =
         numa_memory_provider_create_from_memspace};
diff --git a/src/memspace.c b/src/memspace.c
@@ -208,3 +208,89 @@ umfMemspaceSortDesc(umf_memspace_handle_t hMemspace,
 
     return UMF_RESULT_SUCCESS;
 }
+
+enum umf_result_t umfMemspaceFilter(
+    umf_memspace_handle_t hMemspace,
+    umf_result_t (*getTarget)(umf_memory_target_handle_t initiator,
+                              umf_memory_target_handle_t *nodes,
+                              size_t numNodes,
+                              umf_memory_target_handle_t *target),
+    umf_memspace_handle_t *filteredMemspace) {
+    if (!hMemspace || !getTarget) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
+    umf_memory_target_handle_t *uniqueBestNodes = umf_ba_global_alloc(
+        hMemspace->size * sizeof(umf_memory_target_handle_t *));
+    if (!uniqueBestNodes) {
+        return UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
+    }
+
+    umf_result_t ret = UMF_RESULT_SUCCESS;
+
+    size_t numUniqueBestNodes = 0;
+    for (size_t nodeIdx = 0; nodeIdx < hMemspace->size; nodeIdx++) {
+        umf_memory_target_handle_t target = NULL;
+        ret = getTarget(hMemspace->nodes[nodeIdx], hMemspace->nodes,
+                        hMemspace->size, &target);
+        if (ret != UMF_RESULT_SUCCESS) {
+            goto err_free_best_targets;
+        }
+
+        // check if the target is already present in the best nodes
+        size_t bestTargetIdx;
+        for (bestTargetIdx = 0; bestTargetIdx < numUniqueBestNodes;
+             bestTargetIdx++) {
+            if (uniqueBestNodes[bestTargetIdx] == target) {
+                break;
+            }
+        }
+
+        // if the target is not present, add it to the best nodes
+        if (bestTargetIdx == numUniqueBestNodes) {
+            uniqueBestNodes[numUniqueBestNodes++] = target;
+        }
+    }
+
+    // copy the unique best nodes into a new memspace
+    umf_memspace_handle_t newMemspace =
+        umf_ba_global_alloc(sizeof(struct umf_memspace_t));
+    if (!newMemspace) {
+        ret = UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
+        goto err_free_best_targets;
+    }
+
+    newMemspace->size = numUniqueBestNodes;
+    newMemspace->nodes = umf_ba_global_alloc(
+        sizeof(umf_memory_target_handle_t) * newMemspace->size);
+    if (!newMemspace->nodes) {
+        ret = UMF_RESULT_ERROR_OUT_OF_HOST_MEMORY;
+        goto err_free_new_memspace;
+    }
+
+    size_t cloneIdx = 0;
+    for (size_t cloneIdx = 0; cloneIdx < newMemspace->size; cloneIdx++) {
+        ret = umfMemoryTargetClone(uniqueBestNodes[cloneIdx],
+                                   &newMemspace->nodes[cloneIdx]);
+        if (ret != UMF_RESULT_SUCCESS) {
+            goto err_free_cloned_nodes;
+        }
+    }
+
+    *filteredMemspace = newMemspace;
+    umf_ba_global_free(uniqueBestNodes);
+
+    return UMF_RESULT_SUCCESS;
+
+err_free_cloned_nodes:
+    while (cloneIdx != 0) {
+        cloneIdx--;
+        umfMemoryTargetDestroy(newMemspace->nodes[cloneIdx]);
+    }
+    umf_ba_global_free(newMemspace->nodes);
+err_free_new_memspace:
+    umf_ba_global_free(newMemspace);
+err_free_best_targets:
+    umf_ba_global_free(uniqueBestNodes);
+    return ret;
+}
diff --git a/src/memspace_internal.h b/src/memspace_internal.h
@@ -39,6 +39,18 @@ typedef umf_result_t (*umfGetPropertyFn)(umf_memory_target_handle_t,
 enum umf_result_t umfMemspaceSortDesc(umf_memspace_handle_t hMemspace,
                                       umfGetPropertyFn getProperty);
 
+typedef umf_result_t (*umfGetTargetFn)(umf_memory_target_handle_t initiator,
+                                       umf_memory_target_handle_t *nodes,
+                                       size_t numNodes,
+                                       umf_memory_target_handle_t *target);
+
+///
+/// \brief Filters the targets using getTarget() to create a new memspace
+///
+enum umf_result_t umfMemspaceFilter(umf_memspace_handle_t hMemspace,
+                                    umfGetTargetFn getTarget,
+                                    umf_memspace_handle_t *filteredMemspace);
+
 ///
 /// \brief Destroys memspace
 /// \param hMemspace handle to memspace
@@ -47,6 +59,7 @@ void umfMemspaceDestroy(umf_memspace_handle_t hMemspace);
 
 void umfMemspaceHostAllDestroy(void);
 void umfMemspaceHighestCapacityDestroy(void);
+void umfMemspaceHighestBandwidthDestroy(void);
 
 #ifdef __cplusplus
 }
diff --git a/src/memspaces/memspace_highest_bandwidth.c b/src/memspaces/memspace_highest_bandwidth.c
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
diff --git a/test/memspaces/memspace_highest_bandwidth.cpp b/test/memspaces/memspace_highest_bandwidth.cpp

Original file line number	Diff line number	Diff line change
`@@ -29,6 +29,7 @@ void __attribute__((destructor)) umfDestroy(void) {`
`29`	`29`	`umfMemoryTrackerDestroy(t);`
`30`	`30`	`umfMemspaceHostAllDestroy();`
`31`	`31`	`umfMemspaceHighestCapacityDestroy();`
	`32`	`+ umfMemspaceHighestBandwidthDestroy();`
`32`	`33`	`umfDestroyTopology();`
`33`	`34`	`}`
`34`	`35`