TheRock/rocm-systems/projects/rccl/bindings/nccl4py/nccl/bindings/rocm_extensions.pyx

# SPDX-FileCopyrightText: Copyright (c) 2025-2026 Advanced Micro Devices, Inc. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# RCCL-only collective wrappers.
#
# This module is a deliberate, hand-written companion to the autogenerated
# NVIDIA `nccl.bindings.nccl` module. It exposes RCCL-specific collectives
# that exist in librccl.so but have no equivalent in NVIDIA's NCCL ABI, so
# they are absent from the upstream nccl4py autogen we vendor.
#
# Currently exposed:
#     ncclAllReduceWithBias  -> all_reduce_with_bias
#     ncclAllToAllv          -> all_to_all_v
#
# The naming and calling conventions mirror the upstream `nccl.bindings.nccl`
# wrappers (snake_case, intptr_t pointers, exception-based error reporting
# via NCCLError) so callers see a uniform surface and don't have to know
# whether a given collective came from upstream NCCL or from RCCL.
#
# Each symbol is resolved lazily via dlsym; if the loaded librccl.so does
# not provide it, the wrapper raises NotImplementedError with the C symbol
# name in the message, matching the contract used by the autogenerated
# loader for missing NCCL >= 2.29 host APIs.

import threading

from libc.stdint cimport intptr_t

from .cynccl cimport (
    ncclResult_t,
    ncclInvalidArgument,
    ncclDataType_t,
    ncclRedOp_t,
    ncclComm_t,
    cudaStream_t,
)
from .nccl import check_status


cdef extern from "<dlfcn.h>" nogil:
    void* dlopen(const char*, int)
    char* dlerror()
    void* dlsym(void*, const char*)
    int dlclose(void*)

    enum:
        RTLD_LAZY
        RTLD_NOW
        RTLD_GLOBAL
        RTLD_LOCAL

    const void* RTLD_DEFAULT 'RTLD_DEFAULT'


###############################################################################
# Lazy loader for RCCL-only symbols
###############################################################################

cdef object __rocm_lock = threading.Lock()
cdef bint __rocm_init = False
cdef void* __ncclAllReduceWithBias = NULL
cdef void* __ncclAllToAllv = NULL


cdef void* _load_rccl() except* with gil:
    # Mirrors nccl_linux.load_library(): same RCCL search order, kept as a
    # local copy so this module stays self-contained and does not depend on
    # autogen internals. Used only when the parent loader has not been
    # triggered yet — once it has, RTLD_DEFAULT resolves the symbols below
    # without a second dlopen.
    import os
    import ctypes.util

    cdef void* handle = NULL
    cdef bytes lib_bytes
    cdef const char* err_ptr

    explicit = os.environ.get("NCCL_LIBRARY")
    if explicit:
        candidates = [explicit]
    else:
        candidates = []
        resolved = ctypes.util.find_library("rccl")
        if resolved is not None:
            candidates.append(resolved)
        for env_var in ("ROCM_PATH", "HIP_PATH"):
            root = os.environ.get(env_var)
            if root:
                candidates.append(os.path.join(root, "lib", "librccl.so"))
        candidates.append("librccl.so")

    for cand in candidates:
        lib_bytes = cand.encode("utf-8")
        handle = dlopen(lib_bytes, RTLD_NOW | RTLD_GLOBAL)
        if handle != NULL:
            return handle

    err_ptr = dlerror()
    err_str = err_ptr.decode() if err_ptr != NULL else "no diagnostic"
    raise RuntimeError(
        f"Could not load RCCL library for ROCm-only extensions. "
        f"Tried: {candidates}. Last dlerror: {err_str}."
    )


