From 80379eba841c1ec1cccfb1925e2657e01011be1e Mon Sep 17 00:00:00 2001 From: Raul Estrada Date: Wed, 23 Sep 2026 01:39:42 -0600 Subject: [PATCH] feat(cuda-core): add borrowed DeviceBuffer subviews Signed-off-by: Raul Estrada --- cuda-core/src/lib.rs | 9 +- cuda-core/src/simt/device_buffer.rs | 371 ++++++++++++++++++++++++++ cuda-core/src/simt/mod.rs | 2 +- cuda-core/tests/simt_device_buffer.rs | 80 +++++- 4 files changed, 456 insertions(+), 6 deletions(-) diff --git a/cuda-core/src/lib.rs b/cuda-core/src/lib.rs index ce2a802e5c..acff57b7fc 100644 --- a/cuda-core/src/lib.rs +++ b/cuda-core/src/lib.rs @@ -32,8 +32,9 @@ pub use simt::{ launch_kernel_ex_cooperative, launch_kernel_ex_cooperative_on_stream, launch_kernel_ex_on_stream, launch_kernel_on_stream, BlockRequirement, ConstantHandle, ContextLimit, CudaContext, CudaEvent, CudaFunction, CudaModule, CudaStream, DeviceBuffer, - DeviceCopy, DeviceLaunchLimits, DynamicSharedMemoryRequirement, EmbeddedModule, - EmbeddedModuleError, KernelLaunchConfig, KernelLaunchContract, LaunchAxis, LaunchConfig1D, - LaunchConfig2D, LaunchConfig3D, LaunchContractError, LaunchContractSpec, LaunchDimension, - PinnedHostBuffer, PreparedLaunch, StreamPriorityRange, SyncPolicy, + DeviceCopy, DeviceLaunchLimits, DeviceSlice, DeviceSliceError, DeviceSliceMut, + DynamicSharedMemoryRequirement, EmbeddedModule, EmbeddedModuleError, KernelLaunchConfig, + KernelLaunchContract, LaunchAxis, LaunchConfig1D, LaunchConfig2D, LaunchConfig3D, + LaunchContractError, LaunchContractSpec, LaunchDimension, PinnedHostBuffer, PreparedLaunch, + StreamPriorityRange, SyncPolicy, }; diff --git a/cuda-core/src/simt/device_buffer.rs b/cuda-core/src/simt/device_buffer.rs index 9057de9be9..a58d4d57e1 100644 --- a/cuda-core/src/simt/device_buffer.rs +++ b/cuda-core/src/simt/device_buffer.rs @@ -31,6 +31,7 @@ use std::marker::PhantomData; use std::mem::MaybeUninit; use std::num::Wrapping; +use std::ops::{Bound, RangeBounds}; use std::sync::Arc; use cuda_bindings::CUdeviceptr; @@ -153,6 +154,220 @@ pub struct DeviceBuffer { _marker: PhantomData, } +/// Error returned when constructing or splitting a borrowed device-buffer view. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum DeviceSliceError { + /// An excluded start bound could not be advanced by one element. + StartBoundOverflow, + /// An included end bound could not be advanced to its exclusive form. + EndBoundOverflow, + /// The normalized start index exceeds the parent view length. + StartOutOfBounds { start: usize, len: usize }, + /// The normalized exclusive end index exceeds the parent view length. + EndOutOfBounds { end: usize, len: usize }, + /// The normalized range has its start after its end. + StartAfterEnd { start: usize, end: usize }, + /// Converting the element offset to a byte offset overflowed. + ByteOffsetOverflow { offset: usize, element_size: usize }, + /// Adding the byte offset to the CUDA device pointer overflowed. + PointerOverflow { ptr: CUdeviceptr, byte_offset: u64 }, +} + +impl std::fmt::Display for DeviceSliceError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match *self { + Self::StartBoundOverflow => write!(f, "device slice start bound overflow"), + Self::EndBoundOverflow => write!(f, "device slice end bound overflow"), + Self::StartOutOfBounds { start, len } => { + write!(f, "device slice start {start} exceeds length {len}") + } + Self::EndOutOfBounds { end, len } => { + write!(f, "device slice end {end} exceeds length {len}") + } + Self::StartAfterEnd { start, end } => { + write!(f, "device slice start {start} exceeds end {end}") + } + Self::ByteOffsetOverflow { + offset, + element_size, + } => write!( + f, + "device slice byte offset overflow: {offset} * {element_size}" + ), + Self::PointerOverflow { ptr, byte_offset } => write!( + f, + "device slice pointer overflow: {ptr:#x} + {byte_offset:#x}" + ), + } + } +} + +impl std::error::Error for DeviceSliceError {} + +/// Immutable non-owning view over a contiguous range of a [`DeviceBuffer`]. +/// +/// The view carries only an adjusted CUDA device pointer and element count. +/// Its lifetime borrows the parent allocation, so the allocation cannot be +/// dropped while the view is live. Dropping the view never frees device memory. +/// +/// ```compile_fail +/// use cuda_core::{CudaStream, DeviceBuffer, DeviceSlice}; +/// +/// fn view_cannot_outlive_buffer(stream: &CudaStream) -> DeviceSlice<'static, u32> { +/// let buffer = DeviceBuffer::::zeroed(stream, 4).unwrap(); +/// buffer.slice(1..3).unwrap() +/// } +/// ``` +pub struct DeviceSlice<'a, T> { + ptr: CUdeviceptr, + len: usize, + _borrow: PhantomData<&'a T>, +} + +impl Copy for DeviceSlice<'_, T> {} + +impl Clone for DeviceSlice<'_, T> { + fn clone(&self) -> Self { + *self + } +} + +impl std::fmt::Debug for DeviceSlice<'_, T> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DeviceSlice") + .field("ptr", &self.ptr) + .field("len", &self.len) + .finish() + } +} + +impl DeviceSlice<'_, T> { + /// Returns the adjusted CUDA device pointer for this view. + #[inline] + pub fn cu_deviceptr(&self) -> CUdeviceptr { + self.ptr + } + + /// Number of elements in this view. + #[inline] + pub fn len(&self) -> usize { + self.len + } + + /// Returns `true` when this view contains no elements. + #[inline] + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + /// Splits this immutable view into two contiguous views at `mid`. + pub fn split_at(self, mid: usize) -> Result<(Self, Self), DeviceSliceError> { + if mid > self.len { + return Err(DeviceSliceError::EndOutOfBounds { + end: mid, + len: self.len, + }); + } + let right_ptr = checked_device_ptr_offset::(self.ptr, mid)?; + Ok(( + Self { + ptr: self.ptr, + len: mid, + _borrow: PhantomData, + }, + Self { + ptr: right_ptr, + len: self.len - mid, + _borrow: PhantomData, + }, + )) + } +} + +/// Exclusive non-owning view over a contiguous range of a [`DeviceBuffer`]. +/// +/// The mutable borrow of the parent allocation is retained for `'a`. The type +/// is intentionally neither `Copy` nor `Clone`; splitting consumes the view and +/// returns two disjoint mutable views. +/// +/// ```compile_fail +/// use cuda_core::{CudaStream, DeviceBuffer}; +/// +/// fn overlapping_mutable_views_are_rejected(stream: &CudaStream) { +/// let mut buffer = DeviceBuffer::::zeroed(stream, 4).unwrap(); +/// let first = buffer.slice_mut(0..3).unwrap(); +/// let second = buffer.slice_mut(1..4).unwrap(); +/// let _ = (first, second); +/// } +/// ``` +pub struct DeviceSliceMut<'a, T> { + ptr: CUdeviceptr, + len: usize, + _borrow: PhantomData<&'a mut T>, +} + +impl std::fmt::Debug for DeviceSliceMut<'_, T> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DeviceSliceMut") + .field("ptr", &self.ptr) + .field("len", &self.len) + .finish() + } +} + +impl<'a, T> DeviceSliceMut<'a, T> { + /// Returns the adjusted CUDA device pointer for this view. + #[inline] + pub fn cu_deviceptr(&self) -> CUdeviceptr { + self.ptr + } + + /// Number of elements in this view. + #[inline] + pub fn len(&self) -> usize { + self.len + } + + /// Returns `true` when this view contains no elements. + #[inline] + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + /// Splits this mutable view into two proven-disjoint mutable views. + pub fn split_at_mut(self, mid: usize) -> Result<(Self, Self), DeviceSliceError> { + if mid > self.len { + return Err(DeviceSliceError::EndOutOfBounds { + end: mid, + len: self.len, + }); + } + let right_ptr = checked_device_ptr_offset::(self.ptr, mid)?; + Ok(( + Self { + ptr: self.ptr, + len: mid, + _borrow: PhantomData, + }, + Self { + ptr: right_ptr, + len: self.len - mid, + _borrow: PhantomData, + }, + )) + } + + /// Consumes the exclusive view and downgrades it to an immutable view. + #[inline] + pub fn into_shared(self) -> DeviceSlice<'a, T> { + DeviceSlice { + ptr: self.ptr, + len: self.len, + _borrow: PhantomData, + } + } +} + // SAFETY: CUdeviceptr is a u64 handle valid across threads when the owning // context is bound. The PhantomData is Send if T is Send. unsafe impl Send for DeviceBuffer {} @@ -212,6 +427,42 @@ impl DeviceBuffer { &self.ctx } + /// Borrows a contiguous immutable subrange of this device allocation. + /// + /// The returned view does not allocate, copy, or own memory. Range + /// normalization, element-to-byte conversion, and pointer addition are all + /// checked before the adjusted pointer is exposed. Empty ranges are valid. + pub fn slice(&self, range: R) -> Result, DeviceSliceError> + where + R: RangeBounds, + { + let (start, end) = normalize_device_range(self.len, range)?; + let ptr = checked_device_ptr_offset::(self.ptr, start)?; + Ok(DeviceSlice { + ptr, + len: end - start, + _borrow: PhantomData, + }) + } + + /// Borrows a contiguous exclusive subrange of this device allocation. + /// + /// The exclusive parent borrow is retained by the returned view, preventing + /// safe creation of overlapping mutable views from the same buffer. Use + /// [`DeviceSliceMut::split_at_mut`] to derive multiple disjoint regions. + pub fn slice_mut(&mut self, range: R) -> Result, DeviceSliceError> + where + R: RangeBounds, + { + let (start, end) = normalize_device_range(self.len, range)?; + let ptr = checked_device_ptr_offset::(self.ptr, start)?; + Ok(DeviceSliceMut { + ptr, + len: end - start, + _borrow: PhantomData, + }) + } + /// Constructs a `DeviceBuffer` from pre-existing raw parts. /// /// # Safety @@ -900,6 +1151,58 @@ impl DeviceBuffer { } } +fn normalize_device_range(len: usize, range: R) -> Result<(usize, usize), DeviceSliceError> +where + R: RangeBounds, +{ + let start = match range.start_bound() { + Bound::Included(&start) => start, + Bound::Excluded(&start) => start + .checked_add(1) + .ok_or(DeviceSliceError::StartBoundOverflow)?, + Bound::Unbounded => 0, + }; + let end = match range.end_bound() { + Bound::Included(&end) => end + .checked_add(1) + .ok_or(DeviceSliceError::EndBoundOverflow)?, + Bound::Excluded(&end) => end, + Bound::Unbounded => len, + }; + + if start > len { + return Err(DeviceSliceError::StartOutOfBounds { start, len }); + } + if end > len { + return Err(DeviceSliceError::EndOutOfBounds { end, len }); + } + if start > end { + return Err(DeviceSliceError::StartAfterEnd { start, end }); + } + Ok((start, end)) +} + +fn checked_device_ptr_offset( + ptr: CUdeviceptr, + offset: usize, +) -> Result { + let element_size = std::mem::size_of::(); + let byte_offset = + offset + .checked_mul(element_size) + .ok_or(DeviceSliceError::ByteOffsetOverflow { + offset, + element_size, + })?; + let byte_offset = + u64::try_from(byte_offset).map_err(|_| DeviceSliceError::ByteOffsetOverflow { + offset, + element_size, + })?; + ptr.checked_add(byte_offset) + .ok_or(DeviceSliceError::PointerOverflow { ptr, byte_offset }) +} + fn allocation_size(len: usize) -> Result { len.checked_mul(std::mem::size_of::()).ok_or(DriverError( cuda_bindings::cudaError_enum_CUDA_ERROR_INVALID_VALUE, @@ -925,6 +1228,74 @@ fn chunk_cast_len(bytes: usize, addr: usize, elem_size: usize, align: usize) -> Some(bytes / elem_size) } +#[cfg(test)] +mod device_slice_tests { + use super::{checked_device_ptr_offset, normalize_device_range, DeviceSliceError}; + use std::ops::Bound; + + #[test] + fn normalizes_supported_range_bounds() { + assert_eq!(normalize_device_range(8, 2..6), Ok((2, 6))); + assert_eq!(normalize_device_range(8, 2..=5), Ok((2, 6))); + assert_eq!(normalize_device_range(8, ..4), Ok((0, 4))); + assert_eq!(normalize_device_range(8, 4..), Ok((4, 8))); + assert_eq!(normalize_device_range(8, ..), Ok((0, 8))); + assert_eq!(normalize_device_range(8, 3..3), Ok((3, 3))); + } + + #[test] + fn rejects_invalid_range_bounds() { + assert_eq!( + normalize_device_range(4, 5..), + Err(DeviceSliceError::StartOutOfBounds { start: 5, len: 4 }) + ); + assert_eq!( + normalize_device_range(4, ..5), + Err(DeviceSliceError::EndOutOfBounds { end: 5, len: 4 }) + ); + let reversed = (Bound::Included(3), Bound::Excluded(2)); + assert_eq!( + normalize_device_range(4, reversed), + Err(DeviceSliceError::StartAfterEnd { start: 3, end: 2 }) + ); + } + + #[test] + fn rejects_bound_overflow() { + assert_eq!( + normalize_device_range(usize::MAX, (Bound::Excluded(usize::MAX), Bound::Unbounded),), + Err(DeviceSliceError::StartBoundOverflow) + ); + assert_eq!( + normalize_device_range(usize::MAX, (Bound::Unbounded, Bound::Included(usize::MAX)),), + Err(DeviceSliceError::EndBoundOverflow) + ); + } + + #[test] + fn checks_device_pointer_arithmetic() { + assert_eq!(checked_device_ptr_offset::(0x1000, 3), Ok(0x100c)); + assert_eq!( + checked_device_ptr_offset::<()>(0x1000, usize::MAX), + Ok(0x1000) + ); + assert_eq!( + checked_device_ptr_offset::(0, usize::MAX), + Err(DeviceSliceError::ByteOffsetOverflow { + offset: usize::MAX, + element_size: 8, + }) + ); + assert_eq!( + checked_device_ptr_offset::(u64::MAX - 3, 1), + Err(DeviceSliceError::PointerOverflow { + ptr: u64::MAX - 3, + byte_offset: 8, + }) + ); + } +} + #[cfg(test)] mod chunk_cast_tests { use super::chunk_cast_len; diff --git a/cuda-core/src/simt/mod.rs b/cuda-core/src/simt/mod.rs index 7f07b4c361..15c49deadf 100644 --- a/cuda-core/src/simt/mod.rs +++ b/cuda-core/src/simt/mod.rs @@ -49,7 +49,7 @@ pub use context::{ContextLimit, CudaContext, StreamPriorityRange, SyncPolicy}; /// `#[derive(DeviceCopy)]`, re-exported next to the trait so /// `use cuda_core::DeviceCopy;` brings both into scope (serde pattern). pub use cuda_core_derive::DeviceCopy; -pub use device_buffer::{DeviceBuffer, DeviceCopy}; +pub use device_buffer::{DeviceBuffer, DeviceCopy, DeviceSlice, DeviceSliceError, DeviceSliceMut}; pub use embedded::{EmbeddedModule, EmbeddedModuleError}; pub use event::CudaEvent; pub use launch::{ diff --git a/cuda-core/tests/simt_device_buffer.rs b/cuda-core/tests/simt_device_buffer.rs index 9d67b0471d..48ee16090c 100644 --- a/cuda-core/tests/simt_device_buffer.rs +++ b/cuda-core/tests/simt_device_buffer.rs @@ -3,7 +3,9 @@ * SPDX-License-Identifier: Apache-2.0 */ -use cuda_core::{CudaContext, CudaStream, DeviceBuffer, DriverError, PinnedHostBuffer}; +use cuda_core::{ + CudaContext, CudaStream, DeviceBuffer, DeviceSliceError, DriverError, PinnedHostBuffer, +}; use std::sync::{mpsc, Mutex, MutexGuard}; use std::thread::JoinHandle; use std::time::Duration; @@ -421,3 +423,79 @@ fn uninitialized_async_cast_elem_implicit_drop_is_stream_ordered() { } stream.synchronize().expect("stream sync failed"); } + +#[test] +fn device_buffer_borrowed_subviews_adjust_pointer_and_length() { + let ctx = CudaContext::new(0).expect("failed to create CUDA context"); + let stream = ctx.new_stream().expect("failed to create CUDA stream"); + let data = [10_u32, 20, 30, 40, 50, 60]; + let buffer = DeviceBuffer::from_host(&stream, &data).expect("failed to allocate device buffer"); + let base = buffer.cu_deviceptr(); + + let middle = buffer.slice(1..5).expect("failed to create middle subview"); + assert_eq!(middle.len(), 4); + assert!(!middle.is_empty()); + assert_eq!( + middle.cu_deviceptr(), + base + std::mem::size_of::() as u64 + ); + + let (left, right) = middle.split_at(2).expect("failed to split shared view"); + assert_eq!(left.len(), 2); + assert_eq!(right.len(), 2); + assert_eq!(left.cu_deviceptr(), base + 4); + assert_eq!(right.cu_deviceptr(), base + 12); + + let empty = buffer.slice(3..3).expect("failed to create empty subview"); + assert!(empty.is_empty()); + assert_eq!(empty.cu_deviceptr(), base + 12); +} + +#[test] +fn device_buffer_mutable_subviews_split_into_disjoint_ranges() { + let ctx = CudaContext::new(0).expect("failed to create CUDA context"); + let stream = ctx.new_stream().expect("failed to create CUDA stream"); + let mut buffer = + DeviceBuffer::::zeroed(&stream, 8).expect("failed to allocate device buffer"); + let base = buffer.cu_deviceptr(); + + let middle = buffer + .slice_mut(1..7) + .expect("failed to create mutable middle subview"); + assert_eq!(middle.len(), 6); + assert_eq!(middle.cu_deviceptr(), base + 4); + + let (left, right) = middle + .split_at_mut(2) + .expect("failed to split mutable subview"); + assert_eq!(left.len(), 2); + assert_eq!(right.len(), 4); + assert_eq!(left.cu_deviceptr(), base + 4); + assert_eq!(right.cu_deviceptr(), base + 12); + + let right = right.into_shared(); + assert_eq!(right.len(), 4); + assert_eq!(right.cu_deviceptr(), base + 12); +} + +#[test] +fn device_buffer_subviews_reject_invalid_ranges() { + let ctx = CudaContext::new(0).expect("failed to create CUDA context"); + let stream = ctx.new_stream().expect("failed to create CUDA stream"); + let buffer = DeviceBuffer::::zeroed(&stream, 4).expect("failed to allocate device buffer"); + + assert!(matches!( + buffer.slice(5..), + Err(DeviceSliceError::StartOutOfBounds { start: 5, len: 4 }) + )); + assert!(matches!( + buffer.slice(..5), + Err(DeviceSliceError::EndOutOfBounds { end: 5, len: 4 }) + )); + let start = 3; + let end = 2; + assert!(matches!( + buffer.slice(start..end), + Err(DeviceSliceError::StartAfterEnd { start: 3, end: 2 }) + )); +}