#include <pthread.h>
#include <sched.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/time.h>
#include <stdio.h>
#include <stdlib.h>

int blockthemall=1;
static void inline cpupause()
{
#if defined(i386)
 asm volatile("rep;nop":::"memory");
#else
 asm volatile("":::"memory");
#endif
}
/*
 * Determines number of cpus
 * Can be overiden by the NR_CPUS environment variable
 */
int number_of_cpus()
{
	char line[1024], *p;
	int cnt = 0;
	FILE *F;

	p = getenv("NR_CPUS");
	if (p)
		return atoi(p);
	F = fopen("/proc/cpuinfo", "r");
	if (F == NULL) {
		perror("/proc/cpuinfo");
		return 1;
	}
	while (fgets(line, sizeof(line), F) != NULL) {
		if (memcmp(line, "processor", 9) == 0)
			cnt++;
	}
	fclose(F);
	return cnt;
}

void compute_elapsed(struct timeval *delta, const struct timeval *t0)
{
	struct timeval t1;

	gettimeofday(&t1, NULL);
	delta->tv_sec = t1.tv_sec - t0->tv_sec;
	delta->tv_usec = t1.tv_usec - t0->tv_usec;
	if (delta->tv_usec < 0) {
		delta->tv_usec += 1000000;
		delta->tv_sec--;
	}
}

int nr_loops = 20*1000000;
double incr=0.3456;
void perform_work()
{
	int i;
	double t = 0.0;
	for (i = 0; i < nr_loops; i++) {
		t += incr;
		}
	if (t < 0.0) printf("well... should not happen\n");
}

void set_affinity(int cpu)
{
	long cpu_mask;
	int res;

	cpu_mask = 1L << cpu;
	res = sched_setaffinity(0, sizeof(cpu_mask), &cpu_mask);
	if (res)
		perror("sched_setaffinity");
}

void *thread_work(void *arg)
{
	int cpu = (int)arg;
	set_affinity(cpu);
	while (blockthemall)
		cpupause();
	perform_work();
	return (void *)0;
}

main(int argc, char *argv[])
{
	struct timeval t0, delta;
	int nr_cpus, i;
	pthread_t *tids;

	gettimeofday(&t0, NULL);
	perform_work();
	compute_elapsed(&delta, &t0);
	printf("Time to perform the unit of work on one thread is %d.%06d s\n", delta.tv_sec, delta.tv_usec);

	nr_cpus = number_of_cpus();
	if (nr_cpus <= 1)
		return 0;
	tids = malloc(nr_cpus * sizeof(pthread_t));
	for (i = 1; i < nr_cpus; i++) {
		pthread_create(tids + i, NULL, thread_work, (void *)i);
	}

	set_affinity(0);
	gettimeofday(&t0, NULL);
	blockthemall=0;
	perform_work();
	for (i = 1; i < nr_cpus; i++)
		pthread_join(tids[i], NULL);
	compute_elapsed(&delta, &t0);
	printf("Time to perform the unit of work on %d threads is %d.%06d s\n", nr_cpus, delta.tv_sec, delta.tv_usec);
	
}