metrale_gpu_sys/
nccl.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: Raw NCCL bindings used by metrale-comm, plus the CUDA driver
4//! stream and event calls it needs. The function declarations follow the
5//! installed nccl.h (2.29.7); `NcclConfig` does not (see its doc).
6//!
7//! The `extern` functions are unchecked: a caller passes a communicator from
8//! `ncclCommInitRank`, device pointers valid for the element count on that
9//! communicator's device, and a stream of that device. The wrappers below
10//! (`create_stream` to `stream_ready`) check only the status code.
11//!
12//! Owner: metrale-gpu-sys.
13//! Invariants: none beyond the types.
14
15use std::ffi::c_void;
16
17/// 2026-09-26: `ncclComm_t`, an opaque communicator handle.
18pub type NcclComm = *mut c_void;
19
20/// 2026-09-26: `ncclUniqueId`: 128 bytes, passed by value.
21#[repr(C)]
22#[derive(Copy, Clone)]
23pub struct NcclUniqueId {
24    pub internal: [u8; 128],
25}
26
27/// 2026-09-26: Meant for `ncclCommInitRankConfig`, which nothing in the
28/// workspace calls. It does not match `ncclConfig_t` in the installed nccl.h
29/// (2.29.7): that struct starts with magic `NCCL_API_MAGIC` (0xcafebeef), holds
30/// `netName` as a pointer, and has more fields.
31#[repr(C)]
32pub struct NcclConfig {
33    pub size: usize,
34    pub magic: u32,
35    pub version: u32,
36    pub blocking: i32,
37    pub cga_cluster_size: i32,
38    pub min_ctas: i32,
39    pub max_ctas: i32,
40    pub net_name: [u8; 8],
41    pub split_share: i32,
42}
43
44impl NcclConfig {
45    pub fn non_blocking() -> Self {
46        Self {
47            size: std::mem::size_of::<Self>(),
48            magic: 0x4e43434c,
49            version: 22907,
50            blocking: 0,
51            cga_cluster_size: -1,
52            min_ctas: -1,
53            max_ctas: -1,
54            net_name: [0; 8],
55            split_share: -1,
56        }
57    }
58}
59
60/// 2026-09-26: `ncclResult_t`, values 0 to 7 as in nccl.h.
61#[repr(C)]
62#[derive(Debug, Copy, Clone, PartialEq, Eq)]
63pub enum NcclResult {
64    Success = 0,
65    UnhandledCudaError = 1,
66    SystemError = 2,
67    InternalError = 3,
68    InvalidArgument = 4,
69    InvalidUsage = 5,
70    RemoteError = 6,
71    InProgress = 7,
72}
73
74/// 2026-09-26: `ncclDataType_t`, the values up to `ncclBfloat16` (9) as in
75/// nccl.h.
76#[repr(C)]
77#[derive(Debug, Copy, Clone)]
78#[allow(dead_code)]
79pub enum NcclDataType {
80    Int8 = 0,
81    Uint8 = 1,
82    Int32 = 2,
83    Uint32 = 3,
84    Int64 = 4,
85    Uint64 = 5,
86    Float16 = 6,
87    Float32 = 7,
88    Float64 = 8,
89    Bfloat16 = 9,
90}
91
92/// 2026-09-26: `ncclRedOp_t`'s built-in operations, values as in nccl.h.
93#[repr(C)]
94#[derive(Debug, Copy, Clone)]
95#[allow(dead_code)]
96pub enum NcclRedOp {
97    Sum = 0,
98    Prod = 1,
99    Max = 2,
100    Min = 3,
101    Avg = 4,
102}
103
104#[link(name = "nccl")]
105unsafe extern "C" {
106    pub fn ncclGetUniqueId(id: *mut NcclUniqueId) -> NcclResult;
107
108    pub fn ncclCommInitRank(
109        comm: *mut NcclComm,
110        nranks: i32,
111        id: NcclUniqueId,
112        rank: i32,
113    ) -> NcclResult;
114
115    pub fn ncclCommInitRankConfig(
116        comm: *mut NcclComm,
117        nranks: i32,
118        id: NcclUniqueId,
119        rank: i32,
120        config: *const NcclConfig,
121    ) -> NcclResult;
122
123    pub fn ncclAllReduce(
124        sendbuf: *const c_void,
125        recvbuf: *mut c_void,
126        count: usize,
127        datatype: NcclDataType,
128        op: NcclRedOp,
129        comm: NcclComm,
130        stream: u64,
131    ) -> NcclResult;
132
133    pub fn ncclBroadcast(
134        sendbuf: *const c_void,
135        recvbuf: *mut c_void,
136        count: usize,
137        datatype: NcclDataType,
138        root: i32,
139        comm: NcclComm,
140        stream: u64,
141    ) -> NcclResult;
142
143    pub fn ncclCommDestroy(comm: NcclComm) -> NcclResult;
144
145    pub fn ncclGetErrorString(result: NcclResult) -> *const std::ffi::c_char;
146
147    pub fn ncclCommRegister(
148        comm: NcclComm,
149        buff: *mut c_void,
150        size: usize,
151        handle: *mut *mut c_void,
152    ) -> NcclResult;
153
154    pub fn ncclCommDeregister(comm: NcclComm, handle: *mut c_void) -> NcclResult;
155
156    pub fn ncclMemAlloc(ptr: *mut *mut c_void, size: usize) -> NcclResult;
157
158    pub fn ncclMemFree(ptr: *mut c_void) -> NcclResult;
159
160    pub fn ncclSend(
161        sendbuf: *const c_void,
162        count: usize,
163        datatype: NcclDataType,
164        peer: i32,
165        comm: NcclComm,
166        stream: u64,
167    ) -> NcclResult;
168
169    pub fn ncclRecv(
170        recvbuf: *mut c_void,
171        count: usize,
172        datatype: NcclDataType,
173        peer: i32,
174        comm: NcclComm,
175        stream: u64,
176    ) -> NcclResult;
177
178    // 2026-09-26: `recvbuf` holds `nranks * sendcount` elements, rank i's at
179    // offset `i * sendcount` (nccl.h).
180    pub fn ncclAllGather(
181        sendbuf: *const c_void,
182        recvbuf: *mut c_void,
183        sendcount: usize,
184        datatype: NcclDataType,
185        comm: NcclComm,
186        stream: u64,
187    ) -> NcclResult;
188
189    // 2026-09-26: `sendbuf` holds `nranks * recvcount` elements; each rank
190    // receives `recvcount` (nccl.h).
191    pub fn ncclReduceScatter(
192        sendbuf: *const c_void,
193        recvbuf: *mut c_void,
194        recvcount: usize,
195        datatype: NcclDataType,
196        op: NcclRedOp,
197        comm: NcclComm,
198        stream: u64,
199    ) -> NcclResult;
200
201    pub fn ncclGroupStart() -> NcclResult;
202    pub fn ncclGroupEnd() -> NcclResult;
203
204    pub fn ncclCommGetAsyncError(comm: NcclComm, async_error: *mut NcclResult) -> NcclResult;
205
206    // 2026-09-26: Unlike `ncclCommDestroy`, also aborts operations that may
207    // still be running on the device (nccl.h).
208    pub fn ncclCommAbort(comm: NcclComm) -> NcclResult;
209}
210
211// 2026-09-26: CUDA driver calls for metrale-comm's streams and events.
212#[link(name = "cuda")]
213unsafe extern "C" {
214    fn cuStreamCreate(phStream: *mut u64, flags: u32) -> i32;
215    fn cuEventCreate(phEvent: *mut u64, flags: u32) -> i32;
216    fn cuEventRecord(hEvent: u64, hStream: u64) -> i32;
217    fn cuStreamWaitEvent(hStream: u64, hEvent: u64, flags: u32) -> i32;
218    fn cuEventDestroy_v2(hEvent: u64) -> i32;
219    fn cuStreamDestroy_v2(hStream: u64) -> i32;
220    fn cuStreamSynchronize(hStream: u64) -> i32;
221    fn cuStreamQuery(hStream: u64) -> i32;
222}
223
224pub fn create_stream() -> anyhow::Result<u64> {
225    let mut stream: u64 = 0;
226    let status = unsafe { cuStreamCreate(&mut stream, 1) }; // 2026-09-26: CU_STREAM_NON_BLOCKING
227    if status != 0 {
228        anyhow::bail!("cuStreamCreate failed: status {status}");
229    }
230    Ok(stream)
231}
232
233pub fn create_event() -> anyhow::Result<u64> {
234    let mut event: u64 = 0;
235    let status = unsafe { cuEventCreate(&mut event, 0x02) }; // 2026-09-26: CU_EVENT_DISABLE_TIMING
236    if status != 0 {
237        anyhow::bail!("cuEventCreate failed: status {status}");
238    }
239    Ok(event)
240}
241
242pub fn record_event(event: u64, stream: u64) -> anyhow::Result<()> {
243    let status = unsafe { cuEventRecord(event, stream) };
244    if status != 0 {
245        anyhow::bail!("cuEventRecord failed: status {status}");
246    }
247    Ok(())
248}
249
250pub fn stream_wait_event(stream: u64, event: u64) -> anyhow::Result<()> {
251    let status = unsafe { cuStreamWaitEvent(stream, event, 0) };
252    if status != 0 {
253        anyhow::bail!("cuStreamWaitEvent failed: status {status}");
254    }
255    Ok(())
256}
257
258pub fn destroy_event(event: u64) {
259    if event != 0 {
260        unsafe { cuEventDestroy_v2(event) };
261    }
262}
263
264pub fn destroy_stream(stream: u64) {
265    if stream != 0 {
266        unsafe { cuStreamDestroy_v2(stream) };
267    }
268}
269
270pub fn sync_stream(stream: u64) -> anyhow::Result<()> {
271    let status = unsafe { cuStreamSynchronize(stream) };
272    if status != 0 {
273        anyhow::bail!("cuStreamSynchronize failed: status {status}");
274    }
275    Ok(())
276}
277
278/// 2026-09-26: `ncclMemAlloc`: `size` bytes of device memory from NCCL's
279/// allocator; nccl.h notes the allocation may be larger than requested.
280///
281/// # Safety
282/// The returned pointer must be freed with [`nccl_mem_free`].
283pub unsafe fn nccl_mem_alloc(size: usize) -> anyhow::Result<*mut c_void> {
284    let mut ptr: *mut c_void = std::ptr::null_mut();
285    let result = unsafe { ncclMemAlloc(&mut ptr, size) };
286    check_nccl(result, "ncclMemAlloc")?;
287    Ok(ptr)
288}
289
290/// 2026-09-26: `ncclMemFree` for a pointer from [`nccl_mem_alloc`].
291///
292/// # Safety
293/// `ptr` must have been returned by [`nccl_mem_alloc`] and not yet freed.
294pub unsafe fn nccl_mem_free(ptr: *mut c_void) -> anyhow::Result<()> {
295    let result = unsafe { ncclMemFree(ptr) };
296    check_nccl(result, "ncclMemFree")
297}
298
299/// 2026-09-26: `Ok` for `Success`; otherwise an error naming `context`, the
300/// text of `ncclGetErrorString` and the code.
301pub fn check_nccl(result: NcclResult, context: &str) -> anyhow::Result<()> {
302    if result == NcclResult::Success {
303        Ok(())
304    } else {
305        let msg = unsafe {
306            let ptr = ncclGetErrorString(result);
307            if ptr.is_null() {
308                format!("{result:?}")
309            } else {
310                std::ffi::CStr::from_ptr(ptr).to_string_lossy().into_owned()
311            }
312        };
313        anyhow::bail!("NCCL error in {context}: {msg} ({result:?})")
314    }
315}
316
317/// 2026-09-26: `cuStreamQuery`, without waiting: `true` when the stream's work
318/// is done, `false` on `CUDA_ERROR_NOT_READY` (600), an error otherwise.
319pub fn stream_ready(stream: u64) -> anyhow::Result<bool> {
320    match unsafe { cuStreamQuery(stream) } {
321        0 => Ok(true),
322        600 => Ok(false),
323        status => anyhow::bail!("cuStreamQuery failed: status {status}"),
324    }
325}
metrale_gpu_sys/nccl.rs

metrale_gpu_sys/
nccl.rs
