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(()) }