From f31e0dcaf76fdf6f728f22a0c4feca79b0e22c55 Mon Sep 17 00:00:00 2001
From: Sourav Chakraborty <chakraborty.52@osu.edu>
Date: Mon, 24 Apr 2017 12:45:09 -0400
Subject: [PATCH] Add Non-blocking and Shared memory PMI Calls

---
 contribs/pmi2/pmi2.h             |   8 ++
 contribs/pmi2/pmi2_api.c         | 172 +++++++++++++++++++++++++++++++++++++++
 contribs/pmi2/pmi2_util.c        |  35 ++++++++
 contribs/pmi2/pmi2_util.h        |   9 ++
 contribs/pmi2/slurm/pmi2.h       |  14 ++++
 src/plugins/mpi/pmi2/Makefile.am |   1 +
 src/plugins/mpi/pmi2/Makefile.in |   4 +-
 src/plugins/mpi/pmi2/client.c    |  56 +++++++++++++
 src/plugins/mpi/pmi2/client.h    |   4 +
 src/plugins/mpi/pmi2/kvs.c       | 157 +++++++++++++++++++++++++++++++++++
 src/plugins/mpi/pmi2/kvs.h       |   7 +-
 src/plugins/mpi/pmi2/pmi.h       |   6 +-
 src/plugins/mpi/pmi2/pmi1.c      |   1 +
 src/plugins/mpi/pmi2/pmi2.c      |  40 ++++++++-
 src/plugins/mpi/pmi2/setup.c     |  13 ++-
 src/plugins/mpi/pmi2/setup.h     |   1 +
 src/plugins/mpi/pmi2/shmem.c     | 115 ++++++++++++++++++++++++++
 src/plugins/mpi/pmi2/shmem.h     |  70 ++++++++++++++++
 src/plugins/mpi/pmi2/tree.c      | 163 +++++++++++++++++++++++++++++++++++++
 src/plugins/mpi/pmi2/tree.h      |   2 +
 20 files changed, 872 insertions(+), 6 deletions(-)
 create mode 100644 src/plugins/mpi/pmi2/shmem.c
 create mode 100644 src/plugins/mpi/pmi2/shmem.h

diff --git a/contribs/pmi2/pmi2.h b/contribs/pmi2/pmi2.h
index ddcc36c..8bee03d 100644
--- a/contribs/pmi2/pmi2.h
+++ b/contribs/pmi2/pmi2.h
@@ -52,6 +52,8 @@ static const char NAMEUNPUBLISH_CMD[]     = "name-unpublish";
 static const char NAMEUNPUBLISHRESP_CMD[] = "name-unpublish-response";
 static const char NAMELOOKUP_CMD[]        = "name-lookup";
 static const char NAMELOOKUPRESP_CMD[]    = "name-lookup-response";
+static const char ALLGATHER_CMD[]         = "allgather";
+static const char ALLGATHERRESP_CMD[]     = "allgather-response";
 
 static const char PMIJOBID_KEY[]          = "pmijobid";
 static const char PMIRANK_KEY[]           = "pmirank";
@@ -451,6 +453,12 @@ int PMI2_KVS_Put(const char key[], const char value[]);
 @*/
 int PMI2_KVS_Fence(void);
 
