uhd: Add thread affinity utility functions

3 files changed, 123 insertions, 11 deletions

diff --git a/host/include/uhd/utils/thread.hpp b/host/include/uhd/utils/thread.hpp
index 05cba0d14..aeb27e467 100644
--- a/host/include/uhd/utils/thread.hpp
+++ b/host/include/uhd/utils/thread.hpp
@@ -56,6 +56,12 @@ UHD_API void set_thread_name(boost::thread* thread, const std::string& name);
  */
 UHD_API void set_thread_name(std::thread* thread, const std::string& name);
 
+/*!
+ * Set the affinity of the current thread to a (set of) CPU(s).
+ * \param cpu_affinity_list list of CPU numbers to affinitize the thread to
+ */
+UHD_API void set_thread_affinity(const std::vector<size_t>& cpu_affinity_list);
+
 } // namespace uhd
 
 #endif /* INCLUDED_UHD_UTILS_THREAD_HPP */
diff --git a/host/lib/utils/CMakeLists.txt b/host/lib/utils/CMakeLists.txt
index 74b292fbb..5c4478909 100644
--- a/host/lib/utils/CMakeLists.txt
+++ b/host/lib/utils/CMakeLists.txt
@@ -55,6 +55,10 @@ else()
     set(THREAD_PRIO_DEFS HAVE_THREAD_PRIO_DUMMY)
 endif()
 
