Gahow Wang
28801fbfe5
Every tape op allocates its output via Tensor::zeros -> GpuBuffer::alloc -> cudaMalloc, a synchronous process-serialized driver call. Under the single- process thread-per-GPU DDP model the rank threads' hundreds of per-step allocs serialize through the driver (KI-5 root cause); it costs single-GPU too. Add a per-device, size-classed caching pool: GpuBuffer::alloc serves from a free-list (request rounded up to a size class so repeating training shapes reuse buffers), only cudaMalloc on a miss; Drop returns the buffer to the pool instead of cudaFree. Thread-safe via a global registry keyed by device id with each device's free-list behind its own Mutex (registry lock held only to clone out the per-device Arc<Mutex<_>>, so rank threads don't contend across devices). The buffer records its alloc-time device so Drop returns to the right pool. Transparent: physical capacity may be rounded up, but len()/memset/copy bounds all use the requested length, so the rounded tail is never read and numerics are unchanged. zeros() still memsets (reused buffers hold stale bytes). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
86 lines
2.9 KiB
Rust
86 lines
2.9 KiB
Rust
use crate::error::{self, Result};
|
|
use crate::ffi;
|
|
use crate::pool;
|
|
|
|
/// RAII wrapper around a GPU memory allocation. Dropping returns the buffer to
|
|
/// the per-device caching pool (see [`crate::pool`]) for reuse instead of
|
|
/// calling `cudaFree`.
|
|
///
|
|
/// `len` is the logical (requested) length used for all copy/memset bounds and
|
|
/// exposed via [`GpuBuffer::len`]; `cap` is the physical size class the pool
|
|
/// rounded up to (>= `len`), used only to bucket the buffer for reuse. The
|
|
/// extra `cap - len` bytes are never exposed to callers, so pooling is
|
|
/// numerically transparent. `device` records which device pool to return to.
|
|
pub struct GpuBuffer {
|
|
ptr: *mut u8,
|
|
len: usize,
|
|
cap: usize,
|
|
device: i32,
|
|
}
|
|
|
|
impl GpuBuffer {
|
|
/// Allocate at least `len` bytes on the calling thread's current device,
|
|
/// reusing a pooled buffer when one of the matching size class is free.
|
|
/// The contents are **uninitialized** (a reused buffer holds stale bytes);
|
|
/// callers that need zeros must memset (see [`crate::Storage::zeros`]).
|
|
pub fn alloc(len: usize) -> Result<Self> {
|
|
assert!(len > 0, "cannot allocate 0 bytes on GPU");
|
|
let a = pool::acquire(len)?;
|
|
Ok(Self {
|
|
ptr: a.ptr,
|
|
len,
|
|
cap: a.cap,
|
|
device: a.device,
|
|
})
|
|
}
|
|
|
|
pub fn len(&self) -> usize {
|
|
self.len
|
|
}
|
|
|
|
pub fn is_empty(&self) -> bool {
|
|
self.len == 0
|
|
}
|
|
|
|
pub fn as_ptr(&self) -> *const u8 {
|
|
self.ptr
|
|
}
|
|
|
|
pub fn as_mut_ptr(&mut self) -> *mut u8 {
|
|
self.ptr
|
|
}
|
|
|
|
/// Copy data from a host (CPU) slice to this GPU buffer.
|
|
pub fn copy_from_host(&mut self, src: &[u8]) -> Result<()> {
|
|
assert!(src.len() <= self.len, "source larger than buffer");
|
|
error::check(unsafe {
|
|
ffi::cudaMemcpy(self.ptr, src.as_ptr(), src.len(), ffi::CUDA_MEMCPY_H2D)
|
|
})
|
|
}
|
|
|
|
/// Copy data from this GPU buffer to a host (CPU) slice.
|
|
pub fn copy_to_host(&self, dst: &mut [u8]) -> Result<()> {
|
|
assert!(dst.len() <= self.len, "destination larger than buffer");
|
|
error::check(unsafe {
|
|
ffi::cudaMemcpy(dst.as_mut_ptr(), self.ptr, dst.len(), ffi::CUDA_MEMCPY_D2H)
|
|
})
|
|
}
|
|
|
|
/// Set every byte of the buffer to `value` on the device (no host copy).
|
|
/// Used to zero op-output buffers without a blocking H2D memcpy of zeros.
|
|
pub fn memset(&mut self, value: u8) -> Result<()> {
|
|
error::check(unsafe { ffi::cudaMemset(self.ptr, value as i32, self.len) })
|
|
}
|
|
}
|
|
|
|
impl Drop for GpuBuffer {
|
|
fn drop(&mut self) {
|
|
// Return to the device pool for reuse (no cudaFree). The pool retains
|
|
// the raw pointer for the process lifetime; on process exit the OS
|
|
// reclaims the device context, so this is not a leak.
|
|
pool::release(self.ptr, self.device, self.cap);
|
|
}
|
|
}
|
|
|
|
unsafe impl Send for GpuBuffer {}
|