Unverified Commit 70822f35 authored by Yan Ru Pei's avatar Yan Ru Pei Committed by GitHub
Browse files

fix: gate dynamo-memory numa module behind cfg(target_os = linux) (#6354)


Signed-off-by: default avatarPeaBrane <yanrpei@gmail.com>
Co-authored-by: default avatarCursor <cursoragent@cursor.com>
parent d38954c7
......@@ -14,6 +14,7 @@
pub mod actions;
pub mod arena;
pub mod nixl;
#[cfg(target_os = "linux")]
pub mod numa;
/// Offset-based buffer views into underlying storage.
......@@ -41,6 +42,7 @@ pub use device::DeviceStorage;
#[cfg(target_os = "linux")]
pub use disk::DiskStorage;
pub use external::ExternalDeviceMemory;
#[cfg(target_os = "linux")]
pub use numa::{NumaNode, is_numa_enabled};
pub use offset::OffsetBuffer;
pub use pinned::PinnedStorage;
......
......@@ -68,8 +68,6 @@ impl PinnedStorage {
/// - CUDA context creation fails
/// - Memory allocation fails
pub fn new_for_device(len: usize, device_id: Option<u32>) -> Result<Self> {
use super::numa;
if len == 0 {
return Err(StorageError::AllocationFailed(
"zero-sized allocations are not supported".into(),
......@@ -80,34 +78,30 @@ impl PinnedStorage {
let ctx = cuda_context(gpu_id)?;
let ptr = match device_id {
Some(gpu_id) if numa::is_numa_enabled() => {
// NUMA-aware allocation via worker pool
#[cfg(target_os = "linux")]
Some(gpu_id) if super::numa::is_numa_enabled() => {
tracing::debug!(
"Using NUMA-aware allocation for {} bytes on GPU {}",
len,
gpu_id
);
numa::worker_pool::NumaWorkerPool::global()
super::numa::worker_pool::NumaWorkerPool::global()
.allocate_pinned_for_gpu(len, gpu_id)
.map_err(StorageError::AllocationFailed)? as usize
}
_ => {
// Direct allocation (no NUMA or device_id not specified)
unsafe {
ctx.bind_to_thread().map_err(StorageError::Cuda)?;
let ptr =
cudarc::driver::result::malloc_host(len, sys::CU_MEMHOSTALLOC_DEVICEMAP)
.map_err(StorageError::Cuda)?;
let ptr = ptr as *mut u8;
assert!(!ptr.is_null(), "Failed to allocate pinned memory");
assert!(ptr.is_aligned(), "Pinned memory is not aligned");
assert!(len < isize::MAX as usize);
ptr as usize
}
}
_ => unsafe {
ctx.bind_to_thread().map_err(StorageError::Cuda)?;
let ptr = cudarc::driver::result::malloc_host(len, sys::CU_MEMHOSTALLOC_DEVICEMAP)
.map_err(StorageError::Cuda)?;
let ptr = ptr as *mut u8;
assert!(!ptr.is_null(), "Failed to allocate pinned memory");
assert!(ptr.is_aligned(), "Pinned memory is not aligned");
assert!(len < isize::MAX as usize);
ptr as usize
},
};
Ok(Self { ptr, len, ctx })
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment