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",
         );
     }

crates/xserv-distributed/tests/allreduce.rs

@@ -2,8 +2,8 @@
 
 use half::bf16;
 use std::thread;
-use xserv_cuda::{device, GpuBuffer};
-use xserv_distributed::{get_unique_id, TpContext};
+use xserv_cuda::{GpuBuffer, device};
+use xserv_distributed::{TpContext, get_unique_id};
 
 #[test]
 fn allreduce_two_gpu_sum() {
@@ -25,9 +25,7 @@ fn allreduce_two_gpu_sum() {
                 // Rank r fills its buffer with (r + 1).
                 let val = bf16::from_f32((rank + 1) as f32);
                 let host = vec![val; n];
-                let src = unsafe {
-                    std::slice::from_raw_parts(host.as_ptr() as *const u8, n * 2)
-                };
+                let src = unsafe { std::slice::from_raw_parts(host.as_ptr() as *const u8, n * 2) };
                 let mut buf = GpuBuffer::alloc(n * 2).unwrap();
                 buf.copy_from_host(src).unwrap();
 

crates/xserv-distributed/tests/sendrecv.rs

@@ -6,8 +6,8 @@
 use half::bf16;
 use std::ffi::c_void;
 use std::thread;
-use xserv_cuda::{device, GpuBuffer};
-use xserv_distributed::{get_unique_id, PpContext};
+use xserv_cuda::{GpuBuffer, device};
+use xserv_distributed::{PpContext, get_unique_id};
 
 #[test]
 fn pp_send_recv_two_stages() {
@@ -30,7 +30,8 @@ fn pp_send_recv_two_stages() {
                 if stage == 0 {
                     // Fill with a known pattern and send to stage 1.
                     let host: Vec<bf16> = (0..n).map(|i| bf16::from_f32((i % 97) as f32)).collect();
-                    let src = unsafe { std::slice::from_raw_parts(host.as_ptr() as *const u8, n * 2) };
+                    let src =
+                        unsafe { std::slice::from_raw_parts(host.as_ptr() as *const u8, n * 2) };
                     buf.copy_from_host(src).unwrap();
                     pp.send_bf16_ptr(buf.as_mut_ptr() as *const c_void, n, 1);
                     device::synchronize().unwrap();