+int PMI2_KVS_Ifence(void);
+int PMI2_KVS_Wait(void);
+
+int PMI2_Iallgather(const char value[]);
+int PMI2_Iallgather_wait(void *result);
+
 /*@
   PMI2_KVS_Get - returns the value associated with key in the key-value
       space associated with the job ID jobid
diff --git a/contribs/pmi2/pmi2_api.c b/contribs/pmi2/pmi2_api.c
index f194ebe..8da5a70 100644
--- a/contribs/pmi2/pmi2_api.c
+++ b/contribs/pmi2/pmi2_api.c
@@ -35,6 +35,7 @@ typedef enum {
 } PMI2State;
 
 static PMI2State PMI2_initialized = PMI2_UNINITIALIZED;
+static PMI2ShmemRegion PMI2_Shmem_allgather  = { .fd=-1, .addr=NULL, .filesize=0 };
 
 static int PMI2_debug = 0;
 static int PMI2_fd = -1;
@@ -92,6 +93,7 @@ static void init_kv_strdup_intsuffix(PMI2_Keyvalpair *kv, const char key_prefix[
 static int getPMIFD(void);
 static int PMIi_ReadCommandExp( int fd, PMI2_Command *cmd, const char *exp, int* rc, const char **errmsg );
 static int PMIi_ReadCommand( int fd, PMI2_Command *cmd );
+static int PMIi_ReadRaw( int fd, char *buf );
 
 static int PMIi_WriteSimpleCommand( int fd, PMI2_Command *resp, const char cmd[], PMI2_Keyvalpair *pairs[], int npairs);
 static int PMIi_WriteSimpleCommandStr( int fd, PMI2_Command *resp, const char cmd[], ...);
@@ -356,6 +358,10 @@ int PMI2_Finalize(void)
         pmi2_errno = PMIi_ReadCommandExp(PMI2_fd, &cmd, FINALIZERESP_CMD, &rc, &errmsg);
         if (pmi2_errno) PMI2U_ERR_SETANDJUMP(1, pmi2_errno, "PMIi_ReadCommandExp");
         PMI2U_ERR_CHKANDJUMP(rc, pmi2_errno, PMI2_ERR_OTHER, "**pmi2_finalize %s", errmsg ? errmsg : "unknown");
+        if (PMI2_Shmem_allgather.fd != -1) {
+            pmi2_errno = PMI2U_Destroy_shmem(&PMI2_Shmem_allgather);
+            if (pmi2_errno) PMI2U_ERR_SETANDJUMP(1, pmi2_errno, "PMI2U_Destroy_shmem");
+        }
 
         free(cmd.command);
         freepairs(cmd.pairs, cmd.nPairs);
@@ -753,6 +759,47 @@ fn_fail:
     goto fn_exit;
 }
 
+int PMI2_KVS_Ifence(void)
+{
+    int pmi2_errno = PMI2_SUCCESS;
+    PMI2_Command cmd = {0};
+
+    PMI2U_printf("[BEGIN]");
+
+    pmi2_errno = PMIi_WriteSimpleCommandStr(PMI2_fd, &cmd, KVSFENCE_CMD, NULL);
+    if (pmi2_errno) PMI2U_ERR_SETANDJUMP(1, pmi2_errno, "PMIi_WriteSimpleCommandStr");
+
+fn_exit:
+    free(cmd.command);
+    freepairs(cmd.pairs, cmd.nPairs);
+    PMI2U_printf("[END]");
+    return pmi2_errno;
+fn_fail:
+    goto fn_exit;
+}
+
+int PMI2_KVS_Wait(void)
+{
+    int pmi2_errno = PMI2_SUCCESS;
+    PMI2_Command cmd = {0};
+    int rc;
+    const char *errmsg;
+
+    PMI2U_printf("[BEGIN]");
+
+    pmi2_errno = PMIi_ReadCommandExp(PMI2_fd, &cmd, KVSFENCERESP_CMD, &rc, &errmsg);
+    if (pmi2_errno) PMI2U_ERR_SETANDJUMP(1, pmi2_errno, "PMIi_ReadCommandExp");
+    PMI2U_ERR_CHKANDJUMP(rc, pmi2_errno, PMI2_ERR_OTHER, "**pmi2_kvsfence %s", errmsg ? errmsg : "unknown");
+
+fn_exit:
+    free(cmd.command);
+    freepairs(cmd.pairs, cmd.nPairs);
+    PMI2U_printf("[END]");
+    return pmi2_errno;
+fn_fail:
+    goto fn_exit;
+}
+
 int PMI2_KVS_Get(const char *jobid, int src_pmi_id, const char key[], char value [], int maxValue, int *valLen)
 {
     int pmi2_errno = PMI2_SUCCESS;
@@ -797,6 +844,98 @@ int PMI2_KVS_Get(const char *jobid, int src_pmi_id, const char key[], char value
     goto fn_exit;
 }
 
+int PMI2_Iallgather(const char value[])
+{
+    int pmi2_errno = PMI2_SUCCESS;
+    PMI2_Command cmd = {0};
+
+    PMI2U_printf("[BEGIN]");
+
+    pmi2_errno = PMIi_WriteSimpleCommandStr(PMI2_fd, &cmd, ALLGATHER_CMD, VALUE_KEY, value, NULL);
+    if (pmi2_errno) PMI2U_ERR_SETANDJUMP(1, pmi2_errno, "PMIi_WriteSimpleCommandStr");
+
+fn_exit:
+    free(cmd.command);
+    freepairs(cmd.pairs, cmd.nPairs);
+    PMI2U_printf("[END]");
+    return pmi2_errno;
+fn_fail:
+    goto fn_exit;
+}
+
+int PMI2_Iallgather_wait(void *results)
+{
+    int pmi2_errno = PMI2_SUCCESS;
+
+    PMI2U_printf("[BEGIN]");
+
+    pmi2_errno = PMIi_ReadRaw(PMI2_fd, results);
+    if (pmi2_errno) PMI2U_ERR_SETANDJUMP(1, pmi2_errno, "PMIi_ReadRaw");
+
+fn_exit:
+    PMI2U_printf("[END]");
+    return pmi2_errno;
+fn_fail:
+    goto fn_exit;
+}
+
+int PMI2_SHMEM_Iallgather(const char value[])
+{
+    int pmi2_errno = PMI2_SUCCESS;
+    PMI2_Command cmd = {0};
+
+    PMI2U_printf("[BEGIN]");
+
+    pmi2_errno = PMIi_WriteSimpleCommandStr(PMI2_fd, &cmd, ALLGATHER_CMD,
+                    VALUE_KEY, value, USE_SHMEM_KEY, TRUE_VAL, NULL);
+    if (pmi2_errno)
+        PMI2U_ERR_SETANDJUMP(1, pmi2_errno, "PMIi_WriteSimpleCommandStr");
+
+fn_exit:
+    free(cmd.command);
+    freepairs(cmd.pairs, cmd.nPairs);
+    PMI2U_printf("[END]");
+    return pmi2_errno;
+fn_fail:
+    goto fn_exit;
+}
+
+int PMI2_SHMEM_Iallgather_wait(void **results)
+{
+    int rc, found;
+    PMI2_Command cmd = {0};
+    int pmi2_errno = PMI2_SUCCESS;
+    int shmem_filename_len;
+    const char *shmem_filename;
+    const char *errmsg;
+
+    PMI2U_printf("[BEGIN]");
+    if (PMI2_Shmem_allgather.fd != -1) PMI2U_Destroy_shmem(&PMI2_Shmem_allgather);
+
+    pmi2_errno = PMIi_ReadCommandExp(PMI2_fd, &cmd, ALLGATHERRESP_CMD, &rc, &errmsg);
+    if (pmi2_errno) PMI2U_ERR_SETANDJUMP(1, pmi2_errno, "PMIi_ReadCommandExp");
+    PMI2U_ERR_CHKANDJUMP(rc, pmi2_errno, PMI2_ERR_OTHER, "**pmi2_allgather %s", errmsg ? errmsg : "unknown");
+
+    found = getval(cmd.pairs, cmd.nPairs, SHMEMFILENAME_KEY, &shmem_filename, &shmem_filename_len);
+    PMI2U_ERR_CHKANDJUMP(found == -1, pmi2_errno, PMI2_ERR_OTHER, "**intern");
+    strncpy(PMI2_Shmem_allgather.filename, shmem_filename, shmem_filename_len);
+
+    found = getvalint(cmd.pairs, cmd.nPairs, SHMEMFILESIZE_KEY, &PMI2_Shmem_allgather.filesize);
+    PMI2U_ERR_CHKANDJUMP(found != 1, pmi2_errno, PMI2_ERR_OTHER, "**intern");
+    found = PMI2U_Create_shmem(&PMI2_Shmem_allgather);
+    PMI2U_ERR_CHKANDJUMP(found != 0, pmi2_errno, PMI2_ERR_OTHER, "**intern");
+
+    *results = PMI2_Shmem_allgather.addr;
+
+fn_exit:
+    free(cmd.command);
+    freepairs(cmd.pairs, cmd.nPairs);
+    PMI2U_printf("[END]");
+    return pmi2_errno;
+fn_fail:
+    goto fn_exit;
+}
+
 int PMI2_Info_GetNodeAttr(const char name[], char value[], int valuelen, int *flag, int waitfor)
 {
     int pmi2_errno = PMI2_SUCCESS;
@@ -1325,6 +1464,39 @@ fn_fail:
 }
 
 
+static int PMIi_ReadRaw( int fd, char *buf )
+{
+    int pmi2_errno = PMI2_SUCCESS;
+	int data_len, nbytes, offset;
+
+    PMI2U_printf("[BEGIN]");
+
+	do {
+		nbytes = read(fd, &data_len, sizeof data_len);
+	} while (nbytes == -1 && errno == EINTR);
+
+    PMI2U_printf("[PMIi_ReadRaw] data_len=%d", data_len);
+
+	offset = 0;
+	do
+	{
+		do {
+			nbytes = read(fd, &buf[offset], data_len - offset);
+		} while (nbytes == -1 && errno == EINTR);
+
+		PMI2U_ERR_CHKANDJUMP(nbytes <= 0, pmi2_errno, PMI2_ERR_OTHER, "**read %s", strerror(errno));
+
+		offset += nbytes;
+	}
+	while (offset < data_len);
+
+fn_exit:
+    PMI2U_printf("[END]");
+    return pmi2_errno;
+fn_fail:
+    goto fn_exit;
+}
+
 /* Note that we fill in the fields in a command that is provided.
    We may want to share these routines with the PMI version 2 server */
 int PMIi_ReadCommand( int fd, PMI2_Command *cmd )
diff --git a/contribs/pmi2/pmi2_util.c b/contribs/pmi2/pmi2_util.c
index 8c9b4ee..4f424e8 100644
--- a/contribs/pmi2/pmi2_util.c
+++ b/contribs/pmi2/pmi2_util.c
@@ -17,6 +17,10 @@
 #include <stdarg.h>
 #include <string.h>
 #include <errno.h>
+#include <sys/types.h>
+#include <sys/stat.h>
+#include <sys/mman.h>
+#include <fcntl.h>
 
 #include "pmi2_util.h"
 
@@ -215,6 +219,36 @@ void PMI2U_chgval(const char *keystr, char *valstr) {
     }
 }
 
