use std::fmt;
use crate::error::{CudaError, CudaResult};
use crate::ffi::CUdeviceptr;
use crate::stream::Stream;
use crate::stream_ordered_model::{ModelLimits, StreamOrderId, StreamOrderModel};
pub const CU_MEMPOOL_ATTR_REUSE_FOLLOW_EVENT_DEPENDENCIES: u32 = 1;
pub const CU_MEMPOOL_ATTR_REUSE_ALLOW_OPPORTUNISTIC: u32 = 2;
pub const CU_MEMPOOL_ATTR_REUSE_ALLOW_INTERNAL_DEPENDENCIES: u32 = 3;
pub const CU_MEMPOOL_ATTR_RELEASE_THRESHOLD: u32 = 4;
pub const CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT: u32 = 5;
pub const CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH: u32 = 6;
pub const CU_MEMPOOL_ATTR_USED_MEM_CURRENT: u32 = 7;
pub const CU_MEMPOOL_ATTR_USED_MEM_HIGH: u32 = 8;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StreamOrderedAllocConfig {
pub initial_pool_size: usize,
pub max_pool_size: usize,
pub release_threshold: usize,
pub device: i32,
}
impl StreamOrderedAllocConfig {
pub fn validate(&self) -> CudaResult<()> {
if self.device < 0 {
return Err(CudaError::InvalidValue);
}
if self.max_pool_size > 0 {
if self.initial_pool_size > self.max_pool_size {
return Err(CudaError::InvalidValue);
}
if self.release_threshold > self.max_pool_size {
return Err(CudaError::InvalidValue);
}
}
Ok(())
}
pub fn default_for_device(device: i32) -> Self {
Self {
initial_pool_size: 0,
max_pool_size: 0,
release_threshold: 0,
device,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PoolAttribute {
ReuseFollowEventDependencies,
ReuseAllowOpportunistic,
ReuseAllowInternalDependencies,
ReleaseThreshold(u64),
ReservedMemCurrent,
ReservedMemHigh,
UsedMemCurrent,
UsedMemHigh,
}
impl PoolAttribute {
pub fn to_raw(self) -> u32 {
match self {
Self::ReuseFollowEventDependencies => CU_MEMPOOL_ATTR_REUSE_FOLLOW_EVENT_DEPENDENCIES,
Self::ReuseAllowOpportunistic => CU_MEMPOOL_ATTR_REUSE_ALLOW_OPPORTUNISTIC,
Self::ReuseAllowInternalDependencies => {
CU_MEMPOOL_ATTR_REUSE_ALLOW_INTERNAL_DEPENDENCIES
}
Self::ReleaseThreshold(_) => CU_MEMPOOL_ATTR_RELEASE_THRESHOLD,
Self::ReservedMemCurrent => CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT,
Self::ReservedMemHigh => CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH,
Self::UsedMemCurrent => CU_MEMPOOL_ATTR_USED_MEM_CURRENT,
Self::UsedMemHigh => CU_MEMPOOL_ATTR_USED_MEM_HIGH,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct PoolUsageStats {
pub reserved_current: u64,
pub reserved_high: u64,
pub used_current: u64,
pub used_high: u64,
pub active_allocations: usize,
pub peak_allocations: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum ShareableHandleType {
#[default]
None,
PosixFileDescriptor,
Win32Handle,
Win32KmtHandle,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PoolExportDescriptor {
pub shareable_handle_type: ShareableHandleType,
pub pool_device: i32,
}
pub struct StreamAllocation {
ptr: CUdeviceptr,
size: usize,
stream: u64,
pool: u64,
ready_seq: u64,
freed: bool,
}
impl StreamAllocation {
#[inline]
pub fn as_ptr(&self) -> u64 {
self.ptr
}
#[inline]
pub fn size(&self) -> usize {
self.size
}
#[inline]
pub fn is_freed(&self) -> bool {
self.freed
}
#[inline]
pub fn stream(&self) -> u64 {
self.stream
}
#[inline]
pub fn stream_id(&self) -> StreamOrderId {
StreamOrderId(self.stream)
}
#[inline]
pub fn ready_seq(&self) -> u64 {
self.ready_seq
}
#[inline]
pub fn pool(&self) -> u64 {
self.pool
}
}
impl fmt::Debug for StreamAllocation {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("StreamAllocation")
.field("ptr", &format_args!("0x{:016x}", self.ptr))
.field("size", &self.size)
.field("stream", &format_args!("0x{:016x}", self.stream))
.field("freed", &self.freed)
.finish()
}
}
pub struct StreamMemoryPool {
handle: u64,
owned: bool,
device: i32,
config: StreamOrderedAllocConfig,
active_allocations: usize,
total_allocated: usize,
peak_allocated: usize,
peak_allocation_count: usize,
model: StreamOrderModel,
}
impl fmt::Debug for StreamMemoryPool {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("StreamMemoryPool")
.field("handle", &format_args!("0x{:016x}", self.handle))
.field("device", &self.device)
.field("active_allocations", &self.active_allocations)
.field("total_allocated", &self.total_allocated)
.field("reserved", &self.model.reserved())
.finish()
}
}
impl StreamMemoryPool {
pub fn new(config: StreamOrderedAllocConfig) -> CudaResult<Self> {
config.validate()?;
#[cfg_attr(target_os = "macos", allow(unused_mut))]
let mut pool = Self::with_model(config);
#[cfg(not(target_os = "macos"))]
{
pool.handle = Self::gpu_create_pool(&pool.config)?;
pool.owned = true;
}
Ok(pool)
}
pub fn cpu_pool(config: StreamOrderedAllocConfig) -> CudaResult<Self> {
config.validate()?;
Ok(Self::with_model(config))
}
fn with_model(config: StreamOrderedAllocConfig) -> Self {
let model = StreamOrderModel::new(Self::model_limits(&config));
Self {
handle: 0,
owned: false,
device: config.device,
config,
active_allocations: 0,
total_allocated: 0,
peak_allocated: 0,
peak_allocation_count: 0,
model,
}
}
fn model_limits(config: &StreamOrderedAllocConfig) -> ModelLimits {
ModelLimits {
max_pool_size: config.max_pool_size,
release_threshold: config.release_threshold,
}
}
#[inline]
pub fn stream_id(stream: &Stream) -> StreamOrderId {
StreamOrderId(stream.raw().0 as usize as u64)
}
pub fn alloc_async(&mut self, size: usize, stream: u64) -> CudaResult<StreamAllocation> {
self.alloc_on(size, StreamOrderId(stream))
}
pub fn alloc_async_on_stream(
&mut self,
size: usize,
stream: &Stream,
) -> CudaResult<StreamAllocation> {
self.alloc_on(size, Self::stream_id(stream))
}
pub fn alloc_on(&mut self, size: usize, stream: StreamOrderId) -> CudaResult<StreamAllocation> {
let model_alloc = self.model.alloc(size, stream)?;
self.sync_mirror_stats();
Ok(StreamAllocation {
ptr: model_alloc.ptr,
size: model_alloc.size,
stream: stream.raw(),
pool: self.handle,
ready_seq: model_alloc.ready_seq,
freed: false,
})
}
pub fn free_async(&mut self, alloc: &mut StreamAllocation) -> CudaResult<()> {
let stream = alloc.stream_id();
self.free_on(alloc, stream)
}
pub fn free_async_on_stream(
&mut self,
alloc: &mut StreamAllocation,
stream: &Stream,
) -> CudaResult<()> {
self.free_on(alloc, Self::stream_id(stream))
}
pub fn free_on(
&mut self,
alloc: &mut StreamAllocation,
stream: StreamOrderId,
) -> CudaResult<()> {
if alloc.freed {
return Err(CudaError::InvalidValue);
}
self.model.free(alloc.ptr, stream)?;
self.sync_mirror_stats();
alloc.freed = true;
Ok(())
}
pub fn synchronize_stream(&mut self, stream: StreamOrderId) {
self.model.synchronize(stream);
self.sync_mirror_stats();
}
pub fn is_ready(&self, alloc: &StreamAllocation) -> bool {
let model_alloc = crate::stream_ordered_model::ModelAllocation {
ptr: alloc.ptr,
size: alloc.size,
capacity: alloc.size,
stream: alloc.stream_id(),
ready_seq: alloc.ready_seq,
};
self.model.is_ready_same_stream(&model_alloc)
}
pub fn is_ready_on(
&self,
alloc: &StreamAllocation,
consumer: StreamOrderId,
wait_seq: u64,
) -> bool {
let model_alloc = crate::stream_ordered_model::ModelAllocation {
ptr: alloc.ptr,
size: alloc.size,
capacity: alloc.size,
stream: alloc.stream_id(),
ready_seq: alloc.ready_seq,
};
self.model
.is_ready_cross_stream(&model_alloc, consumer, wait_seq)
}
pub fn record_event(&mut self, stream: StreamOrderId) -> u64 {
self.model.record_event(stream)
}
pub fn model_trim(&mut self, min_bytes_to_keep: usize) {
self.model.trim_to(min_bytes_to_keep);
self.sync_mirror_stats();
}
pub fn trim(&mut self, min_bytes_to_keep: usize) -> CudaResult<()> {
self.platform_trim(min_bytes_to_keep)
}
pub fn stats(&self) -> PoolUsageStats {
PoolUsageStats {
reserved_current: self.model.reserved() as u64,
reserved_high: self.model.reserved_high() as u64,
used_current: self.model.used() as u64,
used_high: self.model.used_high() as u64,
active_allocations: self.model.active(),
peak_allocations: self.model.peak_active(),
}
}
pub fn set_attribute(&mut self, attr: PoolAttribute) -> CudaResult<()> {
match attr {
PoolAttribute::ReservedMemCurrent
| PoolAttribute::UsedMemCurrent
| PoolAttribute::ReservedMemHigh
| PoolAttribute::UsedMemHigh => {
return Err(CudaError::InvalidValue);
}
_ => {}
}
if let PoolAttribute::ReleaseThreshold(val) = attr {
self.config.release_threshold = val as usize;
self.model.set_release_threshold(val as usize);
}
self.platform_set_attribute(attr)
}
pub fn enable_peer_access(&self, peer_device: i32) -> CudaResult<()> {
if peer_device == self.device {
return Err(CudaError::InvalidDevice);
}
self.platform_enable_peer_access(peer_device)
}
pub fn disable_peer_access(&self, peer_device: i32) -> CudaResult<()> {
if peer_device == self.device {
return Err(CudaError::InvalidDevice);
}
self.platform_disable_peer_access(peer_device)
}
pub fn reset_peak_stats(&mut self) {
self.model.reset_peaks();
self.sync_mirror_stats();
}
fn sync_mirror_stats(&mut self) {
self.active_allocations = self.model.active();
self.total_allocated = self.model.used();
self.peak_allocated = self.model.used_high();
self.peak_allocation_count = self.model.peak_active();
}
pub fn default_pool(device: i32) -> CudaResult<Self> {
if device < 0 {
return Err(CudaError::InvalidValue);
}
let config = StreamOrderedAllocConfig::default_for_device(device);
#[cfg(target_os = "macos")]
{
Ok(Self::with_model(config))
}
#[cfg(not(target_os = "macos"))]
{
let handle = Self::gpu_default_pool(device)?;
let mut pool = Self::with_model(config);
pool.handle = handle;
Ok(pool)
}
}
#[inline]
pub fn handle(&self) -> u64 {
self.handle
}
#[inline]
pub fn device(&self) -> i32 {
self.device
}
#[inline]
pub fn config(&self) -> &StreamOrderedAllocConfig {
&self.config
}
fn platform_trim(&mut self, min_bytes_to_keep: usize) -> CudaResult<()> {
#[cfg(target_os = "macos")]
{
let _ = min_bytes_to_keep;
Err(CudaError::NotSupported)
}
#[cfg(not(target_os = "macos"))]
{
Self::gpu_trim(self.handle, min_bytes_to_keep)
}
}
fn platform_set_attribute(&self, attr: PoolAttribute) -> CudaResult<()> {
#[cfg(target_os = "macos")]
{
match attr {
PoolAttribute::ReleaseThreshold(_) => Ok(()),
_ => Err(CudaError::NotSupported),
}
}
#[cfg(not(target_os = "macos"))]
{
Self::gpu_set_attribute(self.handle, attr)
}
}
fn platform_enable_peer_access(&self, peer_device: i32) -> CudaResult<()> {
#[cfg(target_os = "macos")]
{
let _ = peer_device;
Err(CudaError::NotSupported)
}
#[cfg(not(target_os = "macos"))]
{
Self::gpu_enable_peer_access(self.handle, peer_device)
}
}
fn platform_disable_peer_access(&self, peer_device: i32) -> CudaResult<()> {
#[cfg(target_os = "macos")]
{
let _ = peer_device;
Err(CudaError::NotSupported)
}
#[cfg(not(target_os = "macos"))]
{
Self::gpu_disable_peer_access(self.handle, peer_device)
}
}
#[cfg(not(target_os = "macos"))]
fn gpu_create_pool(config: &StreamOrderedAllocConfig) -> CudaResult<u64> {
use crate::ffi::{
CUmemAllocationType, CUmemLocation, CUmemLocationType, CUmemPoolProps, CUmemoryPool,
};
let api = crate::loader::try_driver()?;
let create = api.cu_mem_pool_create.ok_or(CudaError::NotSupported)?;
let props = CUmemPoolProps {
alloc_type: CUmemAllocationType::Pinned as u32,
handle_types: 0,
location: CUmemLocation {
loc_type: CUmemLocationType::Device as u32,
id: config.device,
},
win32_security_attributes: std::ptr::null_mut(),
max_size: config.max_pool_size,
reserved: [0u8; 56],
};
let mut pool = CUmemoryPool::default();
let rc = unsafe { create(&mut pool, &props) };
crate::error::check(rc)?;
Ok(pool.0 as usize as u64)
}
#[cfg(not(target_os = "macos"))]
fn gpu_default_pool(device: i32) -> CudaResult<u64> {
use crate::ffi::CUmemoryPool;
let api = crate::loader::try_driver()?;
let get_default = api
.cu_device_get_default_mem_pool
.ok_or(CudaError::NotSupported)?;
let mut pool = CUmemoryPool::default();
let rc = unsafe { get_default(&mut pool, device) };
crate::error::check(rc)?;
Ok(pool.0 as usize as u64)
}
#[cfg(not(target_os = "macos"))]
#[cfg_attr(not(test), allow(dead_code))]
fn gpu_alloc_async(pool_handle: u64, size: usize, stream: u64) -> CudaResult<CUdeviceptr> {
use crate::ffi::{CUmemoryPool, CUstream};
let api = crate::loader::try_driver()?;
let cu_stream = CUstream(stream as usize as *mut std::ffi::c_void);
let mut dptr: CUdeviceptr = 0;
if pool_handle != 0 {
let alloc_from_pool = api
.cu_mem_alloc_from_pool_async
.ok_or(CudaError::NotSupported)?;
let pool = CUmemoryPool(pool_handle as usize as *mut std::ffi::c_void);
let rc = unsafe { alloc_from_pool(&mut dptr, size, pool, cu_stream) };
crate::error::check(rc)?;
} else {
let alloc_async = api.cu_mem_alloc_async.ok_or(CudaError::NotSupported)?;
let rc = unsafe { alloc_async(&mut dptr, size, cu_stream) };
crate::error::check(rc)?;
}
Ok(dptr)
}
#[cfg(not(target_os = "macos"))]
#[cfg_attr(not(test), allow(dead_code))]
fn gpu_free_async(ptr: CUdeviceptr, stream: u64) -> CudaResult<()> {
use crate::ffi::CUstream;
let api = crate::loader::try_driver()?;
let free_async = api.cu_mem_free_async.ok_or(CudaError::NotSupported)?;
let cu_stream = CUstream(stream as usize as *mut std::ffi::c_void);
crate::error::check(unsafe { free_async(ptr, cu_stream) })
}
#[cfg(not(target_os = "macos"))]
fn gpu_trim(pool_handle: u64, min_bytes_to_keep: usize) -> CudaResult<()> {
use crate::ffi::CUmemoryPool;
let api = crate::loader::try_driver()?;
let trim = api.cu_mem_pool_trim_to.ok_or(CudaError::NotSupported)?;
let pool = CUmemoryPool(pool_handle as usize as *mut std::ffi::c_void);
crate::error::check(unsafe { trim(pool, min_bytes_to_keep) })
}
#[cfg(not(target_os = "macos"))]
fn gpu_set_attribute(pool_handle: u64, attr: PoolAttribute) -> CudaResult<()> {
use crate::ffi::CUmemoryPool;
let api = crate::loader::try_driver()?;
let set_attr = api
.cu_mem_pool_set_attribute
.ok_or(CudaError::NotSupported)?;
let pool = CUmemoryPool(pool_handle as usize as *mut std::ffi::c_void);
let raw_attr = Self::map_pool_attribute(attr)?;
match attr {
PoolAttribute::ReuseFollowEventDependencies
| PoolAttribute::ReuseAllowOpportunistic
| PoolAttribute::ReuseAllowInternalDependencies => {
let mut value: std::ffi::c_int = 1;
let rc = unsafe {
set_attr(pool, raw_attr, (&mut value as *mut std::ffi::c_int).cast())
};
crate::error::check(rc)
}
PoolAttribute::ReleaseThreshold(threshold) => {
let mut value: u64 = threshold;
let rc = unsafe { set_attr(pool, raw_attr, (&mut value as *mut u64).cast()) };
crate::error::check(rc)
}
PoolAttribute::ReservedMemCurrent
| PoolAttribute::ReservedMemHigh
| PoolAttribute::UsedMemCurrent
| PoolAttribute::UsedMemHigh => Err(CudaError::InvalidValue),
}
}
#[cfg(not(target_os = "macos"))]
fn map_pool_attribute(attr: PoolAttribute) -> CudaResult<crate::ffi::CUmemPoolAttribute> {
use crate::ffi::CUmemPoolAttribute;
Ok(match attr {
PoolAttribute::ReuseFollowEventDependencies => {
CUmemPoolAttribute::ReuseFollowEventDependencies
}
PoolAttribute::ReuseAllowOpportunistic => CUmemPoolAttribute::ReuseAllowOpportunistic,
PoolAttribute::ReuseAllowInternalDependencies => {
CUmemPoolAttribute::ReuseAllowInternalDependencies
}
PoolAttribute::ReleaseThreshold(_) => CUmemPoolAttribute::ReleaseThreshold,
PoolAttribute::ReservedMemCurrent => CUmemPoolAttribute::ReservedMemCurrent,
PoolAttribute::ReservedMemHigh => CUmemPoolAttribute::ReservedMemHigh,
PoolAttribute::UsedMemCurrent => CUmemPoolAttribute::UsedMemCurrent,
PoolAttribute::UsedMemHigh => CUmemPoolAttribute::UsedMemHigh,
})
}
#[cfg(not(target_os = "macos"))]
fn gpu_enable_peer_access(pool_handle: u64, peer_device: i32) -> CudaResult<()> {
Self::gpu_set_pool_access(pool_handle, peer_device, true)
}
#[cfg(not(target_os = "macos"))]
fn gpu_disable_peer_access(pool_handle: u64, peer_device: i32) -> CudaResult<()> {
Self::gpu_set_pool_access(pool_handle, peer_device, false)
}
#[cfg(not(target_os = "macos"))]
fn gpu_set_pool_access(pool_handle: u64, peer_device: i32, enable: bool) -> CudaResult<()> {
use crate::ffi::{
CUmemAccessDesc, CUmemAccessFlags, CUmemLocation, CUmemLocationType, CUmemoryPool,
};
let api = crate::loader::try_driver()?;
let set_access = api.cu_mem_pool_set_access.ok_or(CudaError::NotSupported)?;
let pool = CUmemoryPool(pool_handle as usize as *mut std::ffi::c_void);
let flags = if enable {
CUmemAccessFlags::ReadWrite
} else {
CUmemAccessFlags::None
};
let desc = CUmemAccessDesc {
location: CUmemLocation {
loc_type: CUmemLocationType::Device as u32,
id: peer_device,
},
flags: flags as u32,
};
let rc = unsafe { set_access(pool, &desc, 1) };
crate::error::check(rc)
}
}
impl Drop for StreamMemoryPool {
fn drop(&mut self) {
if !self.owned || self.handle == 0 {
return;
}
if let Ok(api) = crate::loader::try_driver() {
if let Some(destroy) = api.cu_mem_pool_destroy {
let pool = crate::ffi::CUmemoryPool(self.handle as usize as *mut std::ffi::c_void);
let rc = unsafe { destroy(pool) };
if rc != 0 {
tracing::warn!(
cuda_error = rc,
pool = format_args!("0x{:016x}", self.handle),
"cuMemPoolDestroy failed during StreamMemoryPool drop"
);
}
}
}
}
}
static DEFAULT_POOL: std::sync::OnceLock<std::sync::Mutex<StreamMemoryPool>> =
std::sync::OnceLock::new();
fn shared_default_pool() -> CudaResult<&'static std::sync::Mutex<StreamMemoryPool>> {
if let Some(pool) = DEFAULT_POOL.get() {
return Ok(pool);
}
let pool = StreamMemoryPool::default_pool(0)?;
Ok(DEFAULT_POOL.get_or_init(|| std::sync::Mutex::new(pool)))
}
pub fn stream_alloc(size: usize, stream: u64) -> CudaResult<StreamAllocation> {
let pool = shared_default_pool()?;
let mut guard = pool.lock().unwrap_or_else(|e| e.into_inner());
guard.alloc_async(size, stream)
}
pub fn stream_free(alloc: &mut StreamAllocation) -> CudaResult<()> {
if alloc.freed {
return Err(CudaError::InvalidValue);
}
if let Some(pool) = DEFAULT_POOL.get() {
let mut guard = pool.lock().unwrap_or_else(|e| e.into_inner());
if guard.handle() == alloc.pool {
let stream = alloc.stream_id();
return guard.free_on(alloc, stream);
}
}
alloc.freed = true;
Ok(())
}
#[cfg(test)]
#[path = "stream_ordered_alloc_tests.rs"]
mod tests;