Module nccl 

Raw NCCL FFI bindings (minimal surface for EP all-reduce + broadcast).

Only the functions Atlas actually calls are bound here — no attempt at complete coverage. Type sizes match NCCL 2.28+ on aarch64 (symmetric memory ncclMemAlloc/ncclMemFree require NCCL ≥ 2.28).

Safety

Every unsafe block wraps a single NCCL FFI call. Invariants:

• The NcclComm/NcclUniqueId arguments come from prior Self constructors that called the matching NCCL init function.
• GPU buffers are valid DevicePtrs alive on the device that owns the comm.
• Counts × dtype-size match the buffer byte count.
• The extern "C" declarations match the NCCL header ABI for the linked library version.

Structs

NcclConfig

NCCL config for ncclCommInitRankConfig. Set blocking=0 for non-blocking mode (ncclGroupEnd returns immediately).

NcclUniqueId

NCCL unique ID for bootstrapping (128 bytes, passed by value in C ABI).

Enums

NcclDataType

NCCL data types (matches nccl.h enum values).

NcclRedOp

NCCL reduction operations.

NcclResult

NCCL result code.

Functions

check_nccl

Convert NCCL result to anyhow::Result.

create_event

create_stream

destroy_event

destroy_stream

ncclAllGather^⚠

ncclAllReduce^⚠

ncclBroadcast^⚠

ncclCommAbort^⚠

ncclCommDeregister^⚠

ncclCommDestroy^⚠

ncclCommGetAsyncError^⚠

ncclCommInitRank^⚠

ncclCommInitRankConfig^⚠

Non-blocking variant: set config.blocking=0 so ncclGroupEnd returns immediately. Poll ncclCommGetAsyncError to check completion and detect hangs with timeout.

ncclCommRegister^⚠

ncclGetErrorString^⚠

ncclGetUniqueId^⚠

ncclGroupEnd^⚠

ncclGroupStart^⚠

ncclMemAlloc^⚠

ncclMemFree^⚠

ncclRecv^⚠

ncclReduceScatter^⚠

ncclSend^⚠

nccl_mem_alloc^⚠

Allocate GPU memory backed by a symmetric memory window across the communicator. NCCL 2.28+ only — older NCCL returns InvalidArgument.

nccl_mem_free^⚠

Free a buffer previously returned by nccl_mem_alloc.

record_event

stream_wait_event

sync_stream

Type Aliases

NcclComm

Opaque NCCL communicator handle.