+int PMI2U_Create_shmem(PMI2ShmemRegion *shmem)
+{
+    struct stat file_status;
+
+    if ((shmem->fd = open(shmem->filename, O_RDONLY)) < 0) {
+        return -1;
+    }
+
+    do {
+        if (fstat(shmem->fd, &file_status) != 0) {
+            return -1;
+        }
+        usleep(1);
+    } while (file_status.st_size != shmem->filesize);
+
+    shmem->addr = mmap(0, shmem->filesize, (PROT_READ), (MAP_SHARED), shmem->fd, 0);
+
+    return 0;
+}
+
+int PMI2U_Destroy_shmem(PMI2ShmemRegion *shmem)
+{
+    munmap(shmem->addr, shmem->filesize);
+    close(shmem->fd);
+    shmem->fd = -1;
+    shmem->addr = NULL;
+    shmem->filesize = 0;
+    return 0;
+}
+
 /* This code is borrowed from mpich2-1.5/src/pm/util/safestr2.c.
    The reason is to keep the save code logic around strncpy() as
    as in the original PMI2 implementation.
@@ -273,3 +307,4 @@ MPIU_Strncpy(char *dest, const char *src, size_t n)
 	    return 1;
     }
 }
+
diff --git a/contribs/pmi2/pmi2_util.h b/contribs/pmi2/pmi2_util.h
index 9e7d8a8..c21dde3 100644
--- a/contribs/pmi2/pmi2_util.h
+++ b/contribs/pmi2/pmi2_util.h
@@ -87,6 +87,13 @@
         free ((void*)( pmi2u_chkmem_stk_[--pmi2u_chkmem_stk_sp_] )); \
 }
 
+typedef struct {
+    int fd;
+    void *addr;
+    char filename[256];
+    int filesize;
+} PMI2ShmemRegion;
+
 /* prototypes for PMIU routines */
 void PMI2U_Set_rank( int PMI_rank );
 void PMI2U_SetServer( void );
@@ -96,6 +103,8 @@ int  PMI2U_parse_keyvals( char *st );
 void PMI2U_dump_keyvals( void );
 char *PMI2U_getval( const char *keystr, char *valstr, int vallen );
 void PMI2U_chgval( const char *keystr, char *valstr );
+int PMI2U_Create_shmem(PMI2ShmemRegion *shmem);
+int PMI2U_Destroy_shmem(PMI2ShmemRegion *shmem);
 int MPIU_Strncpy(char *, const char *, size_t);
 
 #endif /* PMI2UTIL_H_INCLUDED */
diff --git a/contribs/pmi2/slurm/pmi2.h b/contribs/pmi2/slurm/pmi2.h
index e2daf72..d0a1e36 100644
--- a/contribs/pmi2/slurm/pmi2.h
+++ b/contribs/pmi2/slurm/pmi2.h
@@ -52,6 +52,8 @@ static const char NAMEUNPUBLISH_CMD[]     = "name-unpublish";
 static const char NAMEUNPUBLISHRESP_CMD[] = "name-unpublish-response";
 static const char NAMELOOKUP_CMD[]        = "name-lookup";
 static const char NAMELOOKUPRESP_CMD[]    = "name-lookup-response";
+static const char ALLGATHER_CMD[]         = "allgather";
+static const char ALLGATHERRESP_CMD[]     = "allgather-response";
 static const char RING_CMD[]              = "ring";
 static const char RINGRESP_CMD[]          = "ring-response";
 
@@ -86,6 +88,9 @@ static const char INFOVAL_KEY[]           = "infoval%d";
 static const char RING_COUNT_KEY[]        = "ring-count";
 static const char RING_LEFT_KEY[]         = "ring-left";
 static const char RING_RIGHT_KEY[]        = "ring-right";
+static const char SHMEMFILENAME_KEY[]     = "shmemfilename";
+static const char SHMEMFILESIZE_KEY[]     = "shmemfilesize";
+static const char USE_SHMEM_KEY[]         = "use-shmem";
 
 static const char TRUE_VAL[]              = "TRUE";
 static const char FALSE_VAL[]             = "FALSE";
@@ -491,6 +496,15 @@ int PMI2_KVS_Put(const char key[], const char value[]);
 @*/
 int PMI2_KVS_Fence(void);
 
+int PMI2_KVS_IFence(void);
+int PMI2_KVS_IWait(void);
+
+int PMI2_Iallgather(const char value[]);
+int PMI2_Iallgather_Wait(void *result);
+
+int PMI2_SHMEM_Iallgather(const char value[]);
+int PMI2_SHMEM_Iallgather_Wait(void **result);
+
 /*@
   PMI2_KVS_Get - returns the value associated with key in the key-value
       space associated with the job ID jobid
diff --git a/src/plugins/mpi/pmi2/Makefile.am b/src/plugins/mpi/pmi2/Makefile.am
index a7df063..b330e46 100644
--- a/src/plugins/mpi/pmi2/Makefile.am
+++ b/src/plugins/mpi/pmi2/Makefile.am
@@ -18,6 +18,7 @@ mpi_pmi2_la_SOURCES = mpi_pmi2.c \
 	spawn.c spawn.h \
 	tree.c tree.h \
 	nameserv.c nameserv.h \
+	shmem.c shmem.h \
 	ring.c ring.h
 
 mpi_pmi2_la_LDFLAGS = $(SO_LDFLAGS) $(PLUGIN_FLAGS)
diff --git a/src/plugins/mpi/pmi2/Makefile.in b/src/plugins/mpi/pmi2/Makefile.in
index 2450015..f88163b 100644
--- a/src/plugins/mpi/pmi2/Makefile.in
+++ b/src/plugins/mpi/pmi2/Makefile.in
@@ -174,7 +174,7 @@ am__installdirs = "$(DESTDIR)$(pkglibdir)"
 LTLIBRARIES = $(pkglib_LTLIBRARIES)
 mpi_pmi2_la_DEPENDENCIES = $(top_builddir)/src/slurmd/common/libslurmd_reverse_tree_math.la
 am_mpi_pmi2_la_OBJECTS = mpi_pmi2.lo agent.lo client.lo kvs.lo info.lo \
-	pmi1.lo pmi2.lo setup.lo spawn.lo tree.lo nameserv.lo ring.lo
+	pmi1.lo pmi2.lo setup.lo spawn.lo tree.lo nameserv.lo ring.lo shmem.lo
 mpi_pmi2_la_OBJECTS = $(am_mpi_pmi2_la_OBJECTS)
 AM_V_lt = $(am__v_lt_@AM_V@)
 am__v_lt_ = $(am__v_lt_@AM_DEFAULT_V@)
@@ -512,6 +512,7 @@ mpi_pmi2_la_SOURCES = mpi_pmi2.c \
 	spawn.c spawn.h \
 	tree.c tree.h \
 	nameserv.c nameserv.h \
+	shmem.c shmem.h \
 	ring.c ring.h
 
 mpi_pmi2_la_LDFLAGS = $(SO_LDFLAGS) $(PLUGIN_FLAGS)
@@ -608,6 +609,7 @@ distclean-compile:
 @AMDEP_TRUE@@am__include@ @am__quote@./$(DEPDIR)/setup.Plo@am__quote@
 @AMDEP_TRUE@@am__include@ @am__quote@./$(DEPDIR)/spawn.Plo@am__quote@
 @AMDEP_TRUE@@am__include@ @am__quote@./$(DEPDIR)/tree.Plo@am__quote@
+@AMDEP_TRUE@@am__include@ @am__quote@./$(DEPDIR)/shmem.Plo@am__quote@
 
 .c.o:
 @am__fastdepCC_TRUE@	$(AM_V_CC)$(COMPILE) -MT $@ -MD -MP -MF $(DEPDIR)/$*.Tpo -c -o $@ $<
diff --git a/src/plugins/mpi/pmi2/client.c b/src/plugins/mpi/pmi2/client.c
index d4f752b..7844fa3 100644
--- a/src/plugins/mpi/pmi2/client.c
+++ b/src/plugins/mpi/pmi2/client.c
@@ -521,6 +521,17 @@ client_req_get_bool(client_req_t *req, const char *key, bool *pval)
 
 /* ************************************************************ */
 
