gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

style: format Rust workspace

Browse Source
This commit is contained in:
Gahow Wang
2026-06-18 18:11:58 +08:00
parent 013465fc06
commit 531cd3fe08
57 changed files with 4045 additions and 1204 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
crates/xserv-cuda/src/allocator.rs
View File

@@ -165,7 +165,10 @@ pub fn begin_retain() {
/// Stop quarantining and hand the quarantined blocks to the caller.
pub fn end_retain() -> RetainedBlocks {
RETAINED.with(|cell| {
let list = cell.borrow_mut().take().expect("end_retain without begin_retain");
let list = cell
.borrow_mut()
.take()
.expect("end_retain without begin_retain");
RetainedBlocks(list)
})
}

4
crates/xserv-cuda/src/device.rs
View File

@@ -48,9 +48,7 @@ pub fn device_info(device: u32) -> Result<DeviceInfo> {
// Heap-allocate oversized buffer for cudaDeviceProp (layout varies by CUDA version).
// CUDA 12.x struct is ~5-6 KB; use 32 KB to guard against future growth.
let mut prop_buf = vec![0u8; 32768];
error::check(unsafe {
ffi::cudaGetDeviceProperties(prop_buf.as_mut_ptr(), device as i32)
})?;
error::check(unsafe { ffi::cudaGetDeviceProperties(prop_buf.as_mut_ptr(), device as i32) })?;
// Name is always the first field: char[256].
let name = unsafe { CStr::from_ptr(prop_buf.as_ptr() as *const c_char) }
.to_string_lossy()

8
crates/xserv-cuda/src/ffi.rs
View File

@@ -64,11 +64,5 @@ unsafe extern "C" {
pub fn cudaGraphExecDestroy(graph_exec: CudaGraphExec) -> i32;
// --- Our test kernel ---
pub fn launch_vecadd_f32(
a: *const f32,
b: *const f32,
c: *mut f32,
n: i32,
stream: CudaStream,
);
pub fn launch_vecadd_f32(a: *const f32, b: *const f32, c: *mut f32, n: i32, stream: CudaStream);
}

17
crates/xserv-cuda/src/graph.rs
View File

@@ -54,30 +54,21 @@ impl CudaGraph {
// made by THIS thread invalidate the capture. With GLOBAL mode, TP rank
// threads capturing concurrently would poison each other's captures.
error::check(unsafe {
ffi::cudaStreamBeginCapture(
stream.as_raw(),
ffi::CUDA_STREAM_CAPTURE_MODE_THREAD_LOCAL,
)
ffi::cudaStreamBeginCapture(stream.as_raw(), ffi::CUDA_STREAM_CAPTURE_MODE_THREAD_LOCAL)
})
}
/// End capture and instantiate the executable graph.
pub fn end_capture(&mut self, stream: &CudaStream) -> Result<()> {
error::check(unsafe {
ffi::cudaStreamEndCapture(stream.as_raw(), &mut self.graph)
})?;
error::check(unsafe {
ffi::cudaGraphInstantiate(&mut self.exec, self.graph, 0)
})
error::check(unsafe { ffi::cudaStreamEndCapture(stream.as_raw(), &mut self.graph) })?;
error::check(unsafe { ffi::cudaGraphInstantiate(&mut self.exec, self.graph, 0) })
}
/// Replay the captured graph on `stream`.
/// Panics if no graph has been captured yet.
pub fn launch(&self, stream: &CudaStream) -> Result<()> {
assert!(self.is_ready(), "CudaGraph::launch called before capture");
error::check(unsafe {
ffi::cudaGraphLaunch(self.exec, stream.as_raw())
})
error::check(unsafe { ffi::cudaGraphLaunch(self.exec, stream.as_raw()) })
}
fn destroy_inner(&mut self) {

2
crates/xserv-cuda/src/lib.rs
View File

@@ -11,4 +11,4 @@ pub use device::DeviceInfo;
pub use error::{CudaError, Result};
pub use graph::CudaGraph;
pub use memory::{GpuBuffer, PinnedBuffer};
pub use stream::{current_stream_raw, push_stream, CudaStream, StreamGuard};
pub use stream::{CudaStream, StreamGuard, current_stream_raw, push_stream};

46
crates/xserv-cuda/src/memory.rs
View File

@@ -22,7 +22,12 @@ impl GpuBuffer {
assert!(len > 0, "cannot allocate 0 bytes on GPU");
let mut ptr = std::ptr::null_mut();
error::check(unsafe { ffi::cudaMalloc(&mut ptr, len) })?;
Ok(Self { ptr, len, owned: true, pooled: false })
Ok(Self {
ptr,
len,
owned: true,
pooled: false,
})
}
/// Mark this buffer as pooled (returned to caching allocator on drop)
@@ -92,9 +97,7 @@ impl GpuBuffer {
/// Copy from another GPU buffer (D2D).
pub fn copy_from_device(&mut self, src: &GpuBuffer) -> Result<()> {
let n = src.len.min(self.len);
error::check(unsafe {
ffi::cudaMemcpy(self.ptr, src.ptr, n, ffi::CUDA_MEMCPY_D2D)
})
error::check(unsafe { ffi::cudaMemcpy(self.ptr, src.ptr, n, ffi::CUDA_MEMCPY_D2D) })
}
/// Fill buffer with zeros.
@@ -103,7 +106,13 @@ impl GpuBuffer {
}
/// Copy `count` bytes from `src` buffer at `src_offset` to this buffer at `dst_offset`.
pub fn copy_from_device_at(&mut self, src: &GpuBuffer, src_offset: usize, dst_offset: usize, count: usize) -> Result<()> {
pub fn copy_from_device_at(
&mut self,
src: &GpuBuffer,
src_offset: usize,
dst_offset: usize,
count: usize,
) -> Result<()> {
assert!(src_offset + count <= src.len);
assert!(dst_offset + count <= self.len);
error::check(unsafe {
@@ -117,7 +126,14 @@ impl GpuBuffer {
}
/// Async copy `count` bytes from `src` at `src_offset` to `self` at `dst_offset` on `stream`.
pub fn copy_from_device_at_async(&mut self, src: &GpuBuffer, src_offset: usize, dst_offset: usize, count: usize, stream: &CudaStream) -> Result<()> {
pub fn copy_from_device_at_async(
&mut self,
src: &GpuBuffer,
src_offset: usize,
dst_offset: usize,
count: usize,
stream: &CudaStream,
) -> Result<()> {
assert!(src_offset + count <= src.len);
assert!(dst_offset + count <= self.len);
error::check(unsafe {
@@ -161,9 +177,7 @@ impl GpuBuffer {
/// Async zero fill on stream.
pub fn zero_async(&mut self, stream: &CudaStream) -> Result<()> {
error::check(unsafe {
ffi::cudaMemsetAsync(self.ptr, 0, self.len, stream.as_raw())
})
error::check(unsafe { ffi::cudaMemsetAsync(self.ptr, 0, self.len, stream.as_raw()) })
}
/// Consume the buffer without freeing GPU memory. Returns the raw pointer and length.
@@ -178,7 +192,12 @@ impl GpuBuffer {
/// Reconstruct a GpuBuffer from a raw pointer + length.
/// Safety: ptr must have been allocated with cudaMalloc, len must be correct.
pub unsafe fn from_raw(ptr: *mut u8, len: usize) -> Self {
Self { ptr, len, owned: true, pooled: false }
Self {
ptr,
len,
owned: true,
pooled: false,
}
}
/// Create a non-owning view of GPU memory. Dropping this buffer does NOT
@@ -189,7 +208,12 @@ impl GpuBuffer {
/// `ptr` must point to a valid GPU allocation of at least `len` bytes that
/// will remain live for the lifetime of the returned `GpuBuffer`.
pub unsafe fn borrow_raw(ptr: *mut u8, len: usize) -> Self {
Self { ptr, len, owned: false, pooled: false }
Self {
ptr,
len,
owned: false,
pooled: false,
}
}
}