cdef int _check_or_init_rocm() except -1 nogil:
    global __rocm_init
    if __rocm_init:
        return 0

    cdef void* handle = NULL

    with gil, __rocm_lock:
        if __rocm_init:
            return 0

        global __ncclAllReduceWithBias
        __ncclAllReduceWithBias = dlsym(RTLD_DEFAULT, 'ncclAllReduceWithBias')
        if __ncclAllReduceWithBias == NULL:
            if handle == NULL:
                handle = _load_rccl()
            __ncclAllReduceWithBias = dlsym(handle, 'ncclAllReduceWithBias')

        global __ncclAllToAllv
        __ncclAllToAllv = dlsym(RTLD_DEFAULT, 'ncclAllToAllv')
        if __ncclAllToAllv == NULL:
            if handle == NULL:
                handle = _load_rccl()
            __ncclAllToAllv = dlsym(handle, 'ncclAllToAllv')

        __rocm_init = True

    return 0


###############################################################################
# Public wrappers
###############################################################################

cpdef all_reduce_with_bias(
    intptr_t sendbuff,
    intptr_t recvbuff,
    size_t count,
    int datatype,
    int op,
    intptr_t comm,
    intptr_t stream,
    intptr_t acc,
):
    """RCCL-only: ncclAllReduceWithBias.

    Reduces ``sendbuff`` across ranks with ``op``, adds the per-rank bias
    buffer ``acc`` on the receive side, and stores the result in
    ``recvbuff``. No upstream NCCL equivalent.
    """
    cdef ncclResult_t status
    _check_or_init_rocm()
    if __ncclAllReduceWithBias == NULL:
        raise NotImplementedError(
            "function ncclAllReduceWithBias is not found in the loaded RCCL library"
        )
    with nogil:
        status = (<ncclResult_t (*)(
            const void*, void*, size_t, ncclDataType_t, ncclRedOp_t,
            ncclComm_t, cudaStream_t, const void*
        ) noexcept nogil>__ncclAllReduceWithBias)(
            <const void*>sendbuff,
            <void*>recvbuff,
            count,
            <ncclDataType_t>datatype,
            <ncclRedOp_t>op,
            <ncclComm_t>comm,
            <cudaStream_t>stream,
            <const void*>acc,
        )
    check_status(<int>status)


cpdef all_to_all_v(
    intptr_t sendbuff,
    intptr_t sendcounts,
    intptr_t sdispls,
    intptr_t recvbuff,
    intptr_t recvcounts,
    intptr_t rdispls,
    int datatype,
    intptr_t comm,
    intptr_t stream,
):
    """RCCL-only: ncclAllToAllv.

    Vector AllToAll with per-rank send/receive counts and byte displacements.
    Unlike the standard ``allto_all`` (uniform counts, exposed by upstream
    NCCL), this collective is provided by RCCL only.

    ``sendcounts``/``recvcounts``/``sdispls``/``rdispls`` are ``intptr_t``
    pointers to contiguous arrays of ``size_t`` of length equal to the
    communicator size. The caller owns the storage and must keep it alive
    until the stream finishes the operation.
    """
    cdef ncclResult_t status
    _check_or_init_rocm()
    if __ncclAllToAllv == NULL:
        raise NotImplementedError(
            "function ncclAllToAllv is not found in the loaded RCCL library"
        )
    # RCCL 2.27.x's `ncclAlltoAllv_impl` derefs `comm->rank`, `sendcounts`,
    # `recvcounts` in its NVTX payload before any null-check (see
    # `projects/rccl/src/collectives.cc:276-278`); a NULL on any required
    # pointer segfaults inside librccl. Pre-validate here for a controlled
    # `NCCLError(InvalidArgument)` and a cross-RCCL-version-safe contract.
    if (comm == 0 or sendbuff == 0 or sendcounts == 0 or sdispls == 0
            or recvbuff == 0 or recvcounts == 0 or rdispls == 0):
        check_status(<int>ncclInvalidArgument)
    with nogil:
        status = (<ncclResult_t (*)(
            const void*, const size_t*, const size_t*,
            void*, const size_t*, const size_t*,
            ncclDataType_t, ncclComm_t, cudaStream_t
        ) noexcept nogil>__ncclAllToAllv)(
            <const void*>sendbuff,
            <const size_t*>sendcounts,
            <const size_t*>sdispls,
            <void*>recvbuff,
            <const size_t*>recvcounts,
            <const size_t*>rdispls,
            <ncclDataType_t>datatype,
            <ncclComm_t>comm,
            <cudaStream_t>stream,
        )
    check_status(<int>status)