feat(rocr): Bind asyncHandler thread to same NUMA domain as main thread

projects/rocr-runtime/runtime/docs/data/env_variables.rst

@@ -89,10 +89,13 @@
         | 1: Enable
 
     * - | ``HSA_OVERRIDE_CPU_AFFINITY_DEBUG``
-        | Controls whether ROCm helper threads inherit the parent process's CPU affinity mask.
+        | Controls CPU affinity behavior for ROCm helper threads.
       - ``1``
-      - | 0: Enable inheritance. Helper threads use the parent process's core affinity mask, which should be set with enough cores for all threads.
-        | 1: Disable inheritance. Helper threads spawn on all available cores, ignoring the parent's affinity settings, which may affect performance in certain environments.
+      - | 0: Helper threads inherit the parent process's exact CPU affinity mask. Use this when you want strict core assignment.
+        | 1: NUMA-aware affinity (default). ROCr checks the main thread's CPU affinity:
+        |    - If the main thread is restricted to specific CPU cores, helper threads are bound to the same NUMA domain(s) as those cores.
+        |    - If the main thread is unrestricted, helper threads remain unrestricted.
+        | This prevents helper threads from being scheduled on CPU cores in a different NUMA domain than the GPU, which can cause performance degradation on NUMA systems (e.g., MI-300A).
 
     * - | ``HSA_ENABLE_DEBUG``
         | Enables additional debug information and validation in the runtime.

projects/rocr-runtime/runtime/hsa-runtime/CMakeLists.txt

@@ -93,6 +93,99 @@ if(NOT LibElf_FOUND)
   find_package(LibElf REQUIRED)
 endif()
 
