Set thread affinity with +RTS -qa (only on Linux so far)

7 files changed, 47 insertions, 5 deletions

diff --git a/configure.ac b/configure.ac
index 85e0617aa0..40673354ea 100644
--- a/configure.ac
+++ b/configure.ac
@@ -1029,7 +1029,7 @@ dnl    off_t, because it will affect the result of that test.
 AC_SYS_LARGEFILE
 
 dnl ** check for specific header (.h) files that we are interested in
-AC_CHECK_HEADERS([bfd.h ctype.h dirent.h dlfcn.h errno.h fcntl.h grp.h limits.h locale.h nlist.h pthread.h pwd.h signal.h sys/mman.h sys/resource.h sys/time.h sys/timeb.h sys/timers.h sys/times.h sys/utsname.h sys/wait.h termios.h time.h utime.h windows.h winsock.h])
+AC_CHECK_HEADERS([bfd.h ctype.h dirent.h dlfcn.h errno.h fcntl.h grp.h limits.h locale.h nlist.h pthread.h pwd.h signal.h sys/mman.h sys/resource.h sys/time.h sys/timeb.h sys/timers.h sys/times.h sys/utsname.h sys/wait.h termios.h time.h utime.h windows.h winsock.h sched.h])
 
 dnl ** check if it is safe to include both <time.h> and <sys/time.h>
 AC_HEADER_TIME
@@ -1107,7 +1107,7 @@ fi
 
 dnl ** check for more functions
 dnl ** The following have been verified to be used in ghc/, but might be used somewhere else, too.
-AC_CHECK_FUNCS([getclock getrusage gettimeofday setitimer siginterrupt sysconf times ctime_r])
+AC_CHECK_FUNCS([getclock getrusage gettimeofday setitimer siginterrupt sysconf times ctime_r, sched_setaffinity])
 
 AC_TRY_RUN([
 #include <sys/types.h>
diff --git a/includes/OSThreads.h b/includes/OSThreads.h
index 73ef4d2b13..f5c434fc28 100644
--- a/includes/OSThreads.h
+++ b/includes/OSThreads.h
@@ -186,8 +186,8 @@ void  setThreadLocalVar (ThreadLocalKey *key, void *value);
 void  freeThreadLocalKey (ThreadLocalKey *key);
 
 // Processors and affinity
-nat getNumberOfProcessors (void);
-
+nat  getNumberOfProcessors (void);
+void setThreadAffinity     (nat n, nat m);
 #endif // !CMINUSMINUS
 
 #else
diff --git a/includes/RtsFlags.h b/includes/RtsFlags.h
index 77097872a2..348a8f5225 100644
--- a/includes/RtsFlags.h
+++ b/includes/RtsFlags.h
@@ -192,6 +192,7 @@ struct PAR_FLAGS {
   rtsBool        parGcGen;       /* do parallel GC in this generation
                                   * and higher only */
   rtsBool        parGcLoadBalancing; /* do load-balancing in parallel GC */
+  rtsBool        setAffinity;    /* force thread affinity with CPUs */
 };
 #endif /* THREADED_RTS */
 
diff --git a/rts/RtsFlags.c b/rts/RtsFlags.c
index 3fac86bf4e..6a8c692bd9 100644
--- a/rts/RtsFlags.c
+++ b/rts/RtsFlags.c
@@ -233,6 +233,7 @@ void initRtsFlagsDefaults(void)
     RtsFlags.ParFlags.parGcEnabled      = 1;
     RtsFlags.ParFlags.parGcGen          = 1;
     RtsFlags.ParFlags.parGcLoadBalancing = 1;
+    RtsFlags.ParFlags.setAffinity       = 0;
 #endif
 
 #ifdef PAR
@@ -476,6 +477,8 @@ usage_text[] = {
 "  -N        Determine the number of processors to use automatically",
 "  -q1       Use one OS thread for GC (turns off parallel GC)",
 "  -qg<n>    Use parallel GC only for generations >= <n> (default: 1)",
+"  -qb       Disable load-balancing in the parallel GC",
+"  -qa       Use the OS to set thread affinity",
 "  -qm       Don't automatically migrate threads between CPUs",
 "  -qw       Migrate a thread to the current CPU when it is woken up",
 #endif
@@ -1233,6 +1236,9 @@ error = rtsTrue;
 		    case 'b':
 			RtsFlags.ParFlags.parGcLoadBalancing = rtsFalse;
 			break;
+		    case 'a':
+			RtsFlags.ParFlags.setAffinity = rtsTrue;
+			break;
 		    case 'm':
 			RtsFlags.ParFlags.migrate = rtsFalse;
 			break;
diff --git a/rts/Schedule.c b/rts/Schedule.c
index 666b59e3e3..636b5179ad 100644
--- a/rts/Schedule.c
+++ b/rts/Schedule.c
@@ -2125,6 +2125,10 @@ workerStart(Task *task)
     cap = task->cap;
     RELEASE_LOCK(&task->lock);
 
+    if (RtsFlags.ParFlags.setAffinity) {
+        setThreadAffinity(cap->no, n_capabilities);
+    }
+
     // set the thread-local pointer to the Task:
     taskEnter(task);
 
diff --git a/rts/posix/OSThreads.c b/rts/posix/OSThreads.c
index 6eb2d2bc8e..f15fc95c3a 100644
--- a/rts/posix/OSThreads.c
+++ b/rts/posix/OSThreads.c
@@ -7,8 +7,9 @@
  *
  * --------------------------------------------------------------------------*/
 
-#if defined(DEBUG) && defined(__linux__)
+#if defined(__linux__)
 /* We want GNU extensions in DEBUG mode for mutex error checking */
+/* We also want the affinity API, which requires _GNU_SOURCE */
 #define _GNU_SOURCE
 #endif
 
@@ -26,6 +27,10 @@
 #error pthreads.h is required for the threaded RTS on Posix platforms
 #endif
 
+#if defined(HAVE_SCHED_H)
+#include <sched.h>
+#endif
+
 /*
  * This (allegedly) OS threads independent layer was initially
  * abstracted away from code that used Pthreads, so the functions
@@ -197,6 +202,26 @@ getNumberOfProcessors (void)
     return nproc;
 }
 
+// Schedules the thread to run on CPU n of m.  m may be less than the
+// number of physical CPUs, in which case, the thread will be allowed
+// to run on CPU n, n+m, n+2m etc.
+void
+setThreadAffinity (nat n, nat m)
+{
+#if defined(HAVE_SCHED_H) && defined(HAVE_SCHED_SETAFFINITY)
+    nat nproc;
+    cpu_set_t cs;
+    nat i;
+
+    nproc = getNumberOfProcessors();
+    CPU_ZERO(&cs);
+    for (i = n; i < nproc; i+=m) {
+        CPU_SET(n, &cs);
+    }
+    sched_setaffinity(0, sizeof(cpu_set_t), &cs);
+#endif
+}
+
 #else /* !defined(THREADED_RTS) */
 
 int
diff --git a/rts/win32/OSThreads.c b/rts/win32/OSThreads.c
index 929b8171e7..c18bab2d09 100644
--- a/rts/win32/OSThreads.c
+++ b/rts/win32/OSThreads.c
@@ -246,6 +246,12 @@ getNumberOfProcessors (void)
     return nproc;
 }
 
+void
+setThreadAffinity (nat n STG_UNUSED, nat m STG_UNUSED)
+{
+    /* ToDo */
+}
+
 #else /* !defined(THREADED_RTS) */
 
 int