Processor affinity: binding threads to cores

affinity.c

#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <pthread.h>
#include <sched.h>

#define MAX_NUM_THREADS 16 

static void set_thread_affinity(pthread_t t, int cpu)
{
	cpu_set_t cpuset;

	CPU_ZERO(&cpuset);
	CPU_SET(cpu, &cpuset);
	pthread_setaffinity_np(t, sizeof(cpu_set_t), &cpuset);
}

static void print_thread_affinity(pthread_t t, int ID)
{
	cpu_set_t cpuset;
	int i;

	pthread_getaffinity_np(t, sizeof(cpu_set_t), &cpuset);

	for (i = 0; i < CPU_SETSIZE; i++) {
		if (CPU_ISSET(i, &cpuset))
			printf("Thread %2d bound to CPU %d\n", ID, i);
	}
}

static void *thread_func(void *args)
{
	int ID = *(int *)args;

	print_thread_affinity(pthread_self(), ID);

	return NULL;
}

int main(void)
{
	int num_threads, num_cpus;
	int IDs[MAX_NUM_THREADS], i;
	pthread_t threads[MAX_NUM_THREADS];
	cpu_set_t cpuset;

	char *value = getenv("NUM_THREADS");
	if (value) {
		num_threads = abs(atoi(value));
		if (num_threads > MAX_NUM_THREADS)
			num_threads = MAX_NUM_THREADS;
		printf("Creating %d threads\n", num_threads);
	} else { 
		// May not be standard
		num_threads = sysconf(_SC_NPROCESSORS_ONLN);
		printf("Creating %d threads (default)\n", num_threads);
	}

	IDs[0] = 0;
	pthread_getaffinity_np(pthread_self(), sizeof(cpu_set_t), &cpuset);
	num_cpus = CPU_COUNT(&cpuset);
	printf("Number of CPUs: %d\n", num_cpus);

	set_thread_affinity(pthread_self(), 0);

	for (i = 1; i < num_threads; i++) {
		IDs[i] = i;
		pthread_create(&threads[i], NULL, thread_func, &IDs[i]);
		set_thread_affinity(threads[i], i % num_cpus);
	}

	thread_func(&IDs[0]);

	for (i = 1; i < num_threads; i++) {
		pthread_join(threads[i], NULL);
	}

	return 0;
}