+########################################################################
+# Setup defines for thread naming
+########################################################################
+
 set(CMAKE_REQUIRED_LIBRARIES "pthread")
 
 CHECK_CXX_SOURCE_COMPILES("
@@ -82,6 +86,45 @@ else()
     list(APPEND THREAD_PRIO_DEFS HAVE_THREAD_SETNAME_DUMMY)
 endif()
 
+########################################################################
+# Setup defines for thread affinitizing
+########################################################################
+
+CHECK_CXX_SOURCE_COMPILES("
+    #include <pthread.h>
+    int main(){
+        pthread_t pt;
+        cpu_set_t cs;
+        pthread_setaffinity_np(pt, sizeof(cpu_set_t), &cs);
+        return 0;
+    }
+    " HAVE_PTHREAD_SETAFFINITY_NP
+)
+
+if(CYGWIN)
+    set(HAVE_PTHREAD_SETAFFINITY_NP False)
+endif(CYGWIN)
+
+CHECK_CXX_SOURCE_COMPILES("
+    #include <windows.h>
+    int main(){
+        SetThreadAffinityMask(GetCurrentThread(), 0);
+        return 0;
+    }
+    " HAVE_WIN_SETTHREADAFFINITYMASK
+)
+if(HAVE_PTHREAD_SETAFFINITY_NP)
+    message(STATUS "  Setting thread affinity is supported through pthread_setaffinity_np.")
+    list(APPEND THREAD_PRIO_DEFS HAVE_PTHREAD_SETAFFINITY_NP)
+    LIBUHD_APPEND_LIBS(pthread)
+elseif(HAVE_WIN_SETTHREADPRIORITY)
+    message(STATUS "  Setting thread affinity is supported through windows SetThreadAffinityMask.")
+    list(APPEND THREAD_PRIO_DEFS HAVE_WIN_SETTHREADAFFINITYMASK)
+else()
+    message(STATUS "  Setting thread affinity is not supported.")
+    list(APPEND THREAD_PRIO_DEFS HAVE_THREAD_SETAFFINITY_DUMMY)
+endif()
+
 
 set_source_files_properties(
     ${CMAKE_CURRENT_SOURCE_DIR}/thread.cpp
diff --git a/host/lib/utils/thread.cpp b/host/lib/utils/thread.cpp
index d8f52c4c8..e5dde06b2 100644
--- a/host/lib/utils/thread.cpp
+++ b/host/lib/utils/thread.cpp
@@ -58,9 +58,37 @@ static void check_priority_range(float priority){
     }
 #endif /* HAVE_PTHREAD_SETSCHEDPARAM */
 
-/***********************************************************************
- * Windows API to set priority
- **********************************************************************/
+    /***********************************************************************
+     * Pthread API to set affinity
+     **********************************************************************/
+#ifdef HAVE_PTHREAD_SETAFFINITY_NP
+#    include <pthread.h>
+    void uhd::set_thread_affinity(const std::vector<size_t>& cpu_affinity_list)
+    {
+        if (cpu_affinity_list.empty()) {
+            return;
+        }
+
+        cpu_set_t cpu_set;
+        CPU_ZERO(&cpu_set);
+        for (auto cpu_num : cpu_affinity_list) {
+            if (cpu_num > CPU_SETSIZE) {
+                UHD_LOG_WARNING(
+                    "UHD", "CPU index " << cpu_num << " in affinity list out of range");
+            }
+            CPU_SET(cpu_num, &cpu_set);
+        }
+
+        int status = pthread_setaffinity_np(pthread_self(), sizeof(cpu_set_t), &cpu_set);
+        if (status != 0) {
+            UHD_LOG_WARNING("UHD", "Failed to set desired affinity for thread");
+        }
+    }
+#endif /* HAVE_PTHREAD_SETAFFINITYNP */
+
+    /***********************************************************************
+     * Windows API to set priority
+     **********************************************************************/
 #ifdef HAVE_WIN_SETTHREADPRIORITY
     #include <windows.h>
 
@@ -89,20 +117,55 @@ static void check_priority_range(float priority){
     }
 #endif /* HAVE_WIN_SETTHREADPRIORITY */
 
-/***********************************************************************
- * Unimplemented API to set priority
- **********************************************************************/
+    /***********************************************************************
+     * Windows API to set affinity
+     **********************************************************************/
+#ifdef HAVE_WIN_SETTHREADAFFINITYMASK
+#    include <windows.h>
+    void uhd::set_thread_affinity(const std::vector<size_t>& cpu_affinity_list)
+    {
+        if (cpu_affinity_list.empty()) {
+            return;
+        }
+
+        DWORD_PTR cpu_set{0};
+        for (auto cpu_num : cpu_affinity_list) {
+            if (cpu_num > 8 * sizeof(DWORD_PTR)) {
+                UHD_LOG_WARNING(
+                    "UHD", "CPU index " << cpu_num << " in affinity list out of range");
+            }
+            cpu_set |= ((DWORD_PTR)1 << cpu_num);
+        }
+
+        DWORD_PTR status = SetThreadAffinityMask(GetCurrentThread(), cpu_set);
+        if (status == 0) {
+            UHD_LOG_WARNING("UHD", "Failed to set desired affinity for thread");
+        }
+    }
+#endif /* HAVE_WIN_SETTHREADAFFINITYMASK */
+
+    /***********************************************************************
+     * Unimplemented API to set priority
+     **********************************************************************/
 #ifdef HAVE_THREAD_PRIO_DUMMY
     void uhd::set_thread_priority(float, bool){
-        throw uhd::not_implemented_error("set thread priority not implemented");
+        UHD_LOG_DEBUG("UHD", "Setting thread priority is not implemented");
     }
 
 #endif /* HAVE_THREAD_PRIO_DUMMY */
 
-void uhd::set_thread_name(
-    boost::thread *thrd,
-    const std::string &name
-) {
+    /***********************************************************************
+     * Unimplemented API to set affinity
+     **********************************************************************/
+#ifdef HAVE_THREAD_SETAFFINITY_DUMMY
+    void uhd::set_thread_affinity(const std::vector<size_t>& cpu_affinity_list)
+    {
+        UHD_LOG_DEBUG("UHD", "Setting thread affinity is not implemented");
+    }
+#endif /* HAVE_THREAD_SETAFFINITY_DUMMY */
+
+    void uhd::set_thread_name(boost::thread* thrd, const std::string& name)
+    {
 #ifdef HAVE_PTHREAD_SETNAME
     pthread_setname_np(thrd->native_handle(), name.substr(0,16).c_str());
 #endif /* HAVE_PTHREAD_SETNAME */