teams: Add SHMEMX_TEAM_HOSTS pre-defined team

mpp/shmemx-def.h

@@ -22,6 +22,15 @@ typedef struct {
     uint64_t target;
 } shmemx_pcntr_t;
 
+#if SHMEM_HAVE_ATTRIBUTE_VISIBILITY == 1
+    __attribute__((visibility("default"))) extern shmem_team_t SHMEMX_TEAM_HOST;
+#else
+    extern shmem_team_t SHMEMX_TEAM_HOST;
+#endif
+
+#define SHMEMX_EXTERNAL_HEAP_ZE 0
+#define SHMEMX_EXTERNAL_HEAP_CUDA 1
+
 #ifdef __cplusplus
 }
 #endif

src/shmem_team.c

@@ -20,7 +20,8 @@
 
 #define SHMEM_TEAM_WORLD_INDEX   0
 #define SHMEM_TEAM_SHARED_INDEX  1
-#define SHMEM_TEAMS_MIN          2
+#define SHMEM_TEAM_HOST_INDEX    2
+#define SHMEM_TEAMS_MIN          3
 
 #define N_PSYNC_BYTES             8
 #define PSYNC_CHUNK_SIZE          (N_PSYNCS_PER_TEAM * SHMEM_SYNC_SIZE)
@@ -32,6 +33,9 @@ shmem_team_t SHMEM_TEAM_WORLD = (shmem_team_t) &shmem_internal_team_world;
 shmem_internal_team_t shmem_internal_team_shared;
 shmem_team_t SHMEM_TEAM_SHARED = (shmem_team_t) &shmem_internal_team_shared;
 
+shmem_internal_team_t shmem_internal_team_host;
+shmem_team_t SHMEMX_TEAM_HOST = (shmem_team_t) &shmem_internal_team_host;
+
 shmem_internal_team_t **shmem_internal_team_pool;
 long *shmem_internal_psync_pool;
 long *shmem_internal_psync_barrier_pool;
@@ -93,6 +97,16 @@ int shmem_internal_team_init(void)
         shmem_internal_team_shared.psync_avail[i] = 1;
     SHMEM_TEAM_SHARED = (shmem_team_t) &shmem_internal_team_shared;
 
+    /* Initialize SHMEM_TEAM_HOST */
+    shmem_internal_team_host.psync_idx       = SHMEM_TEAM_HOST_INDEX;
+    shmem_internal_team_host.my_pe           = shmem_internal_my_pe;
+    shmem_internal_team_host.config_mask     = 0;
+    shmem_internal_team_host.contexts_len    = 0;
+    memset(&shmem_internal_team_host.config, 0, sizeof(shmem_team_config_t));
+    for (size_t i = 0; i < N_PSYNCS_PER_TEAM; i++)
+        shmem_internal_team_host.psync_avail[i] = 1;
+    SHMEMX_TEAM_HOST = (shmem_team_t) &shmem_internal_team_host;
+
     if (shmem_internal_params.TEAM_SHARED_ONLY_SELF) {
         shmem_internal_team_shared.start         = shmem_internal_my_pe;
         shmem_internal_team_shared.stride        = 1;
@@ -123,6 +137,25 @@ int shmem_internal_team_init(void)
                   shmem_internal_team_shared.size);
     }
 
+    int start = -1, stride = -1, size = 0;
+
+    for (int pe = 0; pe < shmem_internal_num_pes; pe++) {
+        int ret = shmem_runtime_get_node_rank(pe);
+        if (ret < 0) continue;
+
+        ret = check_for_linear_stride(pe, &start, &stride, &size);
+        if (ret < 0) return ret;
+    }
+    shmem_internal_assert(size > 0 && size == shmem_runtime_get_node_size());
+
+    shmem_internal_team_host.start = start;
+    shmem_internal_team_host.stride = (stride == -1) ? 1 : stride;
+    shmem_internal_team_host.size = size;
+
+    DEBUG_MSG("SHMEMX_TEAM_HOST: start=%d, stride=%d, size=%d\n",
+              shmem_internal_team_host.start, shmem_internal_team_host.stride,
+              shmem_internal_team_host.size);
+
     if (shmem_internal_params.TEAMS_MAX > N_PSYNC_BYTES * CHAR_BIT) {
         RETURN_ERROR_MSG("Requested %ld teams, but only %d are supported\n",
                          shmem_internal_params.TEAMS_MAX, N_PSYNC_BYTES * CHAR_BIT);
@@ -140,14 +173,15 @@ int shmem_internal_team_init(void)
     }
     shmem_internal_team_pool[SHMEM_TEAM_WORLD_INDEX] = &shmem_internal_team_world;
     shmem_internal_team_pool[SHMEM_TEAM_SHARED_INDEX] = &shmem_internal_team_shared;
