/******************  com_mpi.c *****************************************
 * 
 *  Communications routines, for MPI interface
 *  Modified from the MILC lattice QCD one
 *  KR 2001
 *



   g_sync() provides a synchronization point for all nodes.
   g_floatsum() sums a floating point number over all nodes.
   g_doublesum() sums a double over all nodes.
   g_vecdoublesum() sums a vector of doubles over all nodes.
   g_floatmax() finds maximum of a floating point number over all nodes.
   g_doublemax() finds maximum of a double over all nodes.
   broadcast_float()  broadcasts a single precision number from
	node 0 to all nodes.
   broadcast_double()  broadcasts a double precision number
   send_integer() sends an integer to one other node
   receive_integer() receives an integer
   terminate() kills the job on all processors

   start_gather() starts asynchronous sends and receives required
   to gather neighbors.
   wait_gather()  waits for receives to finish, insuring that the
   data has actually arrived.

   send_field() sends a field to one other node.
   receive_field() receives a field from some other node.
*/

/* load in definitions, variables etc. */

#include "comdefs.h"
#include "generic.h"
#include "timers.h"   /* includes comm timer calculators */

#ifdef TIMERS
static timer_type start_gather_timer, wait_send_timer, wait_receive_timer, g_sync_timer,
  g_sum_timer, broadcast_timer, send_timer, total_time;
#endif
static double mpi_sent_data = 0.0;


#define MIN_GATHER_INDEX 100
#define MAX_GATHER_INDEX 7100    /* allows 7000 concurrent gathers */
#define FIELD_TYPE 11            /* used in send/receive field */
#define SEND_INTEGER_TYPE 12     /* used in send/receive int */

#define SEND_FLAG 1              /* flags for marking msg_tags */
#define RECEIVE_FLAG 2

extern comlist_struct *comlist;   /* the comlist variables in layout.c */
      
/************************************************************************/

/* get all msg_tags in a single array, avoid allocating
 *   small bits and pieces of messages
 */

#define N_MSG_TAG 40

static msg_tag  msg_tag_arr[N_MSG_TAG];
static msg_tag *msg_tag_free, *msg_tag_last;

void init_msg_tags()
{
  int i;

  for (i=0; i<N_MSG_TAG-1; i++) msg_tag_arr[i].next = &msg_tag_arr[i+1];

  msg_tag_arr[N_MSG_TAG-1].next = NULL;

  msg_tag_free = msg_tag_arr;
}


msg_tag *get_msg_tags(int ntags)
{
  msg_tag *r,*p;
  int i;
  
  r = p = msg_tag_free;
  for (i=0; i<ntags; i++) {
    if (p == NULL) {
      printf("last msg_tag used from array, node %d\n",mynode()); 
      terminate(1212);
    }
    if (i<ntags-1) p = p->next;
  }

  /* p points to the last in the list */
  msg_tag_free = p->next;
  p->next = NULL;
  return(r);
}


void release_msg_tags(msg_tag *tp)
{
  msg_tag *p;

  for (p=tp; p->next != NULL; p=p->next) ;
  p->next = msg_tag_free;
  msg_tag_free = tp;
}

/************************************************************************/


/* Machine initialization */
#include <sys/types.h>
void initialize_machine() {
/*   MPI_Init(&argc,&argv); */
  
#ifdef TIMERS
  timer_reset( &total_time );
  timer_start( &total_time );

  timer_reset( &start_gather_timer );
  timer_reset( &wait_send_timer );
  timer_reset( &wait_receive_timer );
  timer_reset( &g_sync_timer );
  timer_reset( &g_sum_timer );
  timer_reset( &broadcast_timer );
  timer_reset( &send_timer );
#endif

  init_msg_tags();

}


/************************************************************************/

/* this formats the wait_array, used by forallsites_waitA()
 * should be made as fast as possible!
 */

