xserv/crates/xserv-cuda/src/allocator.rs

use crate::error::Result;
use crate::ffi;
use crate::memory::GpuBuffer;
use std::cell::RefCell;
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();
    }
}

thread_local! {
    static ALLOCATOR: RefCell<CachingAllocator> = RefCell::new(CachingAllocator::new());
}

/// Allocate a GPU buffer through the caching allocator.
/// The returned buffer has `pooled = true` so it will be returned
/// to the pool on drop instead of calling cudaFree.
pub fn cached_alloc(size: usize) -> Result<GpuBuffer> {
    ALLOCATOR.with(|cell| {
        let mut buf = cell.borrow_mut().alloc(size)?;
        buf.set_pooled(true);
        Ok(buf)
    })
}

/// Free all cached (unused) GPU buffers back to the driver.
pub fn cached_trim() {
    ALLOCATOR.with(|cell| {
        cell.borrow_mut().trim();
    });
}

/// Return a raw GPU pointer to the caching allocator's free list.
/// Called from `GpuBuffer::Drop` for pooled buffers. Takes raw pointer
/// and size to avoid re-triggering Drop.
pub fn return_to_pool(ptr: *mut u8, len: usize) {
    ALLOCATOR.with(|cell| {
        let mut alloc = cell.borrow_mut();
        let bucket = bucket_size(len);
        alloc.stats.current_allocated = alloc.stats.current_allocated.saturating_sub(len);
        alloc.free_lists.entry(bucket).or_default().push((ptr, len));
    });
}

/// 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);
    }
}