+extern int
+client_resp_send_raw(int fd, const void *buf, int len_buf)
+{
+	safe_write(fd, &len_buf, sizeof len_buf);
+	safe_write(fd, buf, len_buf);
+	return SLURM_SUCCESS;
+
+rwfail:
+	return SLURM_ERROR;
+}
+
 extern client_resp_t *
 client_resp_new(void)
 {
@@ -617,3 +628,48 @@ send_kvs_fence_resp_to_clients(int rc, char *errmsg)
 	client_resp_free(resp);
 	return rc;
 }
+
+extern int
+send_allgather_resp_to_clients(const void *buf, int len_buf)
+{
+	int i;
+	int rc = SLURM_SUCCESS;
+	debug2("send_allgather_resp_to_clients: %d bytes", len_buf);
+
+	for (i = 0; i < job_info.ltasks; i ++) {
+		rc = client_resp_send_raw(STEPD_PMI_SOCK(i), buf, len_buf);
+	}
+
+	return rc;
+}
+
+extern int
+send_shmem_allgather_resp_to_clients(int rc, char *errmsg, PMI2ShmemRegion *shmem)
+{
+    int i;
+    client_resp_t *resp;
+    char *msg;
+
+    resp = client_resp_new();
+    if (rc != 0 && errmsg != NULL) {
+        msg = _str_replace(errmsg, ';', '_');
+        client_resp_append(resp, CMD_KEY"="ALLGATHERRESP_CMD";"
+                RC_KEY"=%d;"ERRMSG_KEY"=%s;",
+                rc, msg);
+        xfree(msg);
+    } else {
+        client_resp_append(resp, CMD_KEY"="ALLGATHERRESP_CMD";"
+                RC_KEY"=%d;", rc);
+        if (shmem != NULL && shmem->fd != -1) {
+            client_resp_append(resp,
+                    SHMEMFILENAME_KEY"=%s;"SHMEMFILESIZE_KEY"=%d;",
+                    shmem->filename, shmem->filesize);
+        }
+    }
+
+    for (i = 0; i < job_info.ltasks; i ++) {
+        rc = client_resp_send(resp, STEPD_PMI_SOCK(i));
+    }
+    client_resp_free(resp);
+    return rc;
+}
diff --git a/src/plugins/mpi/pmi2/client.h b/src/plugins/mpi/pmi2/client.h
index 4d0b427..35a7e82 100644
--- a/src/plugins/mpi/pmi2/client.h
+++ b/src/plugins/mpi/pmi2/client.h
@@ -41,6 +41,7 @@
 #include <inttypes.h>
 
 #include "spawn.h"
+#include "shmem.h"
 
 #define PMI11_VERSION 1
 #define PMI11_SUBVERSION 1
@@ -82,6 +83,7 @@ extern spawn_subcmd_t *client_req_parse_spawn_subcmd(client_req_t *req);
 
 extern client_resp_t *client_resp_new(void);
 extern int  client_resp_send(client_resp_t *req, int fd);
+extern int  client_resp_send_raw(int fd, const void *buf, int len_buf);
 extern void client_resp_free(client_resp_t *resp);
 /* XXX: this requires CPP */
 #define client_resp_append(msg, fmt, ...) do { \
@@ -90,5 +92,7 @@ extern void client_resp_free(client_resp_t *resp);
 
 
 extern int send_kvs_fence_resp_to_clients(int rc, char *errmsg);
+extern int send_allgather_resp_to_clients(const void *buf, int len_buf);
+extern int send_shmem_allgather_resp_to_clients(int rc, char *errmsg, PMI2ShmemRegion *shmem);
 
 #endif	/* _CLIENT_H */
diff --git a/src/plugins/mpi/pmi2/kvs.c b/src/plugins/mpi/pmi2/kvs.c
index 1a99324..4760ffa 100644
--- a/src/plugins/mpi/pmi2/kvs.c
+++ b/src/plugins/mpi/pmi2/kvs.c
@@ -51,6 +51,8 @@ int children_to_wait = 0;
 int kvs_seq = 1; /* starting from 1 */
 int waiting_kvs_resp = 0;
 
+/* for allgather */
+int allg_seq = 1; /* starting from 1 */
 
 /* bucket of key-value pairs */
 typedef struct kvs_bucket {
@@ -66,6 +68,10 @@ static char *temp_kvs_buf = NULL;
 static int temp_kvs_cnt = 0;
 static int temp_kvs_size = 0;
 
+static char *allgather_buf = NULL;
+static int allgather_cnt = 0;
+static int allgather_size = 0;
+
 static int no_dup_keys = 0;
 
 #define TASKS_PER_BUCKET 8
@@ -325,3 +331,154 @@ kvs_clear(void)
 
 	return SLURM_SUCCESS;
 }
