From 7b4773b24d83ef81de6291da3d350ab98045d4a0 Mon Sep 17 00:00:00 2001
From: Sharvil Nanavati <sharvil.nanavati@gmail.com>
Date: Wed, 8 Apr 2020 12:47:41 -0700
Subject: Add API to set thread affinity on Linux.

Issue: #2545
---
 cblas.h                     |  5 +++++
 driver/others/blas_server.c | 18 ++++++++++++++++++
 openblas_config_template.h  |  5 +++++
 3 files changed, 28 insertions(+)

diff --git a/cblas.h b/cblas.h
index 1a87074d6..4bc5588d8 100644
--- a/cblas.h
+++ b/cblas.h
@@ -25,6 +25,11 @@ char* openblas_get_config(void);
 /*Get the CPU corename on runtime.*/
 char* openblas_get_corename(void);
 
+#ifdef OPENBLAS_OS_LINUX
+/* Sets thread affinity for OpenBLAS threads. `thread_idx` is in [0, openblas_get_num_threads()-1]. */
+int openblas_setaffinity(int thread_idx, size_t cpusetsize, cpu_set_t* cpu_set);
+#endif
+
 /* Get the parallelization type which is used by OpenBLAS */
 int openblas_get_parallel(void);
 /* OpenBLAS is compiled for sequential use  */
diff --git a/driver/others/blas_server.c b/driver/others/blas_server.c
index aa0644845..f13b83dd4 100644
--- a/driver/others/blas_server.c
+++ b/driver/others/blas_server.c
@@ -72,6 +72,7 @@ USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include "common.h"
 #if defined(OS_LINUX) || defined(OS_NETBSD) || defined(OS_DARWIN) || defined(OS_ANDROID) || defined(OS_SUNOS) || defined(OS_FREEBSD) || defined(OS_OPENBSD) || defined(OS_DRAGONFLY) || defined(OS_HAIKU)
 #include <dlfcn.h>
+#include <errno.h>
 #include <signal.h>
 #include <sys/resource.h>
 #include <sys/time.h>
@@ -279,6 +280,23 @@ int get_node(void);
 
 static int increased_threads = 0;
 
+#ifdef OS_LINUX
+int openblas_setaffinity(int thread_idx, size_t cpusetsize, cpu_set_t* cpu_set) {
+  const int active_threads = openblas_get_num_threads();
+
+  if (thread_idx < 0 || thread_idx >= active_threads) {
+    errno = EINVAL;
+    return -1;
+  }
+
+  pthread_t thread = (thread_idx == active_threads - 1)
+      ? pthread_self()
+      : blas_threads[thread_idx];
+
+  return pthread_setaffinity_np(thread, cpusetsize, cpu_set);
+}
+#endif
+
 static void* blas_thread_server(void *arg){
 
   /* Thread identifier */
diff --git a/openblas_config_template.h b/openblas_config_template.h
index 52dd49da2..49aea1cab 100644
--- a/openblas_config_template.h
+++ b/openblas_config_template.h
@@ -91,3 +91,8 @@ typedef int blasint;
   #define openblas_complex_xdouble_real(z)           ((z).real)
   #define openblas_complex_xdouble_imag(z)           ((z).imag)
 #endif
+
+/* Inclusion of Linux-specific header is needed for definition of cpu_set_t. */
+#ifdef OPENBLAS_OS_LINUX
+#include <sched.h>
+#endif
-- 
cgit v1.2.3