TheRock/rocm-systems/projects/rocshmem/python/tests/test_smoke.py

"""Torch-free smoke tests for rocshmem4py.

These tests exercise the native host API through ``ctypes`` + HIP, so they
run in any CI environment even when ``torch`` is not installed.

Scope:

  * ``test_create_buffer`` -- validates the framework-agnostic collective
    allocation API (:func:`rocshmem4py.rocshmem_create_buffer`).
  * ``test_symmetric_buffer_lifecycle`` -- SymmetricBuffer RAII lifecycle;
    torch integration path is covered by ``test_basic.py::test_symmetric_buffer``.
  * ``test_host_device_roundtrip`` -- validates the symmetric heap is
    accessible as ordinary HIP device memory via ctypes H2D/D2H transfers.
  * ``test_get_peer_buffer`` -- validates :func:`rocshmem4py.rocshmem_get_peer_buffer`
    (torch-free equivalent of ``symm_rocshmem_tensor``).
  * Multi-PE stream-based transfers using ``rocshmem_create_buffer`` for
    allocation (``*_on_stream`` portable APIs only -- see test_collective.py
    docstring for rationale).

"""

import ctypes
import pytest

from conftest import requires_multi_pe  # noqa: E402
from hip_test_utils import HipStream, d2h, h2d, hip_sync  # noqa: E402

import rocshmem4py  # noqa: E402


# ---------------------------------------------------------------------------
# Single-PE tests unique to this file (torch-free but HIP-specific)
# ---------------------------------------------------------------------------


def test_create_buffer():
    """rocshmem_create_buffer is the framework-agnostic collective allocator.

    Validates that the returned SymmetricBuffer has the correct shape, pointer,
    __cuda_array_interface__ fields, and a valid symmetric address for self.
    """
    nbytes = 4096
    buf = rocshmem4py.rocshmem_create_buffer(nbytes)
    try:
        assert isinstance(buf, rocshmem4py.SymmetricBuffer)
        assert buf.ptr > 0
        assert buf.nbytes == nbytes
        assert buf.size == nbytes
        assert buf.own_data is True
        assert int(buf) == buf.ptr
        # __cuda_array_interface__ must be present for framework interop (e.g. torch)
        cai = buf.__cuda_array_interface__
        assert cai["data"] == (buf.ptr, False)
        assert cai["shape"] == (nbytes,)
        assert cai["version"] == 3
        # Symmetric pointer to self must be valid
        my_pe = rocshmem4py.rocshmem_my_pe()
        assert buf.get_remote_ptr(my_pe) >= 0
    finally:
        buf.free()
        assert buf._freed is True


def test_symmetric_buffer_lifecycle():
    """SymmetricBuffer RAII lifecycle: allocation, pointer validity, remote ptr, and free.

    Covers the torch-free API surface only.  The torch integration path
    (``__cuda_array_interface__``, ``torch.as_tensor``, non-owning view,
    ``_device``) is covered by ``test_basic.py::test_symmetric_buffer``.
    """
    size = 2048
    buf = rocshmem4py.SymmetricBuffer(size)
    try:
        assert buf.size == size and buf.nbytes == size and buf.ptr > 0
        assert int(buf) == buf.ptr
        my_pe = rocshmem4py.rocshmem_my_pe()
        assert buf.get_remote_ptr(my_pe) >= 0
    finally:
        buf.free()
        assert buf._freed is True


def test_host_device_roundtrip():
    """Symmetric heap memory is accessible as ordinary HIP device memory."""
    nelems = 16
    nbytes = nelems * ctypes.sizeof(ctypes.c_int32)
    buf = rocshmem4py.rocshmem_create_buffer(nbytes)
    try:
        tx = (ctypes.c_int32 * nelems)(*range(nelems))
        h2d(buf.ptr, tx)
        hip_sync()

        rx = (ctypes.c_int32 * nelems)()
        d2h(rx, buf.ptr)
        hip_sync()

        assert list(rx) == list(range(nelems))
    finally:
        buf.free()


# ---------------------------------------------------------------------------
# Multi-PE tests — stream-based portable APIs only.
# Host-side blocking APIs are not exercised here; see test_collective.py
# module docstring for the rationale.
# ---------------------------------------------------------------------------


