TheRock/rocm-systems/projects/rccl/bindings/nccl4py/tests/test_rocm_extensions.py

# SPDX-FileCopyrightText: Copyright (c) 2025-2026 Advanced Micro Devices, Inc. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Smoke tests for the RCCL-only collective wrappers shipped in ``nccl.bindings.rocm_extensions``.

These wrappers (``all_reduce_with_bias``, ``all_to_all_v``) are
hand-written on top of the autogenerated upstream binding because the
underlying symbols (``ncclAllReduceWithBias``, ``ncclAllToAllv``) are
RCCL-only and have no NCCL equivalent.

Each test verifies that:

* the wrapper is importable and reachable via ``nccl.bindings``,
* calling it with all-zero arguments raises a controlled exception
  (:class:`NotImplementedError` if the symbol is missing in the loaded
  ``librccl.so``, or :class:`NCCLError` if the symbol is present and
  RCCL rejects the null arguments),
* no segfault reaches Python.

The wrapper checks the ``dlsym`` slot before calling into C, so dummy
zero values are safe when the symbol is missing.
"""

import pytest


def test_import_rocm_extensions():
    from nccl.bindings import all_reduce_with_bias, all_to_all_v  # noqa: F401


def test_all_reduce_with_bias_call_is_controlled():
    from nccl.bindings import all_reduce_with_bias
    from nccl.bindings.nccl import NCCLError

    with pytest.raises((NotImplementedError, NCCLError)):
        all_reduce_with_bias(0, 0, 0, 0, 0, 0, 0, 0)


def test_all_to_all_v_call_is_controlled():
    from nccl.bindings import all_to_all_v
    from nccl.bindings.nccl import NCCLError

    with pytest.raises((NotImplementedError, NCCLError)):
        all_to_all_v(0, 0, 0, 0, 0, 0, 0, 0, 0)