int setup_wait_arr( unsigned char *wait_arr, msg_tag* tag_out[], 
		   msg_tag* tag_in[], int ntag )
{
  register int i,j;
  int nt,dir[NA_MAX];

  if (ntag > NA_MAX) halt("Error in forallsites_waitA: too many gathers");
  for (nt=i=0; i<ntag; i++) if (tag_in[i] != NULL) {
    tag_out[nt] = tag_in[i];
    dir[nt] = tag_in[i]->dir;
    nt++;
  }
  tag_out[nt] = (msg_tag *)NULL;  /* needed for the last loop */

  forallsites(i) {
    wait_arr[i] = 0;    /* basic, no wait */
    for (j=0; j<nt; j++) if ( nb(dir[j],i) >= node.sites ) wait_arr[i] = j+1;
  }
  return( nt );
}


/************************************************************************/
/* GATHER ROUTINES */
/* start_gather() returns a pointer to a list of msg_tag's, which
   be used as input to subsequent wait_gather().

   This list contains msg_tags for all receive buffers, followed by
   end flag.

   If no messages at all are required, the routine will return NULL.
   msg_buf==NULL should be a reliable indicator of no message.

   usage:  tag = start_gather( source, size, direction, parity )
*/



msg_tag* start_gather( field, size, dir, parity )
     /* arguments */
     char * field;      /* pointer to some latfield */
     int size;		/* size in bytes of the field (eg sizeof(su3_vector))*/
     int dir;		/* direction to gather from. eg XUP - index into
			   neighbor tables */
     int parity;	/* parity of sites whose neighbors we gather.
			   one of EVEN, ODD or EVENODD (EVEN+ODD). */
{
  /* local variables */
  int i,j,k,nodepar;	        /* scratch */
  int offset;		/* number of sites in this receive or send */
  int *idx;             /* index array pointer */
  int nsites;
  char *tpt;            /* temp ptr to buffer */
  msg_tag *mbuf, *mp;	/* list of message tags, to be returned */
  comlist_struct *cp;   
  send_struct *sp;
  receive_struct *rp;
  static int index = MIN_GATHER_INDEX;  /* index to identify the operation */


  if (dir < 0 || dir >= NDIRS) {
    printf("No such gather %d, node %d\n",dir,mynode());
    terminate(1212);
  }

  /* First, get the rolling index for the operation 
   * This MUST BE HERE even if there's nothing to do, because
   * somebody else might be doing this!
   */
  ++index ;  if (index > MAX_GATHER_INDEX) index = MIN_GATHER_INDEX;

  cp = &comlist[dir];

  /* Now if there's nothing to do, return - CHECK IF GATHERED */
  if( ( cp->n_send == 0 && cp->n_receive == 0 ) 
      || is_already_gathered( field, size, dir, parity) ) 
    return( (msg_tag *) NULL );

  /* mark gathered, if needed */
  gather_mark_gathered( field, size, dir, parity );

  timer_start( &start_gather_timer );
  
  /* allocate a buffer for the msg_tags.  This is dynamically allocated
     because there may be an arbitrary number of gathers in progress
     in any direction. SIZE = n_msgs sent+received, for end flag */

  /* mbuf = (msg_tag *)malloc((cp->n_send + cp->n_receive + 1)*sizeof(msg_tag) );
     if(mbuf==NULL){
     printf("No room for mbuf, node %d\n",mynode()); 
     terminate(1212);
     }
  */

  mbuf = get_msg_tags(cp->n_send + cp->n_receive);

  /* Loop over node parity */
  forbothparities(nodepar) {

  mp=mbuf;
  /* HANDLE RECEIVES: loop over nodes which will send here */
  for (i=0, rp=cp->from_node; i<cp->n_receive;i++, rp=rp->next, mp=mp->next) {

    if (nodepar == node.parity) {
      MPI_Status status;

    /* note--neighbors of EVEN sites are always first in the list!
     * Thus, for ODD sites we must change the offset 
     */
      switch (parity) {
      case EVEN: nsites = rp->n_even; offset = rp->offset; break;
      case ODD:  nsites = rp->n_odd;  offset = rp->offset + rp->n_even; break;
      case EVENODD: nsites = rp->n;   offset = rp->offset; break;
      }

      mp->flag = RECEIVE_FLAG; /* flag as normal receive */
      mp->dir  = dir;
      /* and post receive -- comes right on spot */
      MPI_Recv( ((char *)field) + offset*size, nsites*size, MPI_BYTE,
		rp->node, index, MPI_COMM_WORLD, &status );
    }
  }

  /* HANDLE SENDS - note: mp automatically correct */
  for(k=0,sp=cp->to_node; k < cp->n_send; k++,sp = sp->next, mp = mp->next) {
    if (nodepar != node.parity) {
      switch (parity) {
      case EVEN: nsites = sp->n_even; offset = 0; break;
      case ODD:  nsites = sp->n_odd;  offset = sp->n_even; break;
      case EVENODD: nsites = sp->n;   offset = 0; break;
      }
    
      /* allocate buffer */
      tpt = (char *)malloc( nsites*size );
      if(tpt==NULL){printf("NO ROOM for tpt, node %d\n",mynode());exit(1);}
      mp->flag = SEND_FLAG; /* flag as send */
      mp->dir  = dir;
      mp->buf  = tpt;
      /* gather data into the buffer */
    
      idx = sp->sitelist + offset;    /* initial offset */
      for (j=0; j<nsites; j++, tpt += size) {
	memcpy( tpt, ((char *)field) + idx[j]*size, size );
      }

      /* start the send -- WHICH TO USE?  Isend/Issend ? */
      MPI_Rsend( mp->buf, nsites*size, MPI_BYTE,
		 sp->node, index, MPI_COMM_WORLD );
    
      free( mp->buf );
    }
  }

  }

  timer_end( &start_gather_timer );

  release_msg_tags( mbuf );
 
  /* return */
  return((msg_tag*)NULL);
}