@requires_multi_pe
def test_putmem_on_stream_peer():
    """Torch-free: push local buffer to a peer via rocshmem_putmem_on_stream."""
    nelems = 32
    nbytes = nelems * ctypes.sizeof(ctypes.c_int32)

    my_pe = rocshmem4py.rocshmem_my_pe()
    n_pes = rocshmem4py.rocshmem_n_pes()
    peer = (my_pe + 1) % n_pes

    src = rocshmem4py.rocshmem_create_buffer(nbytes)
    dst = rocshmem4py.rocshmem_create_buffer(nbytes)
    try:
        src_host = (ctypes.c_int32 * nelems)(*([my_pe] * nelems))
        dst_host = (ctypes.c_int32 * nelems)(*([-1] * nelems))
        h2d(src.ptr, src_host)
        h2d(dst.ptr, dst_host)
        hip_sync()
        rocshmem4py.rocshmem_barrier_all()

        with HipStream() as s:
            rocshmem4py.rocshmem_putmem_on_stream(
                dst.ptr, src.ptr, nbytes, peer, s.handle
            )
            rocshmem4py.rocshmem_barrier_all_on_stream(s.handle)
            s.synchronize()

        rx = (ctypes.c_int32 * nelems)()
        d2h(rx, dst.ptr)
        hip_sync()

        sender = (my_pe - 1 + n_pes) % n_pes
        assert all(v == sender for v in rx), (
            f"PE {my_pe}: expected all {sender}, got {list(rx)[:8]}..."
        )
    finally:
        src.free()
        dst.free()


@requires_multi_pe
def test_getmem_on_stream_peer():
    """Torch-free: pull peer buffer into local buffer via rocshmem_getmem_on_stream."""
    nelems = 32
    nbytes = nelems * ctypes.sizeof(ctypes.c_int32)

    my_pe = rocshmem4py.rocshmem_my_pe()
    n_pes = rocshmem4py.rocshmem_n_pes()
    peer = (my_pe + 1) % n_pes

    src = rocshmem4py.rocshmem_create_buffer(nbytes)
    dst = rocshmem4py.rocshmem_create_buffer(nbytes)
    try:
        src_host = (ctypes.c_int32 * nelems)(*([my_pe] * nelems))
        dst_host = (ctypes.c_int32 * nelems)(*([-1] * nelems))
        h2d(src.ptr, src_host)
        h2d(dst.ptr, dst_host)
        hip_sync()

        with HipStream() as s:
            rocshmem4py.rocshmem_barrier_all_on_stream(s.handle)
            rocshmem4py.rocshmem_getmem_on_stream(
                dst.ptr, src.ptr, nbytes, peer, s.handle
            )
            rocshmem4py.rocshmem_barrier_all_on_stream(s.handle)
            s.synchronize()

        rx = (ctypes.c_int32 * nelems)()
        d2h(rx, dst.ptr)
        hip_sync()

        assert all(v == peer for v in rx), (
            f"PE {my_pe}: expected all {peer}, got {list(rx)[:8]}..."
        )
    finally:
        src.free()
        dst.free()


@requires_multi_pe
def test_get_peer_buffer():
    """``rocshmem_get_peer_buffer`` returns a zero-copy view on IPC backends.

    On backends without direct remote memory access ``rocshmem_ptr`` returns
    NULL and the function raises ``RuntimeError`` (skipped here).  Use
    ``rocshmem_getmem_on_stream`` to copy peer data on those backends —
    that path is exercised by ``test_getmem_on_stream_peer`` above.
    """
    my_pe = rocshmem4py.rocshmem_my_pe()
    n_pes = rocshmem4py.rocshmem_n_pes()
    peer = (my_pe + 1) % n_pes

    nbytes = 256
    buf = rocshmem4py.rocshmem_create_buffer(nbytes)
    try:
        rocshmem4py.rocshmem_barrier_all()
        try:
            peer_buf = rocshmem4py.rocshmem_get_peer_buffer(buf, peer)
        except RuntimeError:
            pytest.skip(
                "rocshmem_ptr returned NULL -- backend lacks direct remote "
                "memory access; copy path covered by test_getmem_on_stream_peer"
            )
        assert isinstance(peer_buf, rocshmem4py.SymmetricBuffer)
        assert peer_buf.ptr > 0
        assert peer_buf.own_data is False
        assert peer_buf.nbytes == nbytes
    finally:
        buf.free()


if __name__ == "__main__":
    pytest.main([__file__, "-v"])