SWDEV-565304 - Pass numa node to migrate pages correctly (#1729)

* SWDEV-565304 - Pass cpuId of the the thread currently running

* SWDEV-565304 - Numa id to be returned

* SWDEV-565304 - Numa id to be returned

projects/clr/hipamd/src/hip_hmm.cpp

@@ -24,6 +24,7 @@
 #include "platform/context.hpp"
 #include "platform/command.hpp"
 #include "platform/memory.hpp"
+#include "os/os.hpp"
 
 namespace hip {
 
@@ -310,9 +311,19 @@ hipError_t ihipMemPrefetchAsync(const void* dev_ptr, size_t count, hipMemLocatio
   const bool cpuAccess = isHost || isHostNuma || isHostCurrent;
 
   // Determine the target device index:
-  //  - for host-prefetch and host-current, always use device 0
+  //  - for host-prefetch, use default CPU agent
+  //  - for host-current, query the current thread's NUMA node ID
   //  - for host-NUMA or device-prefetch, use the provided id
-  int targetDevice = (isHost || isHostCurrent) ? hipCpuDeviceId : location.id;
+  int targetDevice;
+  if (isHost) {
+    targetDevice = hipCpuDeviceId;
+  } else if (isHostCurrent) {
+    uint32_t numa_node = amd::numa::getCurrentNumaNode();
+    targetDevice =
+        (numa_node == static_cast<uint32_t>(-1)) ? hipCpuDeviceId : static_cast<int>(numa_node);
+  } else {
+    targetDevice = location.id;
+  }
 
   amd::Device* dev = nullptr;
   if (cpuAccess == false) {
@@ -378,10 +389,16 @@ hipError_t ihipMemAdvise(const void* dev_ptr, size_t count, hipMemoryAdvise advi
       use_cpu = true;
       break;
     case hipMemLocationTypeHost:
-    case hipMemLocationTypeHostNumaCurrent:
       targetDevice = hipCpuDeviceId;
       use_cpu = true;
       break;
+    case hipMemLocationTypeHostNumaCurrent: {
+      uint32_t numa_node = amd::numa::getCurrentNumaNode();
+      targetDevice =
+          (numa_node == static_cast<uint32_t>(-1)) ? hipCpuDeviceId : static_cast<int>(numa_node);
+      use_cpu = true;
+      break;
+    }
     default:
       return hipErrorInvalidValue;
   }

projects/clr/rocclr/os/os.hpp

@@ -299,6 +299,9 @@ namespace numa {
 
 static constexpr uint32_t kBitsPerUInt64 = 8 * sizeof(uint64_t);
 
+//! Get the NUMA node ID of the current thread
+uint32_t getCurrentNumaNode();
+
 /*! \brief Manage Numa policy.
  *
  *  \note Works in Linux only, dummy in Windows.

projects/clr/rocclr/os/os_posix.cpp

@@ -954,6 +954,15 @@ void Os::CxaDemangle(const std::string& name, std::string* result) {
 
 namespace numa {
 
+// ================================================================================================
+uint32_t getCurrentNumaNode() {
+  unsigned cpu, node;
+  if (syscall(__NR_getcpu, &cpu, &node, nullptr) < 0) {
+    return static_cast<uint32_t>(-1);
+  }
+  return static_cast<uint32_t>(node);
+}
+
 // ================================================================================================
 NumaPolicy::NumaPolicy(const uint32_t numa_node_count) :
   node_map_((numa_node_count + kBitsPerUInt64 - 1) / kBitsPerUInt64, 0) { }

projects/clr/rocclr/os/os_win32.cpp

@@ -704,6 +704,19 @@ void Os::CxaDemangle(const std::string& name, std::string* result) { *result = n
 
 namespace numa {
 
+// ================================================================================================
+uint32_t getCurrentNumaNode() {
+  PROCESSOR_NUMBER procNumber{};
+  GetCurrentProcessorNumberEx(&procNumber);
+
+  USHORT numa_node = static_cast<USHORT>(-1);
+  if (!GetNumaProcessorNodeEx(&procNumber, &numa_node)) {
+    return static_cast<uint32_t>(-1);
+  }
+
+  return static_cast<uint32_t>(numa_node);
+}
+
 // ================================================================================================
 NumaPolicy::NumaPolicy(const uint32_t numa_node_count) {
 }