Merge pull request #22 from edgargabriel/topic/ipc_collectives

add support for sync_all, barrier_all, broadcast, fcollect, and alltoall

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
 )

src/ipc/context_ipc_device.cpp

@@ -43,6 +43,7 @@ __host__ IPCContext::IPCContext(Backend *b)
 
   auto *bp{backend->ipc_backend_proxy.get()};
 
+  barrier_sync = backend->barrier_sync;
   g_ret = bp->g_ret;
   atomic_base_ptr = bp->atomic_ret->atomic_base_ptr;
 }
@@ -100,18 +101,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)

src/ipc/context_ipc_device.hpp

@@ -194,53 +194,44 @@ class IPCContext : public Context {
   template <typename T>
   __device__ void get_nbi_wave(T *dest, const T *source, size_t nelems, int pe);
 
-  // Wait / Test functions
-  template <typename T>
-  __device__ void wait_until(T* ptr, roc_shmem_cmps cmp, T val);
-
-  template <typename T>
-  __device__ void wait_until_all(T* ptr, size_t nelems,
-                                 const int *status,
-                                 roc_shmem_cmps cmp, T val);
-
-  template <typename T>
-  __device__ size_t wait_until_any(T* ptr, size_t nelems,
-                                   const int *status,
-                                   roc_shmem_cmps cmp, T val);
-
-  template <typename T>
-  __device__ size_t wait_until_some(T* ptr, size_t nelems,
-                                    size_t* indices,
-                                    const int *status,
-                                    roc_shmem_cmps cmp, T val);
-
-  template <typename T>
-  __device__ void wait_until_all_vector(T* ptr, size_t nelems,
-                                        const int *status,
-                                        roc_shmem_cmps cmp, T* vals);
-
-  template <typename T>
-  __device__ size_t wait_until_any_vector(T* ptr, size_t nelems,
-                                          const int *status,
-                                          roc_shmem_cmps cmp, T* vals);
- template <typename T>
-  __device__ size_t wait_until_some_vector(T* ptr, size_t nelems,
-                                           size_t* indices,
-                                           const int *status,
-                                           roc_shmem_cmps cmp, T* vals);
-
-  template <typename T>
-  __device__ int test(T* ptr, roc_shmem_cmps cmp, T val);
-  
  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);  // NOLINT(runtime/int)
+
+  template <typename T>
+  __device__ void internal_get_broadcast(T *dst, const T *src, int nelems,
+                                         int pe_root);  // 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

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

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,46 @@ __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) {  // 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) {  // 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 +198,82 @@ __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);
+  } else {
+    internal_get_broadcast(dst, src, nelems, pe_root);
+  }
+
+  // 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;
+
+  // 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;
+
+  // 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
@@ -223,58 +325,6 @@ __device__ void IPCContext::get_nbi_wave(T *dest, const T *source,
   getmem_nbi_wave(dest, source, nelems * sizeof(T), pe);
 }
 
-
-//Wait/test functions
-template <typename T>
-__device__ void wait_until(T* ptr, roc_shmem_cmps cmp, T val) {
-}
-
-template <typename T>
-__device__ void wait_until_all(T* ptr, size_t nelems,
-			       const int *status,
-			       roc_shmem_cmps cmp, T val) {
-}
-
-template <typename T>
-__device__ size_t wait_until_any(T* ptr, size_t nelems,
-				 const int *status,
-				 roc_shmem_cmps cmp, T val) {
-  return 0;
-}
-
-template <typename T>
-__device__ size_t wait_until_some(T* ptr, size_t nelems,
-				  size_t* indices,
-				  const int *status,
-				  roc_shmem_cmps cmp, T val){
-  return 0;
-}
-
-template <typename T>
-__device__ void wait_until_all_vector(T* ptr, size_t nelems,
-				      const int *status,
-				      roc_shmem_cmps cmp, T* vals) {
-}
-
-template <typename T>
-__device__ size_t wait_until_any_vector(T* ptr, size_t nelems,
-					const int *status,
-					roc_shmem_cmps cmp, T* vals){
-  return 0;
-}
-
-template <typename T>
-__device__ size_t wait_until_some_vector(T* ptr, size_t nelems,
-					 size_t* indices,
-					 const int *status,
-					 roc_shmem_cmps cmp, T* vals) {
-}
-
-template <typename T>
-__device__ int test(T* ptr, roc_shmem_cmps cmp, T val) {
-  return 0;
-}
-  
 }  // namespace rocshmem
 
 #endif  // LIBRARY_SRC_IPC_CONTEXT_TMPL_DEVICE_HPP_

src/util.hpp

@@ -212,7 +212,7 @@ __device__ void gpu_dprintf(const char* fmt, const Args&... args) {
       while (atomicCAS(print_lock, 0, 1) == 1) {
       }
 
-      printf("WG (%lu, %lu, %lu) TH (%lu, %lu, %lu) ", hipBlockIdx_x,
+      printf("WG (%u, %u, %u) TH (%u, %u, %u) ", hipBlockIdx_x,
              hipBlockIdx_y, hipBlockIdx_z, hipThreadIdx_x, hipThreadIdx_y,
              hipThreadIdx_z);
       printf(fmt, args...);