msg_tag * wait_gather( msg_tag *mbuf ) 
{
  MPI_Status status;
  msg_tag *mp;

  if (mbuf == NULL) return((msg_tag *)NULL);
  timer_start( &wait_receive_timer );
  /* wait for all receive messages */
  for(mp=mbuf; mp != NULL && mp->flag == RECEIVE_FLAG; mp=mp->next) {
    MPI_Wait( &(mp->mpi), &status );
  }
  timer_end( &wait_receive_timer );
  /* wait for all send messages */
  timer_start( &wait_send_timer );
#ifndef SEND_TEST
  for( ; mp != NULL && mp->flag == SEND_FLAG; mp=mp->next) {
    MPI_Wait( &(mp->mpi), &status );
    /* release the buffer */
    free( mp->buf );
  }
#endif
  timer_end( &wait_send_timer );
  /* and free the mbuf */
  release_msg_tags( mbuf );
  return((msg_tag *)NULL);
}


/****************************************************************
 */


/* SEND AND RECEIVE FIELD */
/* send_field is to be called only by the node doing the sending */
/* get_field is to be called only by the node to which the field was sent */
void send_field(buf,size,tonode) 
     void *buf; int size,tonode; 
{
  timer_start( &send_timer );
  MPI_Send(buf,size,MPI_BYTE,tonode,FIELD_TYPE,MPI_COMM_WORLD);
  timer_end( &send_timer );
  mpi_sent_data += size;
}
void receive_field(buf,size) 
     void *buf; int size; 
{
  MPI_Status status;

  timer_start( &send_timer );
  MPI_Recv(buf,size,MPI_BYTE,MPI_ANY_SOURCE,FIELD_TYPE,
	   MPI_COMM_WORLD,&status);
  timer_end( &send_timer );
  mpi_sent_data += size;
}

/* BASIC COMMUNICATIONS FUNCTIONS */

/* Tell what kind of machine we are on */
static char name[]="MPI (portable)";
char * machine_type(){
  return(name);
}

