AI/rocm-systems
2
0
フォーク 0
コード イシュー プルリクエスト Actions パッケージ プロジェクト リリース Wiki アクティビティ

Add host API for *_on_stream operations (#340)

ソースを参照 
* Add functional test for barrier_all_on_stream

* Add rocshmem_barrier_all_on_stream support for GDA and RO backends

Implements rocshmem_barrier_all_on_stream operation for
GPU Direct Access and Reverse Offload backends.

Previously, rocshmem_barrier_all_on_stream was only supported for IPC backend.

* Add functional test for rocshmem_broadcastmem_on_stream

* Add host-side rocshmem_broadcastmem_on_stream API

Implement stream-based broadcast collective operation

- Add rocshmem_broadcastmem_on_stream host API and kernel implementation
- Add functional test TeamBroadcastmemOnStreamTester with multi-stream
  support and correctness verification
- Use per-workgroup contexts to avoid contention across parallel streams

API:
rocshmem_broadcastmem_on_stream(team, dest, source, nelems, pe_root, stream)

* Add functional test for rocshmem_getmem_on_stream

* Add host-side rocshmem_getmem_on_stream API

Implement stream-based point-to-point RMA get operation

- Add rocshmem_getmem_on_stream host API and kernel implementation
- Support for asynchronous getmem operations on HIP streams
- Add backend support for GDA, RO, and IPC contexts
- Use work-group collective getmem for efficient memory transfer

API:
rocshmem_getmem_on_stream(dest, source, nelems, pe, stream)

(AI Assist)

* Add host-side rocshmem_putmem_on_stream API

- Add rocshmem_putmem_on_stream for asynchronous remote writes
- Support for concurrent RMA operations on HIP streams
- Add backend support for GDA, RO, and IPC contexts
- Use work-group device collective operation

API:
rocshmem_putmem_on_stream(dest, source, bytes, pe, stream)

(AI Assist)

* Add functional test for rocshmem_putmem_on_stream

* Add host-side rocshmem_putmem_signal_on_stream API

Enables asynchronous putmem operations with signaling on HIP streams.

The implementation includes:
- Kernel wrapper rocshmem_putmem_signal_kernel
- Host interface putmem_signal_on_stream method
- Context layer support across all backends (IPC, GDA, RO)
- Public API

Function signature:
void rocshmem_putmem_signal_on_stream(void *dest, const void *source,
                                      size_t bytes, uint64_t *sig_addr,
                                      uint64_t signal, int sig_op,
                                      int pe, hipStream_t stream);

* Add functional test for rocshmem_putmem_signal_on_stream

* Add host-side rocshmem_signal_wait_until_on_stream API

Enables asynchronous signal wait operations on HIP streams.

The implementation includes:
- Kernel wrapper rocshmem_signal_wait_until_kernel
- Host interface signal_wait_until_on_stream method
- Context layer support across all backends (IPC, GDA, RO)
- Native uint64_t support in wait_until API (generated from P2P_SYNC.py)

Function signature:
void rocshmem_signal_wait_until_on_stream(uint64_t *sig_addr, int cmp,
                                          uint64_t cmp_value,
                                          hipStream_t stream);

(AI Assist)

* Add functional test for rocshmem_signal_wait_until_on_stream

* Add documentation for stream API functions

This commit adds API documentation for the following host-side
stream functions:

- rocshmem_barrier_all_on_stream (collective routines)
- rocshmem_broadcastmem_on_stream (collective routines)
- rocshmem_getmem_on_stream (RMA operations)
- rocshmem_putmem_on_stream (RMA operations)
- rocshmem_putmem_signal_on_stream (signaling operations)
- rocshmem_signal_wait_until_on_stream (point-to-point sync)

The documentation includes function signatures, parameter descriptions,
and detailed explanations of asynchronous behavior and stream handling.

(AI Assist)

* Rename "bytes" -> "nelems"

* Add "_TEST_" to the variables used in tests

* Remove incorrect hipStreamDefault usage

hipStreamDefault is not a default stream. This is a flag.

If stream == nullptr, then just pass it to kernel. It will launch the kernel on the default stream

[ROCm/rocshmem commit: d0c8380650]
このコミットが含まれているのは:
Anatolii Rozanov
2025-12-09 15:55:46 +01:00
committed by GitHub
親 b9c172de16
コミット f98c72d627
39個のファイルの変更、2649行の追加、49行の削除
Patchファイルをダウンロード Diffファイルをダウンロード すべてのファイルを展開 すべてのファイルを折り畳む
projects/rocshmem/src/context.hpp
+19
ファイルの表示
@@ -400,6 +400,25 @@ class Context {
const void *source, size_t size,
hipStream_t stream);
__host__ void broadcastmem_on_stream(rocshmem_team_t team, void *dest,
const void *source, size_t nelems,
int pe_root, hipStream_t stream);
__host__ void getmem_on_stream(void *dest, const void *source, size_t nelems,
int pe, hipStream_t stream);
__host__ void putmem_on_stream(void *dest, const void *source, size_t nelems,
int pe, hipStream_t stream);
__host__ void putmem_signal_on_stream(void *dest, const void *source,
size_t nelems, uint64_t *sig_addr,
uint64_t signal, int sig_op, int pe,
hipStream_t stream);
__host__ void signal_wait_until_on_stream(uint64_t *sig_addr, int cmp,
uint64_t cmp_value,
hipStream_t stream);
__host__ void sync_all();
template <typename T>
projects/rocshmem/src/context_host.cpp
+42
ファイルの表示
@@ -129,4 +129,46 @@ __host__ void Context::alltoallmem_on_stream(rocshmem_team_t team, void *dest,
HOST_DISPATCH(alltoallmem_on_stream(team, dest, source, size, stream));
}
__host__ void Context::broadcastmem_on_stream(rocshmem_team_t team, void *dest,
const void *source, size_t nelems,
int pe_root, hipStream_t stream) {
ctxHostStats.incStat(NUM_HOST_BROADCAST);
HOST_DISPATCH(
broadcastmem_on_stream(team, dest, source, nelems, pe_root, stream));
}
__host__ void Context::getmem_on_stream(void *dest, const void *source,
size_t nelems, int pe,
hipStream_t stream) {
ctxHostStats.incStat(NUM_HOST_GET);
HOST_DISPATCH(getmem_on_stream(dest, source, nelems, pe, stream));
}
__host__ void Context::putmem_on_stream(void *dest, const void *source,
size_t nelems, int pe,
hipStream_t stream) {
ctxHostStats.incStat(NUM_HOST_PUT);
HOST_DISPATCH(putmem_on_stream(dest, source, nelems, pe, stream));
}
__host__ void Context::putmem_signal_on_stream(void *dest, const void *source,
size_t nelems,
uint64_t *sig_addr,
uint64_t signal, int sig_op,
int pe, hipStream_t stream) {
ctxHostStats.incStat(NUM_HOST_PUT);
HOST_DISPATCH(putmem_signal_on_stream(dest, source, nelems, sig_addr, signal,
sig_op, pe, stream));
}
__host__ void Context::signal_wait_until_on_stream(uint64_t *sig_addr, int cmp,
uint64_t cmp_value,
hipStream_t stream) {
HOST_DISPATCH(signal_wait_until_on_stream(sig_addr, cmp, cmp_value, stream));
}
} // namespace rocshmem
projects/rocshmem/src/gda/context_gda_host.cpp
+39
ファイルの表示
@@ -113,6 +113,10 @@ __host__ void GDAHostContext::barrier_all() {
host_interface->barrier_all(context_window_info);
}
__host__ void GDAHostContext::barrier_all_on_stream(hipStream_t stream) {
host_interface->barrier_all_on_stream(stream);
}
__host__ void GDAHostContext::alltoallmem_on_stream(rocshmem_team_t team,
void *dest,
const void *source,
@@ -121,4 +125,39 @@ __host__ void GDAHostContext::alltoallmem_on_stream(rocshmem_team_t team,
host_interface->alltoallmem_on_stream(team, dest, source, size, stream);
}
__host__ void GDAHostContext::broadcastmem_on_stream(rocshmem_team_t team,
void *dest,
const void *source,
size_t nelems, int pe_root,
hipStream_t stream) {
host_interface->broadcastmem_on_stream(team, dest, source, nelems, pe_root,
stream);
}
__host__ void GDAHostContext::getmem_on_stream(void *dest, const void *source,
size_t nelems, int pe,
hipStream_t stream) {
host_interface->getmem_on_stream(dest, source, nelems, pe, stream);
}
__host__ void GDAHostContext::putmem_on_stream(void *dest, const void *source,
size_t nelems, int pe,
hipStream_t stream) {
host_interface->putmem_on_stream(dest, source, nelems, pe, stream);
}
__host__ void GDAHostContext::putmem_signal_on_stream(
void *dest, const void *source, size_t nelems, uint64_t *sig_addr,
uint64_t signal, int sig_op, int pe, hipStream_t stream) {
host_interface->putmem_signal_on_stream(dest, source, nelems, sig_addr,
signal, sig_op, pe, stream);
}
__host__ void GDAHostContext::signal_wait_until_on_stream(uint64_t *sig_addr,
int cmp,
uint64_t cmp_value,
hipStream_t stream) {
host_interface->signal_wait_until_on_stream(sig_addr, cmp, cmp_value, stream);
}
} // namespace rocshmem
projects/rocshmem/src/gda/context_gda_host.hpp
+21
ファイルの表示
@@ -82,10 +82,31 @@ class GDAHostContext : public Context {
__host__ void barrier_all();
__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 broadcastmem_on_stream(rocshmem_team_t team, void *dest,
const void *source, size_t nelems,
int pe_root, hipStream_t stream);
__host__ void getmem_on_stream(void *dest, const void *source, size_t nelems,
int pe, hipStream_t stream);
__host__ void putmem_on_stream(void *dest, const void *source, size_t nelems,
int pe, hipStream_t stream);
__host__ void putmem_signal_on_stream(void *dest, const void *source,
size_t nelems, uint64_t *sig_addr,
uint64_t signal, int sig_op, int pe,
hipStream_t stream);
__host__ void signal_wait_until_on_stream(uint64_t *sig_addr, int cmp,
uint64_t cmp_value,
hipStream_t stream);
__host__ void sync_all();
template <typename T>
projects/rocshmem/src/host/host.cpp
+108 -12
ファイルの表示
@@ -25,6 +25,7 @@
#include "host.hpp"
#include "rocshmem/rocshmem_config.h" // NOLINT(build/include_subdir)
#include "rocshmem/rocshmem_SIG_OP.hpp"
#include "envvar.hpp"
#include "host_helpers.hpp"
#include "memory/window_info.hpp"
@@ -325,12 +326,8 @@ __host__ void HostInterface::barrier_all(WindowInfo* window_info) {
}
__host__ void HostInterface::barrier_all_on_stream(hipStream_t stream) {
// launch kernel to do barrier with given stream, if non, use default stream
if (stream == nullptr) {
stream = hipStreamDefault;
}
rocshmem_barrier_all_kernel<<<1, 1, 0, stream>>>();
// Launch kernel to do barrier with given stream
rocshmem_barrier_all_kernel<<<1, 1, 0, stream>>>();
}
__host__ void HostInterface::alltoallmem_on_stream(rocshmem_team_t team,
@@ -338,11 +335,6 @@ __host__ void HostInterface::alltoallmem_on_stream(rocshmem_team_t team,
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
@@ -357,13 +349,117 @@ __host__ void HostInterface::alltoallmem_on_stream(rocshmem_team_t team,
int num_threads_per_block = (optimal_block_size > static_cast<int>(size))
? static_cast<int>(size)
: optimal_block_size;
// Launch kernel to do alltoall with given stream
dim3 gridSize(1);
dim3 blockSize(num_threads_per_block);
rocshmem_alltoallmem_kernel<<<gridSize, blockSize, 0, stream>>>(team, dest,
source, size);
}
__host__ void HostInterface::broadcastmem_on_stream(rocshmem_team_t team,
void *dest,
const void *source,
size_t nelems, int pe_root,
hipStream_t stream) {
// Use dynamic block size determination:
// - Query optimal block size using occupancy API
// - Limit block size to nelems (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_broadcastmem_kernel,
0,
0));
// Limit block size to nelems (bytes) to avoid over-subscription
int num_threads_per_block = (optimal_block_size > static_cast<int>(nelems))
? static_cast<int>(nelems)
: optimal_block_size;
// Launch kernel to do broadcast with given stream
dim3 gridSize(1);
dim3 blockSize(num_threads_per_block);
rocshmem_broadcastmem_kernel<<<gridSize, blockSize, 0, stream>>>(team,
dest,
source,
nelems,
pe_root);
}
__host__ void HostInterface::getmem_on_stream(void *dest, const void *source,
size_t nelems, int pe,
hipStream_t stream) {
int optimal_block_size = 0;
int grid_size = 0;
CHECK_HIP(hipOccupancyMaxPotentialBlockSize(&grid_size, &optimal_block_size,
rocshmem_getmem_kernel, 0, 0));
// Limit block size to nelems to avoid over-subscription
int num_threads_per_block = (optimal_block_size > static_cast<int>(nelems))
? static_cast<int>(nelems)
: optimal_block_size;
// Launch kernel to do getmem with given stream
dim3 gridSize(1);
dim3 blockSize(num_threads_per_block);
rocshmem_getmem_kernel<<<gridSize, blockSize, 0, stream>>>(dest, source,
nelems, pe);
}
__host__ void HostInterface::putmem_on_stream(void *dest, const void *source,
size_t nelems, int pe,
hipStream_t stream) {
int optimal_block_size = 0;
int grid_size = 0;
CHECK_HIP(hipOccupancyMaxPotentialBlockSize(&grid_size, &optimal_block_size,
rocshmem_putmem_kernel, 0, 0));
// Limit block size to nelems to avoid over-subscription
int num_threads_per_block = (optimal_block_size > static_cast<int>(nelems))
? static_cast<int>(nelems)
: optimal_block_size;
// Launch kernel to do putmem with given stream
dim3 gridSize(1);
dim3 blockSize(num_threads_per_block);
rocshmem_putmem_kernel<<<gridSize, blockSize, 0, stream>>>(dest, source,
nelems, pe);
}
__host__ void HostInterface::putmem_signal_on_stream(
void *dest, const void *source, size_t nelems, uint64_t *sig_addr,
uint64_t signal, int sig_op, int pe, hipStream_t stream) {
int optimal_block_size = 0;
int grid_size = 0;
CHECK_HIP(hipOccupancyMaxPotentialBlockSize(
&grid_size, &optimal_block_size, rocshmem_putmem_signal_kernel, 0, 0));
// Limit block size to nelems to avoid over-subscription
int num_threads_per_block = (optimal_block_size > static_cast<int>(nelems))
? static_cast<int>(nelems)
: optimal_block_size;
// Launch kernel to do putmem_signal with given stream
dim3 gridSize(1);
dim3 blockSize(num_threads_per_block);
rocshmem_putmem_signal_kernel<<<gridSize, blockSize, 0, stream>>>(
dest, source, nelems, sig_addr, signal, sig_op, pe);
}
__host__ void HostInterface::signal_wait_until_on_stream(uint64_t *sig_addr,
int cmp,
uint64_t cmp_value,
hipStream_t stream) {
// Use a single thread to wait on the signal
dim3 gridSize(1);
dim3 blockSize(1);
rocshmem_signal_wait_until_kernel<<<gridSize, blockSize, 0, stream>>>(
sig_addr, cmp, cmp_value);
}
__host__ void HostInterface::barrier_for_sync() {
if (host_comm_world_ != MPI_COMM_NULL) {
mpilib_ftable_.Barrier(host_comm_world_);
projects/rocshmem/src/host/host.hpp
+19
ファイルの表示
@@ -200,6 +200,25 @@ class HostInterface {
const void *source, size_t size,
hipStream_t stream);
__host__ void broadcastmem_on_stream(rocshmem_team_t team, void *dest,
const void *source, size_t nelems,
int pe_root, hipStream_t stream);
__host__ void getmem_on_stream(void *dest, const void *source, size_t nelems,
int pe, hipStream_t stream);
__host__ void putmem_on_stream(void *dest, const void *source, size_t nelems,
int pe, hipStream_t stream);
__host__ void putmem_signal_on_stream(void *dest, const void *source,
size_t nelems, uint64_t *sig_addr,
uint64_t signal, int sig_op, int pe,
hipStream_t stream);
__host__ void signal_wait_until_on_stream(uint64_t *sig_addr, int cmp,
uint64_t cmp_value,
hipStream_t stream);
__host__ void barrier_for_sync();
__host__ void sync_all(WindowInfo* window_info);
projects/rocshmem/src/ipc/context_ipc_host.cpp
+35
ファイルの表示
@@ -113,4 +113,39 @@ __host__ void IPCHostContext::alltoallmem_on_stream(rocshmem_team_t team,
host_interface->alltoallmem_on_stream(team, dest, source, size, stream);
}
__host__ void IPCHostContext::broadcastmem_on_stream(rocshmem_team_t team,
void *dest,
const void *source,
size_t nelems, int pe_root,
hipStream_t stream) {
host_interface->broadcastmem_on_stream(team, dest, source, nelems, pe_root,
stream);
}
__host__ void IPCHostContext::getmem_on_stream(void *dest, const void *source,
size_t nelems, int pe,
hipStream_t stream) {
host_interface->getmem_on_stream(dest, source, nelems, pe, stream);
}
__host__ void IPCHostContext::putmem_on_stream(void *dest, const void *source,
size_t nelems, int pe,
hipStream_t stream) {
host_interface->putmem_on_stream(dest, source, nelems, pe, stream);
}
__host__ void IPCHostContext::putmem_signal_on_stream(
void *dest, const void *source, size_t nelems, uint64_t *sig_addr,
uint64_t signal, int sig_op, int pe, hipStream_t stream) {
host_interface->putmem_signal_on_stream(dest, source, nelems, sig_addr,
signal, sig_op, pe, stream);
}
__host__ void IPCHostContext::signal_wait_until_on_stream(uint64_t *sig_addr,
int cmp,
uint64_t cmp_value,
hipStream_t stream) {
host_interface->signal_wait_until_on_stream(sig_addr, cmp, cmp_value, stream);
}
} // namespace rocshmem
projects/rocshmem/src/ipc/context_ipc_host.hpp
+19
ファイルの表示
@@ -88,6 +88,25 @@ class IPCHostContext : public Context {
const void *source, size_t size,
hipStream_t stream);
__host__ void broadcastmem_on_stream(rocshmem_team_t team, void *dest,
const void *source, size_t nelems,
int pe_root, hipStream_t stream);
__host__ void getmem_on_stream(void *dest, const void *source, size_t nelems,
int pe, hipStream_t stream);
__host__ void putmem_on_stream(void *dest, const void *source, size_t nelems,
int pe, hipStream_t stream);
__host__ void putmem_signal_on_stream(void *dest, const void *source,
size_t nelems, uint64_t *sig_addr,
uint64_t signal, int sig_op, int pe,
hipStream_t stream);
__host__ void signal_wait_until_on_stream(uint64_t *sig_addr, int cmp,
uint64_t cmp_value,
hipStream_t stream);
__host__ void sync_all();
template <typename T>
projects/rocshmem/src/reverse_offload/context_ro_host.cpp
+51
ファイルの表示
@@ -133,6 +133,12 @@ __host__ void ROHostContext::barrier_all() {
host_interface->barrier_for_sync();
}
__host__ void ROHostContext::barrier_all_on_stream(hipStream_t stream) {
DPRINTF("Function: ro_net_host_barrier_all_on_stream\n");
host_interface->barrier_all_on_stream(stream);
}
__host__ void ROHostContext::alltoallmem_on_stream(rocshmem_team_t team,
void *dest,
const void *source,
@@ -143,4 +149,49 @@ __host__ void ROHostContext::alltoallmem_on_stream(rocshmem_team_t team,
host_interface->alltoallmem_on_stream(team, dest, source, size, stream);
}
__host__ void ROHostContext::broadcastmem_on_stream(rocshmem_team_t team,
void *dest,
const void *source,
size_t nelems, int pe_root,
hipStream_t stream) {
DPRINTF("Function: ro_net_host_broadcastmem_on_stream\n");
host_interface->broadcastmem_on_stream(team, dest, source, nelems, pe_root,
stream);
}
__host__ void ROHostContext::getmem_on_stream(void *dest, const void *source,
size_t nelems, int pe,
hipStream_t stream) {
DPRINTF("Function: ro_net_host_getmem_on_stream\n");
host_interface->getmem_on_stream(dest, source, nelems, pe, stream);
}
__host__ void ROHostContext::putmem_on_stream(void *dest, const void *source,
size_t nelems, int pe,
hipStream_t stream) {
DPRINTF("Function: ro_net_host_putmem_on_stream\n");
host_interface->putmem_on_stream(dest, source, nelems, pe, stream);
}
__host__ void ROHostContext::putmem_signal_on_stream(
void *dest, const void *source, size_t nelems, uint64_t *sig_addr,
uint64_t signal, int sig_op, int pe, hipStream_t stream) {
DPRINTF("Function: ro_net_host_putmem_signal_on_stream\n");
host_interface->putmem_signal_on_stream(dest, source, nelems, sig_addr,
signal, sig_op, pe, stream);
}
__host__ void ROHostContext::signal_wait_until_on_stream(uint64_t *sig_addr,
int cmp,
uint64_t cmp_value,
hipStream_t stream) {
DPRINTF("Function: ro_net_host_signal_wait_until_on_stream\n");
host_interface->signal_wait_until_on_stream(sig_addr, cmp, cmp_value, stream);
}
} // namespace rocshmem
projects/rocshmem/src/reverse_offload/context_ro_host.hpp
+21
ファイルの表示
@@ -131,10 +131,31 @@ class ROHostContext : public Context {
__host__ void barrier_all();
__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 broadcastmem_on_stream(rocshmem_team_t team, void *dest,
const void *source, size_t nelems,
int pe_root, hipStream_t stream);
__host__ void getmem_on_stream(void *dest, const void *source, size_t nelems,
int pe, hipStream_t stream);
__host__ void putmem_on_stream(void *dest, const void *source, size_t nelems,
int pe, hipStream_t stream);
__host__ void putmem_signal_on_stream(void *dest, const void *source,
size_t nelems, uint64_t *sig_addr,
uint64_t signal, int sig_op, int pe,
hipStream_t stream);
__host__ void signal_wait_until_on_stream(uint64_t *sig_addr, int cmp,
uint64_t cmp_value,
hipStream_t stream);
__host__ void sync_all();
template <typename T>
projects/rocshmem/src/rocshmem.cpp
+49
ファイルの表示
@@ -1007,6 +1007,54 @@ __host__ void rocshmem_alltoallmem_on_stream(rocshmem_team_t team, void *dest,
->alltoallmem_on_stream(team, dest, source, size, stream);
}
__host__ void rocshmem_broadcastmem_on_stream(rocshmem_team_t team, void *dest,
const void *source, size_t nelems,
int pe_root, hipStream_t stream) {
DPRINTF("Host function: rocshmem_broadcastmem_on_stream\n");
get_internal_ctx(ROCSHMEM_HOST_CTX_DEFAULT)
->broadcastmem_on_stream(team, dest, source, nelems, pe_root, stream);
}
__host__ void rocshmem_getmem_on_stream(void *dest, const void *source,
size_t nelems, int pe,
hipStream_t stream) {
DPRINTF("Host function: rocshmem_getmem_on_stream\n");
get_internal_ctx(ROCSHMEM_HOST_CTX_DEFAULT)
->getmem_on_stream(dest, source, nelems, pe, stream);
}
__host__ void rocshmem_putmem_on_stream(void *dest, const void *source,
size_t nelems, int pe,
hipStream_t stream) {
DPRINTF("Host function: rocshmem_putmem_on_stream\n");
get_internal_ctx(ROCSHMEM_HOST_CTX_DEFAULT)
->putmem_on_stream(dest, source, nelems, pe, stream);
}
__host__ void rocshmem_putmem_signal_on_stream(void *dest, const void *source,
size_t nelems,
uint64_t *sig_addr,
uint64_t signal, int sig_op,
int pe, hipStream_t stream) {
DPRINTF("Host function: rocshmem_putmem_signal_on_stream\n");
get_internal_ctx(ROCSHMEM_HOST_CTX_DEFAULT)
->putmem_signal_on_stream(dest, source, nelems, sig_addr, signal, sig_op,
pe, stream);
}
__host__ void rocshmem_signal_wait_until_on_stream(uint64_t *sig_addr, int cmp,
uint64_t cmp_value,
hipStream_t stream) {
DPRINTF("Host function: rocshmem_signal_wait_until_on_stream\n");
get_internal_ctx(ROCSHMEM_HOST_CTX_DEFAULT)
->signal_wait_until_on_stream(sig_addr, cmp, cmp_value, stream);
}
__host__ void rocshmem_sync_all() {
DPRINTF("Host function: rocshmem_sync_all\n");
@@ -1681,6 +1729,7 @@ WAIT_DEF_GEN(unsigned short, ushort)
WAIT_DEF_GEN(unsigned int, uint)
WAIT_DEF_GEN(unsigned long, ulong)
WAIT_DEF_GEN(unsigned long long, ulonglong)
WAIT_DEF_GEN(uint64_t, uint64)
// clang-format on
} // namespace rocshmem
projects/rocshmem/src/rocshmem_gpu.cpp
+52
ファイルの表示
@@ -674,6 +674,57 @@ __global__ ATTR_NO_INLINE void rocshmem_alltoallmem_kernel(rocshmem_team_t team,
}
}
__global__ ATTR_NO_INLINE void rocshmem_broadcastmem_kernel(
rocshmem_team_t team, void *dest, const void *source, size_t nelems,
int pe_root) {
__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 broadcast function with created context and provided team
// Using char type since nelems is in bytes (1 byte per element)
rocshmem_broadcast_wg<char>(ctx, team, (char *) dest, (const char *) source,
(int) nelems, pe_root);
if (ctx_result == 0) {
rocshmem_wg_ctx_destroy(&ctx);
}
}
__global__ ATTR_NO_INLINE void rocshmem_getmem_kernel(void *dest,
const void *source,
size_t nelems, int pe) {
// Use work-group collective getmem with default context
rocshmem_getmem_wg(dest, source, nelems, pe);
}
__global__ ATTR_NO_INLINE void rocshmem_putmem_kernel(void *dest,
const void *source,
size_t nelems, int pe) {
// Use work-group collective putmem with default context
rocshmem_putmem_wg(dest, source, nelems, pe);
}
__global__ ATTR_NO_INLINE void rocshmem_putmem_signal_kernel(
void *dest, const void *source, size_t nelems, uint64_t *sig_addr,
uint64_t signal, int sig_op, int pe) {
// Use work-group collective putmem_signal with default context
rocshmem_putmem_signal_wg(dest, source, nelems, sig_addr, signal, sig_op, pe);
}
__global__ ATTR_NO_INLINE void rocshmem_signal_wait_until_kernel(
uint64_t *sig_addr, int cmp, uint64_t cmp_value) {
// Use default context to wait on signal
rocshmem_uint64_wait_until(sig_addr, cmp, cmp_value);
}
__device__ void rocshmem_barrier_all() {
GPU_DPRINTF("Function: rocshmem_barrier_all (ctx=%zd)\n",
get_internal_ctx(ROCSHMEM_CTX_DEFAULT));
@@ -1867,6 +1918,7 @@ WAIT_DEF_GEN(unsigned short, ushort)
WAIT_DEF_GEN(unsigned int, uint)
WAIT_DEF_GEN(unsigned long, ulong)
WAIT_DEF_GEN(unsigned long long, ulonglong)
WAIT_DEF_GEN(uint64_t, uint64)
// clang-format on
} // namespace rocshmem