phase 0+1: project scaffold + xserv-cuda crate
- Cargo workspace with xserv-cuda crate - CUDA FFI bindings (cudart: memory, stream, device, error) - GpuBuffer RAII wrapper with H2D/D2H/D2D copy - CudaStream wrapper with RAII Drop - CachingAllocator with size-bucketed free lists - PinnedBuffer for page-locked host memory - Device info query via cudaDeviceGetAttribute - Vector-add CUDA kernel smoke test - Integration test suite (11 tests) - build.rs: cc crate compiles .cu for SM 12.0 - sync-and-build.sh for remote build on dash5 - Roadmap doc (docs/00-roadmap.md) and Phase 0+1 design doc Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
109
crates/xserv-cuda/src/allocator.rs
Normal file
109
crates/xserv-cuda/src/allocator.rs
Normal file
@@ -0,0 +1,109 @@
|
||||
use crate::error::Result;
|
||||
use crate::ffi;
|
||||
use crate::memory::GpuBuffer;
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Caching allocator that reuses freed GPU buffers instead of calling
|
||||
/// cudaMalloc/cudaFree on every allocation.
|
||||
///
|
||||
/// Freed buffers are kept in a per-size-bucket free list. On allocation,
|
||||
/// we first check the free list for a buffer of matching (rounded) size.
|
||||
pub struct CachingAllocator {
|
||||
free_lists: HashMap<usize, Vec<(*mut u8, usize)>>,
|
||||
stats: AllocStats,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Clone)]
|
||||
pub struct AllocStats {
|
||||
pub alloc_count: u64,
|
||||
pub cache_hit_count: u64,
|
||||
pub cuda_malloc_count: u64,
|
||||
pub cuda_free_count: u64,
|
||||
pub current_allocated: usize,
|
||||
pub peak_allocated: usize,
|
||||
}
|
||||
|
||||
impl CachingAllocator {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
free_lists: HashMap::new(),
|
||||
stats: AllocStats::default(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn alloc(&mut self, size: usize) -> Result<GpuBuffer> {
|
||||
let bucket = bucket_size(size);
|
||||
self.stats.alloc_count += 1;
|
||||
|
||||
if let Some(list) = self.free_lists.get_mut(&bucket) {
|
||||
if let Some((ptr, actual_len)) = list.pop() {
|
||||
self.stats.cache_hit_count += 1;
|
||||
self.stats.current_allocated += actual_len;
|
||||
if self.stats.current_allocated > self.stats.peak_allocated {
|
||||
self.stats.peak_allocated = self.stats.current_allocated;
|
||||
}
|
||||
return Ok(unsafe { GpuBuffer::from_raw(ptr, actual_len) });
|
||||
}
|
||||
}
|
||||
|
||||
self.stats.cuda_malloc_count += 1;
|
||||
let buf = GpuBuffer::alloc(bucket)?;
|
||||
self.stats.current_allocated += bucket;
|
||||
if self.stats.current_allocated > self.stats.peak_allocated {
|
||||
self.stats.peak_allocated = self.stats.current_allocated;
|
||||
}
|
||||
Ok(buf)
|
||||
}
|
||||
|
||||
/// Return a buffer to the cache instead of freeing it.
|
||||
pub fn dealloc(&mut self, buf: GpuBuffer) {
|
||||
let (ptr, len) = buf.into_raw();
|
||||
let bucket = bucket_size(len);
|
||||
self.stats.current_allocated = self.stats.current_allocated.saturating_sub(len);
|
||||
self.free_lists.entry(bucket).or_default().push((ptr, len));
|
||||
}
|
||||
|
||||
/// Actually free all cached buffers.
|
||||
pub fn trim(&mut self) {
|
||||
for (_bucket, list) in self.free_lists.drain() {
|
||||
for (ptr, _len) in list {
|
||||
unsafe { ffi::cudaFree(ptr) };
|
||||
self.stats.cuda_free_count += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn stats(&self) -> &AllocStats {
|
||||
&self.stats
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for CachingAllocator {
|
||||
fn drop(&mut self) {
|
||||
self.trim();
|
||||
}
|
||||
}
|
||||
|
||||
/// Round up to next power-of-2, minimum 512 bytes.
|
||||
fn bucket_size(size: usize) -> usize {
|
||||
let min = 512;
|
||||
if size <= min {
|
||||
return min;
|
||||
}
|
||||
size.next_power_of_two()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_bucket_size() {
|
||||
assert_eq!(bucket_size(1), 512);
|
||||
assert_eq!(bucket_size(512), 512);
|
||||
assert_eq!(bucket_size(513), 1024);
|
||||
assert_eq!(bucket_size(1024), 1024);
|
||||
assert_eq!(bucket_size(1025), 2048);
|
||||
assert_eq!(bucket_size(1 << 20), 1 << 20);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user