Add thread affinity stuff

4 files changed, 120 insertions, 0 deletions

diff --git a/config.h.in b/config.h.in
index 1393c5a6..1dd5de50 100644
--- a/config.h.in
+++ b/config.h.in
@@ -21,6 +21,9 @@
 /* Define to 1 if CL/cl.h should be used */
 #undef HAVE_CL_CL_H
 
+/* Define to 1 if sched_setaffinity is available. */
+#undef HAVE_CPU_AFFINITY
+
 /* Define to 1 if you have the declaration of `strndup', and to 0 if you
    don't. */
 #undef HAVE_DECL_STRNDUP
diff --git a/configure b/configure
index 6ffb56bc..edd953c3 100755
--- a/configure
+++ b/configure
@@ -7413,6 +7413,52 @@ $as_echo "no" >&6; }
 fi
 
 
+{ $as_echo "$as_me:${as_lineno-$LINENO}: checking for sched_setaffinity in -lpthread" >&5
+$as_echo_n "checking for sched_setaffinity in -lpthread... " >&6; }
+if test "${ac_cv_lib_pthread_sched_setaffinity+set}" = set; then :
+  $as_echo_n "(cached) " >&6
+else
+  ac_check_lib_save_LIBS=$LIBS
+LIBS="-lpthread  $LIBS"
+cat confdefs.h - <<_ACEOF >conftest.$ac_ext
+/* end confdefs.h.  */
+
+/* Override any GCC internal prototype to avoid an error.
+   Use char because int might match the return type of a GCC
+   builtin and then its argument prototype would still apply.  */
+#ifdef __cplusplus
+extern "C"
+#endif
+char sched_setaffinity ();
+int
+main ()
+{
+return sched_setaffinity ();
+  ;
+  return 0;
+}
+_ACEOF
+if ac_fn_c_try_link "$LINENO"; then :
+  ac_cv_lib_pthread_sched_setaffinity=yes
+else
+  ac_cv_lib_pthread_sched_setaffinity=no
+fi
+rm -f core conftest.err conftest.$ac_objext \
+    conftest$ac_exeext conftest.$ac_ext
+LIBS=$ac_check_lib_save_LIBS
+fi
+{ $as_echo "$as_me:${as_lineno-$LINENO}: result: $ac_cv_lib_pthread_sched_setaffinity" >&5
+$as_echo "$ac_cv_lib_pthread_sched_setaffinity" >&6; }
+if test "x$ac_cv_lib_pthread_sched_setaffinity" = x""yes; then :
+
+
+$as_echo "#define HAVE_CPU_AFFINITY 1" >>confdefs.h
+
+
+fi
+
+
+
  if test x$have_opencl = xtrue; then
   HAVE_OPENCL_TRUE=
   HAVE_OPENCL_FALSE='#'
diff --git a/configure.ac b/configure.ac
index e89dd096..f7cbe188 100644
--- a/configure.ac
+++ b/configure.ac
@@ -213,6 +213,12 @@ AS_IF([test "x$enable_gdk_pixbuf" != "xno"],
 ])
 
 
+AC_CHECK_LIB([pthread], [sched_setaffinity], [
+   AC_DEFINE([HAVE_CPU_AFFINITY], [1],
+             [Define to 1 if sched_setaffinity is available.])
+])
+
+
 dnl Conditionals...
 AM_CONDITIONAL([HAVE_OPENCL], test x$have_opencl = xtrue)
 
diff --git a/src/thread-pool.c b/src/thread-pool.c
index 070e4fa0..d672062c 100644
--- a/src/thread-pool.c
+++ b/src/thread-pool.c
@@ -23,9 +23,63 @@
 #include <pthread.h>
 #include <assert.h>
 
+#ifdef HAVE_CPU_AFFINITY
+#include <sched.h>
+#endif
+
+
 #include "utils.h"
 
 
+/* ------------------------------ CPU affinity ------------------------------ */
+
+#ifdef HAVE_CPU_AFFINITY
+
+static int next_cpu(int cur)
+{
+	cur++;
+
+	if ( cur == 73 ) cur = 0;
+
+	return cur;
+}
+
+
+static void set_affinity(int cpu)
+{
+	cpu_set_t c;
+
+	CPU_ZERO(&c);
+	CPU_SET(cpu, &c);
+	if ( sched_setaffinity(0, sizeof(cpu_set_t), &c) ) {
+
+		/* Cannot use ERROR() just yet */
+		fprintf(stderr, "Failed to set CPU affinity.\n");
+
+	} else {
+
+		fprintf(stderr, "Successfully set CPU affinity to %i\n", cpu);
+
+	}
+}
+
+#else /* HAVE_CPU_AFFINITY */
+
+static int next_cpu(int cur)
+{
+	return 0;
+}
+
+
+static void set_affinity(int cpu)
+{
+}
+
+#endif /* HAVE_CPU_AFFINITY */
+
+
+/* --------------------------- Status label stuff --------------------------- */
+
 static int use_status_labels = 0;
 static pthread_key_t status_label_key;
 pthread_mutex_t stderr_lock = PTHREAD_MUTEX_INITIALIZER;
@@ -35,6 +89,7 @@ struct worker_args
 	struct task_queue_range *tqr;
 	struct task_queue *tq;
 	int id;
+	int cpu;
 };
 
 
@@ -81,6 +136,8 @@ static void *range_worker(void *pargsv)
 	struct task_queue_range *q = w->tqr;
 	int *cookie;
 
+	set_affinity(w->cpu);
+
 	cookie = malloc(sizeof(int));
 	*cookie = w->id;
 	pthread_setspecific(status_label_key, cookie);
@@ -132,6 +189,7 @@ void run_thread_range(int n_tasks, int n_threads, const char *text,
 {
 	pthread_t *workers;
 	int i;
+	int cpu = 0;
 	struct task_queue_range q;
 
 	/* The nation of CrystFEL prides itself on having 0% unemployment. */
@@ -166,6 +224,8 @@ void run_thread_range(int n_tasks, int n_threads, const char *text,
 		w->tqr = &q;
 		w->tq = NULL;
 		w->id = i;
+		w->cpu = cpu;
+		cpu = next_cpu(cpu);
 
 		if ( pthread_create(&workers[i], NULL, range_worker, w) ) {
 			/* Not ERROR() here */
@@ -211,6 +271,8 @@ static void *task_worker(void *pargsv)
 	struct task_queue *q = w->tq;
 	int *cookie;
 
+	set_affinity(w->cpu);
+
 	cookie = malloc(sizeof(int));
 	*cookie = w->id;
 	pthread_setspecific(status_label_key, cookie);
@@ -265,6 +327,7 @@ int run_threads(int n_threads, void (*work)(void *, int),
 	pthread_t *workers;
 	int i;
 	struct task_queue q;
+	int cpu = 0;
 
 	pthread_key_create(&status_label_key, NULL);
 
@@ -292,6 +355,8 @@ int run_threads(int n_threads, void (*work)(void *, int),
 		w->tq = &q;
 		w->tqr = NULL;
 		w->id = i;
+		w->cpu = cpu;
+		cpu = next_cpu(cpu);
 
 		if ( pthread_create(&workers[i], NULL, task_worker, w) ) {
 			/* Not ERROR() here */