+    shmem_internal_team_pool[SHMEM_TEAM_HOST_INDEX] = &shmem_internal_team_host;
 
     /* Allocate pSync pool, each with the maximum possible size requirement */
     /* Create two pSyncs per team for back-to-back collectives and one for barriers.
      * Array organization:
      *
-     * [ (world) (shared) (team 1) (team 2) ...  (world) (shared) (team 1) (team 2) ... ]
-     *  <----------- groups 1 & 2-------------->|<------------- group 3 ---------------->
-     *  <--- (bcast, collect, reduce, etc.) --->|<------ (barriers and syncs) ---------->
+     * [ (world) (shared) (host) (team 1) (team 2) ...  (world) (shared) (host) (team 1) (team 2) ... ]
+     *  <----------- groups 1 & 2------------------->|<------------- group 3 ------------------------->
+     *  <--- (bcast, collect, reduce, etc.) -------->|<------ (barriers and syncs) ------------------->
      * */
     long psync_len = shmem_internal_params.TEAMS_MAX * (PSYNC_CHUNK_SIZE + SHMEM_SYNC_SIZE);
     shmem_internal_psync_pool = shmem_internal_shmalloc(sizeof(long) * psync_len);
@@ -171,9 +205,10 @@ int shmem_internal_team_init(void)
         shmem_internal_bit_set(psync_pool_avail, N_PSYNC_BYTES, i);
     }
 
-    /* Set the bits for SHMEM_TEAM_WORLD and SHMEM_TEAM_SHARED to 0: */
+    /* Set the bits for SHMEM_TEAM_WORLD, SHMEM_TEAM_SHARED, and SHMEMX_TEAM_HOST to 0: */
     shmem_internal_bit_clear(psync_pool_avail, N_PSYNC_BYTES, SHMEM_TEAM_WORLD_INDEX);
     shmem_internal_bit_clear(psync_pool_avail, N_PSYNC_BYTES, SHMEM_TEAM_SHARED_INDEX);
+    shmem_internal_bit_clear(psync_pool_avail, N_PSYNC_BYTES, SHMEM_TEAM_HOST_INDEX);
 
     /* Initialize an integer used to agree on an equal return value across PEs in team creation: */
     team_ret_val = shmem_internal_shmalloc(sizeof(int) * 2);
@@ -456,7 +491,8 @@ int shmem_internal_team_destroy(shmem_internal_team_t *team)
     shmem_internal_team_pool[team->psync_idx] = NULL;
     free(team->contexts);
 
-    if (team != &shmem_internal_team_world && team != &shmem_internal_team_shared) {
+    if (team != &shmem_internal_team_world && team != &shmem_internal_team_shared &&
+        team != &shmem_internal_team_host) {
         free(team);
     }
 

src/shmem_team.h

@@ -31,6 +31,7 @@ typedef struct shmem_internal_team_t shmem_internal_team_t;
 
 extern shmem_internal_team_t shmem_internal_team_world;
 extern shmem_internal_team_t shmem_internal_team_shared;
+extern shmem_internal_team_t shmem_internal_team_host;
 
 enum shmem_internal_team_op_t {
     SYNC = 0,

test/shmemx/Makefile.am

@@ -20,8 +20,9 @@ endif
 
 if SHMEMX_TESTS
 check_PROGRAMS += \
-    perf_counter \
-    shmem_malloc_with_hints
+	perf_counter \
+	shmemx_team_host \
+	shmem_malloc_with_hints
 
 if HAVE_PTHREADS
 check_PROGRAMS += \

test/shmemx/shmemx_team_host.c

@@ -0,0 +1,85 @@
+/*
+ * Copyright (c) 2023 Intel Corporation. All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *
+ * - Redistributions of source code must retain the above copyright
+ *   notice, this list of conditions and the following disclaimer.
+ *
+ * - Redistributions in binary form must reproduce the above copyright
+ *   notice, this list of conditions and the following disclaimer listed
+ *   in this license in the documentation and/or other materials
+ *   provided with the distribution.
+ *
+ * - Neither the name of the copyright holders nor the names of its
+ *   contributors may be used to endorse or promote products derived from
+ *   this software without specific prior written permission.
+ *
+ * The copyright holders provide no reassurances that the source code
+ * provided does not infringe any patent, copyright, or any other
+ * intellectual property rights of third parties.  The copyright holders
+ * disclaim any liability to any recipient for claims brought against
+ * recipient by any third party for infringement of that parties
+ * intellectual property rights.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <shmem.h>
+#include <shmemx.h>
+
+
+int main(void)
+{
+    static long lock = 0;
+
+    shmem_init();
+    int me = shmem_my_pe();
+    int npes = shmem_n_pes();
+
+    int team_shared_npes = shmem_team_n_pes(SHMEMX_TEAM_HOST);
+
+    int *peers = malloc(team_shared_npes * sizeof(int));
+    int num_peers = 0;
+
+    /* Print the team members on SHMEM_TEAM_HOST */
+    /* Use a lock for cleaner output */
+    shmem_set_lock(&lock);
+
+    printf("[PE: %d] SHMEM_TEAM_HOST peers: { ", me);
+    for (int i = 0; i < npes; i++) {
+        if (shmem_team_translate_pe(SHMEM_TEAM_WORLD, i,
+                                    SHMEMX_TEAM_HOST) != -1) {
+            peers[num_peers++] = i;
+            printf("%d ", i);
+        }
+    }
+
+    printf("} (num_peers: %d)\n", num_peers);
+
+    fflush(NULL);
+
+    shmem_clear_lock(&lock);
+
+    if (num_peers != team_shared_npes) {
+        shmem_global_exit(1);
+    }
+
+    free(peers);
+    shmem_finalize();
+    return 0;
+}