/* Return my node number */
int mynode()
{
  int node;
  MPI_Comm_rank( MPI_COMM_WORLD, &node );
  return(node);
}

/* Return number of nodes */
int numnodes()
{
  int nodes;
  MPI_Comm_size( MPI_COMM_WORLD, &nodes );
  return(nodes);
}

/* Synchronize all nodes */
void g_sync()
{
  timer_start( &g_sync_timer );
  MPI_Barrier( MPI_COMM_WORLD );
  timer_end( &g_sync_timer );
}

/* Sum float over all nodes to node 0 */
void g_floatsum( float * fpt ) 
{
  float work;

  timer_start( &g_sum_timer );
  if (this_node == 0) {
    MPI_Reduce( fpt, &work, 1, MPI_FLOAT, MPI_SUM, 0 , MPI_COMM_WORLD );
    *fpt = work;
  } else {
    MPI_Reduce( fpt, &work, 1, MPI_FLOAT, MPI_SUM, 0 , MPI_COMM_WORLD );
  }
  timer_end( &g_sum_timer );
}

/* Sum double over all nodes */
void g_doublesum( double * dpt ) 
{
  double work;

  timer_start( &g_sum_timer );
  if (this_node == 0) {
    MPI_Reduce( dpt, &work, 1, MPI_DOUBLE, MPI_SUM, 0 , MPI_COMM_WORLD );
    *dpt = work;
  } else {
    MPI_Reduce( dpt, &work, 1, MPI_DOUBLE, MPI_SUM, 0 , MPI_COMM_WORLD );
  }
  timer_end( &g_sum_timer );
}

/* Sum double over all nodes, and scatter the result */
void  g_doublesum_scatter( double * dpt )
{
  double work;

  timer_start( &g_sum_timer );
  MPI_Allreduce( dpt, &work, 1, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD );
  *dpt = work;
  timer_end( &g_sum_timer );
}


 
/* Sum a vector of doubles over all nodes */
void g_vecintsum( dpt, n ) 
     int *dpt; int n; 
{
  register int *work;
  register int i;

  timer_start( &g_sum_timer );
  work = (int *)malloc(n*sizeof(int));
  if (this_node == 0) {
    MPI_Reduce( dpt, work, n, MPI_INT, MPI_SUM, 0 , MPI_COMM_WORLD );
    for (i=0; i<n; i++) dpt[i] = work[i];
  } else {
    MPI_Reduce( dpt, work, n, MPI_INT, MPI_SUM, 0 , MPI_COMM_WORLD );
  }
  free(work);
  timer_end( &g_sum_timer );
}

 
/* Sum a vector of doubles over all nodes */
void g_vecfloatsum( dpt, nfloats ) 
     float *dpt; int nfloats; 
{
  register float *work;
  register int i;

  timer_start( &g_sum_timer );
  work = (float *)malloc(nfloats*sizeof(float));
  if (this_node == 0) {
    MPI_Reduce( dpt, work, nfloats, MPI_FLOAT, MPI_SUM, 0 , MPI_COMM_WORLD );
    for (i=0; i<nfloats; i++) dpt[i] = work[i];
  } else {
    MPI_Reduce( dpt, work, nfloats, MPI_FLOAT, MPI_SUM, 0 , MPI_COMM_WORLD );
  }
  free(work);
  timer_end( &g_sum_timer );
  /**
  work = (float *)malloc(nfloats*sizeof(float));
  MPI_Allreduce( dpt, work, nfloats, MPI_FLOAT, MPI_SUM, MPI_COMM_WORLD );
  for(i=0;i<nfloats;i++) dpt[i]=work[i];
  free(work);
  **/
}

