/* NUMA test for multicore processors
 * Tyler Simon
 * 7/4/2009
 * Notes:
 * Dynamically selects number of cores
 * Modify MEM_MB for your per core memory size
 *
 * I recommend keeping ITER at 1, but if it makes you feel better increase
 * it, you should get the same BW either way. Check by running with ITER=1,
 * then ITER=5.
 *
 * Enable DEBUG to get per core process stats, I am lazy so execute with numa not ./numa if you can
 *
 * To Compile: c++ -o numa numa.C
 * To Run: numa
 * 
 * Output:
 * [from core] [to core] [BW in MB/sec]
 *
 * You should get an output file <hostname>_numa.txt for plotting
 */


#include <sys/resource.h>
#include <sched.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/time.h>
#include <string.h>
#include <unistd.h>

#define MEM_MB   200
#define MEM_SIZE MEM_MB*1024L*1024L
#define ITER 1 
#define DEBUG 1

char execmd[255];

double getTime() {
  struct timeval tp;
  static double last_timestamp = 0.0;
  double timestamp;
  gettimeofday (&tp, 0);
  timestamp = (double) tp.tv_sec * 1e6 + tp.tv_usec;
  return timestamp;
}

int getNumCPU() {
  cpu_set_t mask;
  if (sched_getaffinity(0,sizeof(cpu_set_t),&mask)) {
    fprintf (stderr, "Unable to retrieve affinity\n");
    exit(1);
  }
  int nproc = 0,i;
  for(i=0; i<CPU_SETSIZE; i++) {
    if( CPU_ISSET(i,&mask) ) {
      nproc++;
    }
  }
  return nproc;
}

void memtest(char *ptr) {
int i;
  for (i=0; i<ITER; i++) {
    memcpy(ptr, ptr+(MEM_SIZE/2),MEM_SIZE/2);
  }
}

void setCPU(int cpu) {
  cpu_set_t mask;
  CPU_ZERO(&mask);
  CPU_SET(cpu, &mask);
  sched_setaffinity(0, sizeof(cpu_set_t), &mask);
}

void test(int cpu, int nproc, char *pscmd, FILE *outfd) {
  long j;
  int i;
  setCPU(cpu);
  char *ptr = (char*) malloc (MEM_SIZE);
  if (!ptr) {
    fprintf (stderr, "failed to malloc\n");
    exit(1);
  }
  if(DEBUG)printf ("\n#%ld MB allocated on cpu %d\n", MEM_MB,cpu);
  
  // make sure it all gets paged in
  for (j=0; j<MEM_SIZE; j++) {
    ptr[j] = j;
  }
  for (i = 0; i < nproc; i++) {
    setCPU(i);


//if(DEBUG)popen(pscmd,"w");
if(DEBUG)system("ps -Cnuma -o pid,pmem,rss,psr,policy,nlwp,pcpu,cmd");


    double start = getTime();
    memtest(ptr);
    double end = getTime();
    fprintf(stderr,"%2.1f\n", (MEM_MB/((end - start) / (1000*1000))*ITER));
    fprintf(outfd,"%d %d %f \n", cpu,i, (MEM_MB/((end - start) / (1000*1000))*ITER));
  }
  free (ptr);
}

int main (int argc, char **argv) {
  int i;
  int nproc = getNumCPU();
  struct rusage usage;
  int ret;
  char pscmd[255];
  char hname[255];
  char out_fname[255];
  FILE *outfd;
  
  gethostname(hname,255);
  sprintf(out_fname,"%s_numa.txt", hname);
  strcpy(execmd,argv[0]);

  sprintf(pscmd,"ps -C%s -o pmem,rss,psr,policy,nlwp,pcpu,sgi_p,ruser,cmd",&execmd);

  outfd=fopen(out_fname,"w");
  if(!outfd){
      fprintf(stderr,"couldn't open output file");
      exit(1);
    }
  

  for (i = 0; i < nproc; i++) {
    test(i,nproc,pscmd,outfd);
  }

  fclose(outfd);
  return 0;
}