1 files changed, 21 insertions, 3 deletions

diff --git a/intern/numaapi/source/numaapi_linux.c b/intern/numaapi/source/numaapi_linux.c
index 62e9dcdfadf..9750f1c17df 100644
--- a/intern/numaapi/source/numaapi_linux.c
+++ b/intern/numaapi/source/numaapi_linux.c
@@ -34,8 +34,6 @@
 #  include <dlfcn.h>
 #endif
 
-#include <stdio.h>
-
 #ifdef WITH_DYNLOAD
 
 // Descriptor numa library.
@@ -64,6 +62,7 @@ typedef void tnuma_free_cpumask(struct bitmask* bitmask);
 typedef void tnuma_free_nodemask(struct bitmask* bitmask);
 typedef int tnuma_run_on_node_mask(struct bitmask *nodemask);
 typedef int tnuma_run_on_node_mask_all(struct bitmask *nodemask);
+typedef struct bitmask *tnuma_get_run_node_mask(void);
 typedef void tnuma_set_interleave_mask(struct bitmask *nodemask);
 typedef void tnuma_set_localalloc(void);
 
@@ -87,6 +86,7 @@ static tnuma_free_nodemask* numa_free_nodemask;
 static tnuma_free_cpumask* numa_free_cpumask;
 static tnuma_run_on_node_mask* numa_run_on_node_mask;
 static tnuma_run_on_node_mask_all* numa_run_on_node_mask_all;
+static tnuma_get_run_node_mask* numa_get_run_node_mask;
 static tnuma_set_interleave_mask* numa_set_interleave_mask;
 static tnuma_set_localalloc* numa_set_localalloc;
 
@@ -162,6 +162,7 @@ static NUMAAPI_Result loadNumaSymbols(void) {
   NUMA_LIBRARY_FIND(numa_free_nodemask);
   NUMA_LIBRARY_FIND(numa_run_on_node_mask);
   NUMA_LIBRARY_FIND(numa_run_on_node_mask_all);
+  NUMA_LIBRARY_FIND(numa_get_run_node_mask);
   NUMA_LIBRARY_FIND(numa_set_interleave_mask);
   NUMA_LIBRARY_FIND(numa_set_localalloc);
 
@@ -204,7 +205,7 @@ int numaAPI_GetNumNodeProcessors(int node) {
   struct bitmask* cpu_mask = numa_allocate_cpumask();
   numa_node_to_cpus(node, cpu_mask);
   const unsigned int num_bytes = numa_bitmask_nbytes(cpu_mask);
-  const unsigned int num_bits = num_bytes  *8;
+  const unsigned int num_bits = num_bytes * 8;
   // TODO(sergey): There might be faster way calculating number of set bits.
   int num_processors = 0;
   for (unsigned int bit = 0; bit < num_bits; ++bit) {
@@ -225,6 +226,23 @@ int numaAPI_GetNumNodeProcessors(int node) {
 }
 
 ////////////////////////////////////////////////////////////////////////////////
+// Topology helpers.
+
+int numaAPI_GetNumCurrentNodesProcessors(void) {
+  struct bitmask* node_mask = numa_get_run_node_mask();
+  const unsigned int num_bytes = numa_bitmask_nbytes(node_mask);
+  const unsigned int num_bits = num_bytes * 8;
+  int num_processors = 0;
+  for (unsigned int bit = 0; bit < num_bits; ++bit) {
+    if (numa_bitmask_isbitset(node_mask, bit)) {
+      num_processors += numaAPI_GetNumNodeProcessors(bit);
+    }
+  }
+  numa_bitmask_free(node_mask);
+  return num_processors;
+}
+
+////////////////////////////////////////////////////////////////////////////////
 // Affinities.
 
 bool numaAPI_RunProcessOnNode(int node) {