style: format Rust workspace

crates/xserv-distributed/src/ffi.rs

@@ -34,7 +34,12 @@ pub const NCCL_SUCCESS: i32 = 0;
 unsafe extern "C" {
     pub fn ncclGetUniqueId(uid: *mut NcclUniqueId) -> i32;
     // ncclUniqueId is passed BY VALUE (a 128-byte struct) per the NCCL ABI.
-    pub fn ncclCommInitRank(comm: *mut NcclComm, nranks: i32, commid: NcclUniqueId, rank: i32) -> i32;
+    pub fn ncclCommInitRank(
+        comm: *mut NcclComm,
+        nranks: i32,
+        commid: NcclUniqueId,
+        rank: i32,
+    ) -> i32;
     pub fn ncclCommDestroy(comm: NcclComm) -> i32;
     pub fn ncclAllReduce(
         sendbuff: *const c_void,
@@ -78,5 +83,10 @@ pub fn err_string(result: i32) -> String {
 }
 
 pub fn check(result: i32, what: &str) {
-    assert_eq!(result, NCCL_SUCCESS, "{what} failed: {}", err_string(result));
+    assert_eq!(
+        result,
+        NCCL_SUCCESS,
+        "{what} failed: {}",
+        err_string(result)
+    );
 }

crates/xserv-distributed/src/lib.rs

@@ -9,8 +9,8 @@ pub mod ffi;
 use std::ffi::c_void;
 
 use ffi::{NcclComm, NcclUniqueId};
-use xserv_cuda::device;
 use xserv_cuda::GpuBuffer;
+use xserv_cuda::device;
 
 pub use ffi::NcclUniqueId as UniqueId;
 
@@ -55,7 +55,12 @@ impl TpContext {
             "ncclCommInitRank",
         );
         ffi::check(unsafe { ffi::ncclGroupEnd() }, "ncclGroupEnd(init)");
-        Self { rank, world, device, comm }
+        Self {
+            rank,
+            world,
+            device,
+            comm,
+        }
     }
 
     /// In-place AllReduce(sum) over `count` BF16 elements in `buf`.
@@ -127,7 +132,12 @@ impl PpContext {
             "ncclCommInitRank",
         );
         ffi::check(unsafe { ffi::ncclGroupEnd() }, "ncclGroupEnd(init)");
-        Self { stage, world, device, comm }
+        Self {
+            stage,
+            world,
+            device,
+            comm,
+        }
     }
 
     /// Send `count` BF16 elements at `ptr` to `peer`, on the null stream so it is
@@ -138,7 +148,16 @@ impl PpContext {
     /// `ptr` must point to at least `count` BF16 elements of valid device memory.
     pub fn send_bf16_ptr(&self, ptr: *const c_void, count: usize, peer: usize) {
         ffi::check(
-            unsafe { ffi::ncclSend(ptr, count, ffi::NCCL_BF16, peer as i32, self.comm, launch_stream()) },
+            unsafe {
+                ffi::ncclSend(
+                    ptr,
+                    count,
+                    ffi::NCCL_BF16,
+                    peer as i32,
+                    self.comm,
+                    launch_stream(),
+                )
+            },
             "ncclSend",
         );
     }
@@ -149,7 +168,16 @@ impl PpContext {
     /// `ptr` must point to at least `count` BF16 elements of valid device memory.
     pub fn recv_bf16_ptr(&self, ptr: *mut c_void, count: usize, peer: usize) {
         ffi::check(
-            unsafe { ffi::ncclRecv(ptr, count, ffi::NCCL_BF16, peer as i32, self.comm, launch_stream()) },
+            unsafe {
+                ffi::ncclRecv(
+                    ptr,
+                    count,
+                    ffi::NCCL_BF16,
+                    peer as i32,
+                    self.comm,
+                    launch_stream(),
+                )
+            },
             "ncclRecv",
         );
     }