+
+/**************************************************************/
+
+extern int
+allgather_init(void)
+{
+	uint16_t cmd;
+	uint32_t nodeid, num_children, size;
+	Buf buf = NULL;
+
+	xfree(allgather_buf);
+	allgather_cnt = 0;
+	allgather_size = TEMP_KVS_SIZE_INC;
+	allgather_buf = xmalloc(allgather_size);
+
+	/* put the tree cmd here to simplify message sending */
+	if (in_stepd()) {
+		cmd = TREE_CMD_ALLGATHER;
+	} else {
+		cmd = TREE_CMD_ALLGATHER_RESP;
+	}
+
+	buf = init_buf(1024);
+	pack16(cmd, buf);
+	if (in_stepd()) {
+		nodeid = job_info.nodeid;
+		/* XXX: TBC */
+		num_children = tree_info.num_children + 1;
+
+		pack32((uint32_t)nodeid, buf); /* from_nodeid */
+		packstr(tree_info.this_node, buf); /* from_node */
+		pack32((uint32_t)num_children, buf); /* num_children */
+		pack32(allg_seq, buf);
+	} else {
+		pack32(allg_seq, buf);
+	}
+	size = get_buf_offset(buf);
+	if (allgather_cnt + size > allgather_size) {
+		allgather_size += TEMP_KVS_SIZE_INC;
+		xrealloc(allgather_buf, allgather_size);
+	}
+	memcpy(&allgather_buf[allgather_cnt], get_buf_data(buf), size);
+	allgather_cnt += size;
+	free_buf(buf);
+
+	tasks_to_wait = 0;
+	children_to_wait = 0;
+
+	return SLURM_SUCCESS;
+}
+
+extern int
+allgather_add(int rank, char *val)
+{
+	Buf buf;
+	uint32_t size;
+
+	if ( val == NULL )
+		return SLURM_SUCCESS;
+
+	buf = init_buf(PMI2_MAX_VALLEN + 3 * sizeof(uint32_t));
+	pack32(rank, buf);
+	packstr(val, buf);
+	size = get_buf_offset(buf);
+	if (allgather_cnt + size > allgather_size) {
+		allgather_size += TEMP_KVS_SIZE_INC;
+		xrealloc(allgather_buf, allgather_size);
+	}
+	memcpy(&allgather_buf[allgather_cnt], get_buf_data(buf), size);
+	allgather_cnt += size;
+	free_buf(buf);
+
+	return SLURM_SUCCESS;
+}
+
+extern int
+allgather_merge(Buf buf)
+{
+	char *data;
+	uint32_t offset, size;
+
+	size = remaining_buf(buf);
+	if (size == 0) {
+		return SLURM_SUCCESS;
+	}
+	data = get_buf_data(buf);
+	offset = get_buf_offset(buf);
+
+	if (allgather_cnt + size > allgather_size) {
+		allgather_size += size;
+		xrealloc(allgather_buf, allgather_size);
+	}
+	memcpy(&allgather_buf[allgather_cnt], &data[offset], size);
+	allgather_cnt += size;
+
+	return SLURM_SUCCESS;
+}
+
+extern int
+allgather_send(void)
+{
+	int rc = SLURM_ERROR, retry = 0;
+	unsigned int delay = 1;
+	hostlist_t hl = NULL;
+	char free_hl = 0;
+
+	if (! in_stepd()) {	/* srun */
+		hl = hostlist_create(job_info.step_nodelist);
+		free_hl = 1;
+	} else if (tree_info.parent_node != NULL) {
+		hl = hostlist_create(tree_info.parent_node);
+		free_hl = 1;
+	}
+
+
+	/* cmd included in allgather_buf */
+	allg_seq ++; /* expecting new kvs after now */
+
+	while (1) {
+		if (retry == 1) {
+			verbose("failed to send allgather kvs, rc=%d, retrying", rc);
+		}
+
+		if (! in_stepd()) {	/* srun */
+			rc = tree_msg_to_stepds(hl,
+						allgather_cnt,
+						allgather_buf);
+		} else if (tree_info.parent_node != NULL) {
+			/* non-first-level stepds */
+			rc = tree_msg_to_stepds(hl,
+						allgather_cnt,
+						allgather_buf);
+		} else {		/* first level stepds */
+			rc = tree_msg_to_srun(allgather_cnt, allgather_buf);
+		}
+		if (rc == SLURM_SUCCESS)
+			break;
+		retry ++;
+		if (retry >= MAX_RETRIES)
+			break;
+		/* wait, in case parent stepd / srun not ready */
+		sleep(delay);
+		delay *= 2;
+	}
+	allgather_init();	/* clear old allgather kvs */
+	if( free_hl ){
+		hostlist_destroy(hl);
+	}
+	return rc;
+}
+
diff --git a/src/plugins/mpi/pmi2/kvs.h b/src/plugins/mpi/pmi2/kvs.h
index de88e9b..e3c4618 100644
--- a/src/plugins/mpi/pmi2/kvs.h
+++ b/src/plugins/mpi/pmi2/kvs.h
@@ -43,10 +43,10 @@
 #include "src/common/xstring.h"
 #include "src/common/pack.h"
 
-
 extern int tasks_to_wait;
 extern int children_to_wait;
 extern int kvs_seq;
+extern int allg_seq;
 extern int waiting_kvs_resp;
 
 extern int   temp_kvs_init(void);
@@ -54,6 +54,11 @@ extern int   temp_kvs_add(char *key, char *val);
 extern int   temp_kvs_merge(Buf buf);
 extern int   temp_kvs_send(void);
 
+extern int   allgather_init(void);
+extern int   allgather_add(int rank, char *val);
+extern int   allgather_merge(Buf buf);
+extern int   allgather_send(void);
+
 extern int   kvs_init(void);
 extern char *kvs_get(char *key);
 extern int   kvs_put(char *key, char *val);
diff --git a/src/plugins/mpi/pmi2/pmi.h b/src/plugins/mpi/pmi2/pmi.h
index 67db227..1db0139 100644
--- a/src/plugins/mpi/pmi2/pmi.h
+++ b/src/plugins/mpi/pmi2/pmi.h
@@ -148,6 +148,8 @@
 #define NAMELOOKUPRESP_CMD     "name-lookup-response"
 #define SPAWN_CMD              "spawn"
 #define SPAWNRESP_CMD          "spawn-response"
+#define ALLGATHER_CMD          "allgather"
+#define ALLGATHERRESP_CMD      "allgather-response"
 #define RING_CMD               "ring"
 #define RINGRESP_CMD           "ring-response"
 
@@ -209,6 +211,9 @@
 #define RING_COUNT_KEY    "ring-count"
 #define RING_LEFT_KEY     "ring-left"
 #define RING_RIGHT_KEY    "ring-right"
+#define SHMEMFILENAME_KEY "shmemfilename"
+#define SHMEMFILESIZE_KEY "shmemfilesize"
+#define USE_SHMEM_KEY     "use-shmem"
 
 #define TRUE_VAL          "TRUE"
 #define FALSE_VAL         "FALSE"
@@ -238,7 +243,6 @@
 #define PMI2_PMI_DEBUGGED_ENV   "PMI_DEBUG"
 #define PMI2_KVS_NO_DUP_KEYS_ENV "SLURM_PMI_KVS_NO_DUP_KEYS"
 
-
 extern int handle_pmi1_cmd(int fd, int lrank);
 extern int handle_pmi2_cmd(int fd, int lrank);
 
diff --git a/src/plugins/mpi/pmi2/pmi1.c b/src/plugins/mpi/pmi2/pmi1.c
index 758441f..ba0b0d8 100644
--- a/src/plugins/mpi/pmi2/pmi1.c
+++ b/src/plugins/mpi/pmi2/pmi1.c
@@ -61,6 +61,7 @@
 #include "kvs.h"
 #include "agent.h"
 #include "nameserv.h"
+#include "shmem.h"
 
 /* client command handlers */
 static int _handle_get_maxes(int fd, int lrank, client_req_t *req);
diff --git a/src/plugins/mpi/pmi2/pmi2.c b/src/plugins/mpi/pmi2/pmi2.c
index 99c61d5..c28b39d 100644
--- a/src/plugins/mpi/pmi2/pmi2.c
+++ b/src/plugins/mpi/pmi2/pmi2.c
@@ -57,6 +57,7 @@
 #include "agent.h"
 #include "nameserv.h"
 #include "ring.h"
+#include "shmem.h"
 
 /* PMI2 command handlers */
 static int _handle_fullinit(int fd, int lrank, client_req_t *req);
@@ -76,6 +77,7 @@ static int _handle_name_publish(int fd, int lrank, client_req_t *req);
 static int _handle_name_unpublish(int fd, int lrank, client_req_t *req);
 static int _handle_name_lookup(int fd, int lrank, client_req_t *req);
 static int _handle_spawn(int fd, int lrank, client_req_t *req);
