first cut on collectives and sync

code is based on the GPUIB implementations of the routines, which seem
however generic enough to work also for the IPC conduit.

Some code is in for broadcast, fcollect, and alltoall.


[ROCm/rocshmem commit: 0de3b5e6fc]

projects/rocshmem/src/ipc/CMakeLists.txt

@@ -30,4 +30,5 @@ target_sources(
     context_ipc_host.cpp
     backend_ipc.cpp
     ipc_team.cpp
+    context_ipc_device_coll.cpp
 )

projects/rocshmem/src/ipc/context_ipc_device.cpp

@@ -100,18 +100,6 @@ __device__ void *IPCContext::shmem_ptr(const void *dest, int pe) {
   return ret;
 }
 
-__device__ void IPCContext::barrier_all() {
-  __syncthreads();
-}
-
-__device__ void IPCContext::sync_all() {
-  __syncthreads();
-}
-
-__device__ void IPCContext::sync(roc_shmem_team_t team) {
-  __syncthreads();
-}
-
 __device__ void IPCContext::putmem_wg(void *dest, const void *source,
                                      size_t nelems, int pe) {
   // TODO (Avinash) check if PE is available for IPC using (isIpcAvailable)

projects/rocshmem/src/ipc/context_ipc_device.hpp

@@ -196,13 +196,44 @@ class IPCContext : public Context {
 
  private:
 
- //context class has IpcImpl object (ipcImpl_)
- IpcImpl *ipcImpl{nullptr};
+  //context class has IpcImpl object (ipcImpl_)
+  IpcImpl *ipcImpl{nullptr};
 
- uint64_t* atomic_base_ptr{nullptr};
+  uint64_t* atomic_base_ptr{nullptr};
 
- char* g_ret;
+  char* g_ret;
 
+  //internal functions used by collective operations
+   template <typename T>
+  __device__ void internal_put_broadcast(T *dst, const T *src, int nelems,
+                                         int pe_root, int PE_start,
+                                         int logPE_stride, int PE_size,
+                                         long *pSync);  // NOLINT(runtime/int)
+
+  template <typename T>
+  __device__ void internal_get_broadcast(T *dst, const T *src, int nelems,
+                                         int pe_root,
+                                         long *pSync);  // NOLINT(runtime/int)
+
+  template <typename T>
+  __device__ void fcollect_linear(roc_shmem_team_t team, T *dest,
+                                  const T *source, int nelems);
+
+  template <typename T>
+  __device__ void alltoall_linear(roc_shmem_team_t team, T *dest,
+                                  const T *source, int nelems);
+
+  __device__ void internal_sync(int pe, int PE_start, int stride, int PE_size,
+                                int64_t *pSync);
+
+  __device__ void internal_direct_barrier(int pe, int PE_start, int stride,
+                                          int n_pes, int64_t *pSync);
+
+  __device__ void internal_atomic_barrier(int pe, int PE_start, int stride,
+                                          int n_pes, int64_t *pSync);
+
+  //Temporary scratchpad memory used by internal barrier algorithms.
+  int64_t *barrier_sync{nullptr};
 };
 
 }  // namespace rocshmem

projects/rocshmem/src/ipc/context_ipc_device_coll.cpp

