diff --git a/python/build_assets/README.md b/python/build_assets/README.md deleted file mode 100644 index 7f568fe2..00000000 --- a/python/build_assets/README.md +++ /dev/null @@ -1,98 +0,0 @@ -# Build Assets for the nccl-extensions Python package - -> **⚠️ Internal Use Only**: This directory is **not released** to the public -> nccl-extensions GitHub repository. - -## Purpose - -The Cython bindings under `python/nccl/_extensions/bindings/` are **generated**, -not hand-written. This directory holds the tooling that generates them: the -cybind config, the Cython templates, and the driver script. The M2N low-level -surface is generated into the shared `nccl._extensions.bindings` package while -its public facade remains under `python/nccl/m2n/`. - -The generated `.pyx`/`.pxd` files are checked into the repository so users can -build the package without internal tooling; the generation tooling here stays -internal. - -## Why not public? - -Generation relies on [cybind](https://gitlab-master.nvidia.com/leof/cybind), an -internal NVIDIA tool. To keep the package buildable without access to it, we: - -1. Run `generate_cython.py` when a bound library's headers change -2. Commit the generated Cython sources -3. Exclude this `build_assets/` directory from public releases - -Ported from nccl4py's `bindings/nccl4py/build_assets/`, trimmed to the targets -this repo owns. nccl4py's `generate_header.py` (which flattens the NCCL *device* -API headers) has no nccl-extensions equivalent and was not ported. - -## Prerequisites - -- CUDA installation, with `CUDA_HOME` or `CUDA_PATH` set (cybind needs `cuda.h`) -- [`uv`](https://docs.astral.sh/uv/) on `PATH` -- SSH access to the cybind repository (unless `--cybind-path` points at a local - checkout) - -## Usage - -Re-run whenever `nccl_ep/include/nccl_ep.h`, `nccl_m2n/src/nccl_m2n.h`, their -templates/configs, or the pinned NCCL header changes: - -```bash -python3 build_assets/generate_cython.py --verbose -``` - -The script clones cybind at the pinned `CYBIND_COMMIT`, stages our configs, -headers and templates into its `assets/`, then regenerates the -complete shared `python/nccl/_extensions/bindings/` package transactionally. -Commit the resulting diff. - -See `generate_cython.py --help` for all options; `--cybind-path` reuses a local -cybind checkout instead of cloning. - -## Headers - -Two different policies, on purpose: - -- **`nccl_ep` headers are not checked in here.** They live in this repo at - `nccl_ep/include/`, and are staged straight from there. The bound version is - read from `NCCL_EP_{MAJOR,MINOR,PATCH}` in `nccl_ep.h` and stamped into - `cybind/configs/nccl_ep.cybind.yaml`, so bindings can never drift from the - header they were generated against. -- **Headers this repo does not own are pinned** under - `cybind/headers///`, matching cybind's own layout. - Currently just `nccl.h`, at the version in `NCCL_PIN`; it fixes the NCCL core - ABI the generated bindings were built against. Bumping it means dropping the - new `nccl.h` in place, updating `NCCL_PIN`, and regenerating. - -## Contents - -- `generate_cython.py` — driver: stages assets, runs cybind, installs output -- `cybind/configs/{nccl_ep,nccl_m2n}.cybind.yaml` — cybind configs for the - bound libraries (including `AUTO_LOWPP_CLASS` struct overrides) -- `cybind/templates/nccl/_extensions/bindings/` — Cython templates, plus the - static files (`_internal/utils.{pxd,pyx}`, `__init__.py`) that cybind does not - process and `generate_cython.py` copies verbatim -- `cybind/headers/` — pinned third-party headers (see above) - -## Generated binding conventions - -All bound libraries share `nccl/_extensions/bindings/`, so common generated -support such as `_internal/utils.pyx` and `_binding_helpers.py` is built and -shipped once. - -A library may provide dedicated templates when its ABI or native-loader -contract differs from the common case. Keep those differences inside the -generated binding layer: - -- add the library configuration under `cybind/configs/`; -- add any library-specific templates under - `cybind/templates/nccl/_extensions/bindings/`; -- keep public, framework-facing APIs in that library's facade package; -- preserve actionable loader errors and ensure native symbols resolve from the - intended library handle. - -After changing a bound header, configuration, or template, regenerate the -complete bindings package and commit the generated diff. diff --git a/python/build_assets/cybind/configs/nccl_ep.cybind.yaml b/python/build_assets/cybind/configs/nccl_ep.cybind.yaml deleted file mode 100644 index b088fffd..00000000 --- a/python/build_assets/cybind/configs/nccl_ep.cybind.yaml +++ /dev/null @@ -1,175 +0,0 @@ -nccl_ep: - module: nccl._extensions.bindings.nccl_ep - data: - versions: - - - 0.1.0 - headers: - - nccl_ep.h - patterns: - function: (^ncclEp)([A-Z].*) - type: (^ncclEp)([A-Z][A-Za-z_]*)_t$|(^nccl)(Comm|Window)_t$ - need_cuda: true - need_driver: false - support_win: false - docstrings: null - need_headers_at_build: false - include_cybind_version: false - status_enum: ncclResult_t - extern: - cython: - - module: nccl.bindings.cynccl - import_type: cimport - imports: - - ncclResult_t - - ncclDataType_t - - _NCCLRESULT_T_INTERNAL_LOADING_ERROR - functions: - ncclEpGetVersion: - return: version - except?: -1 - - # Group lifecycle. - # Allocator hooks live inside ncclEpGroupConfig_t::alloc (ncclEpAllocConfig_t); - # see the Pythonic wrapper in nccl/ep/allocator.py for how alloc_fn / free_fn - # / context get plumbed through. - ncclEpCreateGroup: - return: ep_group - except?: 0 - ncclEpGroupDestroy: {} - - # Tensor lifecycle. - ncclEpTensorAlloc: {} - ncclEpTensorDestroy: {} - - # Handle lifecycle. - # Named-struct ABI: ncclEp*Config_t / ncclEpLayoutInfo_t are declared - # below as AUTO_LOWPP_CLASS so cybind exposes them as Python classes; - # the Pythonic facade in nccl/ep/handle.py constructs them and forwards - # the underlying struct pointer via .ptr. - ncclEpCreateHandle: - return: handle - except?: 0 - ncclEpInitHandle: - return: handle - except?: 0 - ncclEpHandleDestroy: {} - ncclEpHandleMemSize: - return: size_out - except?: -1 - ncclEpUpdateHandle: {} - - # Dispatch / Combine / Complete (named-struct ABI; see note above). - ncclEpDispatch: {} - ncclEpCombine: {} - ncclEpComplete: {} - - types: - ncclEpTensor_t: - AUTO_LOWPP_CLASS: - __config__: - __init__: - post_alloc: |- - self._ptr[0].size = sizeof(ncclEpTensor_t) - self._ptr[0].magic = 0xCAFECAFE - ncclEpTensorAllocConfig_t: - AUTO_LOWPP_CLASS: - __config__: - __init__: - post_alloc: |- - self._ptr[0].size = sizeof(ncclEpTensorAllocConfig_t) - self._ptr[0].magic = 0xC00FFFEE - # ncclEpAllocConfig_t has C function-pointer fields; cybind can't auto- - # convert ncclEpAlloc/FreeFn_t to/from Python, so route through intptr_t - # via per-member getter/setter overrides. Users pass and read raw - # function addresses (e.g. from ctypes.cast(fn, ctypes.c_void_p).value). - ncclEpAllocConfig_t: - AUTO_LOWPP_CLASS: - alloc_fn: - getter: |- - return (self._ptr[0].alloc_fn) - setter: |- - self._ptr[0].alloc_fn = val - free_fn: - getter: |- - return (self._ptr[0].free_fn) - setter: |- - self._ptr[0].free_fn = val - ncclEpGroupConfig_t: - AUTO_LOWPP_CLASS: - __config__: - __init__: - post_alloc: |- - self._ptr[0].size = sizeof(ncclEpGroupConfig_t) - self._ptr[0].magic = 0xC00FFFEE - self._ptr[0].version = 1 - ncclEpLayoutInfo_t: - AUTO_LOWPP_CLASS: - __config__: - __init__: - post_alloc: |- - self._ptr[0].size = sizeof(ncclEpLayoutInfo_t) - self._ptr[0].magic = 0xC00FFFEE - ncclEpDispatchInputs_t: - AUTO_LOWPP_CLASS: - __config__: - __init__: - post_alloc: |- - self._ptr[0].size = sizeof(ncclEpDispatchInputs_t) - self._ptr[0].magic = 0xC00FFFEE - ncclEpDispatchOutputs_t: - AUTO_LOWPP_CLASS: - __config__: - __init__: - post_alloc: |- - self._ptr[0].size = sizeof(ncclEpDispatchOutputs_t) - self._ptr[0].magic = 0xC00FFFEE - ncclEpCombineInputs_t: - AUTO_LOWPP_CLASS: - __config__: - __init__: - post_alloc: |- - self._ptr[0].size = sizeof(ncclEpCombineInputs_t) - self._ptr[0].magic = 0xC00FFFEE - ncclEpCombineOutputs_t: - AUTO_LOWPP_CLASS: - __config__: - __init__: - post_alloc: |- - self._ptr[0].size = sizeof(ncclEpCombineOutputs_t) - self._ptr[0].magic = 0xC00FFFEE - ncclEpHandleConfig_t: - AUTO_LOWPP_CLASS: - __config__: - __init__: - post_alloc: |- - self._ptr[0].size = sizeof(ncclEpHandleConfig_t) - self._ptr[0].magic = 0xC00FFFEE - ncclEpDispatchConfig_t: - AUTO_LOWPP_CLASS: - __config__: - __init__: - post_alloc: |- - self._ptr[0].size = sizeof(ncclEpDispatchConfig_t) - self._ptr[0].magic = 0xC00FFFEE - ncclEpCombineConfig_t: - AUTO_LOWPP_CLASS: - __config__: - __init__: - post_alloc: |- - self._ptr[0].size = sizeof(ncclEpCombineConfig_t) - self._ptr[0].magic = 0xC00FFFEE - ncclEpCompleteConfig_t: - AUTO_LOWPP_CLASS: - __config__: - __init__: - post_alloc: |- - self._ptr[0].size = sizeof(ncclEpCompleteConfig_t) - self._ptr[0].magic = 0xC00FFFEE - - attrs: {} - enums: - ^nccl: '' - ^NCCL_EP_TENSOR_FLAG_: '' - ^ALGO_: '' - ^LAYOUT_: '' - ^TENSOR_TAG_: '' diff --git a/python/build_assets/cybind/configs/nccl_m2n.cybind.yaml b/python/build_assets/cybind/configs/nccl_m2n.cybind.yaml deleted file mode 100644 index 51e2fa91..00000000 --- a/python/build_assets/cybind/configs/nccl_m2n.cybind.yaml +++ /dev/null @@ -1,96 +0,0 @@ -nccl_m2n: - module: nccl._extensions.bindings.nccl_m2n - data: - versions: - - - 0.2.0 - headers: - - nccl_m2n.h - patterns: - function: (^nccl)(M2n.*|Reshard.*) - type: (^ncclM2n)([A-Z].*)_t$|(^nccl)(Mesh|DistTensor)_t$ - need_cuda: true - need_driver: false - support_win: false - docstrings: null - # The generated declarations encode this public ABI. Wheel builds must not - # need a separately installed M2N development header. - need_headers_at_build: false - include_cybind_version: false - status_enum: ncclResult_t - extern: - cython: - - module: nccl.bindings.cynccl - import_type: cimport - imports: - - ncclResult_t - - ncclDataType_t - - ncclComm_t - - ncclWindow_t - - _NCCLRESULT_T_INTERNAL_LOADING_ERROR - functions: - ncclM2nInit: - return: handle - except?: 0 - ncclM2nFinalize: {} - ncclM2nGroupStart: {} - ncclM2nGroupEnd: {} - ncclM2nGroupAbort: {} - ncclReshardWithWindow: {} - ncclReshard: {} - types: - ncclM2nConfig_t: - AUTO_LOWPP_CLASS: - __config__: - __init__: - post_alloc: |- - self._ptr[0].size = sizeof(ncclM2nConfig_t) - self._ptr[0].magic = NCCL_M2N_API_MAGIC - self._ptr[0].version = NCCL_M2N_API_VERSION - self._ptr[0].maxCta = NCCL_M2N_CONFIG_UNDEF_INT - member_aliases: - maxCta: maxCta - ncclMesh_t: - AUTO_LOWPP_CLASS: - dims: - getter: |- - return (self._ptr[0].dims[0], self._ptr[0].dims[1]) - setter: |- - if len(val) != NCCL_RESHARD_MESH_NDIMS: - raise ValueError(f"dims must have length {NCCL_RESHARD_MESH_NDIMS}") - self._ptr[0].dims[0] = val[0] - self._ptr[0].dims[1] = val[1] - member_aliases: - startRank: startRank - ncclDistTensor_t: - AUTO_LOWPP_CLASS: - localShape: - getter: |- - return (self._ptr[0].localShape[0], self._ptr[0].localShape[1], self._ptr[0].localShape[2]) - setter: |- - if len(val) != NCCL_RESHARD_MAX_TENSOR_DIMS: - raise ValueError(f"localShape must have length {NCCL_RESHARD_MAX_TENSOR_DIMS}") - self._ptr[0].localShape[0] = val[0] - self._ptr[0].localShape[1] = val[1] - self._ptr[0].localShape[2] = val[2] - mesh: - getter: |- - return (self._ptr[0].mesh) - setter: |- - self._ptr[0].mesh = val - placements: - getter: |- - return (self._ptr[0].placements[0], self._ptr[0].placements[1]) - setter: |- - if len(val) != NCCL_RESHARD_MESH_NDIMS: - raise ValueError(f"placements must have length {NCCL_RESHARD_MESH_NDIMS}") - self._ptr[0].placements[0] = val[0] - self._ptr[0].placements[1] = val[1] - member_aliases: - dataPtr: dataPtr - localShape: localShape - dtype: dtype - attrs: {} - enums: - ^nccl: '' - ^NCCL_RESHARD_: '' - ^NCCL_M2N_: '' diff --git a/python/build_assets/cybind/headers/nccl/2.30.4/nccl.h b/python/build_assets/cybind/headers/nccl/2.30.4/nccl.h deleted file mode 100644 index 551872c3..00000000 --- a/python/build_assets/cybind/headers/nccl/2.30.4/nccl.h +++ /dev/null @@ -1,842 +0,0 @@ -/************************************************************************* - * SPDX-FileCopyrightText: Copyright (c) 2015-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - * - * See LICENSE.txt for more license information - *************************************************************************/ - -#ifndef NCCL_H_ -#define NCCL_H_ - -#include -#include -#if CUDART_VERSION >= 11000 -#include -#endif -#if __cplusplus && CUDART_VERSION >= 11080 -#include -#endif - -#define NCCL_MAJOR 2 -#define NCCL_MINOR 30 -#define NCCL_PATCH 4 -#define NCCL_SUFFIX "" - -#define NCCL_VERSION_CODE 23004 -#define NCCL_VERSION(X,Y,Z) (((X) <= 2 && (Y) <= 8) ? (X) * 1000 + (Y) * 100 + (Z) : (X) * 10000 + (Y) * 100 + (Z)) - -#ifdef __cplusplus -extern "C" { -#endif - -#include -#include - -/* Opaque handle to communicator */ -typedef struct ncclComm* ncclComm_t; -typedef struct ncclWindow_vidmem* ncclWindow_t; -#define NCCL_COMM_NULL NULL - -#define NCCL_UNIQUE_ID_BYTES 128 -typedef struct { char internal[NCCL_UNIQUE_ID_BYTES]; } ncclUniqueId; - -/* Error type */ -typedef enum { ncclSuccess = 0, - ncclUnhandledCudaError = 1, - ncclSystemError = 2, - ncclInternalError = 3, - ncclInvalidArgument = 4, - ncclInvalidUsage = 5, - ncclRemoteError = 6, - ncclInProgress = 7, - ncclTimeout = 8, - ncclNumResults = 9 } ncclResult_t; - -#define NCCL_CONFIG_UNDEF_INT INT_MIN -#define NCCL_CONFIG_UNDEF_PTR NULL -#define NCCL_SPLIT_NOCOLOR -1 -#define NCCL_UNDEF_FLOAT -1.0f - -/* Internal use only */ -#define NCCL_API_MAGIC 0xcafebeef - -/* Window Registration flags */ -#define NCCL_WIN_DEFAULT 0x00 -#define NCCL_WIN_COLL_SYMMETRIC 0x01 -#define NCCL_WIN_STRICT_ORDERING 0x02 - -#define NCCL_WIN_REQUIRED_ALIGNMENT 4096 - -/* NCCL performance policy */ -#define NCCL_CTA_POLICY_DEFAULT 0x00 -#define NCCL_CTA_POLICY_EFFICIENCY 0x01 -#define NCCL_CTA_POLICY_ZERO 0x02 - -/* ncclCommShrink flags*/ -#define NCCL_SHRINK_DEFAULT 0x00 /* shrink the parent communicator */ -#define NCCL_SHRINK_ABORT 0x01 /* First, terminate ongoing parent operations, and then shrink the parent communicator */ - -/* ncclCommRevoke flags */ -#define NCCL_REVOKE_DEFAULT 0x00 /* reserved for future use; must be 0 */ - -/* Communicator configuration. Users can assign value to attributes to specify the - * behavior of a communicator. */ -typedef struct ncclConfig_v22800 { - /* attributes that users should never touch. */ - size_t size; - unsigned int magic; - unsigned int version; - /* attributes that users are able to customize. */ - int blocking; - int cgaClusterSize; - int minCTAs; - int maxCTAs; - const char *netName; - int splitShare; - int trafficClass; - const char *commName; - int collnetEnable; - int CTAPolicy; - int shrinkShare; - int nvlsCTAs; - int nChannelsPerNetPeer; - int nvlinkCentricSched; - int graphUsageMode; - int numRmaCtx; - int maxP2pPeers; -} ncclConfig_t; - -/* Config initializer must be assigned to initialize config structure when it is created. - * Not initialized config will result in NCCL error. */ -#define NCCL_CONFIG_INITIALIZER { \ - sizeof(ncclConfig_t), /* size */ \ - NCCL_API_MAGIC, /* magic */ \ - NCCL_VERSION_CODE, /* version */ \ - NCCL_CONFIG_UNDEF_INT, /* blocking */ \ - NCCL_CONFIG_UNDEF_INT, /* cgaClusterSize */ \ - NCCL_CONFIG_UNDEF_INT, /* minCTAs */ \ - NCCL_CONFIG_UNDEF_INT, /* maxCTAs */ \ - NCCL_CONFIG_UNDEF_PTR, /* netName */ \ - NCCL_CONFIG_UNDEF_INT, /* splitShare */ \ - NCCL_CONFIG_UNDEF_INT, /* trafficClass */ \ - NCCL_CONFIG_UNDEF_PTR, /* commName */ \ - NCCL_CONFIG_UNDEF_INT, /* collnetEnable */ \ - NCCL_CONFIG_UNDEF_INT, /* CTAPolicy */ \ - NCCL_CONFIG_UNDEF_INT, /* shrinkShare */ \ - NCCL_CONFIG_UNDEF_INT, /* nvlsCTAs */ \ - NCCL_CONFIG_UNDEF_INT, /* nChannelsPerNetPeer */ \ - NCCL_CONFIG_UNDEF_INT, /* nvlinkCentricSched */ \ - NCCL_CONFIG_UNDEF_INT, /* graphUsageMode */ \ - NCCL_CONFIG_UNDEF_INT, /* numRmaCtx */ \ - NCCL_CONFIG_UNDEF_INT, /* maxP2pPeers */ \ -} - -/* This struct will be used by ncclGroupSimulateEnd() API to query information about simulation. */ -typedef struct ncclSimInfo_v22200 { - size_t size; - unsigned int magic; - unsigned int version; - float estimatedTime; -} ncclSimInfo_t; - -/* NCCL_SIM_INFO_INITIALIZER must be assigned to initialize simInfo structure when it is created. - * Not initialized simInfo will result in NCCL error. */ -#define NCCL_SIM_INFO_INITIALIZER { \ - sizeof(ncclSimInfo_t), /* size */ \ - 0x74685283, /* magic */ \ - NCCL_VERSION_CODE, /* version */ \ - NCCL_UNDEF_FLOAT /* estimated time */ \ -} - -/* NCCL malloc and free function for all types of NCCL optimizations - * (e.g. user buffer registration). The actual allocated size might - * be larger than requested due to granularity requirement. */ -ncclResult_t ncclMemAlloc(void** ptr, size_t size); -ncclResult_t pncclMemAlloc(void** ptr, size_t size); - -ncclResult_t ncclMemFree(void *ptr); -ncclResult_t pncclMemFree(void *ptr); - -/* Return the NCCL_VERSION_CODE of the NCCL library in the supplied integer. - * This integer is coded with the MAJOR, MINOR and PATCH level of the - * NCCL library - */ -ncclResult_t ncclGetVersion(int *version); -ncclResult_t pncclGetVersion(int *version); - -/* Generates an Id to be used in ncclCommInitRank. ncclGetUniqueId should be - * called once and the Id should be distributed to all ranks in the - * communicator before calling ncclCommInitRank. */ -ncclResult_t ncclGetUniqueId(ncclUniqueId* uniqueId); -ncclResult_t pncclGetUniqueId(ncclUniqueId* uniqueId); - -/* Create a new communicator (multi thread/process version) with a configuration - * set by users. */ -ncclResult_t ncclCommInitRankConfig(ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank, ncclConfig_t* config); -ncclResult_t pncclCommInitRankConfig(ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank, ncclConfig_t* config); - -/* Creates a new communicator (multi thread/process version). - * rank must be between 0 and nranks-1 and unique within a communicator clique. - * Each rank is associated to a CUDA device, which has to be set before calling - * ncclCommInitRank. - * ncclCommInitRank implicitly syncronizes with other ranks, so it must be - * called by different threads/processes or use ncclGroupStart/ncclGroupEnd. */ -ncclResult_t ncclCommInitRank(ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank); -ncclResult_t pncclCommInitRank(ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank); - -/* Creates a clique of communicators (single process version). - * This is a convenience function to create a single-process communicator clique. - * Returns an array of ndev newly initialized communicators in comm. - * comm should be pre-allocated with size at least ndev*sizeof(ncclComm_t). - * If devlist is NULL, the first ndev CUDA devices are used. - * Order of devlist defines user-order of processors within the communicator. */ -ncclResult_t ncclCommInitAll(ncclComm_t* comm, int ndev, const int* devlist); -ncclResult_t pncclCommInitAll(ncclComm_t* comm, int ndev, const int* devlist); - -/* Finalize a communicator. ncclCommFinalize flushes all issued communications, - * and marks communicator state as ncclInProgress. The state will change to ncclSuccess - * when the communicator is globally quiescent and related resources are freed; then, - * calling ncclCommDestroy can locally free the rest of the resources (e.g. communicator - * itself) without blocking. */ -ncclResult_t ncclCommFinalize(ncclComm_t comm); -ncclResult_t pncclCommFinalize(ncclComm_t comm); - -/* Frees local resources associated with communicator object. */ -ncclResult_t ncclCommDestroy(ncclComm_t comm); -ncclResult_t pncclCommDestroy(ncclComm_t comm); - -/* Frees resources associated with communicator object and aborts any operations - * that might still be running on the device. */ -ncclResult_t ncclCommAbort(ncclComm_t comm); -ncclResult_t pncclCommAbort(ncclComm_t comm); - -/* Revoke a communicator. ncclCommRevoke stops all in-flight operations - * and marks communicator state as ncclInProgress. The state will change to ncclSuccess - * when the communicator is quiescent; then, management operations (destroy, split, - * shrink) can proceed safely. Calling ncclCommFinalize after revoke is invalid. - * Additionally, resource sharing via splitShare/shrinkShare is disabled while revoked. - * revokeFlags must be NCCL_REVOKE_DEFAULT (0). */ -ncclResult_t ncclCommRevoke(ncclComm_t comm, int revokeFlags); -ncclResult_t pncclCommRevoke(ncclComm_t comm, int revokeFlags); - -/* Creates one or more communicators from an existing one. - * Ranks with the same color will end up in the same communicator. - * Within the new communicator, key will be used to order ranks. - * NCCL_SPLIT_NOCOLOR as color will indicate the rank will not be part of any group - * and will therefore return a NULL communicator. - * If config is NULL, the new communicator will inherit the original communicator's - * configuration*/ -ncclResult_t ncclCommSplit(ncclComm_t comm, int color, int key, ncclComm_t *newcomm, ncclConfig_t* config); -ncclResult_t pncclCommSplit(ncclComm_t comm, int color, int key, ncclComm_t *newcomm, ncclConfig_t* config); - -/* Shrink existing communicator. - * Ranks in excludeRanksList will be removed form the existing communicator. - * Within the new communicator, ranks will be re-ordered to fill the gap of removed ones. - * If config is NULL, the new communicator will inherit the original communicator's configuration - * The flag enables NCCL to adapt to various states of the parent communicator, see NCCL_SHRINK flags.*/ -ncclResult_t ncclCommShrink(ncclComm_t comm, int* excludeRanksList, int excludeRanksCount, ncclComm_t* newcomm, ncclConfig_t* config, int shrinkFlags); -ncclResult_t pncclCommShrink(ncclComm_t comm, int* excludeRanksList, int excludeRanksCount, ncclComm_t* newcomm, ncclConfig_t* config, int shrinkFlags); - -/* Generate per-communicator unique ID for grow. - * Constraints: - * - Cannot generate a new UID while a previous UID is unconsumed - * - Each UID can only be used once (no reuse after consumption) - * - Must wait for grow operation to complete before calling again */ -ncclResult_t ncclCommGetUniqueId(ncclComm_t comm, ncclUniqueId* uniqueId); -ncclResult_t pncclCommGetUniqueId(ncclComm_t comm, ncclUniqueId* uniqueId); - -/* Grow communicator by adding new ranks. - * Parameter usage: - * - Existing non-root: comm, uniqueId=NULL, rank=-1 - * - Existing root: comm, uniqueId=&id, rank=-1 - * - New ranks: comm=NULL, uniqueId=&id, rank=assigned - * The UID is consumed upon successful grow and cannot be reused. */ -ncclResult_t ncclCommGrow(ncclComm_t comm, int nRanks, const ncclUniqueId* uniqueId, int rank, ncclComm_t* newcomm, ncclConfig_t* config); -ncclResult_t pncclCommGrow(ncclComm_t comm, int nRanks, const ncclUniqueId* uniqueId, int rank, ncclComm_t* newcomm, ncclConfig_t* config); - -/* Creates a new communicator (multi thread/process version), similar to ncclCommInitRankConfig. - * Allows to use more than one ncclUniqueId (up to one per rank), indicated by nId, to accelerate the init operation. - * The number of ncclUniqueIds and their order must be the same for every rank. - */ -ncclResult_t ncclCommInitRankScalable(ncclComm_t* newcomm, int nranks, int myrank, int nId, ncclUniqueId* commIds, ncclConfig_t* config); -ncclResult_t pncclCommInitRankScalable(ncclComm_t* newcomm, int nranks, int myrank, int nId, ncclUniqueId* commIds, ncclConfig_t* config); - -/* Returns a string for each error code. */ -const char* ncclGetErrorString(ncclResult_t result); -const char* pncclGetErrorString(ncclResult_t result); - -/* Returns a human-readable message of the last error that occurred. */ -const char* ncclGetLastError(ncclComm_t comm); -const char* pncclGetLastError(ncclComm_t comm); - -#ifdef NCCL_OS_LINUX - /* Reload environment variables that determine logging. */ - __attribute__ ((deprecated("ncclResetDebugInit is not supported as part of the NCCL API and will be removed in the future"))) - void ncclResetDebugInit(); - __attribute__ ((deprecated("pncclResetDebugInit is not supported as part of the NCCL API and will be removed in the future"))) - void pncclResetDebugInit(); -#else - #define ncclResetDebugInit() - #define pncclResetDebugInit() -#endif - -/* Checks whether the comm has encountered any asynchronous errors */ -ncclResult_t ncclCommGetAsyncError(ncclComm_t comm, ncclResult_t *asyncError); -ncclResult_t pncclCommGetAsyncError(ncclComm_t comm, ncclResult_t *asyncError); - -/* Gets the number of ranks in the communicator clique. */ -ncclResult_t ncclCommCount(const ncclComm_t comm, int* count); -ncclResult_t pncclCommCount(const ncclComm_t comm, int* count); - -/* Returns the cuda device number associated with the communicator. */ -ncclResult_t ncclCommCuDevice(const ncclComm_t comm, int* device); -ncclResult_t pncclCommCuDevice(const ncclComm_t comm, int* device); - -/* Returns the user-ordered "rank" associated with the communicator. */ -ncclResult_t ncclCommUserRank(const ncclComm_t comm, int* rank); -ncclResult_t pncclCommUserRank(const ncclComm_t comm, int* rank); - -/* Register CUDA buffer for zero-copy operation */ -ncclResult_t ncclCommRegister(const ncclComm_t comm, void* buff, size_t size, void** handle); -ncclResult_t pncclCommRegister(const ncclComm_t comm, void* buff, size_t size, void** handle); - -/* Deregister CUDA buffer */ -ncclResult_t ncclCommDeregister(const ncclComm_t comm, void* handle); -ncclResult_t pncclCommDeregister(const ncclComm_t comm, void* handle); - -/* Communicator suspend flags */ -#define NCCL_SUSPEND_MEM 0x01 // Suspend memory (release dynamic allocations) - -/* - * ncclCommSuspend - * - * Suspend communicator operations to free resources. - * The communicator cannot be used while suspended. - * - * flags: NCCL_SUSPEND_MEM - Release dynamic GPU memory allocations - */ -ncclResult_t ncclCommSuspend(ncclComm_t comm, int flags); -ncclResult_t pncclCommSuspend(ncclComm_t comm, int flags); - -/* - * ncclCommResume - * - * Resume all previously suspended communicator resources. - */ -ncclResult_t ncclCommResume(ncclComm_t comm); -ncclResult_t pncclCommResume(ncclComm_t comm); - -/* Communicator memory statistics */ -typedef enum { - ncclStatGpuMemSuspend = 0, // Allocated GPU memory that can be suspended (bytes) - ncclStatGpuMemSuspended = 1, // GPU memory suspended? (0=active, 1=suspended) - ncclStatGpuMemPersist = 2, // Allocated GPU memory that cannot be suspended (bytes) - ncclStatGpuMemTotal = 3 // Total allocated GPU memory tracked by NCCL (bytes) -} ncclCommMemStat_t; - -/* - * ncclCommMemStats - * - * Query communicator memory statistics. - * - * stat: One of ncclCommMemStat_t values - * value: Output pointer to receive the memory statistic value - */ -ncclResult_t ncclCommMemStats(ncclComm_t comm, ncclCommMemStat_t stat, uint64_t* value); -ncclResult_t pncclCommMemStats(ncclComm_t comm, ncclCommMemStat_t stat, uint64_t* value); - -/* Register memory window */ -ncclResult_t ncclCommWindowRegister(ncclComm_t comm, void* buff, size_t size, ncclWindow_t* win, int winFlags); -ncclResult_t pncclCommWindowRegister(ncclComm_t comm, void* buff, size_t size, ncclWindow_t* win, int winFlags); - -/* Deregister symmetric memory */ -ncclResult_t ncclCommWindowDeregister(ncclComm_t comm, ncclWindow_t win); -ncclResult_t pncclCommWindowDeregister(ncclComm_t comm, ncclWindow_t win); - -/* Get the user pointer from the window */ -ncclResult_t ncclWinGetUserPtr(ncclComm_t comm, ncclWindow_t win, void** outUserPtr); -ncclResult_t pncclWinGetUserPtr(ncclComm_t comm, ncclWindow_t win, void** outUserPtr); - -/* Reduction operation selector */ -typedef enum { ncclNumOps_dummy = 5 } ncclRedOp_dummy_t; -typedef enum { ncclSum = 0, - ncclProd = 1, - ncclMax = 2, - ncclMin = 3, - ncclAvg = 4, - /* ncclNumOps: The number of built-in ncclRedOp_t values. Also - * serves as the least possible value for dynamic ncclRedOp_t's - * as constructed by ncclRedOpCreate*** functions. */ - ncclNumOps = 5, - /* ncclMaxRedOp: The largest valid value for ncclRedOp_t. - * It is defined to be the largest signed value (since compilers - * are permitted to use signed enums) that won't grow - * sizeof(ncclRedOp_t) when compared to previous NCCL versions to - * maintain ABI compatibility. */ - ncclMaxRedOp = 0x7fffffff>>(32-8*sizeof(ncclRedOp_dummy_t)) - } ncclRedOp_t; - -/* Data types */ -typedef enum { ncclInt8 = 0, ncclChar = 0, - ncclUint8 = 1, - ncclInt32 = 2, ncclInt = 2, - ncclUint32 = 3, - ncclInt64 = 4, - ncclUint64 = 5, - ncclFloat16 = 6, ncclHalf = 6, - ncclFloat32 = 7, ncclFloat = 7, - ncclFloat64 = 8, ncclDouble = 8, - ncclBfloat16 = 9, - ncclFloat8e4m3 = 10, - ncclFloat8e5m2 = 11, - ncclNumTypes = 12 -} ncclDataType_t; - -/* ncclScalarResidence_t: Location and dereferencing logic for scalar arguments. */ -typedef enum { - /* ncclScalarDevice: The scalar is in device-visible memory and will be - * dereferenced while the collective is running. */ - ncclScalarDevice = 0, - - /* ncclScalarHostImmediate: The scalar is in host-visible memory and will be - * dereferenced before the ncclRedOpCreate***() function returns. */ - ncclScalarHostImmediate = 1 -} ncclScalarResidence_t; - -/* - * ncclRedOpCreatePreMulSum - * - * Creates a new reduction operator which pre-multiplies input values by a given - * scalar locally before reducing them with peer values via summation. For use - * only with collectives launched against *comm* and *datatype*. The - * *residence* argument indicates how/when the memory pointed to by *scalar* - * will be dereferenced. Upon return, the newly created operator's handle - * is stored in *op*. - */ -ncclResult_t ncclRedOpCreatePreMulSum(ncclRedOp_t *op, void *scalar, ncclDataType_t datatype, ncclScalarResidence_t residence, ncclComm_t comm); -ncclResult_t pncclRedOpCreatePreMulSum(ncclRedOp_t *op, void *scalar, ncclDataType_t datatype, ncclScalarResidence_t residence, ncclComm_t comm); - -/* - * ncclRedOpDestroy - * - * Destroys the reduction operator *op*. The operator must have been created by - * ncclRedOpCreatePreMul with the matching communicator *comm*. An operator may be - * destroyed as soon as the last NCCL function which is given that operator returns. - */ -ncclResult_t ncclRedOpDestroy(ncclRedOp_t op, ncclComm_t comm); -ncclResult_t pncclRedOpDestroy(ncclRedOp_t op, ncclComm_t comm); - -/* - * Collective communication operations - * - * Collective communication operations must be called separately for each - * communicator in a communicator clique. - * - * They return when operations have been enqueued on the CUDA stream. - * - * Since they may perform inter-CPU synchronization, each call has to be done - * from a different thread or process, or need to use Group Semantics (see - * below). - */ - -/* - * Reduce - * - * Reduces data arrays of length count in sendbuff into recvbuff using op - * operation. - * recvbuff may be NULL on all calls except for root device. - * root is the rank (not the CUDA device) where data will reside after the - * operation is complete. - * - * In-place operation will happen if sendbuff == recvbuff. - */ -ncclResult_t ncclReduce(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype, - ncclRedOp_t op, int root, ncclComm_t comm, cudaStream_t stream); -ncclResult_t pncclReduce(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype, - ncclRedOp_t op, int root, ncclComm_t comm, cudaStream_t stream); - -/* - * (deprecated) Broadcast (in-place) - * - * Copies count values from root to all other devices. - * root is the rank (not the CUDA device) where data resides before the - * operation is started. - * - * This operation is implicitely in place. - */ -ncclResult_t ncclBcast(void* buff, size_t count, ncclDataType_t datatype, int root, - ncclComm_t comm, cudaStream_t stream); -ncclResult_t pncclBcast(void* buff, size_t count, ncclDataType_t datatype, int root, - ncclComm_t comm, cudaStream_t stream); - -/* - * Broadcast - * - * Copies count values from root to all other devices. - * root is the rank (not the CUDA device) where data resides before the - * operation is started. - * - * In-place operation will happen if sendbuff == recvbuff. - */ -ncclResult_t ncclBroadcast(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype, int root, - ncclComm_t comm, cudaStream_t stream); -ncclResult_t pncclBroadcast(const void* sendbuff, void* recvbuff, size_t count, ncclDataType_t datatype, int root, - ncclComm_t comm, cudaStream_t stream); - -/* - * All-Reduce - * - * Reduces data arrays of length count in sendbuff using op operation, and - * leaves identical copies of result on each recvbuff. - * - * In-place operation will happen if sendbuff == recvbuff. - */ -ncclResult_t ncclAllReduce(const void* sendbuff, void* recvbuff, size_t count, - ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm, cudaStream_t stream); -ncclResult_t pncclAllReduce(const void* sendbuff, void* recvbuff, size_t count, - ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm, cudaStream_t stream); - -/* - * Reduce-Scatter - * - * Reduces data in sendbuff using op operation and leaves reduced result - * scattered over the devices so that recvbuff on rank i will contain the i-th - * block of the result. - * Assumes sendcount is equal to nranks*recvcount, which means that sendbuff - * should have a size of at least nranks*recvcount elements. - * - * In-place operations will happen if recvbuff == sendbuff + rank * recvcount. - */ -ncclResult_t ncclReduceScatter(const void* sendbuff, void* recvbuff, - size_t recvcount, ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm, - cudaStream_t stream); -ncclResult_t pncclReduceScatter(const void* sendbuff, void* recvbuff, - size_t recvcount, ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm, - cudaStream_t stream); - -/* - * All-Gather - * - * Each device gathers sendcount values from other GPUs into recvbuff, - * receiving data from rank i at offset i*sendcount. - * Assumes recvcount is equal to nranks*sendcount, which means that recvbuff - * should have a size of at least nranks*sendcount elements. - * - * In-place operations will happen if sendbuff == recvbuff + rank * sendcount. - */ -ncclResult_t ncclAllGather(const void* sendbuff, void* recvbuff, size_t sendcount, - ncclDataType_t datatype, ncclComm_t comm, cudaStream_t stream); -ncclResult_t pncclAllGather(const void* sendbuff, void* recvbuff, size_t sendcount, - ncclDataType_t datatype, ncclComm_t comm, cudaStream_t stream); - -/* - * All-to-All - * - * Each device sends count values to all other devices and receives count values - * from all other devices. Data to send to destination rank j is taken from - * sendbuff+j*count and data received from source rank i is placed at - * recvbuff+i*count. - */ -ncclResult_t ncclAlltoAll(const void* sendbuff, void* recvbuff, size_t count, - ncclDataType_t datatype, ncclComm_t comm, cudaStream_t stream); -ncclResult_t pncclAlltoAll(const void* sendbuff, void* recvbuff, size_t count, - ncclDataType_t datatype, ncclComm_t comm, cudaStream_t stream); - -/* - * Gather - * - * Each rank sends count elements from sendbuff to the root rank. - * On the root rank, data from rank i is placed at recvbuff + i*count. - * On non-root ranks, recvbuff is not used. - * root is the rank where data will be gathered. - * - * In-place operations will happen if sendbuff == recvbuff + root * count. - */ -ncclResult_t ncclGather(const void* sendbuff, void* recvbuff, size_t count, - ncclDataType_t datatype, int root, ncclComm_t comm, cudaStream_t stream); -ncclResult_t pncclGather(const void* sendbuff, void* recvbuff, size_t count, - ncclDataType_t datatype, int root, ncclComm_t comm, cudaStream_t stream); - -/* - * Scatter - * - * On the root rank, count elements from sendbuff+i*count are sent to rank i. - * On non-root ranks, sendbuff is not used. - * Each rank receives count elements into recvbuff. - * root is the rank that will distribute the data. - * - * In-place operations will happen if recvbuff == sendbuff + root * count. - */ -ncclResult_t ncclScatter(const void* sendbuff, void* recvbuff, size_t count, - ncclDataType_t datatype, int root, ncclComm_t comm, cudaStream_t stream); -ncclResult_t pncclScatter(const void* sendbuff, void* recvbuff, size_t count, - ncclDataType_t datatype, int root, ncclComm_t comm, cudaStream_t stream); - -/* - * Send - * - * Send data from sendbuff to rank peer. - * - * Rank peer needs to call ncclRecv with the same datatype and the same count from this - * rank. - * - * This operation is blocking for the GPU. If multiple ncclSend and ncclRecv operations - * need to progress concurrently to complete, they must be fused within a ncclGroupStart/ - * ncclGroupEnd section. - */ -ncclResult_t ncclSend(const void* sendbuff, size_t count, ncclDataType_t datatype, int peer, - ncclComm_t comm, cudaStream_t stream); -ncclResult_t pncclSend(const void* sendbuff, size_t count, ncclDataType_t datatype, int peer, - ncclComm_t comm, cudaStream_t stream); - -/* - * Receive - * - * Receive data from rank peer into recvbuff. - * - * Rank peer needs to call ncclSend with the same datatype and the same count to this - * rank. - * - * This operation is blocking for the GPU. If multiple ncclSend and ncclRecv operations - * need to progress concurrently to complete, they must be fused within a ncclGroupStart/ - * ncclGroupEnd section. - */ -ncclResult_t pncclRecv(void* recvbuff, size_t count, ncclDataType_t datatype, int peer, - ncclComm_t comm, cudaStream_t stream); -ncclResult_t ncclRecv(void* recvbuff, size_t count, ncclDataType_t datatype, int peer, - ncclComm_t comm, cudaStream_t stream); - - -/* - * Put - * - * One-sided communication operation that writes data from the local buffer to a - * remote peer's registered memory window without explicit participation from the - * target process. - * - * Parameters: - * localbuff - Local source buffer containing data to be transferred - * count - Number of elements to transfer - * datatype - NCCL data type of each element - * peer - Target rank to write data to - * peerWin - Memory window object registered by the target peer - * peerWinOffset- Offset in bytes from the start of peer's registered window - * sigIdx - Signal index identifier for the operation - * ctx - Context identifier for the operation - * flags - Reserved for future use - * comm - NCCL communicator - * stream - CUDA stream to enqueue the operation on - * - * Returns: - * ncclSuccess on successful enqueue, error code otherwise - */ -ncclResult_t ncclPutSignal(const void* localbuff, size_t count, ncclDataType_t datatype, - int peer, ncclWindow_t peerWin, size_t peerWinOffset, - int sigIdx, int ctx, unsigned int flags, ncclComm_t comm, cudaStream_t stream); - -ncclResult_t pncclPutSignal(const void* localbuff, size_t count, ncclDataType_t datatype, - int peer, ncclWindow_t peerWin, size_t peerWinOffset, - int sigIdx, int ctx, unsigned int flags, ncclComm_t comm, cudaStream_t stream); - -/* - * Signal - * - * Sends a signal to the specified peer without transferring data. - * - * Parameters: - * peer - Target rank to send signal to - * sigIdx - Signal index identifier for the operation - * ctx - Context identifier for the operation - * flags - Reserved for future use - * comm - NCCL communicator - * stream - CUDA stream to enqueue the operation on - * - * Returns: - * ncclSuccess on successful signal enqueue, error code otherwise - */ -ncclResult_t ncclSignal(int peer, int sigIdx, int ctx, unsigned int flags, ncclComm_t comm, cudaStream_t stream); -ncclResult_t pncclSignal(int peer, int sigIdx, int ctx, unsigned int flags, ncclComm_t comm, cudaStream_t stream); - -/* - * Wait Signal Descriptor - * - * Describes how many signal operations to wait for - * from a particular rank on a given signal index and context. - */ -typedef struct { - int opCnt; // Number of signal operations to wait for - int peer; // Target peer to wait for signals from - int sigIdx; // Signal index identifier - int ctx; // Context identifier -} ncclWaitSignalDesc_t; - -/* - * Wait Signal - * - * Waits for signals as described in the signal descriptor array. - * - * Parameters: - * nDesc - Number of signal descriptors in the array - * signalDescs - Array of descriptors specifying the signals to wait for. - * Each descriptor indicates how many signals to expect from - * a specific peer on a particular signal index and context. - * comm - NCCL communicator - * stream - CUDA stream to enqueue the operation on - * - * Returns: - * ncclSuccess when all required signals received, error code otherwise - */ -ncclResult_t ncclWaitSignal(int nDesc, ncclWaitSignalDesc_t* signalDescs, ncclComm_t comm, cudaStream_t stream); -ncclResult_t pncclWaitSignal(int nDesc, ncclWaitSignalDesc_t* signalDescs, ncclComm_t comm, cudaStream_t stream); - -/* - * Group semantics - * - * When managing multiple GPUs from a single thread, and since NCCL collective - * calls may perform inter-CPU synchronization, we need to "group" calls for - * different ranks/devices into a single call. - * - * Grouping NCCL calls as being part of the same collective operation is done - * using ncclGroupStart and ncclGroupEnd. ncclGroupStart will enqueue all - * collective calls until the ncclGroupEnd call, which will wait for all calls - * to be complete. Note that for collective communication, ncclGroupEnd only - * guarantees that the operations are enqueued on the streams, not that - * the operation is effectively done. - * - * Both collective communication and ncclCommInitRank can be used in conjunction - * of ncclGroupStart/ncclGroupEnd, but not together. - * - * Group semantics also allow to fuse multiple operations on the same device - * to improve performance (for aggregated collective calls), or to permit - * concurrent progress of multiple send/receive operations. - */ - -/* - * Group Start - * - * Start a group call. All calls to NCCL until ncclGroupEnd will be fused into - * a single NCCL operation. Nothing will be started on the CUDA stream until - * ncclGroupEnd. - */ -ncclResult_t ncclGroupStart(); -ncclResult_t pncclGroupStart(); - -/* - * Group End - * - * End a group call. Start a fused NCCL operation consisting of all calls since - * ncclGroupStart. Operations on the CUDA stream depending on the NCCL operations - * need to be called after ncclGroupEnd. - */ -ncclResult_t ncclGroupEnd(); -ncclResult_t pncclGroupEnd(); - -/* - * Group Simulate End - * - * Simulate a ncclGroupEnd() call and return NCCL's simulation info in a struct. - */ -ncclResult_t ncclGroupSimulateEnd(ncclSimInfo_t* simInfo); -ncclResult_t pncclGroupSimulateEnd(ncclSimInfo_t* simInfo); - -/* - * Parameter access - * - * For accessing NCCL runtime parameters. Parameters are identified by string keys - * and can be read as typed values or as strings. NCCL provides two styles of parameter - * APIs: handle-based and key-based. - */ - -/* - * Handle-based API (allows typed access) - * - * Handles represents parameters and are opaque (internal details hidden). - */ -typedef struct ncclParamHandle* ncclParamHandle_t; - -/* - * Look up the parameter identified by key and store a handle to it in - * out. The returned handle is owned by the parameter system and must not - * be freed by the caller. - */ -ncclResult_t ncclParamBind(ncclParamHandle_t* out, const char* key); -ncclResult_t pncclParamBind(ncclParamHandle_t* out, const char* key); - -/* - * Read the value of the parameter bound to h as the type of out. - * Function names are suffixed with I/U and 8/16/32/64 for 8-, 16-, 32- and 64-bit - * signed and unsigned integers. - */ -ncclResult_t ncclParamGetI8(ncclParamHandle_t h, int8_t* out); -ncclResult_t pncclParamGetI8(ncclParamHandle_t h, int8_t* out); - -ncclResult_t ncclParamGetI16(ncclParamHandle_t h, int16_t* out); -ncclResult_t pncclParamGetI16(ncclParamHandle_t h, int16_t* out); - -ncclResult_t ncclParamGetI32(ncclParamHandle_t h, int32_t* out); -ncclResult_t pncclParamGetI32(ncclParamHandle_t h, int32_t* out); - -ncclResult_t ncclParamGetI64(ncclParamHandle_t h, int64_t* out); -ncclResult_t pncclParamGetI64(ncclParamHandle_t h, int64_t* out); - -ncclResult_t ncclParamGetU8(ncclParamHandle_t h, uint8_t* out); -ncclResult_t pncclParamGetU8(ncclParamHandle_t h, uint8_t* out); - -ncclResult_t ncclParamGetU16(ncclParamHandle_t h, uint16_t* out); -ncclResult_t pncclParamGetU16(ncclParamHandle_t h, uint16_t* out); - -ncclResult_t ncclParamGetU32(ncclParamHandle_t h, uint32_t* out); -ncclResult_t pncclParamGetU32(ncclParamHandle_t h, uint32_t* out); - -ncclResult_t ncclParamGetU64(ncclParamHandle_t h, uint64_t* out); -ncclResult_t pncclParamGetU64(ncclParamHandle_t h, uint64_t* out); - -/* - * Read the value of the parameter bound to h as a string. - * Returned pointer is owned by the parameter system and is valid until the - * next ncclParamGetStr() call on the same thread. - */ -ncclResult_t ncclParamGetStr(ncclParamHandle_t h, const char** out); -ncclResult_t pncclParamGetStr(ncclParamHandle_t h, const char** out); - -/* - * Read the value of the parameter bound to h as raw binary data. - * The user needs to allocate a buffer for the result and the parameter value is copied - * into user buffer as bytes. - */ -ncclResult_t ncclParamGet(ncclParamHandle_t h, void* out, int maxLen, int* len); -ncclResult_t pncclParamGet(ncclParamHandle_t h, void* out, int maxLen, int* len); - -/* - * Key-based API (no handle required, typeless access, return value as string) - */ - -/* - * Get parameter value as string by key. Returned pointer is owned by the - * parameter system and is valid until the next ncclParamGetParameter() call - * on the same thread. - */ -ncclResult_t ncclParamGetParameter(const char* key, const char** value, int* valueLen); -ncclResult_t pncclParamGetParameter(const char* key, const char** value, int* valueLen); - -/* - * Get all registered parameter keys. Returned pointer table is owned by the - * parameter system and is valid until the next ncclParamGetAllParameterKeys() - * call on the same thread. By default, the results include only parameters published - * in NCCL documentation. Setting NCCL_PARAM_DUMP_ALL=true will include all parameters. - */ -ncclResult_t ncclParamGetAllParameterKeys(const char*** table, int* tableLen); -ncclResult_t pncclParamGetAllParameterKeys(const char*** table, int* tableLen); - -/* - * Dump all parameters to log output. By default, the result includes only parameters published - * in NCCL documentation. Setting NCCL_PARAM_DUMP_ALL=true will include all parameters. - */ -void ncclParamDumpAll(void); -void pncclParamDumpAll(void); - -#ifdef __cplusplus -} // end extern "C" -#endif - -#endif // end include guard diff --git a/python/build_assets/cybind/templates/nccl/_extensions/bindings/__init__.py b/python/build_assets/cybind/templates/nccl/_extensions/bindings/__init__.py deleted file mode 100644 index 65e4f643..00000000 --- a/python/build_assets/cybind/templates/nccl/_extensions/bindings/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# See LICENSE.txt for more license information diff --git a/python/build_assets/cybind/templates/nccl/_extensions/bindings/_internal/__init__.py b/python/build_assets/cybind/templates/nccl/_extensions/bindings/_internal/__init__.py deleted file mode 100644 index 65e4f643..00000000 --- a/python/build_assets/cybind/templates/nccl/_extensions/bindings/_internal/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# See LICENSE.txt for more license information diff --git a/python/build_assets/cybind/templates/nccl/_extensions/bindings/_internal/nccl_ep.pxd b/python/build_assets/cybind/templates/nccl/_extensions/bindings/_internal/nccl_ep.pxd deleted file mode 100644 index c16fb4f5..00000000 --- a/python/build_assets/cybind/templates/nccl/_extensions/bindings/_internal/nccl_ep.pxd +++ /dev/null @@ -1,14 +0,0 @@ -# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 -# -# This code was automatically generated $version_span. Do not modify it directly. - -from ..cy$libname cimport * - - -############################################################################### -# Wrapper functions -############################################################################### - -$wrapper_decls diff --git a/python/build_assets/cybind/templates/nccl/_extensions/bindings/_internal/nccl_ep_linux.pyx b/python/build_assets/cybind/templates/nccl/_extensions/bindings/_internal/nccl_ep_linux.pyx deleted file mode 100644 index 7a624f77..00000000 --- a/python/build_assets/cybind/templates/nccl/_extensions/bindings/_internal/nccl_ep_linux.pyx +++ /dev/null @@ -1,176 +0,0 @@ -# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 -# -# This code was automatically generated $version_span. Do not modify it directly. - -from libc.stdint cimport intptr_t, uint64_t, uintptr_t - -import os -import threading - -from .utils import FunctionNotFoundError, NotSupportedError - -from cuda.pathfinder import load_nvidia_dynamic_lib - -${snippet_linux_externs_pxd} - -cdef extern from "" nogil: - ctypedef struct Dl_info: - const char* dli_fname - void* dli_fbase - const char* dli_sname - void* dli_saddr - int dladdr(const void*, Dl_info*) - - -############################################################################### -# Library resolution (mirrors cuda.pathfinder.load_nvidia_dynamic_lib precedence, -# adapted for libnccl_ep.so which is not registered as an NVIDIA pip wheel.) -############################################################################### - -# Resolved at first import via _resolve_library_path() below. Path lookup runs -# once, then dlopen handle is cached in the lowpp ${libname} init guard. -# -# Each library's .so ships under its own facade package -- nccl_ep -> nccl/ep/, -# nccl_m2n -> nccl/m2n/ -- so derive that directory from the library name rather -# than hardcoding one library's. -_PACKAGE_LIB_RELPATH = os.path.join( - "${libname}".removeprefix("nccl_"), "lib", "lib${libname}.so" -) - - -def _resolve_library_path() -> str: - # 1. nccl-extensions package path (replaces cuda.pathfinder's NVIDIA-pip-wheel - # step). lib${libname}.so is at nccl//lib/; this file lives in - # nccl/_extensions/bindings/_internal/, so go up three dirs to reach nccl/. - pkg_lib = os.path.normpath(os.path.join( - os.path.dirname(__file__), "..", "..", "..", _PACKAGE_LIB_RELPATH - )) - if os.path.exists(pkg_lib): - return pkg_lib - - # 2. CONDA_PREFIX/lib[64] - conda_prefix = os.environ.get("CONDA_PREFIX") - if conda_prefix: - for sub in ("lib", "lib64"): - candidate = os.path.join(conda_prefix, sub, "lib${libname}.so") - if os.path.exists(candidate): - return candidate - - # 3. CUDA_HOME / CUDA_PATH lib[64] - for env_var in ("CUDA_HOME", "CUDA_PATH"): - root = os.environ.get(env_var) - if root: - for sub in ("lib", "lib64"): - candidate = os.path.join(root, sub, "lib${libname}.so") - if os.path.exists(candidate): - return candidate - - # 4. SONAME fallback — let dlopen perform its own search across - # LD_LIBRARY_PATH, /etc/ld.so.cache, and /lib, /usr/lib, /lib64, - # /usr/lib64. If it fails the caller surfaces a clear error. - return "lib${libname}.so" - - -############################################################################### -# Wrapper init -############################################################################### - -cdef object __symbol_lock = threading.Lock() -cdef bint __py_${libname}_init = False - -$wrapper_init - - -cdef void* load_library() except* with gil: - # libnccl_ep.so has NEEDED libnccl.so.2. Pre-load it with RTLD_GLOBAL so the - # SONAME is already mapped when libnccl_ep.so's NEEDED is resolved, - # without depending on filesystem search. - load_nvidia_dynamic_lib("nccl") - - cdef bytes path_bytes = _resolve_library_path().encode() - cdef void* handle = dlopen(path_bytes, RTLD_NOW | RTLD_GLOBAL) - if handle == NULL: - err_msg = dlerror() - raise RuntimeError( - f'Failed to dlopen lib${libname} ({err_msg.decode()}); ' - f'tried path {path_bytes.decode()!r}' - ) - return handle - - -cdef int _check_or_init_${libname}() except -1 nogil: - global __py_${libname}_init - if __py_${libname}_init: - return 0 - - cdef void* handle = NULL - - with gil, __symbol_lock: - # Recheck the flag after obtaining the locks - if __py_${libname}_init: - return 0 - - # Load function -${set_wrapper} - __py_${libname}_init = True - return 0 - - -cdef dict func_ptrs = None - - -cpdef dict _inspect_function_pointers(): - global func_ptrs - if func_ptrs is not None: - return func_ptrs - - _check_or_init_${libname}() - cdef dict data = {} - -${set_functor} - - func_ptrs = data - return data - - -cpdef _inspect_function_pointer(str name): - global func_ptrs - if func_ptrs is None: - func_ptrs = _inspect_function_pointers() - return func_ptrs[name] - - -cdef object __${libname}_loaded_so_path = None - - -cpdef object _inspect_loaded_library_path(): - # Path of the .so backing the loaded symbols, via dladdr() on a - # resolved entry point. None if it cannot be determined. - global __${libname}_loaded_so_path - if __${libname}_loaded_so_path is not None: - return __${libname}_loaded_so_path - - cdef dict ptrs = _inspect_function_pointers() - # Any resolved symbol maps to the same .so. - cdef intptr_t addr = 0 - for value in ptrs.values(): - if value: - addr = value - break - - cdef Dl_info info - if addr == 0: - return None - if dladdr(addr, &info) == 0 or info.dli_fname == NULL: - return None - __${libname}_loaded_so_path = os.fsdecode(info.dli_fname) - return __${libname}_loaded_so_path - - -############################################################################### -# Wrapper functions -############################################################################### - -$wrapper_defs diff --git a/python/build_assets/cybind/templates/nccl/_extensions/bindings/_internal/nccl_m2n.pxd b/python/build_assets/cybind/templates/nccl/_extensions/bindings/_internal/nccl_m2n.pxd deleted file mode 100644 index 7c159586..00000000 --- a/python/build_assets/cybind/templates/nccl/_extensions/bindings/_internal/nccl_m2n.pxd +++ /dev/null @@ -1,20 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# See LICENSE.txt for more license information. - -from ..cynccl_m2n cimport * - - -############################################################################### -# Wrapper functions -############################################################################### - -cdef ncclResult_t _ncclM2nInit(ncclM2nHandle_t* handle, const ncclM2nConfig_t* config) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil -cdef ncclResult_t _ncclM2nFinalize(ncclM2nHandle_t handle) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil -cdef ncclResult_t _ncclM2nGroupStart() except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil -cdef ncclResult_t _ncclM2nGroupEnd() except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil -cdef ncclResult_t _ncclM2nGroupAbort() except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil -cdef const char* _ncclM2nGetLastError() noexcept nogil -cdef ncclResult_t _ncclReshardWithWindow(ncclM2nHandle_t handle, ncclComm_t comm, ncclWindow_t window, const ncclDistTensor_t* src, const ncclDistTensor_t* dst, cudaStream_t stream) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil -cdef ncclResult_t _ncclReshard(ncclM2nHandle_t handle, ncclComm_t comm, const ncclDistTensor_t* src, const ncclDistTensor_t* dst, cudaStream_t stream) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil diff --git a/python/build_assets/cybind/templates/nccl/_extensions/bindings/_internal/nccl_m2n_linux.pyx b/python/build_assets/cybind/templates/nccl/_extensions/bindings/_internal/nccl_m2n_linux.pyx deleted file mode 100644 index a046756d..00000000 --- a/python/build_assets/cybind/templates/nccl/_extensions/bindings/_internal/nccl_m2n_linux.pyx +++ /dev/null @@ -1,307 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# See LICENSE.txt for more license information. - -from libc.stdint cimport intptr_t - -import os -import threading - -from cuda.pathfinder import load_nvidia_dynamic_lib -from .utils import FunctionNotFoundError - - -############################################################################### -# Extern -############################################################################### - -cdef extern from "" nogil: - void* dlopen(const char*, int) - char* dlerror() - void* dlsym(void*, const char*) - int dlclose(void*) - - enum: - RTLD_NOW - RTLD_GLOBAL - -############################################################################### -# Library resolution -############################################################################### - -def _candidate_library_paths() -> list[str]: - explicit = os.environ.get("NCCL_M2N_LIBRARY") - if explicit: - return [explicit] - - # With no explicit override, prefer the native library bundled with this - # facade before environment and SONAME fallbacks. - candidates = [os.path.normpath(os.path.join( - os.path.dirname(__file__), "..", "..", "..", "m2n", "lib", "libnccl_m2n.so" - ))] - - home = os.environ.get("NCCL_M2N_HOME") - if home: - candidates.append(os.path.join(home, "lib", "libnccl_m2n.so")) - - conda_prefix = os.environ.get("CONDA_PREFIX") - if conda_prefix: - for subdir in ("lib", "lib64"): - candidates.append(os.path.join(conda_prefix, subdir, "libnccl_m2n.so")) - - for env_var in ("CUDA_HOME", "CUDA_PATH"): - root = os.environ.get(env_var) - if root: - for subdir in ("lib", "lib64"): - candidates.append(os.path.join(root, subdir, "libnccl_m2n.so")) - - candidates.append("libnccl_m2n.so") - return candidates - - -cdef void* load_library() except* with gil: - load_nvidia_dynamic_lib("nccl") - - cdef void* handle = NULL - cdef bytes path_bytes - cdef char* err_msg - errors = [] - - for path in _candidate_library_paths(): - if path != "libnccl_m2n.so" and not os.path.exists(path): - errors.append(f"{path}: not found") - continue - path_bytes = path.encode() - handle = dlopen(path_bytes, RTLD_NOW | RTLD_GLOBAL) - if handle != NULL: - return handle - err_msg = dlerror() - if err_msg != NULL: - errors.append(f"{path}: {err_msg.decode()}") - else: - errors.append(f"{path}: dlopen failed") - - raise RuntimeError( - "Failed to dlopen libnccl_m2n.so. Set NCCL_M2N_LIBRARY to the " - "shared library path or NCCL_M2N_HOME to an install prefix. Tried: " - + "; ".join(errors) - ) - - -############################################################################### -# Wrapper init -############################################################################### - -cdef object __symbol_lock = threading.Lock() -cdef bint __py_nccl_m2n_init = False -cdef void* __library_handle = NULL - -cdef void* __ncclM2nInit = NULL -cdef void* __ncclM2nFinalize = NULL -cdef void* __ncclM2nGroupStart = NULL -cdef void* __ncclM2nGroupEnd = NULL -cdef void* __ncclM2nGroupAbort = NULL -cdef void* __ncclM2nGetLastError = NULL -cdef void* __ncclReshardWithWindow = NULL -cdef void* __ncclReshard = NULL - - -cdef int _check_or_init_nccl_m2n() except -1 nogil: - global __py_nccl_m2n_init - if __py_nccl_m2n_init: - return 0 - - cdef void* handle = NULL - cdef void* init_fn = NULL - cdef void* finalize_fn = NULL - cdef void* group_start_fn = NULL - cdef void* group_end_fn = NULL - cdef void* group_abort_fn = NULL - cdef void* get_last_error_fn = NULL - cdef void* reshard_with_window_fn = NULL - cdef void* reshard_fn = NULL - - with gil, __symbol_lock: - if __py_nccl_m2n_init: - return 0 - - global __ncclM2nInit - global __ncclM2nFinalize - global __ncclM2nGroupStart - global __ncclM2nGroupEnd - global __ncclM2nGroupAbort - global __ncclM2nGetLastError - global __ncclReshardWithWindow - global __ncclReshard - - handle = load_library() - init_fn = dlsym(handle, 'ncclM2nInit') - finalize_fn = dlsym(handle, 'ncclM2nFinalize') - group_start_fn = dlsym(handle, 'ncclM2nGroupStart') - group_end_fn = dlsym(handle, 'ncclM2nGroupEnd') - group_abort_fn = dlsym(handle, 'ncclM2nGroupAbort') - get_last_error_fn = dlsym(handle, 'ncclM2nGetLastError') - reshard_with_window_fn = dlsym(handle, 'ncclReshardWithWindow') - reshard_fn = dlsym(handle, 'ncclReshard') - - missing = [] - if init_fn == NULL: - missing.append("ncclM2nInit") - if finalize_fn == NULL: - missing.append("ncclM2nFinalize") - if group_start_fn == NULL: - missing.append("ncclM2nGroupStart") - if group_end_fn == NULL: - missing.append("ncclM2nGroupEnd") - if group_abort_fn == NULL: - missing.append("ncclM2nGroupAbort") - if get_last_error_fn == NULL: - missing.append("ncclM2nGetLastError") - if reshard_with_window_fn == NULL: - missing.append("ncclReshardWithWindow") - if reshard_fn == NULL: - missing.append("ncclReshard") - if missing: - dlclose(handle) - raise FunctionNotFoundError( - "libnccl_m2n.so does not provide the complete M2N v2 API; " - "missing: " + ", ".join(missing) - ) - - global __library_handle - __library_handle = handle - __ncclM2nInit = init_fn - __ncclM2nFinalize = finalize_fn - __ncclM2nGroupStart = group_start_fn - __ncclM2nGroupEnd = group_end_fn - __ncclM2nGroupAbort = group_abort_fn - __ncclM2nGetLastError = get_last_error_fn - __ncclReshardWithWindow = reshard_with_window_fn - __ncclReshard = reshard_fn - - __py_nccl_m2n_init = True - return 0 - - -cdef dict func_ptrs = None - - -cpdef dict _inspect_function_pointers(): - global func_ptrs - if func_ptrs is not None: - return func_ptrs - - _check_or_init_nccl_m2n() - cdef dict data = {} - - global __ncclM2nInit - data["__ncclM2nInit"] = __ncclM2nInit - - global __ncclM2nFinalize - data["__ncclM2nFinalize"] = __ncclM2nFinalize - - global __ncclM2nGroupStart - data["__ncclM2nGroupStart"] = __ncclM2nGroupStart - - global __ncclM2nGroupEnd - data["__ncclM2nGroupEnd"] = __ncclM2nGroupEnd - - global __ncclM2nGroupAbort - data["__ncclM2nGroupAbort"] = __ncclM2nGroupAbort - - global __ncclM2nGetLastError - data["__ncclM2nGetLastError"] = __ncclM2nGetLastError - - global __ncclReshardWithWindow - data["__ncclReshardWithWindow"] = __ncclReshardWithWindow - - global __ncclReshard - data["__ncclReshard"] = __ncclReshard - - func_ptrs = data - return data - - -cpdef _inspect_function_pointer(str name): - global func_ptrs - if func_ptrs is None: - func_ptrs = _inspect_function_pointers() - return func_ptrs[name] - - -############################################################################### -# Wrapper functions -############################################################################### - -cdef ncclResult_t _ncclM2nInit(ncclM2nHandle_t* handle, const ncclM2nConfig_t* config) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil: - global __ncclM2nInit - _check_or_init_nccl_m2n() - if __ncclM2nInit == NULL: - with gil: - raise FunctionNotFoundError("function ncclM2nInit is not found") - return (__ncclM2nInit)(handle, config) - - -cdef ncclResult_t _ncclM2nFinalize(ncclM2nHandle_t handle) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil: - global __ncclM2nFinalize - _check_or_init_nccl_m2n() - if __ncclM2nFinalize == NULL: - with gil: - raise FunctionNotFoundError("function ncclM2nFinalize is not found") - return (__ncclM2nFinalize)(handle) - - -cdef ncclResult_t _ncclM2nGroupStart() except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil: - global __ncclM2nGroupStart - _check_or_init_nccl_m2n() - if __ncclM2nGroupStart == NULL: - with gil: - raise FunctionNotFoundError("function ncclM2nGroupStart is not found") - return (__ncclM2nGroupStart)() - - -cdef ncclResult_t _ncclM2nGroupEnd() except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil: - global __ncclM2nGroupEnd - _check_or_init_nccl_m2n() - if __ncclM2nGroupEnd == NULL: - with gil: - raise FunctionNotFoundError("function ncclM2nGroupEnd is not found") - return (__ncclM2nGroupEnd)() - - -cdef ncclResult_t _ncclM2nGroupAbort() except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil: - global __ncclM2nGroupAbort - _check_or_init_nccl_m2n() - if __ncclM2nGroupAbort == NULL: - with gil: - raise FunctionNotFoundError("function ncclM2nGroupAbort is not found") - return (__ncclM2nGroupAbort)() - - -cdef const char* _ncclM2nGetLastError() noexcept nogil: - global __ncclM2nGetLastError - _check_or_init_nccl_m2n() - if __ncclM2nGetLastError == NULL: - return NULL - return (__ncclM2nGetLastError)() - - -cdef ncclResult_t _ncclReshardWithWindow(ncclM2nHandle_t handle, ncclComm_t comm, ncclWindow_t window, const ncclDistTensor_t* src, const ncclDistTensor_t* dst, cudaStream_t stream) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil: - global __ncclReshardWithWindow - _check_or_init_nccl_m2n() - if __ncclReshardWithWindow == NULL: - with gil: - raise FunctionNotFoundError("function ncclReshardWithWindow is not found") - return (__ncclReshardWithWindow)( - handle, comm, window, src, dst, stream) - - -cdef ncclResult_t _ncclReshard(ncclM2nHandle_t handle, ncclComm_t comm, const ncclDistTensor_t* src, const ncclDistTensor_t* dst, cudaStream_t stream) except?_NCCLRESULT_T_INTERNAL_LOADING_ERROR nogil: - global __ncclReshard - _check_or_init_nccl_m2n() - if __ncclReshard == NULL: - with gil: - raise FunctionNotFoundError("function ncclReshard is not found") - return (__ncclReshard)( - handle, comm, src, dst, stream) diff --git a/python/build_assets/cybind/templates/nccl/_extensions/bindings/_internal/utils.pxd b/python/build_assets/cybind/templates/nccl/_extensions/bindings/_internal/utils.pxd deleted file mode 100644 index cd48ca15..00000000 --- a/python/build_assets/cybind/templates/nccl/_extensions/bindings/_internal/utils.pxd +++ /dev/null @@ -1,167 +0,0 @@ -# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 - -from libc.stdint cimport int32_t, int64_t, intptr_t -from libcpp.vector cimport vector -from libcpp cimport bool as cppbool -from libcpp cimport nullptr_t, nullptr -from libcpp.memory cimport unique_ptr - - -cdef extern from * nogil: - """ - template - class nullable_unique_ptr { - public: - nullable_unique_ptr() noexcept = default; - - nullable_unique_ptr(std::nullptr_t) noexcept = delete; - - explicit nullable_unique_ptr(T* data, bool own_data): - own_data_(own_data) - { - if (own_data) - manager_.reset(data); - else - raw_data_ = data; - } - - nullable_unique_ptr(const nullable_unique_ptr&) = delete; - - nullable_unique_ptr& operator=(const nullable_unique_ptr&) = delete; - - nullable_unique_ptr(nullable_unique_ptr&& other) noexcept - { - own_data_ = other.own_data_; - other.own_data_ = false; // ownership is transferred - if (own_data_) - { - manager_ = std::move(other.manager_); - raw_data_ = nullptr; // just in case - } - else - { - manager_.reset(nullptr); // just in case - raw_data_ = other.raw_data_; - } - } - - nullable_unique_ptr& operator=(nullable_unique_ptr&& other) noexcept - { - own_data_ = other.own_data_; - other.own_data_ = false; // ownership is transferred - if (own_data_) - { - manager_ = std::move(other.manager_); - raw_data_ = nullptr; // just in case - } - else - { - manager_.reset(nullptr); // just in case - raw_data_ = other.raw_data_; - } - return *this; - } - - ~nullable_unique_ptr() = default; - - void reset(T* data, bool own_data) - { - own_data_ = own_data; - if (own_data_) - { - manager_.reset(data); - raw_data_ = nullptr; - } - else - { - manager_.reset(nullptr); - raw_data_ = data; - } - } - - void swap(nullable_unique_ptr& other) noexcept - { - std::swap(manager_, other.manager_); - std::swap(raw_data_, other.raw_data_); - std::swap(own_data_, other.own_data_); - } - - /* - * Get the pointer to the underlying object (this is different from data()!). - */ - T* get() const noexcept - { - if (own_data_) - return manager_.get(); - else - return raw_data_; - } - - /* - * Get the pointer to the underlying buffer (this is different from get()!). - */ - void* data() noexcept - { - if (own_data_) - return manager_.get()->data(); - else - return raw_data_; - } - - T& operator*() - { - if (own_data_) - return *manager_; - else - return *raw_data_; - } - - private: - std::unique_ptr manager_{}; - T* raw_data_{nullptr}; - bool own_data_{false}; - }; - """ - # xref: cython/Cython/Includes/libcpp/memory.pxd - cdef cppclass nullable_unique_ptr[T]: - nullable_unique_ptr() - nullable_unique_ptr(T*, cppbool) - nullable_unique_ptr(nullable_unique_ptr[T]&) - - # Modifiers - void reset(T*, cppbool) - void swap(nullable_unique_ptr&) - - # Observers - T* get() - T& operator*() - void* data() - - -ctypedef fused ResT: - int - int32_t - int64_t - char - float - double - - -ctypedef fused PtrT: - void - - -cdef cppclass nested_resource[T]: - nullable_unique_ptr[ vector[intptr_t] ] ptrs - nullable_unique_ptr[ vector[vector[T]] ] nested_resource_ptr - - -# accepts the output pointer as input to use the return value for exception propagation -cdef int get_resource_ptr(nullable_unique_ptr[vector[ResT]] &in_out_ptr, object obj, ResT* __unused) except 1 -cdef int get_resource_ptrs(nullable_unique_ptr[ vector[PtrT*] ] &in_out_ptr, object obj, PtrT* __unused) except 1 -cdef int get_nested_resource_ptr(nested_resource[ResT] &in_out_ptr, object obj, ResT* __unused) except 1 - -cdef bint is_nested_sequence(data) -cdef void* get_buffer_pointer(buf, Py_ssize_t size, readonly=*) except* diff --git a/python/build_assets/cybind/templates/nccl/_extensions/bindings/_internal/utils.pyx b/python/build_assets/cybind/templates/nccl/_extensions/bindings/_internal/utils.pyx deleted file mode 100644 index 3e2ef8d3..00000000 --- a/python/build_assets/cybind/templates/nccl/_extensions/bindings/_internal/utils.pyx +++ /dev/null @@ -1,129 +0,0 @@ -# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 - -cimport cpython -from libc.stdint cimport intptr_t -from libcpp.utility cimport move -from cython.operator cimport dereference as deref - - -cdef bint is_nested_sequence(data): - if not cpython.PySequence_Check(data): - return False - else: - for i in data: - if not cpython.PySequence_Check(i): - return False - else: - return True - - -cdef void* get_buffer_pointer(buf, Py_ssize_t size, readonly=True) except*: - """The caller must ensure ``buf`` is alive when the returned pointer is in use.""" - cdef void* bufPtr - cdef int flags = cpython.PyBUF_ANY_CONTIGUOUS - if not readonly: - flags |= cpython.PyBUF_WRITABLE - cdef int status = -1 - cdef cpython.Py_buffer view - - if isinstance(buf, int): - bufPtr = buf - else: # try buffer protocol - try: - status = cpython.PyObject_GetBuffer(buf, &view, flags) - # when the caller does not provide a size, it is set to -1 at generate-time by cybind - if size != -1: - assert view.len == size - assert view.ndim == 1 - except Exception as e: - adj = "writable " if not readonly else "" - raise ValueError( - "buf must be either a Python int representing the pointer " - f"address to a valid buffer, or a 1D contiguous {adj}" - "buffer, of size bytes") from e - else: - bufPtr = view.buf - finally: - if status == 0: - cpython.PyBuffer_Release(&view) - - return bufPtr - - -# Cython can't infer the ResT overload when it is wrapped in nullable_unique_ptr, -# so we need a dummy (__unused) input argument to help it -cdef int get_resource_ptr(nullable_unique_ptr[vector[ResT]] &in_out_ptr, object obj, ResT* __unused) except 1: - if cpython.PySequence_Check(obj): - vec = new vector[ResT](len(obj)) - # set the ownership immediately to avoid leaking the `vec` memory in - # case of exception in the following loop - in_out_ptr.reset(vec, True) - for i in range(len(obj)): - deref(vec)[i] = obj[i] - else: - in_out_ptr.reset(obj, False) - return 0 - - -cdef int get_resource_ptrs(nullable_unique_ptr[ vector[PtrT*] ] &in_out_ptr, object obj, PtrT* __unused) except 1: - if cpython.PySequence_Check(obj): - vec = new vector[PtrT*](len(obj)) - # set the ownership immediately to avoid leaking the `vec` memory in - # case of exception in the following loop - in_out_ptr.reset(vec, True) - for i in range(len(obj)): - deref(vec)[i] = (obj[i]) - else: - in_out_ptr.reset(obj, False) - return 0 - - -cdef int get_nested_resource_ptr(nested_resource[ResT] &in_out_ptr, object obj, ResT* __unused) except 1: - cdef nullable_unique_ptr[ vector[intptr_t] ] nested_ptr - cdef nullable_unique_ptr[ vector[vector[ResT]] ] nested_res_ptr - cdef vector[intptr_t]* nested_vec = NULL - cdef vector[vector[ResT]]* nested_res_vec = NULL - cdef size_t i = 0, length = 0 - cdef intptr_t addr - - if is_nested_sequence(obj): - length = len(obj) - nested_res_vec = new vector[vector[ResT]](length) - nested_vec = new vector[intptr_t](length) - # set the ownership immediately to avoid leaking memory in case of - # exception in the following loop - nested_res_ptr.reset(nested_res_vec, True) - nested_ptr.reset(nested_vec, True) - for i, obj_i in enumerate(obj): - if ResT is char: - obj_i_bytes = ((obj_i)).encode() - str_len = (len(obj_i_bytes)) + 1 # including null termination - deref(nested_res_vec)[i].resize(str_len) - obj_i_ptr = (obj_i_bytes) - # cast to size_t explicitly to work around a potentially Cython bug - deref(nested_res_vec)[i].assign(obj_i_ptr, obj_i_ptr + str_len) - else: - deref(nested_res_vec)[i] = obj_i - deref(nested_vec)[i] = (deref(nested_res_vec)[i].data()) - elif cpython.PySequence_Check(obj): - length = len(obj) - nested_vec = new vector[intptr_t](length) - nested_ptr.reset(nested_vec, True) - for i, addr in enumerate(obj): - deref(nested_vec)[i] = addr - nested_res_ptr.reset(NULL, False) - else: - # obj is an int (ResT**) - nested_res_ptr.reset(NULL, False) - nested_ptr.reset(obj, False) - - in_out_ptr.ptrs = move(nested_ptr) - in_out_ptr.nested_resource_ptr = move(nested_res_ptr) - return 0 - - -class FunctionNotFoundError(RuntimeError): pass - -class NotSupportedError(RuntimeError): pass diff --git a/python/build_assets/cybind/templates/nccl/_extensions/bindings/cynccl_ep.pxd b/python/build_assets/cybind/templates/nccl/_extensions/bindings/cynccl_ep.pxd deleted file mode 100644 index d165461d..00000000 --- a/python/build_assets/cybind/templates/nccl/_extensions/bindings/cynccl_ep.pxd +++ /dev/null @@ -1,39 +0,0 @@ -# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 -# -# This code was automatically generated $version_span. Do not modify it directly. - - -from libc.stdint cimport uint64_t - -$external_imports - - -############################################################################### -# Types (structs, enums, ...) -############################################################################### - -# enums -$enum_decls - - -# types -cdef extern from *: - """ - #include - #include - #include - """ - ctypedef void* cudaStream_t 'cudaStream_t' - ctypedef int cudaError_t 'cudaError_t' - - -$type_decls - - -############################################################################### -# Functions -############################################################################### - -$func_decls diff --git a/python/build_assets/cybind/templates/nccl/_extensions/bindings/cynccl_ep.pyx b/python/build_assets/cybind/templates/nccl/_extensions/bindings/cynccl_ep.pyx deleted file mode 100644 index 3c8756c3..00000000 --- a/python/build_assets/cybind/templates/nccl/_extensions/bindings/cynccl_ep.pyx +++ /dev/null @@ -1,14 +0,0 @@ -# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 -# -# This code was automatically generated $version_span. Do not modify it directly. - -from ._internal cimport $libname as _$libname - - -############################################################################### -# Wrapper functions -############################################################################### - -$wrapper_defs diff --git a/python/build_assets/cybind/templates/nccl/_extensions/bindings/cynccl_m2n.pxd b/python/build_assets/cybind/templates/nccl/_extensions/bindings/cynccl_m2n.pxd deleted file mode 100644 index 9dc9c6f4..00000000 --- a/python/build_assets/cybind/templates/nccl/_extensions/bindings/cynccl_m2n.pxd +++ /dev/null @@ -1,58 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# This code was automatically generated $version_span. Do not modify it directly. - -$external_imports - - -############################################################################### -# Types and constants -############################################################################### - -# The generated declarations encode the M2N ABI, including these public -# constants, so wheel builds do not require nccl_m2n.h at build time. -cdef extern from *: - """ - #include - enum { - NCCL_RESHARD_MESH_NDIMS = 2, - NCCL_RESHARD_MAX_TENSOR_DIMS = 3, - NCCL_RESHARD_REPLICATE = -1, - NCCL_M2N_CONFIG_UNDEF_INT = INT_MIN, - NCCL_M2N_API_MAGIC = 0x4d324e32u, - NCCL_M2N_API_VERSION = 2u - }; - """ - enum: - NCCL_RESHARD_MESH_NDIMS - NCCL_RESHARD_MAX_TENSOR_DIMS - NCCL_RESHARD_REPLICATE - NCCL_M2N_CONFIG_UNDEF_INT - NCCL_M2N_API_MAGIC - NCCL_M2N_API_VERSION - -# enums -$enum_decls - -# types -cdef extern from *: - """ - #include - #include - #include - """ - ctypedef void* cudaStream_t 'cudaStream_t' - ctypedef int cudaError_t 'cudaError_t' - -$type_decls - - -############################################################################### -# Functions -############################################################################### - -$func_decls - -# Keep the error-detail query on the original no-throw Cython contract. -cdef const char* ncclM2nGetLastError() noexcept nogil diff --git a/python/build_assets/cybind/templates/nccl/_extensions/bindings/cynccl_m2n.pyx b/python/build_assets/cybind/templates/nccl/_extensions/bindings/cynccl_m2n.pyx deleted file mode 100644 index 31725266..00000000 --- a/python/build_assets/cybind/templates/nccl/_extensions/bindings/cynccl_m2n.pyx +++ /dev/null @@ -1,17 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# This code was automatically generated $version_span. Do not modify it directly. - -from ._internal cimport $libname as _$libname - - -############################################################################### -# Wrapper functions -############################################################################### - -$wrapper_defs - - -cdef const char* ncclM2nGetLastError() noexcept nogil: - return _$libname._ncclM2nGetLastError() diff --git a/python/build_assets/cybind/templates/nccl/_extensions/bindings/nccl_ep.pxd b/python/build_assets/cybind/templates/nccl/_extensions/bindings/nccl_ep.pxd deleted file mode 100644 index f256b7ff..00000000 --- a/python/build_assets/cybind/templates/nccl/_extensions/bindings/nccl_ep.pxd +++ /dev/null @@ -1,33 +0,0 @@ -# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 -# -# This code was automatically generated $version_span. Do not modify it directly. - -from libc.stdint cimport intptr_t - -from .cy${libname} cimport * - - -############################################################################### -# Types -############################################################################### - -$type_decls - -ctypedef cudaStream_t Stream - - -############################################################################### -# Enum -############################################################################### - -$enum_decls - - -############################################################################### -# Functions -############################################################################### - -$func_decls -cpdef object get_library_path() diff --git a/python/build_assets/cybind/templates/nccl/_extensions/bindings/nccl_ep.pyx b/python/build_assets/cybind/templates/nccl/_extensions/bindings/nccl_ep.pyx deleted file mode 100644 index 5a6808db..00000000 --- a/python/build_assets/cybind/templates/nccl/_extensions/bindings/nccl_ep.pyx +++ /dev/null @@ -1,64 +0,0 @@ -# Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. -# -# SPDX-License-Identifier: Apache-2.0 -# -# This code was automatically generated $version_span. Do not modify it directly. - -cimport cython # NOQA -from libc.stdint cimport uint64_t -from libcpp.vector cimport vector - -from ._internal.utils cimport (nested_resource, nullable_unique_ptr, get_buffer_pointer, - get_resource_ptr, get_nested_resource_ptr) - -from enum import IntEnum as _IntEnum - - -$snippet_auto_lowpp_imports_pyx - - -############################################################################### -# POD -############################################################################### - -$pod_defs - - -############################################################################### -# Enum -############################################################################### - -$enum_defs - - -############################################################################### -# Error handling -############################################################################### - -class NCCLEpError(Exception): - - def __init__(self, status): - self.status = status - cdef str err = f"NCCL EP error code {status}" - super(NCCLEpError, self).__init__(err) - - def __reduce__(self): - return (type(self), (self.status,)) - - -@cython.profile(False) -cpdef inline check_status(int status): - if status != 0: - raise NCCLEpError(status) - - -############################################################################### -# Wrapper functions -############################################################################### - -$wrapper_defs - - -cpdef object get_library_path(): - from ._internal.nccl_ep import _inspect_loaded_library_path - return _inspect_loaded_library_path() diff --git a/python/build_assets/cybind/templates/nccl/_extensions/bindings/nccl_m2n.pxd b/python/build_assets/cybind/templates/nccl/_extensions/bindings/nccl_m2n.pxd deleted file mode 100644 index 1460b0f2..00000000 --- a/python/build_assets/cybind/templates/nccl/_extensions/bindings/nccl_m2n.pxd +++ /dev/null @@ -1,32 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# This code was automatically generated $version_span. Do not modify it directly. - -from libc.stdint cimport intptr_t - -from .cy${libname} cimport * - - -############################################################################### -# Types -############################################################################### - -$type_decls - -ctypedef ncclComm_t Comm -ctypedef ncclWindow_t Window -ctypedef cudaStream_t Stream - - -############################################################################### -# Functions -############################################################################### - -cpdef intptr_t init(intptr_t config) except? 0 -cpdef finalize(intptr_t handle) -cpdef group_start() -cpdef group_end() -cpdef group_abort() -cpdef reshard_with_window(intptr_t handle, intptr_t comm, intptr_t window, intptr_t src, intptr_t dst, intptr_t stream) -cpdef reshard(intptr_t handle, intptr_t comm, intptr_t src, intptr_t dst, intptr_t stream) diff --git a/python/build_assets/cybind/templates/nccl/_extensions/bindings/nccl_m2n.pyx b/python/build_assets/cybind/templates/nccl/_extensions/bindings/nccl_m2n.pyx deleted file mode 100644 index 855180e8..00000000 --- a/python/build_assets/cybind/templates/nccl/_extensions/bindings/nccl_m2n.pyx +++ /dev/null @@ -1,178 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# This code was automatically generated $version_span. Do not modify it directly. - -cimport cython # NOQA -from libc.stdint cimport uint64_t -from libcpp.vector cimport vector - -from ._internal.utils cimport (nested_resource, nullable_unique_ptr, get_buffer_pointer, - get_resource_ptr, get_nested_resource_ptr) - -from enum import IntEnum as _IntEnum - -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# See LICENSE.txt for more license information. - -from libc.stdlib cimport calloc, free, malloc -from cython cimport view -cimport cpython.buffer -cimport cpython.memoryview -cimport cpython -from libc.string cimport memcmp, memcpy -import numpy as _numpy - - -cdef __from_data(data, dtype_name, expected_dtype, lowpp_type): - # _numpy.recarray is a subclass of _numpy.ndarray, so implicitly handled here. - if isinstance(data, lowpp_type): - return data - if not isinstance(data, _numpy.ndarray): - raise TypeError("data argument must be a NumPy ndarray") - if data.size != 1: - raise ValueError("data array must have a size of 1") - if data.dtype != expected_dtype: - raise ValueError(f"data array must be of dtype {dtype_name}") - return lowpp_type.from_ptr(data.ctypes.data, not data.flags.writeable, data) - - -cdef __from_buffer(buffer, size, lowpp_type): - cdef Py_buffer view - if cpython.PyObject_GetBuffer(buffer, &view, cpython.PyBUF_SIMPLE) != 0: - raise TypeError("buffer argument does not support the buffer protocol") - try: - if view.itemsize != 1: - raise ValueError("buffer itemsize must be 1 byte") - if view.len != size: - raise ValueError(f"buffer length must be {size} bytes") - return lowpp_type.from_ptr(view.buf, view.readonly, buffer) - finally: - cpython.PyBuffer_Release(&view) - - -cdef __getbuffer(object self, cpython.Py_buffer *buffer, void *ptr, int size, bint readonly): - buffer.buf = ptr - buffer.format = 'b' - buffer.internal = NULL - buffer.itemsize = 1 - buffer.len = size - buffer.ndim = 1 - buffer.obj = self - buffer.readonly = readonly - buffer.shape = &buffer.len - buffer.strides = &buffer.itemsize - buffer.suboffsets = NULL - - -############################################################################### -# POD -############################################################################### - -$pod_defs - - -MESH_NDIMS = NCCL_RESHARD_MESH_NDIMS -MAX_TENSOR_DIMS = NCCL_RESHARD_MAX_TENSOR_DIMS -REPLICATE = NCCL_RESHARD_REPLICATE - - -############################################################################### -# Error handling -############################################################################### - -from nccl.bindings.nccl import NCCLError as _NCCLError -from nccl._extensions._runtime import NATIVE_CALL_LOCK as _NATIVE_CALL_LOCK -from ._internal.utils import FunctionNotFoundError - - -class NCCLReshardError(_NCCLError): - - def __init__(self, status, detail=None): - self.status = int(status) - self.detail = detail - message = f"NCCL Reshard error code {self.status}" - if detail: - message += f": {detail}" - Exception.__init__(self, message) - - def __reduce__(self): - return (type(self), (self.status, self.detail)) - - -@cython.profile(False) -cpdef inline check_status(int status): - cdef const char* detail = NULL - cdef bytes detail_bytes - cdef object detail_text = None - if status != 0: - detail = ncclM2nGetLastError() - if detail != NULL: - detail_bytes = detail - if detail_bytes: - detail_text = detail_bytes.decode("utf-8", "replace") - raise NCCLReshardError(status, detail_text) - - -############################################################################### -# Wrapper functions -############################################################################### - -cpdef intptr_t init(intptr_t config) except? 0: - cdef Handle handle - cdef ncclResult_t status - with _NATIVE_CALL_LOCK: - with nogil: - status = ncclM2nInit(&handle, config) - check_status(status) - return handle - - -cpdef finalize(intptr_t handle): - cdef ncclResult_t status - with _NATIVE_CALL_LOCK: - with nogil: - status = ncclM2nFinalize(handle) - check_status(status) - - -cpdef group_start(): - cdef ncclResult_t status - with _NATIVE_CALL_LOCK: - with nogil: - status = ncclM2nGroupStart() - check_status(status) - - -cpdef group_end(): - cdef ncclResult_t status - with _NATIVE_CALL_LOCK: - with nogil: - status = ncclM2nGroupEnd() - check_status(status) - - -cpdef group_abort(): - cdef ncclResult_t status - with _NATIVE_CALL_LOCK: - with nogil: - status = ncclM2nGroupAbort() - check_status(status) - - -cpdef reshard_with_window(intptr_t handle, intptr_t comm, intptr_t window, intptr_t src, intptr_t dst, intptr_t stream): - cdef ncclResult_t status - with _NATIVE_CALL_LOCK: - with nogil: - status = ncclReshardWithWindow(handle, comm, window, src, dst, stream) - check_status(status) - - -cpdef reshard(intptr_t handle, intptr_t comm, intptr_t src, intptr_t dst, intptr_t stream): - cdef ncclResult_t status - with _NATIVE_CALL_LOCK: - with nogil: - status = ncclReshard(handle, comm, src, dst, stream) - check_status(status) diff --git a/python/build_assets/generate_cython.py b/python/build_assets/generate_cython.py deleted file mode 100644 index a2ba80e3..00000000 --- a/python/build_assets/generate_cython.py +++ /dev/null @@ -1,452 +0,0 @@ -#!/usr/bin/env python3 -# -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# See LICENSE.txt for more license information -# - -""" -Generate Cython bindings for the nccl-extensions libraries using cybind. - -Ported from nccl4py's ``build_assets/generate_cython.py``, trimmed to the -generated targets this repo owns (``nccl_ep`` and ``nccl_m2n``). Output goes -to ``python/nccl/_extensions/bindings/`` as flat sibling modules. - -Unlike nccl4py, the bound headers are *not* checked in under -``cybind/headers//``: nccl_ep's public header lives in this repo, so -it is staged straight from ``nccl_ep/include/`` and the version is read from -its ``NCCL_EP_{MAJOR,MINOR,PATCH}`` macros. Only headers this repo does not -own -- currently just ``nccl.h`` -- are pinned under ``cybind/headers/``. -""" - -from __future__ import annotations - -import argparse -import logging -import os -import re -import shutil -import subprocess -import sys -import tempfile -from contextlib import contextmanager, nullcontext -from dataclasses import dataclass -from pathlib import Path - -from packaging.version import Version - - -# cybind repository configuration -CYBIND_COMMIT = "cde8bae486ff7ff88accf3cfff4e62527fd06199" -CYBIND_SSH_URL = "ssh://git@gitlab-master.nvidia.com:12051/xiakunl/cybind.git" - -# Script directory for resolving default paths -SCRIPT_DIR = Path(__file__).resolve().parent -PYTHON_DIR = SCRIPT_DIR.parent -REPO_ROOT = PYTHON_DIR.parent - -# Shared paths: one cybind/ assets dir feeds every target; cybind emits into -# one bindings package (flat sibling layout). -ASSETS_DIR = SCRIPT_DIR / "cybind" -BINDINGS_DIR = PYTHON_DIR / "nccl" / "_extensions" / "bindings" - -# Pinned copies of headers owned by other repos, laid out the way cybind's own -# assets/headers/ does: //. Bumping NCCL_PIN means dropping -# the matching nccl.h here and regenerating. -HEADERS_DIR = ASSETS_DIR / "headers" -NCCL_PIN = Version("2.30.4") - -# Static files in templates/ that cybind doesn't process -- copied verbatim -# into BINDINGS_DIR after cybind finishes. -STATIC_FILES = ( - "__init__.py", - "_internal/__init__.py", - "_internal/utils.pxd", - "_internal/utils.pyx", -) - -# Every target emits into ``nccl._extensions.bindings.*``, so cybind looks up -# their templates under one shared subtree. -_TEMPLATES_RELPATH = Path("nccl", "_extensions", "bindings") - -# Global logger - will be configured in main() -logger = logging.getLogger(__name__) - - -@dataclass(frozen=True) -class Target: - """Per-binding-target configuration for cybind.""" - - name: str # cybind library name (top-level key in YAML) - version: Version - # Headers staged into cybind's assets/headers///, as - # {path relative to that dir: source file}. Must cover the transitive - # includes of the config's ``data.headers`` entry, laid out the way - # those #include directives spell them. - headers: dict[str, Path] - - -def _read_version(header: Path, prefix: str) -> Version: - """Read ``_{MAJOR,MINOR,PATCH}`` #defines out of a C header.""" - text = header.read_text() - parts = [] - for component in ("MAJOR", "MINOR", "PATCH"): - match = re.search(rf"^#define\s+{prefix}_{component}\s+(\d+)", text, re.MULTILINE) - if match is None: - raise RuntimeError(f"No {prefix}_{component} #define found in {header}") - parts.append(match.group(1)) - return Version(".".join(parts)) - - -def _nccl_ep_target() -> Target: - """Bind nccl_ep against this repo's own headers. - - The staged layout mirrors what ``nccl_ep/CMakeLists.txt`` installs (public - header at the root, everything else under ``nccl_ep/``), which is how - ``nccl_ep.h``'s own #include directives spell them. - """ - include_dir = REPO_ROOT / "nccl_ep" / "include" - public_header = include_dir / "nccl_ep.h" - if not public_header.is_file(): - raise RuntimeError(f"nccl_ep public header not found: {public_header}") - return Target( - name="nccl_ep", - version=_read_version(public_header, "NCCL_EP"), - headers={ - "nccl_ep.h": public_header, - "nccl_ep/ep_enums.h": include_dir / "ep_enums.h", - "nccl.h": HEADERS_DIR / "nccl" / str(NCCL_PIN) / "nccl.h", - }, - ) - - -def _nccl_m2n_target() -> Target: - """Bind nccl_m2n against its public header. - - The generated declarations encode the public ABI, so downstream wheel - builds do not need a separate M2N header installation. ``nccl.h`` stays - pinned with the other third-party headers because M2N imports its NCCL - result and datatype definitions. - """ - public_header = REPO_ROOT / "nccl_m2n" / "src" / "nccl_m2n.h" - if not public_header.is_file(): - raise RuntimeError(f"nccl_m2n public header not found: {public_header}") - return Target( - name="nccl_m2n", - version=_read_version(public_header, "NCCL_M2N"), - headers={ - "nccl_m2n.h": public_header, - "nccl.h": HEADERS_DIR / "nccl" / str(NCCL_PIN) / "nccl.h", - }, - ) - - -def _stamp_asset_yaml(yaml_path: Path, version: Version) -> None: - """Stamp ``data.versions: - - X.Y.Z`` in the asset YAML in place. - - Text-level edits (vs. parse -> mutate -> dump) preserve the file's - comments, ordering, and formatting that ``yaml.safe_dump`` would strip. - The substitution is idempotent: the regex matches any prior value, so - reruns at the same version are no-ops. - """ - text = yaml_path.read_text() - text, n_versions = re.subn(r"(versions:\n\s*- - )\S+", rf"\g<1>{version}", text) - if n_versions != 1: - raise RuntimeError( - f"Expected exactly one `versions:` block in {yaml_path}, found {n_versions}" - ) - yaml_path.write_text(text) - - -def clone_cybind(cybind_dir: Path) -> None: - try: - subprocess.run( - ["git", "clone", CYBIND_SSH_URL, str(cybind_dir)], - check=True, - capture_output=True, - text=True, - ) - subprocess.run( - ["git", "switch", "--detach", CYBIND_COMMIT], - cwd=cybind_dir, - check=True, - capture_output=True, - text=True, - ) - logger.debug(f"Cloned -> {cybind_dir}") - except subprocess.CalledProcessError as e: - raise RuntimeError(f"Failed to clone cybind: {e.stderr}") from e - - -def prepare_assets(cybind_dir: Path, targets: list[Target]) -> None: - """Stage our targets' configs, headers, and templates into cybind's - shared ``assets/`` -- only the slots we own (one ``configs/*.cybind.yaml`` - per target, per-target ``headers///``, and the shared - ``templates/nccl/_extensions/bindings/`` subtree). Sibling files for - other libs are left untouched.""" - cybind_assets = cybind_dir / "cybind" / "assets" - - for target in targets: - config_filename = f"{target.name}.cybind.yaml" - config_src = ASSETS_DIR / "configs" / config_filename - if not config_src.exists(): - raise FileNotFoundError(f"Config file not found: {config_src}") - _stamp_asset_yaml(config_src, target.version) - config_dst = cybind_assets / "configs" / config_filename - config_dst.parent.mkdir(parents=True, exist_ok=True) - shutil.copy2(config_src, config_dst) - logger.debug(f"Staged {config_filename} (version={target.version})") - - headers_dst = cybind_assets / "headers" / target.name / str(target.version) - if headers_dst.exists(): - shutil.rmtree(headers_dst) - for relpath, src in target.headers.items(): - if not src.is_file(): - raise FileNotFoundError(f"Header not found: {src}") - dst = headers_dst / relpath - dst.parent.mkdir(parents=True, exist_ok=True) - shutil.copy2(src, dst) - logger.debug(f"Staged headers/{target.name}/{target.version}/{relpath} <- {src}") - - templates_src = ASSETS_DIR / "templates" / _TEMPLATES_RELPATH - if not templates_src.is_dir(): - raise FileNotFoundError(f"Templates dir not found: {templates_src}") - templates_dst = cybind_assets / "templates" / _TEMPLATES_RELPATH - if templates_dst.exists(): - shutil.rmtree(templates_dst) - templates_dst.parent.mkdir(parents=True, exist_ok=True) - shutil.copytree(templates_src, templates_dst) - logger.debug(f"Staged templates/{_TEMPLATES_RELPATH}/") - -def run_cybind(cybind_dir: Path, libnames: list[str], output_dir: Path) -> None: - """Run cybind once for all libraries in libnames. - - Cybind's ``--generate`` is multi-valued; passing all libnames in one - invocation amortizes the venv setup. Each library's headers must live at - ``cybind/assets/headers///`` (cybind's default lookup - when no ``--input-dir`` is passed); ``prepare_assets`` handles that. - - Args: - - cybind_dir: Path to cybind repository. - - libnames: Library keys to generate (top-level YAML keys). - - output_dir: Cybind's ``--output`` value. Cybind emits each library - into ``output_dir//``. - - Note: - - Requires CUDA_PATH for CUDA header resolution. - - Uses uv to create an isolated venv and install cybind. - """ - output_dir.mkdir(parents=True, exist_ok=True) - - cmd = [ - "uv", - "run", - "--isolated", - "--with", - str(cybind_dir), - "-m", - "cybind", - "--generate", - *libnames, - "--output", - str(output_dir), - ] - - logger.debug(f"Command: {' '.join(cmd)}") - logger.debug(f"Working directory: {cybind_dir}") - logger.debug(f"CUDA_PATH: {os.environ.get('CUDA_PATH', 'not set')}") - - try: - result = subprocess.run(cmd, cwd=cybind_dir, check=True, capture_output=True, text=True) - if result.stdout: - logger.debug("cybind execution stdout:") - for line in result.stdout.splitlines(): - logger.debug(f" {line}") - if result.stderr: - logger.debug("cybind execution stderr:") - for line in result.stderr.splitlines(): - logger.debug(f" {line}") - logger.debug("Successfully generated bindings") - except subprocess.CalledProcessError as e: - logger.error(f"cybind failed with exit code {e.returncode}") - if e.stdout: - logger.error("=== stdout ===") - for line in e.stdout.splitlines(): - logger.error(line) - if e.stderr: - logger.error("=== stderr ===") - for line in e.stderr.splitlines(): - logger.error(line) - raise RuntimeError("cybind execution failed") from e - - -@contextmanager -def backup_and_restore_on_failure(path: Path): - """Restore the prior binding package if generation does not complete.""" - backup = path.with_name(path.name + ".bak") - if backup.exists(): - shutil.rmtree(backup) - if path.exists(): - logger.info(f">>> Backing up {path} -> {backup}") - path.rename(backup) - try: - yield - except BaseException: - if path.exists(): - shutil.rmtree(path) - if backup.exists(): - logger.error(f"Bindings generation failed; restoring {path}") - backup.rename(path) - raise - else: - if backup.exists(): - logger.info(f">>> Removing backup {backup}") - shutil.rmtree(backup) - - -def verify_generated_m2n_loader_contract() -> None: - """Check the generated M2N loader's atomic loading diagnostics contract.""" - loader = BINDINGS_DIR / "_internal" / "nccl_m2n_linux.pyx" - text = loader.read_text() - required = ( - 'errors.append(f"{path}: not found")', - 'missing.append("ncclM2nGetLastError")', - ) - missing = [snippet for snippet in required if snippet not in text] - if missing: - raise RuntimeError( - f"Generated {loader} does not preserve the M2N loader contract: {missing}" - ) - - -def verify_generated_m2n_import_layering() -> None: - """Keep low-level bindings independent from the public M2N facade.""" - binding = BINDINGS_DIR / "nccl_m2n.pyx" - text = binding.read_text() - required = "from nccl._extensions._runtime import NATIVE_CALL_LOCK as _NATIVE_CALL_LOCK" - if required not in text or "nccl.m2n" in text: - raise RuntimeError( - f"Generated {binding} imports public nccl.m2n state and can create an import cycle" - ) - - -# One entry per bound library. Each is a zero-arg factory so a missing header -# for one target raises with that target's own error message. -TARGETS = (_nccl_ep_target, _nccl_m2n_target) - - -def main() -> int: - parser = argparse.ArgumentParser( - description="Generate Cython bindings for nccl-extensions using cybind" - ) - parser.add_argument( - "--cuda-home", - type=Path, - default=None, - help="Path to CUDA installation (default: $CUDA_PATH or $CUDA_HOME)", - ) - parser.add_argument( - "--cybind-path", - type=Path, - default=None, - help=( - "Use a local cybind checkout at this path instead of cloning. " - "Note: our targets' slots under the checkout's assets/ " - "(configs/.cybind.yaml, headers///, " - "templates/nccl/_extensions/bindings/) are overwritten; sibling " - "files for other libs are left alone. Default: clone " - "CYBIND_SSH_URL at the pinned CYBIND_COMMIT into a temp dir." - ), - ) - parser.add_argument("--verbose", "-v", action="store_true", help="Enable verbose (debug) output") - args = parser.parse_args() - - log_level = logging.DEBUG if args.verbose else logging.INFO - logging.basicConfig(level=log_level, format="%(message)s") - - # Fail fast on missing tools rather than after the cybind clone. - required_tools = ["uv"] - if args.cybind_path is None: - required_tools.append("git") - for tool in required_tools: - if shutil.which(tool) is None: - logger.error(f"{tool} not found on PATH") - return 1 - - if args.cybind_path is not None and not args.cybind_path.is_dir(): - logger.error(f"--cybind-path is not a directory: {args.cybind_path}") - return 1 - - # CUDA_PATH (cybind uses it to find cuda.h) - if args.cuda_home: - if not args.cuda_home.exists(): - logger.error(f"CUDA home not found at {args.cuda_home}") - return 1 - cuda_path = str(args.cuda_home) - else: - cuda_path = os.environ.get("CUDA_PATH") or os.environ.get("CUDA_HOME") - if not cuda_path: - logger.error("Provide --cuda-home or set CUDA_PATH or CUDA_HOME") - return 1 - if not Path(cuda_path).is_dir(): - logger.error(f"CUDA path is not a directory: {cuda_path}") - return 1 - - os.environ["CUDA_PATH"] = cuda_path - logger.debug(f"CUDA: {cuda_path}") - - try: - targets = [make_target() for make_target in TARGETS] - except (RuntimeError, FileNotFoundError) as e: - logger.error(str(e)) - return 1 - - # cybind tree: use --cybind-path in place, or clone into a temp dir. - if args.cybind_path is not None: - logger.info(f">>> Using local cybind from {args.cybind_path} (in place)") - cybind_ctx = nullcontext(args.cybind_path) - else: - cybind_ctx = tempfile.TemporaryDirectory(prefix="nccl_extensions_cybind_") - - with cybind_ctx as cybind_dir_: - cybind_dir = Path(cybind_dir_) - if args.cybind_path is None: - logger.info(">>> Cloning cybind...") - clone_cybind(cybind_dir) - - logger.info(">>> Preparing cybind assets") - prepare_assets(cybind_dir, targets) - - # Back up the real bindings dir, then run cybind with --output pointing - # at PYTHON_DIR so the emitted // path lands directly on BINDINGS_DIR. The context manager - # removes the backup on success and restores it on failure. - with backup_and_restore_on_failure(BINDINGS_DIR): - logger.info(">>> Running cybind for all targets") - run_cybind(cybind_dir, [t.name for t in targets], PYTHON_DIR) - - # Copy shared static template files cybind doesn't process. - logger.debug("Copying static files from templates...") - templates_root = ASSETS_DIR / "templates" / _TEMPLATES_RELPATH - for rel in STATIC_FILES: - src = templates_root / rel - if not src.exists(): - continue - dst = BINDINGS_DIR / rel - dst.parent.mkdir(parents=True, exist_ok=True) - logger.debug(f" {rel}") - shutil.copy2(src, dst) - - verify_generated_m2n_loader_contract() - verify_generated_m2n_import_layering() - - logger.info("=" * 60) - logger.info(f"Bindings location: {BINDINGS_DIR}") - logger.info(f"Generated: {', '.join(f'{t.name} {t.version}' for t in targets)}") - logger.info("=" * 60) - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/python/ci/verify_m2n_package.sh b/python/ci/verify_m2n_package.sh deleted file mode 100644 index 2f0727dd..00000000 --- a/python/ci/verify_m2n_package.sh +++ /dev/null @@ -1,103 +0,0 @@ -#!/usr/bin/env bash -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -# Build and inspect the Python distribution with the M2N native artifact from -# the preceding child-pipeline build job. This intentionally tests the public -# package boundary rather than the source tree through PYTHONPATH. - -set -euo pipefail - -if [[ $# -ne 1 ]]; then - echo "usage: $0 " >&2 - exit 2 -fi - -: "${CUDA_HOME:?CUDA_HOME must be set}" -: "${NCCL_HOME:?NCCL_HOME must be set}" - -ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" -BUILD_DIR="$(cd "$1" && pwd)" -PACKAGE_DIR="${ROOT}/python" -M2N_PACKAGE_DIR="${PACKAGE_DIR}/nccl/m2n" -DIST_DIR="${PACKAGE_DIR}/dist" -VENV_DIR="${PACKAGE_DIR}/.ci-m2n-wheel-venv" -BOOTSTRAP_DIR="${PACKAGE_DIR}/.ci-m2n-bootstrap" -TOOLS_VENV_DIR="${PACKAGE_DIR}/.ci-m2n-tools" -UV_CACHE_DIR="${PACKAGE_DIR}/.ci-m2n-uv-cache" -UV_PYTHON_INSTALL_DIR="${PACKAGE_DIR}/.ci-m2n-uv-python" - -LIBRARY="${BUILD_DIR}/lib/libnccl_m2n.so" -HEADER="${BUILD_DIR}/include/nccl_m2n.h" -test -s "${LIBRARY}" -test -s "${HEADER}" -test -d "${NCCL_HOME}/lib" - -cleanup() { - rm -rf "${M2N_PACKAGE_DIR}/lib" "${M2N_PACKAGE_DIR}/include" "${VENV_DIR}" \ - "${BOOTSTRAP_DIR}" "${TOOLS_VENV_DIR}" "${UV_CACHE_DIR}" \ - "${UV_PYTHON_INSTALL_DIR}" -} -trap cleanup EXIT -rm -rf "${DIST_DIR}" -cleanup - -mkdir -p "${M2N_PACKAGE_DIR}/lib" "${M2N_PACKAGE_DIR}/include" "${DIST_DIR}" -cp "${LIBRARY}" "${M2N_PACKAGE_DIR}/lib/libnccl_m2n.so" -cp "${HEADER}" "${M2N_PACKAGE_DIR}/include/nccl_m2n.h" - -# The wheel build below Cython-compiles the checked-in generated sources for -# both EP and M2N. Regeneration is intentionally out of scope for this package -# gate: it needs private cybind source access unavailable in the build image. -# The build-tools image has Python 3.8 and a read-only home directory, while -# this package requires Python 3.10+. Bootstrap uv into the writable checkout, -# then let it provision an isolated supported interpreter and tool environment. -python3 -m pip install --disable-pip-version-check --target "${BOOTSTRAP_DIR}" uv -export UV_CACHE_DIR UV_PYTHON_INSTALL_DIR -PYTHONPATH="${BOOTSTRAP_DIR}" python3 -m uv venv --python 3.12 --seed "${TOOLS_VENV_DIR}" -"${TOOLS_VENV_DIR}/bin/python" -m pip install --disable-pip-version-check build uv -export PATH="${TOOLS_VENV_DIR}/bin:${PATH}" -PYTHON="${TOOLS_VENV_DIR}/bin/python" - -# Build the wheel from the staged checkout, then build the source-only sdist -# separately. The default `build` sequence derives its wheel from the sdist, -# which intentionally excludes the staged shared library. -"${PYTHON}" -m build --wheel --outdir "${DIST_DIR}" "${PACKAGE_DIR}" -"${PYTHON}" -m build --sdist --outdir "${DIST_DIR}" "${PACKAGE_DIR}" - -WHEEL="$(find "${DIST_DIR}" -maxdepth 1 -name '*.whl' -print -quit)" -SDIST="$(find "${DIST_DIR}" -maxdepth 1 -name '*.tar.gz' -print -quit)" -test -n "${WHEEL}" -test -n "${SDIST}" - -"${PYTHON}" - "${WHEEL}" <<'PY' -import sys -import zipfile - -with zipfile.ZipFile(sys.argv[1]) as wheel: - names = wheel.namelist() - for expected in ( - "nccl/m2n/lib/libnccl_m2n.so", - "nccl/m2n/include/nccl_m2n.h", - ): - if expected not in names: - raise SystemExit(f"wheel is missing {expected}") - metadata = next(name for name in names if name.endswith(".dist-info/METADATA")) - if "Provides-Extra: bench" not in wheel.read(metadata).decode(): - raise SystemExit("wheel metadata is missing the bench extra") -PY -tar -tzf "${SDIST}" | grep -E '/nccl/_extensions/bindings/(nccl_m2n|cynccl_m2n)\.pyx$' -tar -tzf "${SDIST}" | grep -E '/nccl/m2n/include/nccl_m2n\.h$' -if tar -tzf "${SDIST}" | grep -E '\.so$'; then - echo "ERROR: the source distribution contains a shared library" >&2 - exit 1 -fi - -"${PYTHON}" -m venv "${VENV_DIR}" -# Match the CUDA 12 build image and install the benchmark dependency through -# the public extras, rather than relying on undeclared transitive packages. -"${VENV_DIR}/bin/pip" install --disable-pip-version-check "${WHEEL}[cu12,bench]" -LD_LIBRARY_PATH="${NCCL_HOME}/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" \ - "${VENV_DIR}/bin/python" -c 'import nccl.m2n; from nccl._extensions.bindings import nccl_m2n' -LD_LIBRARY_PATH="${NCCL_HOME}/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" \ - "${VENV_DIR}/bin/python" -m nccl.m2n.benchmarks.reshard_bench --help