+static int _handle_allgather(int fd, int lrank, client_req_t *req);
 
 
 static struct {
@@ -99,6 +101,7 @@ static struct {
 	{ NAMEUNPUBLISH_CMD,     _handle_name_unpublish },
 	{ NAMELOOKUP_CMD,        _handle_name_lookup },
 	{ SPAWN_CMD,             _handle_spawn },
+	{ ALLGATHER_CMD,         _handle_allgather },
 	{ NULL, NULL},
 };
 
@@ -163,6 +166,7 @@ static int
 _handle_finalize(int fd, int lrank, client_req_t *req)
 {
 	client_resp_t *resp;
+    static int count = 0;
 	int rc = 0;
 
 	resp = client_resp_new();
@@ -174,6 +178,9 @@ _handle_finalize(int fd, int lrank, client_req_t *req)
 	shutdown(fd, SHUT_RDWR);
 	close(fd);
 	task_finalize(lrank);
+    if (++count == job_info.ltasks) {
+        kvs_destroy_shmem(&PMI2_Shmem_allgather);
+    }
 	return rc;
 }
 
@@ -186,8 +193,9 @@ _handle_abort(int fd, int lrank, client_req_t *req)
 	debug3("mpi/pmi2: in _handle_abort");
 	client_req_parse_body(req);
 	client_req_get_bool(req, ISWORLD_KEY, &is_world);
-	/* no response needed. just cancel the job step if required */
+    /* no response needed. just cancel the job step if required */
 	if (is_world) {
+        kvs_destroy_shmem(&PMI2_Shmem_allgather);
 		slurm_kill_job_step(job_info.jobid, job_info.stepid, SIGKILL);
 	}
 	return rc;
@@ -226,6 +234,36 @@ _handle_job_disconnect(int fd, int lrank, client_req_t *req)
 }
 
 static int
+_handle_allgather(int fd, int lrank, client_req_t *req)
+{
+	int rc = SLURM_SUCCESS;
+	char *val = NULL;
+    bool use_shmem = false;
+
+	debug2("mpi/pmi2: in _handle_allgather");
+	client_req_parse_body(req);
+	client_req_get_str(req, VALUE_KEY, &val);
+	client_req_get_bool(req, USE_SHMEM_KEY, &use_shmem);
+    use_shmem_allgather = use_shmem;
+
+	rc = allgather_add(job_info.gtids[lrank], val);
+
+	if (tasks_to_wait == 0 && children_to_wait == 0) {
+		tasks_to_wait = job_info.ltasks;
+		children_to_wait = tree_info.num_children;
+	}
+	tasks_to_wait --;
+
+	if (tasks_to_wait == 0 && children_to_wait == 0) {
+		rc = allgather_send();
+		waiting_kvs_resp = 1;
+	}
+
+	debug2("mpi/pmi2: out _handle_allgather");
+	return rc;
+}
+
+static int
 _handle_ring(int fd, int lrank, client_req_t *req)
 {
 	int rc = SLURM_SUCCESS;
diff --git a/src/plugins/mpi/pmi2/setup.c b/src/plugins/mpi/pmi2/setup.c
index 21740de..6ba958f 100644
--- a/src/plugins/mpi/pmi2/setup.c
+++ b/src/plugins/mpi/pmi2/setup.c
@@ -249,7 +249,9 @@ _setup_stepd_tree_info(const stepd_step_rec_t *job, char ***env)
 	/* init kvs seq to 0. TODO: reduce array size */
 	tree_info.children_kvs_seq = xmalloc(sizeof(uint32_t) *
 					     job_info.nnodes);
-
+	tree_info.children_allg_seq = xmalloc(sizeof(uint32_t) *
+					     job_info.nnodes);
+	
 	return SLURM_SUCCESS;
 }
 
@@ -318,6 +320,10 @@ _setup_stepd_kvs(const stepd_step_rec_t *job, char ***env)
 	if (rc != SLURM_SUCCESS)
 		return rc;
 
+	rc = allgather_init();
+	if (rc != SLURM_SUCCESS)
+		return rc;
+
 	rc = kvs_init();
 	if (rc != SLURM_SUCCESS)
 		return rc;
@@ -644,7 +650,9 @@ _setup_srun_tree_info(const mpi_plugin_client_info_t *job)
 	/* init kvs seq to 0. TODO: reduce array size */
 	tree_info.children_kvs_seq = xmalloc(sizeof(uint32_t) *
 					     job_info.nnodes);
-
+	tree_info.children_allg_seq = xmalloc(sizeof(uint32_t) *
+					     job_info.nnodes);
+	
 	return SLURM_SUCCESS;
 }
 
@@ -668,6 +676,7 @@ _setup_srun_kvs(const mpi_plugin_client_info_t *job)
 
 	kvs_seq = 1;
 	rc = temp_kvs_init();
+	rc = allgather_init();
 	return rc;
 }
 
diff --git a/src/plugins/mpi/pmi2/setup.h b/src/plugins/mpi/pmi2/setup.h
index cd46da6..e13c6f3 100644
--- a/src/plugins/mpi/pmi2/setup.h
+++ b/src/plugins/mpi/pmi2/setup.h
@@ -88,6 +88,7 @@ typedef struct pmi2_tree_info {
 	uint16_t pmi_port;	 /* PMI2 comm port of this srun */
 	slurm_addr_t *srun_addr; /* PMI2 comm address parent srun */
 	uint32_t *children_kvs_seq; /* sequence number of children nodes */
+	uint32_t *children_allg_seq; /* sequence number of children nodes */
 } pmi2_tree_info_t;
 
 