@@ -0,0 +1,120 @@
+/******************************************************************************
+ * Copyright (c) 2024 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to
+ * deal in the Software without restriction, including without limitation the
+ * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+ * sell copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+ * IN THE SOFTWARE.
+ *****************************************************************************/
+
+#include "roc_shmem/roc_shmem.hpp"
+#include "../context_incl.hpp"
+#include "context_ipc_tmpl_device.hpp"
+#include "../util.hpp"
+#include "ipc_team.hpp"
+
+namespace rocshmem {
+
+__device__ void IPCContext::internal_direct_barrier(int pe, int PE_start,
+                                                      int stride, int n_pes,
+                                                      int64_t *pSync) {
+  int64_t flag_val = 1;
+  if (pe == PE_start) {
+    // Go through all PE offsets (except current offset = 0)
+    // and wait until they all reach
+    for (size_t i = 1; i < n_pes; i++) {
+      wait_until(&pSync[i], ROC_SHMEM_CMP_EQ, flag_val);
+      pSync[i] = ROC_SHMEM_SYNC_VALUE;
+    }
+    threadfence_system();
+
+    // Announce to other PEs that all have reached
+    for (size_t i = 1, j = PE_start + stride; i < n_pes; ++i, j += stride) {
+      put_nbi(&pSync[0], &flag_val, 1, j);
+    }
+  } else {
+    // Mark current PE offset as reached
+    size_t pe_offset = (pe - PE_start) / stride;
+    put_nbi(&pSync[pe_offset], &flag_val, 1, PE_start);
+    wait_until(&pSync[0], ROC_SHMEM_CMP_EQ, flag_val);
+    pSync[0] = ROC_SHMEM_SYNC_VALUE;
+    threadfence_system();
+  }
+}
+
+__device__ void IPCContext::internal_atomic_barrier(int pe, int PE_start,
+                                                      int stride, int n_pes,
+                                                      int64_t *pSync) {
+  int64_t flag_val = 1;
+  if (pe == PE_start) {
+    wait_until(&pSync[0], ROC_SHMEM_CMP_EQ, (int64_t)(n_pes - 1));
+    pSync[0] = ROC_SHMEM_SYNC_VALUE;
+    threadfence_system();
+
+    for (size_t i = 1, j = PE_start + stride; i < n_pes; ++i, j += stride) {
+      put_nbi(&pSync[0], &flag_val, 1, j);
+    }
+  } else {
+    amo_add<int64_t>(&pSync[0], flag_val, PE_start);
+    wait_until(&pSync[0], ROC_SHMEM_CMP_EQ, flag_val);
+    pSync[0] = ROC_SHMEM_SYNC_VALUE;
+    threadfence_system();
+  }
+}
+
+// Uses PE values that are relative to world
+__device__ void IPCContext::internal_sync(int pe, int PE_start, int stride,
+                                          int PE_size, int64_t *pSync) {
+  __syncthreads();
+  if (is_thread_zero_in_block()) {
+    if (PE_size < 64) {
+      internal_direct_barrier(pe, PE_start, stride, PE_size, pSync);
+    } else {
+      internal_atomic_barrier(pe, PE_start, stride, PE_size, pSync);
+    }
+  }
+  __threadfence();
+  __syncthreads();
+}
+
+__device__ void IPCContext::sync(roc_shmem_team_t team) {
+  IPCTeam *team_obj = reinterpret_cast<IPCTeam *>(team);
+
+  /**
+   * Ensure that the stride is a multiple of 2.
+   */
+  int log_pe_stride = static_cast<int>(team_obj->tinfo_wrt_world->log_stride);
+  int pe = team_obj->my_pe_in_world;
+  int pe_start = team_obj->tinfo_wrt_world->pe_start;
+  int pe_stride = (1 << log_pe_stride);
+  int pe_size = team_obj->num_pes;
+
+  internal_sync(pe, pe_start, pe_stride, pe_size, barrier_sync);
+}
+
+__device__ void IPCContext::sync_all() {
+  internal_sync(my_pe, 0, 1, num_pes, barrier_sync);
+}
+
+__device__ void IPCContext::barrier_all() {
+  if (is_thread_zero_in_block()) {
+    quiet();
+  }
+  sync_all();
+  __syncthreads();
+}
+
+}  // namespace rocshmem

projects/rocshmem/src/ipc/context_ipc_tmpl_device.hpp

@@ -27,6 +27,7 @@
 #include "roc_shmem/roc_shmem.hpp"
 #include "context_ipc_device.hpp"
 #include "../util.hpp"
