From 7c6657610640c5d0f0aa2128da0eef4873f72738 Mon Sep 17 00:00:00 2001
From: Abhishek Bagusetty <59661409+abagusetty@users.noreply.github.com>
Date: Thu, 6 Jun 2024 18:14:23 -0500
Subject: [PATCH] Numa fix (#125)

* Fix cpuinfo names for Arm CPUs

* Fix a bug leading to segfaults from previous commit for the arm

* fix to choose appropriate NUMA nodes in RMM memory pool
---
 src/tamm/rmm_memory_pool.hpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/tamm/rmm_memory_pool.hpp b/src/tamm/rmm_memory_pool.hpp
index fd3172753..bf13e9c29 100644
--- a/src/tamm/rmm_memory_pool.hpp
+++ b/src/tamm/rmm_memory_pool.hpp
@@ -196,6 +196,7 @@ class RMMMemoryManager {
     EXPECTS_STR((numa_available() != -1), "[TAMM ERROR]: numa APIs are not available!");
 
     numa_set_bind_policy(1);
+    numa_set_strict(1);
     unsigned numNumaNodes = numa_num_task_nodes();
 
     // for ranks_pn_=1, there is no need to check the mapping to numa-nodes (mostly used for CI)
@@ -208,7 +209,7 @@ class RMMMemoryManager {
     struct bitmask* numaNodes = numa_get_mems_allowed();
     numa_bind(numaNodes);
 
-    int  numa_id         = numa_preferred();
+    int  numa_id         = numa_node_of_cpu(sched_getcpu());
     long numa_total_size = numa_node_size(numa_id, &max_host_bytes);
     max_host_bytes *= 0.40; // reserve 40% only of the free numa-node memory (reserving rest of
                             // GA, non-pool allocations)