#ifndef _GNU_SOURCE
#define _GNU_SOURCE 1
#endif

#include <stdio.h>
//#include <sched.h>
//#include <hwloc.h>
#include <stdlib.h>
#include <pthread.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/syscall.h>

#define NR_CPUS 8
#define NR_THREADS 5
#define USE_THREADAFFINITY 1

pid_t tid[NR_THREADS];

pthread_barrier_t init_start, barrier2;

void print_affinity(pthread_t pid, int threadno) {
	int i;
	cpu_set_t set;
	
	CPU_ZERO(&set);
	
	sched_getaffinity( pid, sizeof(set), &set);  

	
	for( i = 0; i < NR_CPUS; i++) {
		if (CPU_ISSET(i, &set))
			printf("Cpu #%d is in %d thread's affinity\n", i, threadno);
	}
}

void set_schedaffinity(pthread_t pid, int cpu) {
	cpu_set_t set;
	
	CPU_ZERO(&set);
	CPU_SET(cpu, &set);
	
	sched_setaffinity(pid, sizeof(set), &set);
}

void set_threadaffinity(pthread_t pid, int cpu) {
	cpu_set_t set;
	
	CPU_ZERO(&set);
	CPU_SET(cpu, &set);
	
	pthread_setaffinity_np(pid, sizeof(set), &set);
}

void * threadCode(void * threadNo) {
	long int thread = (long int) threadNo;
	
	if (!USE_THREADAFFINITY) {
		//Set current thread id
		tid[thread] = (pid_t) syscall (SYS_gettid);
	
		//Wait for all the threads to set their thread id
		pthread_barrier_wait(&init_start);
	}
	
	//Wait for main to set the affinity
	pthread_barrier_wait(&barrier2);
	//sleep(1);
	
	//Print affinity
	print_affinity(0,thread);
	
	//Do something so that we can see cpu being used
	while (1);

	pthread_exit(0);
}



int main(int argc , char *argv[]) {
	long int i;
	pthread_t pid[NR_THREADS];
	
	// Initialize barriers
	if(pthread_barrier_init(&init_start, NULL, NR_THREADS+1)) {
		printf("Could not create a barrier\n");
		exit(1);
	}
	if(pthread_barrier_init(&barrier2, NULL, NR_THREADS+1)) {
		printf("Could not create a barrier\n");
		exit(1);
	}
	
	//Create threads
	for (i = 0; i < NR_THREADS; i++)
		pthread_create(&pid[i], NULL, threadCode, (void *) i);
	
	//Wait for threads to set their thread id
	if (!USE_THREADAFFINITY)
		pthread_barrier_wait(&init_start);
	
	//Set affinity
	for (i = 0; i < NR_THREADS; i++) {
		if (USE_THREADAFFINITY)
			set_threadaffinity(pid[i],0);
		else
			set_schedaffinity(tid[i],0);
	}
	
	//Tell the threads to continue after setting affinity
	pthread_barrier_wait(&barrier2);
	
	getchar(); 
	
    return 0;
}