+#include "ipc_team.hpp"
 
 namespace rocshmem {
 
@@ -150,11 +151,48 @@ __device__ void IPCContext::to_all(T *dest, const T *source, int nreduce,
 				   long *pSync) {  // NOLINT(runtime/int)
 }
 
+template <typename T>
+__device__ void IPCContext::internal_put_broadcast(
+    T *dst, const T *src, int nelems, int pe_root, int pe_start,
+    int log_pe_stride, int pe_size,
+    [[maybe_unused]] long *p_sync) {  // NOLINT(runtime/int)
+  if (my_pe == pe_root) {
+    int stride = 1 << log_pe_stride;
+    int finish = pe_start + stride * pe_size;
+    for (int i = pe_start; i < finish; i += stride) {
+      if (i != my_pe) {
+        put_nbi_wg(dst, src, nelems, i);
+      }
+    }
+  }
+}
+
+template <typename T>
+__device__ void IPCContext::internal_get_broadcast(
+    T *dst, const T *src, int nelems, int pe_root,
+    [[maybe_unused]] long *pSync) {  // NOLINT(runtime/int)
+  if (my_pe != pe_root) {
+    get_wg(dst, src, nelems, pe_root);
+  }
+}
+
 template <typename T>
 __device__ void IPCContext::broadcast(roc_shmem_team_t team, T *dst,
 				      const T *src, int nelems, int pe_root) {
-  //broadcast<T>(dst, src, nelems, pe_root_world, pe_start, log_pe_stride,
-  //             pe_size, p_sync);
+  IPCTeam *team_obj = reinterpret_cast<IPCTeam *>(team);
+
+  /**
+   * Ensure that the stride is a multiple of 2 .
+   */
+  int log_pe_stride = static_cast<int>(team_obj->tinfo_wrt_world->log_stride);
+  int pe_start = team_obj->tinfo_wrt_world->pe_start;
+  int pe_size = team_obj->tinfo_wrt_world->size;
+  long *p_sync = team_obj->bcast_pSync;
+
+  // Passed pe_root is relative to team, convert to world root
+  int pe_root_world = team_obj->get_pe_in_world(pe_root);
+  broadcast<T>(dst, src, nelems, pe_root_world, pe_start, log_pe_stride,
+               pe_size, p_sync);
 }
 
 template <typename T>
@@ -162,16 +200,88 @@ __device__ void IPCContext::broadcast(T *dst, const T *src, int nelems,
 				      int pe_root, int pe_start,
 				      int log_pe_stride, int pe_size,
 				      long *p_sync) {  // NOLINT(runtime/int)
+  if (num_pes < 4) {
+    internal_put_broadcast(dst, src, nelems, pe_root, pe_start, log_pe_stride,
+                           pe_size, p_sync);
+  } else {
+    internal_get_broadcast(dst, src, nelems, pe_root, p_sync);
+  }
+
+  // Synchronize on completion of broadcast
+  internal_sync(my_pe, pe_start, (1 << log_pe_stride), pe_size, p_sync);
 }
 
 template <typename T>
 __device__ void IPCContext::alltoall(roc_shmem_team_t team, T *dst,
 				     const T *src, int nelems) {
+  alltoall_linear(team, dst, src, nelems);
+}
+
+template <typename T>
+__device__ void IPCContext::alltoall_linear(roc_shmem_team_t team, T *dst,
+                                            const T *src, int nelems) {
+  IPCTeam *team_obj = reinterpret_cast<IPCTeam *>(team);
+
+  /**
+   * Ensure that the stride is a multiple of 2
+   */
+  int log_pe_stride = static_cast<int>(team_obj->tinfo_wrt_world->log_stride);
+  int pe_start = team_obj->tinfo_wrt_world->pe_start;
+  int pe_size = team_obj->num_pes;
+  int stride = 1 << log_pe_stride;
+  long *pSync = team_obj->alltoall_pSync;
+  int my_pe_in_team = team_obj->my_pe;
+
+  /* TODO (egabriel): shouldn't the loop handle pe_start, pe_size, and stride?
+   * otherwise what is the point of the few lines above?
+   */
+  // Have each PE put their designated data to the other PEs
+  for (int j = 0; j < pe_size; j++) {
+    int dest_pe = team_obj->get_pe_in_world(j);
+    put_nbi_wg(&dst[my_pe_in_team * nelems], &src[j * nelems], nelems, dest_pe);
+  }
+  if (is_thread_zero_in_block()) {
+    quiet();
+  }
+  // wait until everyone has obtained their designated data
+  internal_sync(my_pe, pe_start, stride, pe_size, pSync);
 }
 
 template <typename T>
 __device__ void IPCContext::fcollect(roc_shmem_team_t team, T *dst,
 				     const T *src, int nelems) {
+  fcollect_linear(team, dst, src, nelems);
+}
+
+template <typename T>
+__device__ void IPCContext::fcollect_linear(roc_shmem_team_t team, T *dst,
+                                            const T *src, int nelems) {
+  IPCTeam *team_obj = reinterpret_cast<IPCTeam *>(team);
+
+  /**
+   * Ensure that the stride is a multiple of 2.
+   */
+  int log_pe_stride = static_cast<int>(team_obj->tinfo_wrt_world->log_stride);
+  int pe_start = team_obj->tinfo_wrt_world->pe_start;
+  int pe_size = team_obj->num_pes;
+  int stride = 1 << log_pe_stride;
+  long *pSync = team_obj->alltoall_pSync;
+  int my_pe_in_team = team_obj->my_pe;
+
+  /* TODO (egabriel): shouldn't the loop handle pe_start, pe_size, and stride?
+   * otherwise what is the point of the few lines above?
+   */
+  // Have each PE put their designated data to the other PEs
+  for (int j = 0; j < pe_size; j++) {
+    int dest_pe = team_obj->get_pe_in_world(j);
+    put_nbi_wg(&dst[my_pe_in_team * nelems], src, nelems, dest_pe);
+  }
+
+  if (is_thread_zero_in_block()) {
+    quiet();
+  }
+  // wait until everyone has obtained their designated data
+  internal_sync(my_pe, pe_start, stride, pe_size, pSync);
 }
 
 // Block/wave functions