diff --git a/src/plugins/mpi/pmi2/shmem.c b/src/plugins/mpi/pmi2/shmem.c
new file mode 100644
index 0000000..fd7aa42
--- /dev/null
+++ b/src/plugins/mpi/pmi2/shmem.c
@@ -0,0 +1,115 @@
+/*****************************************************************************\
+ **  shmem.c - Shared Memory manipulation functions
+ *****************************************************************************
+ *  Copyright (C) 2016 The Ohio State University
+ *  This file was developed by the team members of The
+ *  Ohio State University's Network-Based Computing Laboratory (NBCL),
+ *  headed by Professor Dhabaleswar K. (DK) Panda.
+ *  Contact:
+ *  Prof. Dhabaleswar K. (DK) Panda
+ *  Dept. of Computer Science and Engineering
+ *  The Ohio State University
+ *  2015 Neil Avenue
+ *  Columbus, OH - 43210-1277
+ *  Tel: (614)-292-5199; Fax: (614)-292-2911
+ *  E-mail:panda@cse.ohio-state.edu
+ *  All rights reserved.
+ *
+ *  This file is part of SLURM, a resource management program.
+ *  For details, see <http://slurm.schedmd.com/>.
+ *  Please also read the included file: DISCLAIMER.
+ *
+ *  SLURM is free software; you can redistribute it and/or modify it under
+ *  the terms of the GNU General Public License as published by the Free
+ *  Software Foundation; either version 2 of the License, or (at your option)
+ *  any later version.
+ *
+ *  In addition, as a special exception, the copyright holders give permission
+ *  to link the code of portions of this program with the OpenSSL library under
+ *  certain conditions as described in each individual source file, and
+ *  distribute linked combinations including the two. You must obey the GNU
+ *  General Public License in all respects for all of the code used other than
+ *  OpenSSL. If you modify file(s) with this exception, you may extend this
+ *  exception to your version of the file(s), but you are not obligated to do
+ *  so. If you do not wish to do so, delete this exception statement from your
+ *  version.  If you delete this exception statement from all source files in
+ *  the program, then also delete it here.
+ *
+ *  SLURM is distributed in the hope that it will be useful, but WITHOUT ANY
+ *  WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
+ *  FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
+ *  details.
+ *
+ *  You should have received a copy of the GNU General Public License along
+ *  with SLURM; if not, write to the Free Software Foundation, Inc.,
+ *  51 Franklin Street, Fifth Floor, Boston, MA 02110-1301  USA.
+\*****************************************************************************/
+
+#include <stdlib.h>
+#include <unistd.h>
+#include <fcntl.h>
+#include <sys/mman.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+
+#include "shmem.h"
+
+#include "src/common/slurm_xlator.h"
+#include "src/common/slurm_protocol_interface.h"
+#include "src/common/slurm_protocol_api.h"
+#include "src/common/xmalloc.h"
+
+int use_shmem_allgather = 0;
+PMI2ShmemRegion PMI2_Shmem_allgather = { .fd=-1, .addr=NULL, .filesize=0 };
+
+extern int
+kvs_create_shmem(PMI2ShmemRegion *shmem)
+{
+    struct stat file_status;
+    debug("in kvs_create_shmem: filename: %s, filesize: %d", shmem->filename, shmem->filesize);
+
+    if ((shmem->fd = open(shmem->filename, O_RDWR | O_CREAT, S_IRWXU | S_IRGRP | S_IROTH)) < 0) {
+        error("Can not open shmem file: %s, error: %m", shmem->filename);
+        return -1;
+    }
+
+    if(ftruncate(shmem->fd, 0)
+            || ftruncate(shmem->fd, shmem->filesize)
+            || lseek(shmem->fd, 0, SEEK_SET)
+      ) {
+        error("Can not ftruncate shmem file: %s, error: %m", shmem->filename);
+        return -1;
+    }
+
+    do {
+        if (fstat(shmem->fd, &file_status) != 0) {
+            error("Can not fstat shmem file: %s, error: %m", shmem->filename);
+            return -1;
+        }
+        usleep(1);
+    } while (file_status.st_size != shmem->filesize);
+
+    shmem->addr = mmap(0, shmem->filesize, (PROT_READ | PROT_WRITE), (MAP_SHARED), shmem->fd, 0);
+    if (shmem->addr == MAP_FAILED) {
+        error("Can not mmap shmem file: %s, size: %d, error: %m", shmem->filename, shmem->filesize);
+    }
+    memset(shmem->addr, 0, shmem->filesize);
+
+    debug("out kvs_create_shmem: filename: %s, filesize: %d, fd: %d", shmem->filename, shmem->filesize, shmem->fd);
+    return 0;
+}
+
+extern int
+kvs_destroy_shmem(PMI2ShmemRegion *shmem)
+{
+    if (shmem->fd != -1) {
+        munmap(shmem->addr, shmem->filesize);
+        close(shmem->fd);
+        unlink(shmem->filename);
+        shmem->addr = NULL;
+        shmem->filesize = 0;
+        shmem->fd = -1;
+    }
+    return 0;
+}
+
diff --git a/src/plugins/mpi/pmi2/shmem.h b/src/plugins/mpi/pmi2/shmem.h
new file mode 100644
index 0000000..b096028
--- /dev/null
+++ b/src/plugins/mpi/pmi2/shmem.h
@@ -0,0 +1,70 @@
+/*****************************************************************************\
+ **  shmem.h - Shared Memory manipulation functions
+ *****************************************************************************
+ *  Copyright (C) 2016 The Ohio State University
+ *  This file was developed by the team members of The
+ *  Ohio State University's Network-Based Computing Laboratory (NBCL),
+ *  headed by Professor Dhabaleswar K. (DK) Panda.
+ *  Contact:
+ *  Prof. Dhabaleswar K. (DK) Panda
+ *  Dept. of Computer Science and Engineering
+ *  The Ohio State University
+ *  2015 Neil Avenue
+ *  Columbus, OH - 43210-1277
+ *  Tel: (614)-292-5199; Fax: (614)-292-2911
+ *  E-mail:panda@cse.ohio-state.edu
+ *  All rights reserved.
+ *
+ *  This file is part of SLURM, a resource management program.
+ *  For details, see <http://slurm.schedmd.com/>.
+ *  Please also read the included file: DISCLAIMER.
+ *
+ *  SLURM is free software; you can redistribute it and/or modify it under
+ *  the terms of the GNU General Public License as published by the Free
+ *  Software Foundation; either version 2 of the License, or (at your option)
+ *  any later version.
+ *
+ *  In addition, as a special exception, the copyright holders give permission
+ *  to link the code of portions of this program with the OpenSSL library under
+ *  certain conditions as described in each individual source file, and
+ *  distribute linked combinations including the two. You must obey the GNU
+ *  General Public License in all respects for all of the code used other than
+ *  OpenSSL. If you modify file(s) with this exception, you may extend this
+ *  exception to your version of the file(s), but you are not obligated to do
+ *  so. If you do not wish to do so, delete this exception statement from your
+ *  version.  If you delete this exception statement from all source files in
+ *  the program, then also delete it here.
+ *
+ *  SLURM is distributed in the hope that it will be useful, but WITHOUT ANY
+ *  WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
+ *  FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
+ *  details.
+ *
+ *  You should have received a copy of the GNU General Public License along
+ *  with SLURM; if not, write to the Free Software Foundation, Inc.,
+ *  51 Franklin Street, Fifth Floor, Boston, MA 02110-1301  USA.
+\*****************************************************************************/
+
+#ifndef _SHMEM_H
+#define _SHMEM_H
+
+#if     HAVE_CONFIG_H
+#  include "config.h"
+#endif
+
+#define PMI2_SHMEM_FILENAME_ALLGATHER  "/tmp/SLURM_PMI2_SHMEM_ALLG_%llu_%llu.tmp"
+
+typedef struct {
+    int fd;
+    void *addr;
+    char filename[256];
+    int filesize;
+} PMI2ShmemRegion;
+
+extern int use_shmem_allgather;
+extern PMI2ShmemRegion PMI2_Shmem_allgather;
+
+extern int   kvs_create_shmem(PMI2ShmemRegion *shmem);
+extern int   kvs_destroy_shmem(PMI2ShmemRegion *shmem);
+
+#endif	/* _SHMEM_H */
diff --git a/src/plugins/mpi/pmi2/tree.c b/src/plugins/mpi/pmi2/tree.c
index 30ad957..4fca342 100644
--- a/src/plugins/mpi/pmi2/tree.c
+++ b/src/plugins/mpi/pmi2/tree.c
@@ -58,6 +58,7 @@
 #include "pmi.h"
 #include "nameserv.h"
 #include "ring.h"
+#include "shmem.h"
 
 static int _handle_kvs_fence(int fd, Buf buf);
 static int _handle_kvs_fence_resp(int fd, Buf buf);
@@ -66,6 +67,8 @@ static int _handle_spawn_resp(int fd, Buf buf);
 static int _handle_name_publish(int fd, Buf buf);
 static int _handle_name_unpublish(int fd, Buf buf);
 static int _handle_name_lookup(int fd, Buf buf);
+static int _handle_allgather(int fd, Buf buf);
+static int _handle_allgather_resp(int fd, Buf buf);
 static int _handle_ring(int fd, Buf buf);
 static int _handle_ring_resp(int fd, Buf buf);
 
