diff --git a/openinfer-comm/crates/openinfer-comm-cuda-lib/Cargo.toml b/openinfer-comm/crates/openinfer-comm-cuda-lib/Cargo.toml index 0a5201f1..e52323b4 100644 --- a/openinfer-comm/crates/openinfer-comm-cuda-lib/Cargo.toml +++ b/openinfer-comm/crates/openinfer-comm-cuda-lib/Cargo.toml @@ -23,6 +23,7 @@ cudart-sys = { path = "../openinfer-comm-cudart-sys", package = "openinfer-comm- gdrapi-sys = { path = "../openinfer-comm-gdrapi-sys", package = "openinfer-comm-gdrapi-sys" } libc = { workspace = true } +log = { workspace = true } thiserror = { workspace = true } bincode = { workspace = true, features = ["derive", "alloc"] } diff --git a/openinfer-comm/crates/openinfer-comm-cuda-lib/src/gdr.rs b/openinfer-comm/crates/openinfer-comm-cuda-lib/src/gdr.rs index d89811bd..3fd49810 100644 --- a/openinfer-comm/crates/openinfer-comm-cuda-lib/src/gdr.rs +++ b/openinfer-comm/crates/openinfer-comm-cuda-lib/src/gdr.rs @@ -25,9 +25,28 @@ impl Drop for GdrContextHandle { } } -/// Public wrapper around the GDRCopy context. +/// Backing for the small CPU-mapped control buffers the a2a proxy writes and the +/// a2a CUDA kernel spins on. Two modes with identical semantics: +/// +/// - `Gdr`: the production path. GDRCopy pins GPU memory and maps its BAR into +/// CPU address space, so the proxy pokes GPU-resident flags directly. +/// - `HostPinned`: a fallback for hosts **without the `gdrdrv` kernel module** +/// (`gdr_open` returns null when `/dev/gdrdrv` is absent and the caller has no +/// `CAP_SYS_MODULE` to `insmod` it). It mirrors GDRCopy the other way round — +/// `cuMemHostAlloc(DEVICEMAP)` + `cuMemHostGetDevicePointer` map host-pinned +/// memory into the GPU's address space, so the kernel spins on it over PCIe +/// (higher small-message latency) while the proxy reads/writes it locally. The +/// RDMA data plane is untouched (it never used GDRCopy — it registers device +/// buffers via `ibv_reg_mr`/dma-buf), so this only degrades control-flag +/// signalling latency and lets P/D benchmarks run without root kernel access. +enum GdrMode { + Gdr(Arc), + HostPinned, +} + +/// Public wrapper around the GDRCopy context (or its host-pinned fallback). pub struct GdrCopyContext { - context: Arc, + mode: GdrMode, } fn align_to(ptr: u64, alignment: usize) -> u64 { @@ -38,12 +57,34 @@ impl GdrCopyContext { pub fn new() -> GdrResult { let handle = unsafe { gdrapi_sys::gdr_open() }; if handle.is_null() { - return Err(CudaError::GdrCopyError("Failed to create GDR copy handle")); + // gdrdrv kernel module not loaded — fall back to host-pinned mapped + // memory rather than failing. Only control-flag latency degrades. + log::warn!( + "gdrdrv unavailable (gdr_open returned null); using host-pinned mapped \ + memory for a2a control flags — higher small-message latency, identical \ + semantics, RDMA data path unaffected" + ); + return Ok(GdrCopyContext { mode: GdrMode::HostPinned }); } - Ok(GdrCopyContext { context: Arc::new(GdrContextHandle { handle }) }) + Ok(GdrCopyContext { mode: GdrMode::Gdr(Arc::new(GdrContextHandle { handle })) }) + } + + /// True when running on the host-pinned fallback (no gdrdrv). + pub fn is_host_pinned_fallback(&self) -> bool { + matches!(self.mode, GdrMode::HostPinned) } fn alloc_buffer(&self, nbytes: usize) -> GdrResult { + match &self.mode { + GdrMode::Gdr(context) => Self::alloc_gdr(context, nbytes), + GdrMode::HostPinned => Self::alloc_host_pinned(nbytes), + } + } + + fn alloc_gdr( + context: &Arc, + nbytes: usize, + ) -> GdrResult { let mut device_ptr: u64 = 0; let page_size: usize = 1 << 16; // 64KB page size let bytesize = nbytes.div_ceil(page_size) * page_size; @@ -55,9 +96,7 @@ impl GdrCopyContext { } let aligned_device_ptr = align_to(device_ptr, page_size); - - let context = self.context.clone(); - + let context = context.clone(); let g = context.handle; let mut mh = gdrapi_sys::gdr_mh_t { h: 0 }; @@ -80,24 +119,78 @@ impl GdrCopyContext { } Ok(GdrBuffer { - device_ptr, - aligned_device_ptr, - mapped_ptr, - bytesize, - mh, - context, + inner: GdrBufferInner::Gdr { + device_ptr, + aligned_device_ptr, + mapped_ptr, + bytesize, + mh, + context, + }, }) } + + fn alloc_host_pinned(nbytes: usize) -> GdrResult { + let page_size: usize = 1 << 16; + let bytesize = nbytes.div_ceil(page_size) * page_size; + + // DEVICEMAP: the allocation also gets a device-accessible address so the + // a2a kernel can poll it over PCIe. WRITECOMBINED is deliberately NOT set + // — the CPU both reads (flag waits) and writes these bytes. + let mut host_ptr: *mut c_void = null_mut(); + if unsafe { + cuda_sys::cuMemHostAlloc( + &mut host_ptr, + bytesize, + cuda_sys::CU_MEMHOSTALLOC_DEVICEMAP, + ) + } != cuda_sys::CUDA_SUCCESS + { + return Err(CudaError::GdrCopyError( + "Failed to allocate host-pinned fallback buffer", + )); + } + + let mut device_ptr: u64 = 0; + if unsafe { + cuda_sys::cuMemHostGetDevicePointer_v2(&mut device_ptr, host_ptr, 0) + } != cuda_sys::CUDA_SUCCESS + { + unsafe { cuda_sys::cuMemFreeHost(host_ptr) }; + return Err(CudaError::GdrCopyError( + "Failed to map host-pinned fallback buffer to the device", + )); + } + + // Flags start unset; the CPU allocation is not zeroed by the driver. + unsafe { std::ptr::write_bytes(host_ptr as *mut u8, 0, bytesize) }; + + Ok(GdrBuffer { + inner: GdrBufferInner::HostPinned { host_ptr, device_ptr, bytesize }, + }) + } +} + +/// Raw control buffer visible to both the GPU (`device_ptr`) and the CPU proxy. +enum GdrBufferInner { + Gdr { + device_ptr: u64, + aligned_device_ptr: u64, + mapped_ptr: *mut c_void, + mh: gdrapi_sys::gdr_mh_t, + bytesize: usize, + context: Arc, + }, + HostPinned { + host_ptr: *mut c_void, + device_ptr: u64, + #[allow(dead_code)] + bytesize: usize, + }, } -/// Raw buffer allocated on the CPU, copied using GDRCopy. struct GdrBuffer { - device_ptr: u64, - aligned_device_ptr: u64, - mapped_ptr: *mut c_void, - mh: gdrapi_sys::gdr_mh_t, - bytesize: usize, - context: Arc, + inner: GdrBufferInner, } unsafe impl Send for GdrBuffer {} @@ -105,62 +198,100 @@ unsafe impl Sync for GdrBuffer {} impl Drop for GdrBuffer { fn drop(&mut self) { - let g = self.context.handle; - unsafe { - gdrapi_sys::gdr_unmap(g, self.mh, self.mapped_ptr, self.bytesize); - gdrapi_sys::gdr_unpin_buffer(g, self.mh); - cuda_sys::cuMemFree(self.device_ptr); - }; + match &self.inner { + GdrBufferInner::Gdr { + device_ptr, + mapped_ptr, + mh, + bytesize, + context, + .. + } => { + let g = context.handle; + unsafe { + gdrapi_sys::gdr_unmap(g, *mh, *mapped_ptr, *bytesize); + gdrapi_sys::gdr_unpin_buffer(g, *mh); + cuda_sys::cuMemFree(*device_ptr); + }; + } + GdrBufferInner::HostPinned { host_ptr, .. } => { + unsafe { cuda_sys::cuMemFreeHost(*host_ptr) }; + } + } } } trait GdrRead { - fn read(mapped_ptr: *mut c_void) -> Self; + fn read(cpu_ptr: *mut c_void) -> Self; } impl GdrRead for u8 { #[inline(always)] - fn read(mapped_ptr: *mut c_void) -> Self { - let flag = unsafe { AtomicU8::from_ptr(mapped_ptr as *mut u8) }; + fn read(cpu_ptr: *mut c_void) -> Self { + let flag = unsafe { AtomicU8::from_ptr(cpu_ptr as *mut u8) }; flag.load(Ordering::Acquire) } } trait GdrWrite { - fn write(mapped_ptr: *mut c_void, value: Self); + fn write(cpu_ptr: *mut c_void, value: Self); } impl GdrWrite for u8 { #[inline(always)] - fn write(mapped_ptr: *mut c_void, value: Self) { - let flag = unsafe { AtomicU8::from_ptr(mapped_ptr as *mut u8) }; + fn write(cpu_ptr: *mut c_void, value: Self) { + let flag = unsafe { AtomicU8::from_ptr(cpu_ptr as *mut u8) }; flag.store(value, Ordering::Release); } } impl GdrBuffer { + /// The GPU-side address the a2a kernel reads/writes. fn get_device_ptr(&self) -> *mut c_void { - self.aligned_device_ptr as *mut c_void + match &self.inner { + GdrBufferInner::Gdr { aligned_device_ptr, .. } => { + *aligned_device_ptr as *mut c_void + } + GdrBufferInner::HostPinned { device_ptr, .. } => *device_ptr as *mut c_void, + } + } + + /// The CPU-side view of the same bytes the proxy thread pokes. + fn cpu_ptr(&self) -> *mut c_void { + match &self.inner { + GdrBufferInner::Gdr { mapped_ptr, .. } => *mapped_ptr, + GdrBufferInner::HostPinned { host_ptr, .. } => *host_ptr, + } } #[inline(always)] fn read(&self) -> T { - T::read(self.mapped_ptr) + T::read(self.cpu_ptr()) } #[inline(always)] fn write(&self, value: T) { - T::write(self.mapped_ptr, value); + T::write(self.cpu_ptr(), value); } fn copy_to(&self, src: *const c_void, nbytes: usize) { - unsafe { - gdrapi_sys::gdr_copy_to_mapping(self.mh, self.mapped_ptr, src, nbytes); + match &self.inner { + GdrBufferInner::Gdr { mh, mapped_ptr, .. } => unsafe { + gdrapi_sys::gdr_copy_to_mapping(*mh, *mapped_ptr, src, nbytes); + }, + GdrBufferInner::HostPinned { host_ptr, .. } => unsafe { + // Host-pinned memory is plain CPU memory — a direct copy suffices. + std::ptr::copy_nonoverlapping( + src as *const u8, + *host_ptr as *mut u8, + nbytes, + ); + }, } } } -/// Byte-flag implemented using GDRCopy. +/// Byte-flag implemented over a GDR (or host-pinned) control buffer. pub struct GdrFlag { buffer: GdrBuffer, }