From 235a0390eec64c54f97e6bcbfc0e24307031812a Mon Sep 17 00:00:00 2001
From: daanx <daan@effp.org>
Date: Wed, 2 Apr 2025 10:34:00 -0700
Subject: [PATCH] refactor numa_node_count

---
 include/mimalloc/internal.h | 22 +++------------------
 src/os.c                    | 39 ++++++++++++++++++++++++-------------
 2 files changed, 29 insertions(+), 32 deletions(-)

diff --git a/include/mimalloc/internal.h b/include/mimalloc/internal.h
index 5b3e7e23..51fad09c 100644
--- a/include/mimalloc/internal.h
+++ b/include/mimalloc/internal.h
@@ -140,9 +140,11 @@ void*       _mi_os_alloc_aligned_at_offset(size_t size, size_t alignment, size_t
 void*       _mi_os_get_aligned_hint(size_t try_alignment, size_t size);
 bool        _mi_os_use_large_page(size_t size, size_t alignment);
 size_t      _mi_os_large_page_size(void);
-
 void*       _mi_os_alloc_huge_os_pages(size_t pages, int numa_node, mi_msecs_t max_secs, size_t* pages_reserved, size_t* psize, mi_memid_t* memid);
 
+int         _mi_os_numa_node_count(void);
+int         _mi_os_numa_node(void);
+
 // arena.c
 mi_arena_id_t _mi_arena_id_none(void);
 void        _mi_arena_free(void* p, size_t size, size_t still_committed_size, mi_memid_t memid);
@@ -813,24 +815,6 @@ static inline uintptr_t _mi_random_shuffle(uintptr_t x) {
   return x;
 }
 
-// -------------------------------------------------------------------
-// Optimize numa node access for the common case (= one node)
-// -------------------------------------------------------------------
-
-int    _mi_os_numa_node_get(void);
-size_t _mi_os_numa_node_count_get(void);
-
-extern mi_decl_hidden _Atomic(size_t) _mi_numa_node_count;
-static inline int _mi_os_numa_node(void) {
-  if mi_likely(mi_atomic_load_relaxed(&_mi_numa_node_count) == 1) { return 0; }
-  else return _mi_os_numa_node_get();
-}
-static inline size_t _mi_os_numa_node_count(void) {
-  const size_t count = mi_atomic_load_relaxed(&_mi_numa_node_count);
-  if mi_likely(count > 0) { return count; }
-  else return _mi_os_numa_node_count_get();
-}
-
 
 
 // -----------------------------------------------------------------------
diff --git a/src/os.c b/src/os.c
index 12cc5da3..894e3a45 100644
--- a/src/os.c
+++ b/src/os.c
@@ -696,34 +696,47 @@ static void mi_os_free_huge_os_pages(void* p, size_t size) {
   }
 }
 
+
 /* ----------------------------------------------------------------------------
 Support NUMA aware allocation
 -----------------------------------------------------------------------------*/
 
-_Atomic(size_t)  _mi_numa_node_count; // = 0   // cache the node count
+static _Atomic(size_t) mi_numa_node_count; // = 0   // cache the node count
 
-size_t _mi_os_numa_node_count_get(void) {
-  size_t count = mi_atomic_load_acquire(&_mi_numa_node_count);
-  if (count <= 0) {
+int _mi_os_numa_node_count(void) {
+  size_t count = mi_atomic_load_acquire(&mi_numa_node_count);
+  if mi_unlikely(count == 0) {
     long ncount = mi_option_get(mi_option_use_numa_nodes); // given explicitly?
-    if (ncount > 0) {
+    if (ncount > 0 && ncount < INT_MAX) {
       count = (size_t)ncount;
     }
     else {
-      count = _mi_prim_numa_node_count(); // or detect dynamically
-      if (count == 0) count = 1;
+      const size_t n = _mi_prim_numa_node_count(); // or detect dynamically
+      if (n == 0 || n > INT_MAX) { count = 1; }
+                            else { count = n; }
     }
-    mi_atomic_store_release(&_mi_numa_node_count, count); // save it
+    mi_atomic_store_release(&mi_numa_node_count, count); // save it
     _mi_verbose_message("using %zd numa regions\n", count);
   }
-  return count;
+  mi_assert_internal(count > 0 && count <= INT_MAX);
+  return (int)count;
 }
 
-int _mi_os_numa_node_get(void) {
-  size_t numa_count = _mi_os_numa_node_count();
+static int mi_os_numa_node_get(void) {
+  int numa_count = _mi_os_numa_node_count();
   if (numa_count<=1) return 0; // optimize on single numa node systems: always node 0
   // never more than the node count and >= 0
-  size_t numa_node = _mi_prim_numa_node();
+  const size_t n = _mi_prim_numa_node();
+  int numa_node = (n < INT_MAX ? (int)n : 0);
   if (numa_node >= numa_count) { numa_node = numa_node % numa_count; }
-  return (int)numa_node;
+  return numa_node;
+}
+
+int _mi_os_numa_node(void) {
+  if mi_likely(mi_atomic_load_relaxed(&mi_numa_node_count) == 1) { 
+    return 0; 
+  }
+  else {
+    return mi_os_numa_node_get();
+  }
 }