context_switch_latency: Add CPU control

1 files changed, 29 insertions, 2 deletions

diff --git a/context_switch_latency.c b/context_switch_latency.c
index db82805..fe39a7a 100644
--- a/context_switch_latency.c
+++ b/context_switch_latency.c
@@ -1,8 +1,11 @@
+#define _GNU_SOURCE
+
 #include <time.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
 #include <unistd.h>
+#include <sched.h>
 #include <pthread.h>
 #include <sys/eventfd.h>
 
@@ -10,6 +13,7 @@ int to_thread, from_thread;
 volatile int state;  // 0 = warmup.  1 = benchmark.  2 = exit.
 
 int use_xstate = 0;
+int thread_cpu;
 
 void maybe_use_xstate(void)
 {
@@ -19,6 +23,12 @@ void maybe_use_xstate(void)
 
 void *threadproc(void *x)
 {
+  cpu_set_t cpuset;
+  CPU_ZERO(&cpuset);
+  CPU_SET(thread_cpu, &cpuset);
+  if (sched_setaffinity(0, sizeof(cpuset), &cpuset) != 0)
+    printf("[WARN]\tsched_setaffinity failed\n");
+
   while(1)
     {
       uint64_t buf;
@@ -62,9 +72,10 @@ int main(int argc, char **argv)
   pthread_t thread;
   uint64_t i, iters = 100000;
   uint64_t ns;
+  int main_cpu;
 
-  if (argc != 2) {
-    printf("Usage: %s <0|1>\n\nSet the parameter to 1 to use xstate\n",
+  if (argc != 3) {
+    printf("Usage: %s <0|1> <same|different>\n\nSet the parameter to 1 to use xstate\nUse 'same' for same CPU or 'different' for cross-CPU\n",
 	   argv[0]);
     return 1;
   } else {
@@ -72,6 +83,16 @@ int main(int argc, char **argv)
       use_xstate = 1;
     else if (strcmp(argv[1], "0"))
       abort();
+
+    if (!strcmp(argv[2], "same")) {
+      main_cpu = 0;
+      thread_cpu = 0;
+    } else if (!strcmp(argv[2], "different")) {
+      main_cpu = 0;
+      thread_cpu = 1;
+    } else {
+      abort();
+    }
   }
 
   printf("use_xstate = %d\n", use_xstate);
@@ -86,6 +107,12 @@ int main(int argc, char **argv)
   if (pthread_create(&thread, 0, threadproc, 0) != 0)
     abort();
 
+  cpu_set_t cpuset;
+  CPU_ZERO(&cpuset);
+  CPU_SET(main_cpu, &cpuset);
+  if (sched_setaffinity(0, sizeof(cpuset), &cpuset) != 0)
+    printf("[WARN]\tsched_setaffinity failed\n");
+
   /* Warm up (and burn the xstate heuristic) */
   for(i = 0; i < 10000; i++) {
     bounce();