emissary-mpi reduce operation

Author: amitamd7

offload/DeviceRTL/include/EmissaryMPI.h

@@ -16,20 +16,32 @@ typedef enum {
   _MPI_INVALID,
   _MPI_Send_idx,
   _MPI_Recv_idx,
+  _MPI_Allreduce_idx,
+  _MPI_Reduce_idx,
 } offload_emis_mpi_t;
 
 ///  Device stubs that call _emissary_exec using identical host API interface
 #if defined(__NVPTX__) || defined(__AMDGCN__)
 extern "C" int MPI_Send(const void *buf, int count, MPI_Datatype datatype,
                         int dest, int tag, MPI_Comm comm) {
-  return (int)_emissary_exec(_PACK_EMIS_IDS(EMIS_ID_MPI, _MPI_Send_idx), buf,
-                             count, datatype, dest, tag, comm);
+  return (int)_emissary_exec(_PACK_EMIS_IDS(EMIS_ID_MPI, _MPI_Send_idx), 
+		         buf, count, datatype, dest, tag, comm);
 }
 extern "C" int MPI_Recv(void *buf, int count, MPI_Datatype datatype, int source,
                         int tag, MPI_Comm comm, MPI_Status *st) {
   return (int)_emissary_exec(_PACK_EMIS_IDS(EMIS_ID_MPI, _MPI_Recv_idx), buf,
                              count, datatype, source, tag, comm, st);
 }
+extern "C" int MPI_Allreduce(const void *sendbuf, void *recvbuf, int count,
+                              MPI_Datatype datatype, MPI_Op op, MPI_Comm comm) {
+  return (int)_emissary_exec(_PACK_EMIS_IDS(EMIS_ID_MPI, _MPI_Allreduce_idx), 
+		  sendbuf, recvbuf, count, datatype, op, comm);
+}
+extern "C" int MPI_Reduce(const void *sendbuf, void *recvbuf, int count, MPI_Datatype datatype,
+                              MPI_Op op, int root, MPI_Comm comm) {
+  return (int)_emissary_exec(_PACK_EMIS_IDS(EMIS_ID_MPI, _MPI_Reduce_idx), 
+                   sendbuf, recvbuf, count, datatype, op, root, comm);
+}
 #endif
 
 /// Host variadic wrapper functions.
@@ -61,6 +73,35 @@ extern int V_MPI_Recv(void *fnptr, ...) {
   int rval = MPI_Recv(v0, v1, v2, v3, v4, v5, v6);
   return rval;
 }
+extern int V_MPI_Allreduce(void *fnptr, ...) {
+  va_list args;
+  va_start(args, fnptr);
+  void *buf             = va_arg(args, void *);
+  void *recvbuf         = va_arg(args, void *);
+  int count             = va_arg(args, int);
+  MPI_Datatype datatype = va_arg(args, MPI_Datatype);
+  MPI_Op op             = va_arg(args, MPI_Op);
+  MPI_Comm comm         = va_arg(args, MPI_Comm);
+  va_end(args);
+  int rval = MPI_Allreduce(
+    buf, recvbuf, count, datatype, op, comm);
+  return rval;
+}
+extern int V_MPI_Reduce(void *fnptr, ...) {
+  va_list args;
+  va_start(args, fnptr);
+  void *buf             = va_arg(args, void *);
+  void *recvbuf         = va_arg(args, void *);
+  int count             = va_arg(args, int);
+  MPI_Datatype datatype = va_arg(args, MPI_Datatype);
+  MPI_Op op             = va_arg(args, MPI_Op);
+  int root              = va_arg(args, int);
+  MPI_Comm comm         = va_arg(args, MPI_Comm);
+  va_end(args);
+  int rval = MPI_Reduce(
+    buf, recvbuf, count, datatype, op, root, comm);
+  return rval;
+}
 
 /// EmissaryMPI function selector
 emis_return_t EmissaryMPI(char *data, emisArgBuf_t *ab, emis_argptr_t *a[]) {
@@ -78,6 +119,18 @@ emis_return_t EmissaryMPI(char *data, emisArgBuf_t *ab, emis_argptr_t *a[]) {
         V_MPI_Recv(fnptr, a[0], a[1], a[2], a[3], a[4], a[5], a[6]);
     return (emis_return_t)return_value_int;
   }
+  case _MPI_Allreduce_idx: {
+    void *fnptr = (void *)V_MPI_Allreduce;
+    int return_value_int =
+        V_MPI_Allreduce(fnptr, a[0], a[1], a[2], a[3], a[4], a[5]);
+    return (emis_return_t) return_value_int;
+  }
+  case _MPI_Reduce_idx: {
+    void *fnptr = (void *)V_MPI_Reduce;
+    int return_value_int =
+        V_MPI_Reduce(fnptr, a[0], a[1], a[2], a[3], a[4], a[5], a[6]);
+    return (emis_return_t) return_value_int;
+  }
   }
   return (emis_return_t)0;
 }