AI/rocm-systems
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add host API for *_on_stream operations (#340)

Browse Source 
* 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]
This commit is contained in:
Anatolii Rozanov
2025-12-09 15:55:46 +01:00
committed by GitHub
parent b9c172de16
commit f98c72d627
39 changed files with 2649 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
projects/rocshmem/docs/api/coll.rst
+36 -1
View File
@@ -22,6 +22,20 @@ This routine performs a collective barrier across all PEs in the system.
The caller is blocked until the barrier is resolved and all updates local and remote are completed.
These APIs should be called from only one thread/wavefront/workgroup within the grid to avoid undefined behavior.
ROCSHMEM_BARRIER_ALL_ON_STREAM
-------------------------------
.. cpp:function:: __host__ void rocshmem_barrier_all_on_stream(hipStream_t stream)
:param stream: HIP stream on which to enqueue the operation.
:returns: None.
**Description:**
This routine enqueues a collective barrier operation on a HIP stream. The barrier is performed
across all PEs in the system. The operation is enqueued on the specified stream and will execute
asynchronously. The caller must synchronize the stream (e.g., using ``hipStreamSynchronize``)
to ensure completion.
ROCSHMEM_BARRIER
----------------
@@ -109,7 +123,6 @@ execute asynchronously. The caller must synchronize the stream (e.g., using
This function creates a separate context for each workgroup to avoid contention on the
default context, allowing parallel execution across multiple streams.
If ``stream`` is ``nullptr``, the operation will use ``hipStreamDefault``.
ROCSHMEM_BROADCAST
------------------
@@ -131,6 +144,28 @@ The caller is blocked until the broadcast completes.
Valid ``TYPENAME`` and ``TYPE`` values are listed in :ref:`RMA_TYPES`.
ROCSHMEM_BROADCASTMEM_ON_STREAM
--------------------------------
.. cpp:function:: __host__ void rocshmem_broadcastmem_on_stream(rocshmem_team_t team, void *dest, const void *source, size_t nelems, int pe_root, hipStream_t stream)
:param team: The team participating in the collective.
:param dest: Destination address. Must be an address on the symmetric heap.
:param source: Source address. Must be an address on the symmetric heap.
:param nelems: Number of bytes to broadcast.
:param pe_root: Root PE (relative to team) from which to broadcast.
:param stream: HIP stream on which to enqueue the operation.
:returns: None.
**Description:**
This routine enqueues a broadcast collective operation on a HIP stream. The function broadcasts
data from the root PE to all other PEs participating in the collective routine. The operation
is enqueued on the specified stream and will execute asynchronously. The caller must synchronize
the stream (e.g., using ``hipStreamSynchronize``) to ensure completion.
This function creates a separate context for each workgroup to avoid contention on the
default context, allowing parallel execution across multiple streams.
ROCSHMEM_FCOLLECT
-----------------
projects/rocshmem/docs/api/pt2pt_sync.rst
+19
View File
@@ -96,6 +96,25 @@ ROCSHMEM_TEST
**Description:**
This routine tests if the condition ``(*ivars cmp val)`` is true.
ROCSHMEM_SIGNAL_WAIT_UNTIL_ON_STREAM
-------------------------------------
.. cpp:function:: __host__ void rocshmem_signal_wait_until_on_stream(uint64_t *sig_addr, int cmp, uint64_t cmp_value, hipStream_t stream)
:param sig_addr: Address of the signal variable on the symmetric heap.
:param cmp: Comparison operator (e.g., ROCSHMEM_CMP_EQ, ROCSHMEM_CMP_GE, etc.).
:param cmp_value: Value to compare against.
:param stream: HIP stream on which to enqueue the operation.
:returns: None.
**Description:**
This routine enqueues a wait operation on a HIP stream. The function blocks the calling thread
until the signal variable at ``sig_addr`` satisfies the comparison condition ``(*sig_addr cmp cmp_value)``.
The wait operation is executed asynchronously on the specified stream. The caller must synchronize
the stream (e.g., using ``hipStreamSynchronize``) to ensure the wait condition has been satisfied.
Valid ``cmp`` values are listed in :ref:`CMP_VALUES`.
.. _CMP_VALUES:
Supported comparisons
projects/rocshmem/docs/api/rma.rst
+38
View File
@@ -67,6 +67,25 @@ ROCSHMEM_PUTMEM
**Description:**
This routine writes contiguous data of ``nelems`` bytes from source on the calling PE to ``dest`` at ``pe``.
ROCSHMEM_PUTMEM_ON_STREAM
--------------------------
.. cpp:function:: __host__ void rocshmem_putmem_on_stream(void *dest, const void *source, size_t nelems, int pe, hipStream_t stream)
:param dest: Destination address. Must be an address on the symmetric heap.
:param source: Source address. Must be an address on the symmetric heap.
:param nelems: Size of the transfer in bytes.
:param pe: PE of the remote process.
:param stream: HIP stream on which to enqueue the operation.
:returns: None.
**Description:**
This routine enqueues a putmem RMA operation on a HIP stream. The function writes contiguous
data of ``nelems`` bytes from source on the calling PE to ``dest`` at ``pe``. The operation
is enqueued on the specified stream and will execute asynchronously. The caller must
synchronize the stream (e.g., using ``hipStreamSynchronize``) to ensure completion.
ROCSHMEM_P
----------
@@ -137,6 +156,25 @@ ROCSHMEM_GETMEM
**Description:**
This routine reads contiguous data of ``nelems`` bytes from source on ``pe`` to ``dest`` on the calling PE.
ROCSHMEM_GETMEM_ON_STREAM
--------------------------
.. cpp:function:: __host__ void rocshmem_getmem_on_stream(void *dest, const void *source, size_t nelems, int pe, hipStream_t stream)
:param dest: Destination address. Must be an address on the symmetric heap.
:param source: Source address. Must be an address on the symmetric heap.
:param nelems: Size of the transfer in bytes.
:param pe: PE of the remote process.
:param stream: HIP stream on which to enqueue the operation.
:returns: None.
**Description:**
This routine enqueues a getmem RMA operation on a HIP stream. The function reads contiguous
data of ``nelems`` bytes from source on ``pe`` to ``dest`` on the calling PE. The operation
is enqueued on the specified stream and will execute asynchronously. The caller must
synchronize the stream (e.g., using ``hipStreamSynchronize``) to ensure completion.
ROCSHMEM_G
----------
.. cpp:function:: __device__ float rocshmem_ctx_float_g(rocshmem_ctx_t ctx, const float *source, int pe)
projects/rocshmem/docs/api/sigops.rst
+24
View File
@@ -71,6 +71,30 @@ then applies ``sig_op`` at ``sig_addr`` with the signal value.
Valid ``sig_op values`` are listed in SIGNAL_OPERATORS_.
Valid ``TYPENAME`` and ``TYPE`` values are listed in :ref:`RMA_TYPES`.
ROCSHMEM_PUTMEM_SIGNAL_ON_STREAM
---------------------------------
.. cpp:function:: __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)
:param dest: Destination address on the remote PE. Must be an address on the symmetric heap.
:param source: Source address on the local PE. Must be an address on the symmetric heap.
:param nelems: Size of the transfer in bytes.
:param sig_addr: Address of signal variable on the remote PE. Must be an address on the symmetric heap.
:param signal: Signal value to be written.
:param sig_op: Signal operation (ROCSHMEM_SIGNAL_SET or ROCSHMEM_SIGNAL_ADD).
:param pe: PE number of the remote PE.
:param stream: HIP stream on which to enqueue the operation.
:returns: None.
**Description:**
This routine enqueues a put-with-signal operation on a HIP stream. The function writes contiguous
data of ``nelems`` bytes from source on the calling PE to ``dest`` at ``pe``, then applies ``sig_op``
at ``sig_addr`` with the signal value. The operation is enqueued on the specified stream and will
execute asynchronously. The caller must synchronize the stream (e.g., using ``hipStreamSynchronize``)
to ensure completion.
Valid ``sig_op`` values are listed in SIGNAL_OPERATORS_.
ROCSHMEM_SIGNAL_FETCH
---------------------