Add host API for alltoallmem_on_stream collective operation (#333)

* Add host-side rocshmem_alltoallmem_on_stream function

Function signature:
  rocshmem_alltoallmem_on_stream(rocshmem_team_t team, void *dest,
                                 const void *source, size_t size,
                                 hipStream_t stream)

- The function launches rocshmem_alltoallmem_kernel which calls
device-side alltoall<char> workgroup collective through default context.
- Uses dynamic block size determination via occupancy API.
- Implemented for all backends.

* Fix incorrect sync buffer size allocation for alltoall in GDA and IPC backends

When allocating memory for alltoall_pSync_pool in setup_teams() and
teams_init() functions, the code incorrectly used ROCSHMEM_BCAST_SYNC_SIZE
instead of ROCSHMEM_ALLTOALL_SYNC_SIZE.

* Add functional test for team_alltoallmem_on_stream

This commit adds a new functional test to verify the correctness of
the host-side rocshmem_team_alltoallmem_on_stream API.

* Add documentation for rocshmem_alltoallmem_on_stream

This commit adds API documentation for the host-side
rocshmem_alltoallmem_on_stream function in the collective routines
section. The documentation includes:

[ROCm/rocshmem commit: 5577feb70d]

projects/rocshmem/src/context.hpp

@@ -397,6 +397,10 @@ class Context {
 
   __host__ void barrier_all_on_stream(hipStream_t stream);
 
+  __host__ void alltoallmem_on_stream(rocshmem_team_t team, void *dest,
+                                      const void *source, size_t size,
+                                      hipStream_t stream);
+
   __host__ void sync_all();
 
   template <typename T>

projects/rocshmem/src/context_host.cpp

@@ -122,4 +122,12 @@ __host__ void Context::barrier_all_on_stream(hipStream_t stream) {
   HOST_DISPATCH(barrier_all_on_stream(stream));
 }
 
+__host__ void Context::alltoallmem_on_stream(rocshmem_team_t team, void *dest,
+                                             const void *source, size_t size,
+                                             hipStream_t stream) {
+  ctxHostStats.incStat(NUM_HOST_ALLTOALL);
+
+  HOST_DISPATCH(alltoallmem_on_stream(team, dest, source, size, stream));
+}
+
 }  // namespace rocshmem

projects/rocshmem/src/gda/backend_gda.cpp

@@ -453,8 +453,8 @@ void GDABackend::setup_teams() {
                             * max_num_teams;
 
   alltoall_pSync_pool = reinterpret_cast<long *>(wrk_sync_pool_top_);
-  wrk_sync_pool_top_ += sizeof(long) * ROCSHMEM_BCAST_SYNC_SIZE
-                            * max_num_teams;
+  wrk_sync_pool_top_ += sizeof(long) * ROCSHMEM_ALLTOALL_SYNC_SIZE *
+                        max_num_teams;
 
   /* Accommodating for largest possible data type for pWrk */
   pWrk_pool = reinterpret_cast<void *>(wrk_sync_pool_top_);

projects/rocshmem/src/gda/context_gda_host.cpp

@@ -113,4 +113,12 @@ __host__ void GDAHostContext::barrier_all() {
   host_interface->barrier_all(context_window_info);
 }
 
+__host__ void GDAHostContext::alltoallmem_on_stream(rocshmem_team_t team,
+                                                    void *dest,
+                                                    const void *source,
+                                                    size_t size,
+                                                    hipStream_t stream) {
+  host_interface->alltoallmem_on_stream(team, dest, source, size, stream);
+}
+
 }  // namespace rocshmem

projects/rocshmem/src/gda/context_gda_host.hpp

@@ -82,6 +82,10 @@ class GDAHostContext : public Context {
 
   __host__ void barrier_all();
 
+  __host__ void alltoallmem_on_stream(rocshmem_team_t team, void *dest,
+                                      const void *source, size_t size,
+                                      hipStream_t stream);
+
   __host__ void sync_all();
 
   template <typename T>

projects/rocshmem/src/host/host.cpp

@@ -333,6 +333,36 @@ __host__ void HostInterface::barrier_all_on_stream(hipStream_t stream) {
   rocshmem_barrier_all_kernel<<<1, 1, 0,  stream>>>();
 }
 
+__host__ void HostInterface::alltoallmem_on_stream(rocshmem_team_t team,
+                                                   void *dest,
+                                                   const void *source,
+                                                   size_t size,
+                                                   hipStream_t stream) {
+  // launch kernel to do alltoall with given stream, if none, use default stream
+  if (stream == nullptr) {
+    stream = hipStreamDefault;
+  }
+
+  // Use dynamic block size determination:
+  // - Query optimal block size using occupancy API
+  // - Limit block size to size (number of bytes) to avoid over-subscription
+  // - Always use 1 block (single workgroup collective)
+  int optimal_block_size = 0;
+  int grid_size = 0;
+  CHECK_HIP(hipOccupancyMaxPotentialBlockSize(&grid_size, &optimal_block_size,
+                                              rocshmem_alltoallmem_kernel, 0,
+                                              0));
+
+  // Limit block size to size (bytes) to avoid over-subscription
+  int num_threads_per_block = (optimal_block_size > static_cast<int>(size))
+                                  ? static_cast<int>(size)
+                                  : optimal_block_size;
+                                  
+  dim3 gridSize(1);
+  dim3 blockSize(num_threads_per_block);
+  rocshmem_alltoallmem_kernel<<<gridSize, blockSize, 0, stream>>>(team, dest,
+                                                                  source, size);
+}
 
 __host__ void HostInterface::barrier_for_sync() {
   if (host_comm_world_ != MPI_COMM_NULL) {

projects/rocshmem/src/host/host.hpp

@@ -196,6 +196,10 @@ class HostInterface {
   
   __host__ void barrier_all_on_stream(hipStream_t stream);
 
+  __host__ void alltoallmem_on_stream(rocshmem_team_t team, void *dest,
+                                      const void *source, size_t size,
+                                      hipStream_t stream);
+
   __host__ void barrier_for_sync();
 
   __host__ void sync_all(WindowInfo* window_info);

projects/rocshmem/src/ipc/backend_ipc.cpp

@@ -482,8 +482,8 @@ void IPCBackend::teams_init() {
                             * max_num_teams;
 
   alltoall_pSync_pool = reinterpret_cast<long *>(wrk_sync_pool_top_);
-  wrk_sync_pool_top_ += sizeof(long) * ROCSHMEM_BCAST_SYNC_SIZE
-                            * max_num_teams;
+  wrk_sync_pool_top_ += sizeof(long) * ROCSHMEM_ALLTOALL_SYNC_SIZE *
+                        max_num_teams;
 
   /* Accommodating for largest possible data type for pWrk */
   pWrk_pool = reinterpret_cast<void *>(wrk_sync_pool_top_);

projects/rocshmem/src/ipc/context_ipc_host.cpp

@@ -105,4 +105,12 @@ __host__ void IPCHostContext::barrier_all_on_stream(hipStream_t stream) {
   host_interface->barrier_all_on_stream(stream);
 }
 
+__host__ void IPCHostContext::alltoallmem_on_stream(rocshmem_team_t team,
+                                                    void *dest,
+                                                    const void *source,
+                                                    size_t size,
+                                                    hipStream_t stream) {
+  host_interface->alltoallmem_on_stream(team, dest, source, size, stream);
+}
+
 }  // namespace rocshmem

projects/rocshmem/src/ipc/context_ipc_host.hpp

@@ -84,6 +84,10 @@ class IPCHostContext : public Context {
 
   __host__ void barrier_all_on_stream(hipStream_t stream);
 
+  __host__ void alltoallmem_on_stream(rocshmem_team_t team, void *dest,
+                                      const void *source, size_t size,
+                                      hipStream_t stream);
+
   __host__ void sync_all();
 
   template <typename T>

projects/rocshmem/src/reverse_offload/context_ro_host.cpp

@@ -133,4 +133,14 @@ __host__ void ROHostContext::barrier_all() {
   host_interface->barrier_for_sync();
 }
 
+__host__ void ROHostContext::alltoallmem_on_stream(rocshmem_team_t team,
+                                                    void *dest,
+                                                    const void *source,
+                                                    size_t size,
+                                                    hipStream_t stream) {
+  DPRINTF("Function: ro_net_host_alltoallmem_on_stream\n");
+
+  host_interface->alltoallmem_on_stream(team, dest, source, size, stream);
+}
+
 }  // namespace rocshmem

projects/rocshmem/src/reverse_offload/context_ro_host.hpp

@@ -131,6 +131,10 @@ class ROHostContext : public Context {
 
   __host__ void barrier_all();
 
+  __host__ void alltoallmem_on_stream(rocshmem_team_t team, void *dest,
+                                      const void *source, size_t size,
+                                      hipStream_t stream);
+
   __host__ void sync_all();
 
   template <typename T>

projects/rocshmem/src/rocshmem.cpp

@@ -999,6 +999,15 @@ __host__ void rocshmem_barrier_all_on_stream(hipStream_t stream) {
   get_internal_ctx(ROCSHMEM_HOST_CTX_DEFAULT)->barrier_all_on_stream(stream);
 }
 
+__host__ void rocshmem_alltoallmem_on_stream(rocshmem_team_t team, void *dest,
+                                             const void *source, size_t size,
+                                             hipStream_t stream) {
+  DPRINTF("Host function: rocshmem_alltoallmem_on_stream\n");
+
+  get_internal_ctx(ROCSHMEM_HOST_CTX_DEFAULT)
+      ->alltoallmem_on_stream(team, dest, source, size, stream);
+}
+
 __host__ void rocshmem_sync_all() {
   DPRINTF("Host function: rocshmem_sync_all\n");
 

projects/rocshmem/src/rocshmem_gpu.cpp

@@ -648,6 +648,32 @@ __global__ ATTR_NO_INLINE void rocshmem_barrier_all_kernel(){
   rocshmem_barrier_all();
 }
 
+__global__ ATTR_NO_INLINE void rocshmem_alltoallmem_kernel(rocshmem_team_t team,
+                                                           void *dest,
+                                                           const void *source,
+                                                           size_t size) {
+  // Create a context for this workgroup to avoid contention on default context
+  // This allows parallel execution across multiple streams without serialization
+  __shared__ rocshmem_ctx_t ctx;
+  __shared__ int ctx_result;
+  
+  ctx_result = rocshmem_wg_team_create_ctx(team, 0, &ctx);
+  
+  // If context creation failed, fall back to default context
+  if (ctx_result != 0) {
+    ctx = ROCSHMEM_CTX_DEFAULT;
+    __syncthreads();
+  }
+
+  // Call device alltoall function with created context and provided team
+  // Using char type since size is in bytes (1 byte per element)
+  rocshmem_alltoall_wg<char>(ctx, team, (char *) dest,
+                             (const char *) source, (int) size);
+
+  if (ctx_result == 0) {
+    rocshmem_wg_ctx_destroy(&ctx);
+  }
+}
 
 __device__ void rocshmem_barrier_all() {
   GPU_DPRINTF("Function: rocshmem_barrier_all (ctx=%zd)\n",

projects/rocshmem/src/stats.hpp

@@ -142,6 +142,7 @@ enum rocshmem_host_stats {
   NUM_HOST_SHMEM_PTR,
   NUM_HOST_SYNC_ALL,
   NUM_HOST_BROADCAST,
+  NUM_HOST_ALLTOALL,
   NUM_HOST_STATS
 };