/* Sum a vector of doubles over all nodes */
void g_vecdoublesum( dpt, ndoubles ) 
     double *dpt; int ndoubles;
{
  register double *work;
  register int i;
  
  timer_start( &g_sum_timer );
  work = (double *)malloc(ndoubles*sizeof(double));
  if (this_node == 0) {
    MPI_Reduce( dpt, work, ndoubles, MPI_DOUBLE, MPI_SUM, 0 , MPI_COMM_WORLD );
    for (i=0; i<ndoubles; i++) dpt[i] = work[i];
  } else {
    MPI_Reduce( dpt, work, ndoubles, MPI_DOUBLE, MPI_SUM, 0 , MPI_COMM_WORLD );
  }
  free(work);
  timer_end( &g_sum_timer );

  /**
  work = (double *)malloc(ndoubles*sizeof(double));
  MPI_Allreduce( dpt, work, ndoubles, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD );
  for(i=0;i<ndoubles;i++)dpt[i]=work[i];
  free(work);
  **/
}

/* Find maximum of float over all nodes */
void g_floatmax( fpt ) 
     float *fpt; 
{
  float work;
  MPI_Allreduce( fpt, &work, 1, MPI_FLOAT, MPI_MAX, MPI_COMM_WORLD );
  *fpt = work;
}

/* Find maximum of double over all nodes */
void g_doublemax( dpt ) 
     double *dpt; 
{
  double work;
  MPI_Allreduce( dpt, &work, 1, MPI_DOUBLE, MPI_MAX, MPI_COMM_WORLD );
  *dpt = work;
}

/* Broadcast a whole field */
void broadcast_field(void *pt, int size) 
{
  timer_start( &broadcast_timer );
  MPI_Bcast( pt, size, MPI_BYTE, 0, MPI_COMM_WORLD );
  timer_end( &broadcast_timer );
}

/* Broadcast floating point number from node zero */
void broadcast_float(float *fpt) 
{
  timer_start( &broadcast_timer );
  MPI_Bcast( fpt, 1, MPI_FLOAT, 0, MPI_COMM_WORLD );
  timer_end( &broadcast_timer );
}

/* Broadcast double precision floating point number from node zero */
void broadcast_double( double *dpt )
{
  timer_start( &broadcast_timer );
  MPI_Bcast( dpt, 1, MPI_DOUBLE, 0, MPI_COMM_WORLD );
  timer_end( &broadcast_timer );
}

/* Broadcast double precision floating point number from node zero */
void broadcast_int( int *dpt )
{
  timer_start( &broadcast_timer );
  MPI_Bcast( dpt, 1, MPI_INT, 0, MPI_COMM_WORLD );
  timer_end( &broadcast_timer );
}

/* Send an integer to one other node */
/* This is to be called only by the node doing the sending */
void send_integer(tonode,address) 
     int tonode; int *address; 
{
  MPI_Send( address, 1, MPI_INT, tonode, SEND_INTEGER_TYPE,
	    MPI_COMM_WORLD);
}

/* Receive an integer from another node */
/* Note we do not check if this was really meant for us */
void receive_integer(address) 
     int *address; 
{
  MPI_Status status;
  MPI_Recv( address, 1, MPI_INT, MPI_ANY_SOURCE, SEND_INTEGER_TYPE,
	    MPI_COMM_WORLD, &status);
}



/****************************************************************
 * Latfield status functions
 * -- these keep track of the gathers, i.e. is the field already
 * gathered so that new fetch is unnecessary.
 */

#define GATHER_FOLLOW 14232
#define GATHER_NOT_FOLLOW 12

int n_gather_done = 0, n_gather_avoided = 0;

int is_already_gathered( char *field, int size, int dir, int parity )
{
  struct gather_status_arr *a;
  /* cast the latfield array */
  a = (struct gather_status_arr *) (field + size*node.latfield_size);
  if ( a->status == GATHER_FOLLOW && (a->gathered[dir] ^ parity) == 0) {
#ifdef CHECK_GATHER_FIELDS
    /* now looks like is gathered, check if the checkup field has changed */
    if (((parity & EVEN) && (*((unsigned int *)field) != a->check_even[dir] )) ||
	((parity & ODD ) && (*((unsigned int *)(field+size*(node.sites-1))) 
			     != a->check_odd[dir] ))) {
      printf(" #### GATHER CHECK: Forgotten mark_changed() somewhere!\n");
      if (size < sizeof(int))
	printf(" Because field size %d < sizeof(int), can be spurious\n",size);
      else {
	printf(" Field size %d chars, dir %d parity %d\n",size, dir, parity);
      }
      halt(" ###### ");
    }
#endif
    n_gather_avoided++;
    return(1);
  } 
  n_gather_done++;
  return(0);
}

