refactor numa_node_count

This commit is contained in:
daanx 2025-04-02 10:34:00 -07:00
parent 77b622511a
commit 235a0390ee
2 changed files with 29 additions and 32 deletions

View file

@ -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();
}
// -----------------------------------------------------------------------

View file

@ -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();
}
}