@@ -81,6 +84,8 @@ static int (*tree_cmd_handlers[]) (int fd, Buf buf) = {
 	_handle_name_publish,
 	_handle_name_unpublish,
 	_handle_name_lookup,
+	_handle_allgather,
+	_handle_allgather_resp,
 	_handle_ring,
 	_handle_ring_resp,
 	NULL
@@ -94,11 +99,15 @@ static char *tree_cmd_names[] = {
 	"TREE_CMD_NAME_PUBLISH",
 	"TREE_CMD_NAME_UNPUBLISH",
 	"TREE_CMD_NAME_LOOKUP",
+	"TREE_CMD_ALLGATHER",
+	"TREE_CMD_ALLGATHER_RESP",
 	"TREE_CMD_RING",
 	"TREE_CMD_RING_RESP",
 	NULL,
 };
 
+/**************************************************************/
+
 static int
 _handle_kvs_fence(int fd, Buf buf)
 {
@@ -223,6 +232,160 @@ unpack_error:
 	goto resp;
 }
 
+/**************************************************************/
+
+static int
+_sort_gathered_values(Buf buf, char *addr, int maxlen)
+{
+    int rc = SLURM_SUCCESS;
+    uint32_t rank, temp32;
+
+    debug3("mpi/pmi2: _sort_gathered_values");
+
+    temp32 = remaining_buf(buf);
+    while (remaining_buf(buf) > 0) {
+        safe_unpack32(&rank, buf);
+        unpackmem(&addr[rank*maxlen], &temp32, buf);
+    }
+
+out:
+    return rc;
+
+unpack_error:
+    error("mpi/pmi2: unpack error in allgather resp");
+    rc = SLURM_ERROR;
+    goto out;
+}
+
+
+static int
+_handle_allgather(int fd, Buf buf)
+{
+	uint32_t from_nodeid, num_children, temp32, seq;
+	char *from_node = NULL;
+	int rc = SLURM_SUCCESS;
+
+	safe_unpack32(&from_nodeid, buf);
+	safe_unpackstr_xmalloc(&from_node, &temp32, buf);
+	safe_unpack32(&num_children, buf);
+	safe_unpack32(&seq, buf);
+
+	debug3("mpi/pmi2: in _handle_allgather, from node %u(%s) representing"
+	       " %u offspring, seq=%u", from_nodeid, from_node, num_children,
+	       seq);
+	if (seq != allg_seq) {
+		error("mpi/pmi2: invalid kvs seq from node %u(%s) ignored, "
+		      "expect %u got %u", 
+		      from_nodeid, from_node, allg_seq, seq);
+		goto out;
+	}
+	if (seq == tree_info.children_allg_seq[from_nodeid]) {
+		info("mpi/pmi2: duplicate ALLGATHER request from node %u(%s) "
+		      "ignored, seq=%u", from_nodeid, from_node, seq);
+		goto out;
+	}
+	tree_info.children_allg_seq[from_nodeid] = seq;
+	
+	if (tasks_to_wait == 0 && children_to_wait == 0) {
+		tasks_to_wait = job_info.ltasks;
+		children_to_wait = tree_info.num_children;
+	}
+	children_to_wait -= num_children;
+
+	allgather_merge(buf);
+
+	if (children_to_wait == 0 && tasks_to_wait == 0) {
+
+		rc = allgather_send();
+		if (rc != SLURM_SUCCESS) {
+			if (in_stepd()) {
+				error("mpi/pmi2: failed to send allgather kvs"
+				      " to %s",
+				      tree_info.parent_node ?: "srun");
+			} else {
+				error("mpi/pmi2: failed to send allgather kvs"
+				      " to compute nodes");
+			}
+			/* cancel the step to avoid tasks hang */
+			slurm_kill_job_step(job_info.jobid, job_info.stepid,
+					    SIGKILL);
+		} else {
+			if (in_stepd())
+				waiting_kvs_resp = 1;
+		}
+	}
+	debug3("mpi/pmi2: out _handle_allgather, tasks_to_wait=%d, "
+	       "children_to_wait=%d", tasks_to_wait, children_to_wait);
+out:
+	xfree(from_node);
+	return rc;
+
+unpack_error:
+	error("mpi/pmi2: failed to unpack allgather message");
+	rc = SLURM_ERROR;
+	goto out;
+}
+
+static int
+_handle_allgather_resp(int fd, Buf buf)
+{
+	int rc = SLURM_SUCCESS;
+	uint32_t seq;
+	int result_len, maxlen = PMI2_MAX_VALLEN;
+	char *result_buf = NULL, *errmsg;
+    PMI2ShmemRegion *shmem = &PMI2_Shmem_allgather;
+
+	debug3("mpi/pmi2: in _handle_allgather_resp");
+
+	safe_unpack32(&seq, buf);
+	debug("_handle_allgather_resp, seq: %u", seq);
+	if (seq != allg_seq - 1) {
+		error("mpi/pmi2: invalid kvs seq from srun, expect %u"
+		      " got %u", allg_seq - 1, seq);
+		return SLURM_ERROR;
+	}
+	if (! waiting_kvs_resp) {
+		debug("mpi/pmi2: duplicate ALLGATHER_RESP from srun ignored");
+		return rc;
+	} else {
+		waiting_kvs_resp = 0;
+	}
+
+    if (use_shmem_allgather) {
+        kvs_destroy_shmem(shmem);
+        sprintf(shmem->filename, PMI2_SHMEM_FILENAME_ALLGATHER, job_info.jobid, job_info.stepid);
+        shmem->filesize = job_info.ntasks * maxlen * sizeof (char);
+        debug("Allgather filename: %s, size: %d", shmem->filename, shmem->filesize);
+        kvs_create_shmem(shmem);
+        _sort_gathered_values(buf, shmem->addr, maxlen);
+    } else {
+        result_len = job_info.ntasks * maxlen * sizeof (char);
+        result_buf = xmalloc(result_len);
+        memset(result_buf, '\0', result_len);
+        _sort_gathered_values(buf, result_buf, maxlen);
+    }
+
+resp:
+    if (use_shmem_allgather) {
+	    rc = send_shmem_allgather_resp_to_clients(rc, errmsg, shmem);
+    } else {
+	    rc = send_allgather_resp_to_clients(result_buf, result_len);
+    }
+	if (rc != SLURM_SUCCESS) {
+		slurm_kill_job_step(job_info.jobid, job_info.stepid, SIGKILL);
+	}
+	xfree(result_buf);
+	return rc;
+
+unpack_error:
+	error("mpi/pmi2: unpack error in allgather resp");
+	errmsg = "mpi/pmi2: unpack error in allgather resp";
+	rc = SLURM_ERROR;
+	goto resp;
+}
+
+/**************************************************************/
+
 /* only called in srun */
 static int
 _handle_spawn(int fd, Buf buf)
diff --git a/src/plugins/mpi/pmi2/tree.h b/src/plugins/mpi/pmi2/tree.h
index 5848158..1420624 100644
--- a/src/plugins/mpi/pmi2/tree.h
+++ b/src/plugins/mpi/pmi2/tree.h
@@ -46,6 +46,8 @@ enum {
 	TREE_CMD_NAME_PUBLISH,
 	TREE_CMD_NAME_UNPUBLISH,
 	TREE_CMD_NAME_LOOKUP,
+	TREE_CMD_ALLGATHER,
+	TREE_CMD_ALLGATHER_RESP,
 	TREE_CMD_RING,
 	TREE_CMD_RING_RESP,
 	TREE_CMD_COUNT
-- 
1.8.3.1

