From cd21af757ed3f2b48718a01240060e2221d05067 Mon Sep 17 00:00:00 2001
From: Saleel Kudchadker <Saleel.Kudchadker@amd.com>
Date: Fri, 2 Jul 2021 09:05:33 -0700
Subject: [PATCH] SWDEV-260448 - Honor NUMACTL for Direct Dispatch

Setting AMD_CPU_AFFINITY=1 will keep Async Handler thread within the
bounds set by numactl.

Change-Id: Id01b30df5127d65c29ac072bf74a04986b7128de
---
 rocclr/device/rocm/rocvirtual.cpp |  3 ++-
 rocclr/os/os.hpp                  |  5 ++++-
 rocclr/os/os_posix.cpp            | 11 ++++++++++-
 rocclr/os/os_win32.cpp            |  3 +++
 4 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/rocclr/device/rocm/rocvirtual.cpp b/rocclr/device/rocm/rocvirtual.cpp
index d31f01982c..2e62b03360 100644
--- a/rocclr/device/rocm/rocvirtual.cpp
+++ b/rocclr/device/rocm/rocvirtual.cpp
@@ -157,7 +157,8 @@ bool HsaAmdSignalHandler(hsa_signal_value_t value, void* arg) {
 
   amd::Thread* thread = amd::Thread::current();
   if (!(thread != nullptr ||
-      ((thread = new amd::HostThread()) != nullptr && thread == amd::Thread::current()))) {
+      ((thread = new amd::HostThread()) != nullptr && thread == amd::Thread::current() &&
+        amd::Os::setThreadAffinityToMainThread()))) {
     return false;
   }
 
diff --git a/rocclr/os/os.hpp b/rocclr/os/os.hpp
index 989363fc04..811bc0691f 100644
--- a/rocclr/os/os.hpp
+++ b/rocclr/os/os.hpp
@@ -121,7 +121,7 @@ class Os : AllStatic {
   // Given a valid file name, returns mmapped memory with the mapped size.
   static bool MemoryMapFile(const char* fname, const void** mmap_ptr, size_t* mmap_size);
 
-  // Given a valid file name amd mapped size, returns ftruncated mmaped memory 
+  // Given a valid file name amd mapped size, returns ftruncated mmaped memory
   static bool MemoryMapFileTruncated(const char* fname, const void** mmap_ptr, size_t mmap_size);
 
   // Given a valid mmaped ptr with correct size, unmaps the ptr from memory
@@ -181,6 +181,9 @@ class Os : AllStatic {
   static void setThreadAffinity(const void* handle, const ThreadAffinityMask& mask);
   //! Set the currently running thread's name.
   static void setCurrentThreadName(const char* name);
+  //! Set current threads affinity to that of main thread
+  static bool setThreadAffinityToMainThread();
+
   //! Check if the thread is alive
   static bool isThreadAlive(const Thread& osThread);
 
diff --git a/rocclr/os/os_posix.cpp b/rocclr/os/os_posix.cpp
index 2d5b89a587..0868dbc152 100644
--- a/rocclr/os/os_posix.cpp
+++ b/rocclr/os/os_posix.cpp
@@ -131,6 +131,7 @@ static pthread_setaffinity_fn pthread_setaffinity_fptr;
 
 static void init() __attribute__((constructor(101)));
 static void init() { Os::init(); }
+static cpu_set_t nativeMask_;
 
 bool Os::installSigfpeHandler() {
   // Install a SIGFPE signal handler @todo: Chain the handlers
@@ -160,6 +161,7 @@ bool Os::init() {
   pageSize_ = (size_t)::sysconf(_SC_PAGESIZE);
   processorCount_ = ::sysconf(_SC_NPROCESSORS_CONF);
 
+  pthread_getaffinity_np(pthread_self(), sizeof(cpu_set_t), &nativeMask_);
   pthread_setaffinity_fptr = (pthread_setaffinity_fn)dlsym(RTLD_NEXT, "pthread_setaffinity_np");
 
   return Thread::init();
@@ -392,13 +394,20 @@ const void* Os::createOsThread(amd::Thread* thread) {
   return reinterpret_cast<const void*>(handle);
 }
 
-
 void Os::setThreadAffinity(const void* handle, const Os::ThreadAffinityMask& mask) {
   if (pthread_setaffinity_fptr != NULL) {
     pthread_setaffinity_fptr((pthread_t)handle, sizeof(cpu_set_t), &mask.mask_);
   }
 }
 
+bool Os::setThreadAffinityToMainThread() {
+  if (AMD_CPU_AFFINITY) {
+    ClPrint(amd::LOG_INFO, amd::LOG_INIT, "Setting Affinity to the main thread's affinity");
+    pthread_setaffinity_np(pthread_self(), sizeof(cpu_set_t), &nativeMask_);
+  }
+  return true;
+}
+
 void Os::yield() { ::sched_yield(); }
 
 uint64_t Os::timeNanos() {
diff --git a/rocclr/os/os_win32.cpp b/rocclr/os/os_win32.cpp
index 327bd0ecf0..6b4f6aac32 100644
--- a/rocclr/os/os_win32.cpp
+++ b/rocclr/os/os_win32.cpp
@@ -341,6 +341,9 @@ void Os::setThreadAffinity(const void* handle, const Os::ThreadAffinityMask& mas
   }
 }
 
+bool Os::setThreadAffinityToMainThread() {
+  return true;
+}
 void Os::yield() { ::SwitchToThread(); }
 
 uint64_t Os::timeNanos() {