add mem info utils

Author: madsbk

cpp/include/rapidsmpf/system_info.hpp

@@ -6,6 +6,17 @@
 
 namespace rapidsmpf {
 
+/**
+ * @brief Get the total amount of host (main) memory.
+ *
+ * @return Total host memory in bytes.
+ *
+ * @note On WSL and in containerized environments, the returned value
+ * reflects the memory visible to the Linux kernel instance, which may
+ * differ from the host system's physical memory.
+ */
+std::uint64_t get_total_host_memory() noexcept;
+
 /**
  * @brief Get the NUMA node ID associated with the calling CPU thread.
  *
@@ -27,4 +38,21 @@ namespace rapidsmpf {
  */
 int get_current_numa_node_id() noexcept;
 
+/**
+ * @brief Get the total amount of host memory for a NUMA node.
+ *
+ * @param numa_id
+ *     NUMA node for which to query the total host memory. Defaults to the
+ *     current NUMA node as returned by `get_current_numa_node_id()`.
+ *
+ * @note If NUMA support is not available or the node size cannot be
+ * determined, this function falls back to returning the total host memory.
+ *
+ * @return Total host memory of the NUMA node in bytes.
+ */
+std::uint64_t get_numa_node_host_memory(
+    int numa_id = get_current_numa_node_id()
+) noexcept;
+
+
 }  // namespace rapidsmpf

cpp/src/system_info.cpp

@@ -12,14 +12,44 @@
 
 namespace rapidsmpf {
 
+std::uint64_t get_total_host_memory() noexcept {
+    static const uint64_t ret = [] {
+        auto const page_size = ::sysconf(_SC_PAGE_SIZE);
+        auto const phys_pages = ::sysconf(_SC_PHYS_PAGES);
+
+        // Because `get_total_host_memory()` is marked `noexcept`, any failure
+        // here results in process termination, which is intentional.
+        RAPIDSMPF_EXPECTS(
+            page_size > 0 && phys_pages > 0,
+            "sysconf(_SC_PAGE_SIZE/_SC_PHYS_PAGES) failed",
+            std::runtime_error
+        );
+        return static_cast<std::uint64_t>(page_size)
+               * static_cast<std::uint64_t>(phys_pages);
+    }();
+    return ret;
+}
+
 int get_current_numa_node_id() noexcept {
-    static const int numa_node_id = [] {
+    static const int ret = [] {
         if (!numa_available()) {
             return 0;
         }
         return numa_node_of_cpu(sched_getcpu());
     }();
-    return numa_node_id;
+    return ret;
+}
+
+std::uint64_t get_numa_node_host_memory(int numa_id) noexcept {
+    if (numa_available() < 0) {
+        return get_total_host_memory();
+    }
+    long long free_ll = 0;  // ignored.
+    long long total_ll = numa_node_size64(numa_id, &free_ll);
+    if (total_ll < 0) {
+        return get_total_host_memory();
+    }
+    return static_cast<std::uint64_t>(total_ll);
 }
 
 }  // namespace rapidsmpf