From cd9ea471535b6bb1d6224ea47a35a926ad3d7182 Mon Sep 17 00:00:00 2001 From: Alexander Droste Date: Fri, 23 Oct 2026 12:18:15 +0100 Subject: [PATCH] fix(cuda): retain pinned sources across transfer failures Validate the transfer context and wait for the previous source fence before replacing it. On submission or fence failure, synchronize the stream before recycling the source, retaining the allocation if completion cannot be established. Include ownership regression tests and bind the context explicitly before pinned allocation. Original-commit: 94897853df Original-commit: a2f6ce3b37 Signed-off-by: Alexander Droste --- vortex-cuda/src/pinned.rs | 336 +++++++++++++++++++++++++++++++++++--- 1 file changed, 314 insertions(+), 22 deletions(-) diff --git a/vortex-cuda/src/pinned.rs b/vortex-cuda/src/pinned.rs index 40371a042af..7fbdf9b4725 100644 --- a/vortex-cuda/src/pinned.rs +++ b/vortex-cuda/src/pinned.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::mem::ManuallyDrop; use std::ops::Range; use std::sync::Arc; use std::sync::atomic::AtomicU64; @@ -62,6 +63,9 @@ impl PinnedByteBuffer { let event = ctx .new_event(Some(CUevent_flags::CU_EVENT_BLOCKING_SYNC)) .map_err(|e| vortex_err!("failed to create pinned host buffer event: {e}"))?; + // Bind explicitly rather than relying on event creation to bind this worker thread. + ctx.bind_to_thread() + .map_err(|e| vortex_err!("failed to bind pinned host buffer context: {e}"))?; // Keep file-I/O staging cacheable: cudarc's allocator uses write-combined memory, // which makes CPU reads expensive. let ptr = unsafe { result::malloc_host(capacity, 0) } @@ -90,7 +94,9 @@ impl PinnedByteBuffer { } } -/// A borrowed range that preserves the pinned buffer's stream synchronization. +/// A borrowed range that records completion on the pinned buffer's event. +/// +/// Construct only after successfully waiting on the old event, and use only with that stream. struct PinnedByteBufferView<'a> { buffer: &'a mut PinnedByteBuffer, range: Range, @@ -105,32 +111,39 @@ impl HostSlice for PinnedByteBufferView<'_> { &'a self, stream: &'a CudaStream, ) -> (&'a [u8], SyncOnDrop<'a>) { - stream.context().record_err(stream.wait(&self.buffer.event)); - // SAFETY: The view retains the allocation, and the caller must use the returned slice - // only with this stream. The guard records completion on the owner's event. - let slice = unsafe { std::slice::from_raw_parts(self.buffer.ptr, self.buffer.logical_len) }; + // SAFETY: The range is initialized and bounds-checked before constructing the view. + // The old fence was successfully waited on before creating any recording guard. + // The caller uses the slice only with that stream, and the owner retains the allocation. + let slice = unsafe { + std::slice::from_raw_parts(self.buffer.ptr.add(self.range.start), self.range.len()) + }; let sync = SyncOnDrop::Record(Some((&self.buffer.event, stream))); - (&slice[self.range.clone()], sync) + (slice, sync) } unsafe fn stream_synced_mut_slice<'a>( &'a mut self, stream: &'a CudaStream, ) -> (&'a mut [u8], SyncOnDrop<'a>) { - stream.context().record_err(stream.wait(&self.buffer.event)); - // SAFETY: The view exclusively borrows the allocation, and the caller must use the - // returned slice only with this stream. The guard tracks writes on the owner's event. - let slice = - unsafe { std::slice::from_raw_parts_mut(self.buffer.ptr, self.buffer.logical_len) }; + // SAFETY: The view exclusively borrows the allocation and covers only the initialized, + // validated range. The old fence was waited on before constructing the view, and the + // caller uses the slice only with that stream. The guard tracks writes on the same event. + let slice = unsafe { + std::slice::from_raw_parts_mut(self.buffer.ptr.add(self.range.start), self.range.len()) + }; let sync = SyncOnDrop::Record(Some((&self.buffer.event, stream))); - (&mut slice[self.range.clone()], sync) + (slice, sync) } } impl Drop for PinnedByteBuffer { fn drop(&mut self) { let context = self.event.context(); - context.record_err(self.event.synchronize()); + if let Err(error) = self.event.synchronize() { + // An old or failed fence cannot justify freeing a possible DMA source. + context.record_err::<()>(Err(error)); + return; + } context.record_err(unsafe { result::free_host(self.ptr.cast()) }); } } @@ -351,6 +364,10 @@ impl PooledPinnedBuffer { range: Range, destination: &mut CudaViewMut<'_, u8>, ) -> VortexResult<()> { + vortex_ensure!( + &self.pool.ctx == stream.context(), + "pinned buffer and transfer stream must use the same CUDA context" + ); let pinned = self.inner.as_mut().vortex_expect("buffer already consumed"); vortex_ensure!( range.start <= range.end && range.end <= pinned.logical_len, @@ -364,21 +381,47 @@ impl PooledPinnedBuffer { "pinned host buffer range length does not match destination length" ); + // Chain the old transfer before this stream replaces the completion fence, even for + // buffers returned directly to a bucket. Check the wait before creating a source guard: + // on failure, preserve the old fence rather than recording over it on this stream. + stream + .wait(&pinned.event) + .map_err(|error| vortex_err!("Failed to wait for pinned source transfer: {error}"))?; let source = PinnedByteBufferView { buffer: pinned, range, }; - // The page-locked source and its completion event keep this asynchronous. - stream + // Keep the source guard's record as well as the explicit completion record. The latter + // also fences the source when cudarc's automatic event tracking is disabled. + let completion = stream .memcpy_htod(&source, destination) - .map_err(|e| vortex_err!("Failed to schedule H2D copy: {}", e))?; + .and_then(|()| source.buffer.event.record(stream)); + self.retain_after_transfer(stream, completion) + } - let event = stream - .record_event(None) - .map_err(|e| vortex_err!("Failed to record CUDA event: {}", e))?; + fn retain_after_transfer( + mut self, + stream: &VortexCudaStream, + completion: Result<(), result::DriverError>, + ) -> VortexResult<()> { + let event = match completion.and_then(|()| stream.record_event(None)) { + Ok(event) => event, + Err(error) => { + // A failed record can leave an old event behind after DMA was already submitted. + // Fence the actual stream before recycling, independently of event tracking. + if let Err(sync_error) = stream.synchronize() { + // Completion is unknown: leaking is safer than freeing an active DMA source. + let _leaked = ManuallyDrop::new( + self.inner.take().vortex_expect("buffer already consumed"), + ); + return Err(vortex_err!( + "H2D transfer fence failed: {error}; stream synchronization failed: {sync_error}" + )); + } + return Err(vortex_err!("H2D transfer fence failed: {error}")); + } + }; - // On earlier errors, Drop returns the buffer to the pool, but the HostSlice event still - // gates access and freeing. On success, the inflight queue retains it until completion. let inner = self.inner.take().vortex_expect("buffer already consumed"); self.pool.inflight.lock().push(InflightPinnedBuffer { event, @@ -399,8 +442,10 @@ impl Drop for PooledPinnedBuffer { #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::mpsc::Sender; use cudarc::driver::CudaContext; + use cudarc::driver::sys::cudaError_enum; use rstest::rstest; use vortex::array::buffer::DeviceBuffer; use vortex::buffer::Alignment; @@ -494,8 +539,16 @@ mod tests { #[rstest] #[crate::test] - fn transfer_retains_and_reuses_source(#[values(false, true)] ranged: bool) -> VortexResult<()> { + fn transfer_retains_and_reuses_source( + #[values(false, true)] ranged: bool, + #[values(false, true)] automatic_tracking: bool, + ) -> VortexResult<()> { let (pool, stream) = setup()?; + if !automatic_tracking { + // SAFETY: Device allocations are used and freed on their allocation stream; + // the pool must independently fence host reuse even without cudarc tracking. + unsafe { stream.context().disable_event_tracking() }; + } let mut pinned = pool.get(1024)?; pinned.as_mut_slice().fill(0xAB); let allocation = pinned.as_mut_slice().as_ptr(); @@ -530,6 +583,245 @@ mod tests { assert_eq!(host.as_ref(), vec![0xAB; if ranged { 256 } else { 1024 }]); pool.reclaim_completed()?; assert_eq!(pool.stats().puts, 1); + + let other_stream = VortexCudaStream( + stream + .context() + .new_stream() + .map_err(|e| vortex_err!("Failed to create stream: {e}"))?, + ); + let second = reused.transfer_to_device(&other_stream)?; + let host = second.copy_to_host_sync(Alignment::of::())?; + assert_eq!(host.as_ref(), vec![0xCD; 1024]); + pool.reclaim_completed()?; + assert_eq!(pool.stats().puts, 2); + Ok(()) + } + + #[crate::test] + fn host_access_waits_when_transfer_returns_directly_to_bucket() -> VortexResult<()> { + let (pool, stream) = setup()?; + let mut pinned = pool.get(1024 * 1024)?; + pinned.as_mut_slice().fill(0x5A); + let mut destination = stream.device_alloc::(1024 * 1024)?; + submit_to_bucket(pinned, &stream, &mut destination)?; + let mut reused = pool.get(1024 * 1024)?; + reused.as_mut_slice().fill(0xA5); + let host = CudaDeviceBuffer::new(destination).copy_to_host_sync(Alignment::of::())?; + assert_eq!(host.as_ref(), vec![0x5A; 1024 * 1024]); + Ok(()) + } + + fn setup_resubmission( + automatic_tracking: bool, + ) -> VortexResult<( + Arc, + VortexCudaStream, + VortexCudaStream, + )> { + // Other tests synchronize the primary context. Keep the gated DMA independent of them. + let ctx = CudaContext::new_non_primary(0, 0) + .map_err(|e| vortex_err!("Failed to initialize CUDA: {e}"))?; + let pool = Arc::new(PinnedByteBufferPool::new(Arc::clone(&ctx))); + let stream = VortexCudaStream( + ctx.new_stream() + .map_err(|e| vortex_err!("Failed to create stream: {e}"))?, + ); + if !automatic_tracking { + // SAFETY: Each destination stays on its allocation stream. The pool fences sources. + unsafe { stream.context().disable_event_tracking() }; + } + let other = VortexCudaStream( + pool.ctx + .new_stream() + .map_err(|error| vortex_err!("{error}"))?, + ); + Ok((pool, stream, other)) + } + + fn submit_to_bucket( + mut pinned: PooledPinnedBuffer, + stream: &VortexCudaStream, + destination: &mut cudarc::driver::CudaSlice, + ) -> VortexResult<()> { + let owner = pinned.inner.as_mut().unwrap(); + // SAFETY: Tests initialize the source. On failure retain_after_transfer synchronizes + // the stream or leaks the allocation; on success the bucket retains the recorded fence. + let source = unsafe { std::slice::from_raw_parts(owner.ptr, owner.logical_len) }; + let completion = stream + .memcpy_htod(source, destination) + .and_then(|()| owner.event.record(stream)); + if completion.is_err() { + return pinned.retain_after_transfer(stream, completion); + } + // Exercise direct bucket return rather than normal inflight reclamation. + drop(pinned); + Ok(()) + } + + fn block_stream(stream: &VortexCudaStream) -> VortexResult> { + use std::sync::mpsc::Receiver; + use std::sync::mpsc::channel; + let (release, receiver) = channel::<()>(); + stream + .context() + .bind_to_thread() + .map_err(|e| vortex_err!("{e}"))?; + unsafe extern "C" fn wait(data: *mut std::ffi::c_void) { + // SAFETY: Registration transfers sole ownership of the boxed receiver here. + let receiver = unsafe { Box::from_raw(data.cast::>()) }; + // Dropping the sender releases the gate, including during unwinding. The timeout + // also bounds cleanup if a failing CUDA operation synchronizes before unwinding. + let _ = receiver.recv_timeout(std::time::Duration::from_secs(10)); + } + let data = Box::into_raw(Box::new(receiver)); + // SAFETY: The stream is live; the callback owns data on successful registration. + unsafe { + result::stream::launch_host_function(stream.cu_stream(), wait, data.cast()).map_err( + |error| { + drop(Box::from_raw(data)); + vortex_err!("Failed to block test stream: {error}") + }, + )?; + } + Ok(release) + } + + fn assert_source_retained(pool: &PinnedByteBufferPool) -> VortexResult<()> { + let deadline = std::time::Instant::now() + std::time::Duration::from_millis(250); + // Give an incorrectly unfenced second transfer time to finish; an immediate query + // could pass just because that stream has not progressed yet. + while std::time::Instant::now() < deadline { + pool.reclaim_completed()?; + assert_eq!( + pool.inflight.lock().len(), + 1, + "source reclaimed before first DMA" + ); + std::thread::sleep(std::time::Duration::from_millis(1)); + } + Ok(()) + } + + fn assert_device_bytes(device: cudarc::driver::CudaSlice, byte: u8) -> VortexResult<()> { + let host = CudaDeviceBuffer::new(device).copy_to_host_sync(Alignment::of::())?; + assert!(host.as_ref().iter().all(|value| *value == byte)); + Ok(()) + } + + #[rstest] + #[case::untracked(false)] + #[case::tracked(true)] + #[crate::test] + fn pending_bucket_source_resubmission(#[case] automatic_tracking: bool) -> VortexResult<()> { + let (pool, stream, other) = setup_resubmission(automatic_tracking)?; + let len = 1024 * 1024; + let mut pinned = pool.get(len)?; + pinned.as_mut_slice().fill(0x5A); + let mut first = stream.device_alloc::(len)?; + let mut second = other.device_alloc::(len)?; + pool.ctx.synchronize().map_err(|e| vortex_err!("{e}"))?; + // Declare the gate last so unwinding releases it before CUDA allocations are dropped. + let gate = block_stream(&stream)?; + submit_to_bucket(pinned, &stream, &mut first)?; + let pinned = pool.get(len)?; + pinned.copy_to_device(&other, 0..len, &mut second.slice_mut(..))?; + assert_source_retained(&pool)?; + drop(gate); + other.synchronize().map_err(|e| vortex_err!("{e}"))?; + pool.reclaim_completed()?; + assert!(pool.inflight.lock().is_empty()); + pool.get(len)?.as_mut_slice().fill(0xA5); + assert_device_bytes(first, 0x5A)?; + assert_device_bytes(second, 0x5A) + } + + #[rstest] + #[case::untracked(false)] + #[case::tracked(true)] + #[crate::test] + fn pending_bucket_source_wait_failure(#[case] automatic_tracking: bool) -> VortexResult<()> { + let (pool, stream, other) = setup_resubmission(automatic_tracking)?; + let len = 1024 * 1024; + let mut pinned = pool.get(len)?; + pinned.as_mut_slice().fill(0x5A); + let mut first = stream.device_alloc::(len)?; + let mut second = other + .clone_htod(&vec![0xCCu8; len]) + .map_err(|e| vortex_err!("{e}"))?; + submit_to_bucket(pinned, &stream, &mut first)?; + let pinned = pool.get(len)?; + // Consumed by wait -> bind_to_thread -> check_err; the error text checks that path. + pool.ctx.record_err::<()>(Err(result::DriverError( + cudaError_enum::CUDA_ERROR_INVALID_VALUE, + ))); + let error = pinned + .copy_to_device(&other, 0..len, &mut second.slice_mut(..)) + .expect_err("pending wait must fail before submission"); + assert!( + error + .to_string() + .contains("Failed to wait for pinned source") + ); + assert!(pool.inflight.lock().is_empty()); + let mut returned = pool.get(len)?; + returned.as_mut_slice().fill(0xA5); + assert_device_bytes(first, 0x5A)?; + assert_device_bytes(second, 0xCC) + } + + #[rstest] + #[crate::test] + fn failed_fence_waits_before_reuse( + #[values(false, true)] automatic_tracking: bool, + ) -> VortexResult<()> { + let (pool, stream) = setup()?; + if !automatic_tracking { + // SAFETY: The destination is used only on its allocation stream; the explicit + // recovery path must fence the source without depending on automatic tracking. + unsafe { stream.context().disable_event_tracking() }; + } + let mut pinned = pool.get(1024 * 1024)?; + pinned.as_mut_slice().fill(0x3C); + let mut destination = stream.device_alloc::(1024 * 1024)?; + let owner = pinned.inner.as_mut().unwrap(); + // SAFETY: The source is initialized above and retained by pinned. + // retain_after_transfer synchronizes the stream or leaks the allocation on failure. + let source = unsafe { std::slice::from_raw_parts(owner.ptr, owner.logical_len) }; + let completion = stream.memcpy_htod(source, &mut destination); + if completion.is_err() { + return pinned.retain_after_transfer(&stream, completion); + } + let result = pinned.retain_after_transfer( + &stream, + Err(result::DriverError( + cudaError_enum::CUDA_ERROR_INVALID_VALUE, + )), + ); + assert!(result.is_err()); + assert!(pool.inflight.lock().is_empty()); + let mut reused = pool.get(1024 * 1024)?; + reused.as_mut_slice().fill(0xC3); + let host = CudaDeviceBuffer::new(destination).copy_to_host_sync(Alignment::of::())?; + assert_eq!(host.as_ref(), vec![0x3C; 1024 * 1024]); + Ok(()) + } + + #[crate::test] + fn rejects_transfer_between_contexts() -> VortexResult<()> { + let (pool, _) = setup()?; + let context = CudaContext::new_non_primary(0, 0) + .map_err(|e| vortex_err!("Failed to create context: {e}"))?; + let stream = VortexCudaStream( + context + .new_stream() + .map_err(|e| vortex_err!("Failed to create stream: {e}"))?, + ); + let mut pinned = pool.get(1024)?; + pinned.as_mut_slice().fill(0xFF); + assert!(pinned.transfer_to_device(&stream).is_err()); + assert!(pool.inflight.lock().is_empty()); + pool.get(1024)?; Ok(()) }