+## Find NUMA library
+if(UNIX)
+  set(NUMA_FOUND FALSE)
+
+  # Try pkg-config first (for TheRock builds which may use rocm_sysdeps_numa)
+  if(PKG_CONFIG_FOUND)
+    pkg_check_modules(NUMA_PC numa)
+    if(NUMA_PC_FOUND)
+      # Resolve actual library file using LIBRARY_DIRS as hints
+      find_library(NUMA_LIBRARY_FILE
+        NAMES numa rocm_sysdeps_numa
+        HINTS ${NUMA_PC_LIBRARY_DIRS}
+        NO_DEFAULT_PATH
+      )
+
+      if(NUMA_LIBRARY_FILE)
+        # Create SHARED IMPORTED target with resolved path
+        add_library(numa SHARED IMPORTED GLOBAL)
+        set_target_properties(numa PROPERTIES
+          IMPORTED_LOCATION "${NUMA_LIBRARY_FILE}"
+          INTERFACE_INCLUDE_DIRECTORIES "${NUMA_PC_INCLUDE_DIRS}"
+        )
+        set(NUMA_FOUND TRUE)
+        message(STATUS "Found NUMA via pkg-config (SHARED): ${NUMA_LIBRARY_FILE}")
+      else()
+        # pkg-config reported numa but no actual library file resolved under its
+        # LIBRARY_DIRS. Do NOT accept a link-flags-only INTERFACE target here: a
+        # bare "numa" item resolves back to this target and CMake silently drops
+        # it, leaving NUMA_FOUND=TRUE while nothing is linked -> unresolved
+        # numa_* symbols at link time. Leave NUMA_FOUND FALSE so the
+        # find_package / manual-search tiers below can resolve a real library
+        # file, or we degrade to the all-CPU path if none exists.
+        message(STATUS "NUMA via pkg-config: no library file resolved under "
+                       "'${NUMA_PC_LIBRARY_DIRS}'; trying find_package/manual search")
+      endif()
+    endif()
+  endif()
+
+  # Try CMake package config. ROCm ships this as numa::numa, while the
+  # link lines below use the target name numa.
+  if(NOT NUMA_FOUND)
+    set(_numa_hints)
+    if(DEFINED ROCM_PATH)
+      list(APPEND _numa_hints "${ROCM_PATH}/lib/rocm_sysdeps")
+    endif()
+    if(DEFINED ENV{ROCM_PATH})
+      list(APPEND _numa_hints "$ENV{ROCM_PATH}/lib/rocm_sysdeps")
+    endif()
+    find_package(NUMA QUIET HINTS ${_numa_hints})
+    if(TARGET numa::numa)
+      set(_numa_package_target numa::numa)
+    elseif(TARGET NUMA::NUMA)
+      set(_numa_package_target NUMA::NUMA)
+    else()
+      set(_numa_package_target "")
+    endif()
+
+    if(_numa_package_target)
+      if(NOT TARGET numa)
+        add_library(numa INTERFACE IMPORTED GLOBAL)
+        target_link_libraries(numa INTERFACE ${_numa_package_target})
+      endif()
+      set(NUMA_FOUND TRUE)
+      message(STATUS "Found NUMA via CMake package: ${_numa_package_target}")
+    else()
+      set(NUMA_FOUND FALSE)
+    endif()
+    unset(_numa_package_target)
+    unset(_numa_hints)
+  endif()
+
+  # Manual detection as fallback
+  if(NOT NUMA_FOUND)
+    find_library(NUMA_LIBRARY NAMES numa rocm_sysdeps_numa)
+    find_path(NUMA_INCLUDE_DIR NAMES numa.h)
+
+    if(NUMA_LIBRARY AND NUMA_INCLUDE_DIR)
+      add_library(numa SHARED IMPORTED GLOBAL)
+      set_target_properties(numa PROPERTIES
+        IMPORTED_LOCATION "${NUMA_LIBRARY}"
+        INTERFACE_INCLUDE_DIRECTORIES "${NUMA_INCLUDE_DIR}")
+      set(NUMA_FOUND TRUE)
+      message(STATUS "Found NUMA via manual search: ${NUMA_LIBRARY}")
+    endif()
+  endif()
+
+  if(NOT NUMA_FOUND)
+    message(WARNING "NUMA library not found; helper threads will use all CPUs. "
+                    "Install libnuma-dev for NUMA-aware CPU affinity (perf optimization).")
+  endif()
+endif()
+
+
 ## Find AMD CUID library.
 set(AMDCUID_ROOT "/opt/amdcuid" CACHE PATH "Path to AMDCUID installation")
 find_library(AMDCUID_LIB
@@ -384,7 +477,14 @@ if(${IMAGE_SUPPORT})
 endif()
 
 if (UNIX)
-  target_link_libraries ( ${CORE_RUNTIME_TARGET} PRIVATE elf::elf dl pthread rt )
+  if(NUMA_FOUND)
+    target_link_libraries ( ${CORE_RUNTIME_TARGET} PRIVATE elf::elf dl pthread rt numa )
+    # Gate the numa_* code in os_linux.cpp on the same condition that links the
+    # library, so the compiled calls and the linked library never disagree.
+    target_compile_definitions ( ${CORE_RUNTIME_TARGET} PRIVATE HSA_HAS_NUMA=1 )
+  else()
+    target_link_libraries ( ${CORE_RUNTIME_TARGET} PRIVATE elf::elf dl pthread rt )
+  endif()
 else()
   target_link_libraries ( ${CORE_RUNTIME_TARGET} PRIVATE Ws2_32 )
   target_link_directories(${CORE_RUNTIME_TARGET} PRIVATE ${DXCORE_LIB_PATH})
@@ -449,7 +549,11 @@ if( NOT ${BUILD_SHARED_LIBS} )
   ## Add external link requirements.
   target_link_libraries ( ${CORE_RUNTIME_NAME} INTERFACE hsakmt-staticdrm::hsakmt-staticdrm )
   if (UNIX)
-    target_link_libraries ( ${CORE_RUNTIME_NAME} INTERFACE elf::elf dl pthread rt )
+    if(NUMA_FOUND)
+      target_link_libraries ( ${CORE_RUNTIME_NAME} INTERFACE elf::elf dl pthread rt numa )
+    else()
+      target_link_libraries ( ${CORE_RUNTIME_NAME} INTERFACE elf::elf dl pthread rt )
+    endif()
   else()
     target_link_libraries ( ${CORE_RUNTIME_NAME} INTERFACE Ws2_32 )
   endif()

projects/rocr-runtime/runtime/hsa-runtime/core/util/lnx/os_linux.cpp

@@ -65,6 +65,17 @@
 #include <sys/mman.h>
 #include <sys/socket.h>
 #include <sys/un.h>
+// HSA_HAS_NUMA is defined by the build system only when it also links libnuma
+// (see NUMA_FOUND in hsa-runtime/CMakeLists.txt). Gating on it rather than on
+// __has_include(<numa.h>) keeps the compiled numa_* calls and the linked
+// library in lockstep, avoiding unresolved numa_* symbols when the header is
+// visible but the library is not linked.
+#if defined(HSA_HAS_NUMA) && HSA_HAS_NUMA
+#include <numa.h>
+#define HAS_NUMA 1
+#else
+#define HAS_NUMA 0
+#endif
 #if defined(__i386__) || defined(__x86_64__)
 #include <cpuid.h>
 #endif
@@ -91,6 +102,128 @@ void* __stdcall ThreadTrampoline(void* arg) {
   return nullptr;
 }
 
+#if HAS_NUMA
+// Helper function to get NUMA node mask for allowed CPUs.
+// Returns a bitmask where each bit represents whether the corresponding NUMA node
+// is allowed based on the current process's CPU affinity.
+// Returns 0 if all NUMA nodes are allowed (unrestricted) or on error.
+// NOTE: This function supports up to 64 NUMA nodes. Systems with more nodes
+// will fall back to allowing all CPUs.
+static unsigned long GetAllowedNumaNodeMask() {
+  if (numa_available() == -1) {
+    return 0;  // NUMA unavailable
+  }
+
+  struct bitmask* node_mask = numa_get_run_node_mask();
+  if (!node_mask) {
+    return 0;  // Error - allow all CPUs
+  }
+
+  int max_node = numa_max_node();
+
+  // Safety check: if system has more than 64 NUMA nodes, we can't represent
+  // them in an unsigned long. Fall back to unrestricted mode.
+  constexpr int max_supported_nodes = static_cast<int>(sizeof(unsigned long) * 8);
+  if (max_node >= max_supported_nodes) {
+    fprintf(stderr,
+            "Warning: System has %d NUMA nodes, but only %d are supported. "
+            "Falling back to unrestricted CPU affinity for helper threads.\n",
+            max_node + 1, max_supported_nodes);
+    numa_bitmask_free(node_mask);
+    return 0;  // Unrestricted fallback for large HPC systems
+  }
+
+  // Track two separate pieces of information:
+  // 1. Whether ALL nodes are allowed (for unrestricted check)
+  // 2. The bitmask of allowed nodes (for restricted case)
+  bool all_nodes_allowed = true;
+  unsigned long allowed_node_bitmask = 0;
+
+  for (int i = 0; i <= max_node; i++) {
+    if (numa_bitmask_isbitset(node_mask, i)) {
+      allowed_node_bitmask |= (1UL << i);
+    } else {
+      all_nodes_allowed = false;
+    }
+  }
+
+  numa_bitmask_free(node_mask);
+
+  // Unrestricted - signal "allow all NUMA nodes" with 0.
+  if (all_nodes_allowed) {
+    return 0;
+  }
+  return allowed_node_bitmask;
+}
+
+// Fill cpuset with CPUs from the allowed NUMA nodes.
+// If allowed_node_mask is 0, allows all CPUs.
+// Otherwise, only sets CPUs from NUMA nodes indicated by the bitmask.
+static void FillCpuSetFromNumaNodes(cpu_set_t* cpuset, int cores, unsigned long allowed_node_mask) {
+  CPU_ZERO_S(CPU_ALLOC_SIZE(cores), cpuset);
+
+  if (allowed_node_mask == 0) {
+    // Unrestricted or error - allow all CPU cores
+    for (int i = 0; i < cores; i++) {
+      CPU_SET_S(i, CPU_ALLOC_SIZE(cores), cpuset);
+    }
+    return;
+  }
+
+  // Restrict to allowed NUMA nodes
+  int max_node = numa_max_node();
+  int num_cpus = numa_num_possible_cpus();
+
+  struct bitmask* node_cpus = numa_allocate_cpumask();
+
+  if (node_cpus == nullptr) {
+    // Allocation failed - allow all CPU cores as fallback
+    for (int i = 0; i < cores; i++) {
+      CPU_SET_S(i, CPU_ALLOC_SIZE(cores), cpuset);
+    }
+    return;
+  }
+
+  bool added_cpu = false;
+  for (int node = 0; node <= max_node; node++) {
+    // Safety check: ensure we don't shift beyond the width of unsigned long
+    if (node >= static_cast<int>(sizeof(unsigned long) * 8)) {
+      break;
+    }
+
+    // Check if this NUMA node is allowed
+    if ((allowed_node_mask & (1UL << node)) == 0) {
+      continue;
+    }
+
+    // Get CPUs for this NUMA node
+    if (numa_node_to_cpus(node, node_cpus) != 0) {
+      fprintf(stderr, "numa_node_to_cpus failed for node %d\n", node);
+      continue;
+    }
+
+    // Add all CPUs from this node to the cpuset
+    for (int cpu = 0; cpu < num_cpus && cpu < cores; cpu++) {
+      if (numa_bitmask_isbitset(node_cpus, cpu)) {
+        CPU_SET_S(cpu, CPU_ALLOC_SIZE(cores), cpuset);
+        added_cpu = true;
+      }
+    }
+  }
+
+  numa_free_cpumask(node_cpus);
+
+  // An empty affinity mask makes pthread_setaffinity_np fail and the helper
+  // thread is never created. Preserve the previous all-CPU behavior instead.
+  if (!added_cpu) {
+    for (int i = 0; i < cores; i++) {
+      CPU_SET_S(i, CPU_ALLOC_SIZE(cores), cpuset);
+    }
+  }
+}
+#endif  // HAS_NUMA
+
+
 // Thread container allows multiple waits and separate close (destroy).
 class os_thread {
  public:
@@ -138,10 +271,18 @@ class os_thread {
         fprintf(stderr, "CPU_ALLOC failed: %s\n", strerror(errno));
         return;
       }
-      CPU_ZERO_S(CPU_ALLOC_SIZE(cores), cpuset);
+
+#if HAS_NUMA
+      // Get the NUMA node mask based on main thread's CPU affinity
+      unsigned long allowed_node_mask = GetAllowedNumaNodeMask();
+      FillCpuSetFromNumaNodes(cpuset, cores, allowed_node_mask);
+#else
+      // NUMA not available - allow all CPU cores
       for (int i = 0; i < cores; i++) {
         CPU_SET_S(i, CPU_ALLOC_SIZE(cores), cpuset);
       }
+#endif
+
 #ifdef HAVE_PTHREAD_ATTR_SETAFFINITY_NP
       err = pthread_attr_setaffinity_np(&attrib, CPU_ALLOC_SIZE(cores), cpuset);
       CPU_FREE(cpuset);
@@ -173,6 +314,7 @@ class os_thread {
     } while (stackSize < 20 * 1024 * 1024);
 
 #ifndef HAVE_PTHREAD_ATTR_SETAFFINITY_NP
+    // Reuse the already-filled cpuset from above (no need to reallocate)
     if (cores && cpuset) {
       err = pthread_setaffinity_np(thread, CPU_ALLOC_SIZE(cores), cpuset);
       CPU_FREE(cpuset);