void gather_status_reset( char *field, int size )
{
  struct gather_status_arr *a;
  /* cast the latfield array */
  a = (struct gather_status_arr *) (field + size*node.latfield_size);
  a->status = GATHER_NOT_FOLLOW;
}

void gather_mark_dirty( char *field, int size, int parity )
{
  int dir,p;
  struct gather_status_arr *a;

  /* cast the latfield array */
  a = (struct gather_status_arr *) (field + size*node.latfield_size);
  a->status = GATHER_FOLLOW;
  p = opp_parity(parity); 
  /* mark opposite parity, because will fetch from there! */
  /* Remember that need to mark opposite directions too! */
  for(dir=0; dir<NDIRS; dir++) a->gathered[dir] &= (!p);
}

void gather_mark_gathered( char *field, int size, int dir, int parity )
{
  struct gather_status_arr *a;

  /* cast the latfield array */
  a = (struct gather_status_arr *) (field + size*node.latfield_size);
  if (a->status == GATHER_FOLLOW) a->gathered[dir] |= parity;
#ifdef CHECK_GATHER_FIELDS
  if (parity & EVEN) a->check_even[dir] = *((unsigned int *)field);
  if (parity & ODD ) a->check_odd[dir]  = 
    *((unsigned int *)(field+size*(node.sites-1)));
#endif
}


/****************************************************************/



/* version of exit for multinode processes -- kill all nodes */
void terminate(int status)
{
  printf("Termination: node %d, status = %d\n",this_node,status);
  fflush(stdout);
  MPI_Abort( MPI_COMM_WORLD, 0);
  exit(status);
}




/* clean exit from all nodes */
void finishrun() 
{
#ifdef TIMERS
  report_comm_timers();
#endif

  if (this_node == 0) {
    extern int n_gather_done,n_gather_avoided; 

    printf(" COMMS from node 0: %d done, %d (%.2g%%) optimized away\n",
	   n_gather_done, n_gather_avoided, 
	   100.0*n_gather_avoided/(n_gather_avoided+n_gather_done));
  }

  fflush(stdout);
  fflush(NULL);  /* Try to flush all open files */
  MPI_Finalize();
  exit(0);
}


void report_comm_timers()
{
#ifdef TIMERS
  double tot;
  if (this_node == 0) {
    printf(" *************************\n");
    printf(" MPI communications timers from node 0:\n");
    
    printf(" start_get:    ");
    timer_report( &start_gather_timer );
    printf(" waiting send: ");
    timer_report( &wait_send_timer );
    printf("      receive: ");
    timer_report( &wait_receive_timer );
    printf(" g_sync:       ");
    timer_report( &g_sync_timer );
    printf(" g_sum:        ");
    timer_report( &g_sum_timer );
    printf(" broadcast:    ");
    timer_report( &broadcast_timer );
    printf(" send/receive: ");
    timer_report( &send_timer );
    printf(" send/receive %g MB, bandwith %g MB/sec\n",
	   mpi_sent_data*1e-6, mpi_sent_data*1e-6/send_timer.total);

    tot = start_gather_timer.total + wait_send_timer.total + wait_receive_timer.total + 
      g_sync_timer.total +
      g_sum_timer.total + broadcast_timer.total + send_timer.total ;

    /* find current time */
    timer_end( &total_time );
    
    printf(" Total comm. time %.3g, total time %.3g, comm %.2g%%\n",
	   tot, total_time.total, 100*tot/total_time.total );
    printf(" ***** \n");
  }
#endif
}


