1use std::{fmt, marker::PhantomData, ptr};
11
12use cugraph_sys::{
13 cugraph_error_t, cugraph_type_erased_device_array_create, cugraph_type_erased_device_array_t,
14 cugraph_type_erased_device_array_view, cugraph_type_erased_device_array_view_copy_from_host,
15 cugraph_type_erased_device_array_view_copy_to_host, cugraph_type_erased_device_array_view_create,
16 cugraph_type_erased_device_array_view_free, cugraph_type_erased_device_array_view_pointer,
17 cugraph_type_erased_device_array_view_size, cugraph_type_erased_device_array_view_t,
18 cugraph_type_erased_device_array_view_type, cugraph_type_erased_host_array_view_create,
19 cugraph_type_erased_host_array_view_free, cugraph_type_erased_host_array_view_t,
20};
21
22use crate::{
23 cuda::{DeviceResourceKind, drop_device_resource},
24 error::{Error, ErrorKind, Result, check_error, require_non_null},
25 handle::{EnteredContext, ResourceHandleLease},
26 types::{CugraphElement, DataType},
27};
28
29pub(crate) struct RawDeviceArray {
31 raw: *mut cugraph_type_erased_device_array_t,
32 raw_view: *mut cugraph_type_erased_device_array_view_t,
33 device: i32,
34 _resource_handle_lease: Option<ResourceHandleLease>,
36 _nosend_nosync: PhantomData<*mut ()>,
37}
38
39fn checked_byte_len(len: usize, dtype: DataType, operation: &'static str) -> Result<usize> {
40 len.checked_mul(dtype.size_of()).ok_or_else(|| {
41 Error::new(ErrorKind::InvalidInput, operation)
42 .with_message("array byte size overflow")
43 .with_context("len", len.to_string())
44 .with_context("dtype", format!("{:?}", dtype))
45 .with_context("element_size", dtype.size_of().to_string())
46 })
47}
48
49impl RawDeviceArray {
50 pub(crate) unsafe fn from_raw(
58 raw: *mut cugraph_type_erased_device_array_t,
59 device: i32,
60 operation: &'static str,
61 pointer_name: &'static str,
62 resource_handle_lease: Option<ResourceHandleLease>,
63 ) -> Result<Self> {
64 let raw = require_non_null(raw, operation, pointer_name)?;
65 let raw_view = unsafe { cugraph_type_erased_device_array_view(raw) };
68 let raw_view = match require_non_null(raw_view, operation, "device_array_view") {
69 Ok(raw_view) => raw_view,
70 Err(error) => {
71 let mut raw = raw;
72 drop_device_resource(&mut raw, device, operation, DeviceResourceKind::DeviceArray);
73 return Err(error);
74 },
75 };
76
77 Ok(Self { raw, raw_view, device, _resource_handle_lease: resource_handle_lease, _nosend_nosync: PhantomData })
78 }
79
80 pub(crate) fn new(handle: &EnteredContext<'_>, len: usize, dtype: DataType) -> Result<Self> {
81 let mut raw: *mut cugraph_type_erased_device_array_t = ptr::null_mut();
82 let mut error: *mut cugraph_error_t = ptr::null_mut();
83
84 unsafe {
87 let code =
88 cugraph_type_erased_device_array_create(handle.handle_ptr(), len, dtype.into(), &mut raw, &mut error);
89 check_error(code, error, handle.device(), "RawDeviceArray::new")?;
90 }
91
92 let raw = require_non_null(raw, "RawDeviceArray::new", "device_array")?;
93 let raw_view = unsafe { cugraph_type_erased_device_array_view(raw) };
96 let raw_view = match require_non_null(raw_view, "RawDeviceArray::new", "device_array_view") {
97 Ok(raw_view) => raw_view,
98 Err(error) => {
99 let mut raw = raw;
100 drop_device_resource(&mut raw, handle.device(), "RawDeviceArray::new", DeviceResourceKind::DeviceArray);
101 return Err(error);
102 },
103 };
104
105 Ok(Self {
106 raw,
107 raw_view,
108 device: handle.device(),
109 _resource_handle_lease: handle.runtime_bound_lease(),
110 _nosend_nosync: PhantomData,
111 })
112 }
113
114 pub(crate) fn raw_view(&self) -> RawDeviceArrayView<'_> {
115 RawDeviceArrayView {
116 raw: self.raw_view,
117 owned: false,
118 device: Some(self.device),
119 _lifetime: PhantomData,
120 _nosend_nosync: PhantomData,
121 }
122 }
123
124 fn copy_from_host(&mut self, handle: &EnteredContext<'_>, data: &[u8]) -> Result<()> {
125 let view = self.raw_view();
126 let mut error: *mut cugraph_error_t = ptr::null_mut();
127
128 unsafe {
132 let code = cugraph_type_erased_device_array_view_copy_from_host(
133 handle.handle_ptr(),
134 view.raw,
135 data.as_ptr() as *const i8,
136 &mut error,
137 );
138 check_error(code, error, handle.device(), "RawDeviceArray::copy_from_host")?;
139 }
140
141 Ok(())
142 }
143}
144
145impl Drop for RawDeviceArray {
146 fn drop(&mut self) {
147 unsafe { cugraph_type_erased_device_array_view_free(self.raw_view) };
152 drop_device_resource(&mut self.raw, self.device, "RawDeviceArray::drop.array", DeviceResourceKind::DeviceArray);
153 }
154}
155
156pub(crate) struct RawDeviceArrayView<'a> {
162 pub(crate) raw: *mut cugraph_type_erased_device_array_view_t,
163 owned: bool,
164 device: Option<i32>,
165 _lifetime: PhantomData<&'a ()>,
166 _nosend_nosync: PhantomData<*mut ()>,
167}
168
169impl fmt::Debug for RawDeviceArrayView<'_> {
170 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
171 f.debug_struct("RawDeviceArrayView")
172 .field("len", &self.len())
173 .field("dtype", &self.dtype().ok())
174 .finish_non_exhaustive()
175 }
176}
177
178impl RawDeviceArrayView<'_> {
179 #[cfg(test)]
190 pub(crate) unsafe fn from_borrowed(
191 raw: *mut cugraph_type_erased_device_array_view_t,
192 device: Option<i32>,
193 operation: &'static str,
194 pointer_name: &'static str,
195 ) -> Result<Self> {
196 let raw = require_non_null(raw, operation, pointer_name)?;
197 Ok(unsafe { Self::from_borrowed_unchecked(raw, device) })
201 }
202
203 pub(crate) unsafe fn from_borrowed_unchecked(
204 raw: *mut cugraph_type_erased_device_array_view_t,
205 device: Option<i32>,
206 ) -> Self {
207 RawDeviceArrayView { raw, owned: false, device, _lifetime: PhantomData, _nosend_nosync: PhantomData }
208 }
209
210 pub(crate) unsafe fn from_owned_view(
220 raw: *mut cugraph_type_erased_device_array_view_t,
221 device: Option<i32>,
222 operation: &'static str,
223 pointer_name: &'static str,
224 ) -> Result<Self> {
225 let raw = require_non_null(raw, operation, pointer_name)?;
226 Ok(RawDeviceArrayView { raw, owned: true, device, _lifetime: PhantomData, _nosend_nosync: PhantomData })
227 }
228
229 pub(crate) unsafe fn try_from_raw(
235 pointer: *mut std::ffi::c_void,
236 len: usize,
237 dtype: DataType,
238 operation: &'static str,
239 ) -> Result<Self> {
240 checked_byte_len(len, dtype, operation)?;
241 let raw = unsafe { cugraph_type_erased_device_array_view_create(pointer, len, dtype.into()) };
244 if raw.is_null() {
245 return Err(Error::new(ErrorKind::OutOfMemory, operation)
246 .with_message("failed to allocate device array view metadata")
247 .set_temporary());
248 }
249 Ok(RawDeviceArrayView { raw, owned: true, device: None, _lifetime: PhantomData, _nosend_nosync: PhantomData })
250 }
251
252 pub(crate) fn device_ptr(&self) -> *mut std::ffi::c_void {
254 unsafe { cugraph_type_erased_device_array_view_pointer(self.raw) as *mut std::ffi::c_void }
256 }
257
258 pub(crate) fn len(&self) -> usize {
260 unsafe { cugraph_type_erased_device_array_view_size(self.raw) }
262 }
263
264 pub(crate) fn is_empty(&self) -> bool {
266 self.len() == 0
267 }
268
269 pub(crate) fn dtype(&self) -> Result<DataType> {
271 let c_type = unsafe { cugraph_type_erased_device_array_view_type(self.raw) };
273 DataType::try_from(c_type).map_err(|err| err.with_operation("RawDeviceArrayView::dtype"))
274 }
275
276 pub(crate) fn as_ptr(&self) -> *mut cugraph_type_erased_device_array_view_t {
277 self.raw
278 }
279
280 pub(crate) fn device(&self) -> Option<i32> {
281 self.device
282 }
283
284 pub(crate) fn raw_view(&self) -> RawDeviceArrayView<'_> {
286 unsafe { RawDeviceArrayView::from_borrowed_unchecked(self.raw, self.device) }
289 }
290
291 pub(crate) fn to_host_vec(&self, handle: &EnteredContext<'_>) -> Result<Vec<u8>> {
296 let dtype = self.dtype()?;
297 let byte_len = checked_byte_len(self.len(), dtype, "RawDeviceArrayView::to_host_vec")?;
298 let mut buf = vec![0u8; byte_len];
299 let mut error: *mut cugraph_error_t = ptr::null_mut();
300
301 unsafe {
305 let code = cugraph_type_erased_device_array_view_copy_to_host(
306 handle.handle_ptr(),
307 buf.as_mut_ptr() as *mut i8,
308 self.raw,
309 &mut error,
310 );
311 check_error(code, error, handle.device(), "RawDeviceArrayView::to_host_vec")?;
312 }
313
314 Ok(buf)
315 }
316}
317
318impl Drop for RawDeviceArrayView<'_> {
319 fn drop(&mut self) {
320 if self.owned {
321 unsafe { cugraph_type_erased_device_array_view_free(self.raw) };
325 }
326 }
327}
328
329pub struct DeviceArray<T: CugraphElement> {
352 inner: RawDeviceArray,
353 _element: PhantomData<T>,
354}
355
356impl<T: CugraphElement> fmt::Debug for DeviceArray<T> {
357 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
358 let mut d = f.debug_struct("DeviceArray");
359 d.field("dtype", &T::data_type());
360 d.field("len", &self.inner.raw_view().len());
361 d.finish_non_exhaustive()
362 }
363}
364
365impl<T: CugraphElement> DeviceArray<T> {
366 pub(crate) unsafe fn from_raw(
373 raw: *mut cugraph_type_erased_device_array_t,
374 device: i32,
375 operation: &'static str,
376 pointer_name: &'static str,
377 resource_handle_lease: Option<ResourceHandleLease>,
378 ) -> Result<Self> {
379 let inner = unsafe { RawDeviceArray::from_raw(raw, device, operation, pointer_name, resource_handle_lease) }?;
381 let dtype = inner.raw_view().dtype()?;
382 if dtype != T::data_type() {
383 return Err(Error::new(ErrorKind::InvalidState, operation).with_message(format!(
384 "{pointer_name} dtype {dtype} does not match requested dtype {}",
385 T::data_type()
386 )));
387 }
388 Ok(Self { inner, _element: PhantomData })
389 }
390
391 pub fn new(handle: &EnteredContext<'_>, len: usize) -> Result<Self> {
393 let inner = RawDeviceArray::new(handle, len, T::data_type())?;
394 Ok(Self { inner, _element: PhantomData })
395 }
396
397 pub fn from_slice(handle: &EnteredContext<'_>, data: &[T]) -> Result<Self> {
414 let mut arr = Self::new(handle, data.len())?;
415 let byte_len = std::mem::size_of_val(data);
416 let bytes = unsafe { std::slice::from_raw_parts(data.as_ptr() as *const u8, byte_len) };
419 arr.inner.copy_from_host(handle, bytes)?;
420 Ok(arr)
421 }
422
423 pub fn view(&self) -> Result<DeviceArrayView<'_, T>> {
425 Ok(DeviceArrayView { inner: self.inner.raw_view(), _element: PhantomData })
426 }
427}
428
429pub struct DeviceArrayView<'a, T: CugraphElement> {
435 inner: RawDeviceArrayView<'a>,
436 _element: PhantomData<T>,
437}
438
439impl<T: CugraphElement> fmt::Debug for DeviceArrayView<'_, T> {
440 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
441 f.debug_struct("DeviceArrayView")
442 .field("len", &self.len())
443 .field("dtype", &T::data_type())
444 .finish_non_exhaustive()
445 }
446}
447
448impl<'a, T: CugraphElement> DeviceArrayView<'a, T> {
449 #[cfg(test)] pub(crate) unsafe fn from_raw(pointer: *mut std::ffi::c_void, len: usize) -> Result<Self> {
459 let inner =
461 unsafe { RawDeviceArrayView::try_from_raw(pointer, len, T::data_type(), "DeviceArrayView::from_raw") }?;
462 Ok(DeviceArrayView { inner, _element: PhantomData })
463 }
464
465 pub(crate) fn from_raw_unchecked(raw: RawDeviceArrayView<'a>) -> Self {
471 DeviceArrayView { inner: raw, _element: PhantomData }
472 }
473
474 pub fn device_ptr(&self) -> *mut std::ffi::c_void {
476 self.inner.device_ptr()
477 }
478
479 pub fn len(&self) -> usize {
481 self.inner.len()
482 }
483
484 pub fn is_empty(&self) -> bool {
486 self.len() == 0
487 }
488
489 pub(crate) fn as_ptr(&self) -> *mut cugraph_type_erased_device_array_view_t {
491 self.inner.as_ptr()
492 }
493
494 pub(crate) fn raw_view(&self) -> RawDeviceArrayView<'_> {
496 self.inner.raw_view()
497 }
498
499 pub fn to_vec(&self, handle: &EnteredContext<'_>) -> Result<Vec<T>> {
514 let type_size = std::mem::size_of::<T>();
515 let bytes = self.inner.to_host_vec(handle)?;
516 let len = bytes.len() / type_size;
517
518 let mut result = Vec::with_capacity(len);
519 for chunk in bytes.chunks_exact(type_size) {
520 let value = unsafe { std::ptr::read_unaligned(chunk.as_ptr() as *const T) };
524 result.push(value);
525 }
526
527 Ok(result)
528 }
529}
530
531pub(crate) struct HostArrayView<'a> {
540 pub(crate) raw: *mut cugraph_type_erased_host_array_view_t,
541 _lifetime: PhantomData<&'a ()>,
542 _nosend_nosync: PhantomData<*mut ()>,
543}
544
545impl fmt::Debug for HostArrayView<'_> {
546 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
547 f.debug_struct("HostArrayView").finish_non_exhaustive()
548 }
549}
550
551impl<'a> HostArrayView<'a> {
552 pub(crate) fn new(data: &'a [u8], n_elems: usize, dtype: DataType) -> Result<Self> {
562 let required = checked_byte_len(n_elems, dtype, "HostArrayView::new")?;
563 if data.len() < required {
564 return Err(Error::new(ErrorKind::InvalidInput, "HostArrayView::new")
565 .with_message("slice is too small for the requested number of elements")
566 .with_context("required_bytes", required.to_string())
567 .with_context("actual_bytes", data.len().to_string())
568 .with_context("n_elems", n_elems.to_string())
569 .with_context("dtype", format!("{:?}", dtype)));
570 }
571
572 let data_ptr = data.as_ptr().cast::<std::ffi::c_void>().cast_mut();
576
577 let raw = unsafe { cugraph_type_erased_host_array_view_create(data_ptr, n_elems, dtype.into()) };
580 if raw.is_null() {
581 return Err(Error::new(ErrorKind::OutOfMemory, "HostArrayView::new")
582 .with_message("failed to allocate host array view metadata")
583 .set_temporary());
584 }
585 Ok(HostArrayView { raw, _lifetime: PhantomData, _nosend_nosync: PhantomData })
586 }
587
588 pub(crate) fn as_ptr(&self) -> *mut cugraph_type_erased_host_array_view_t {
589 self.raw
590 }
591}
592
593impl Drop for HostArrayView<'_> {
594 fn drop(&mut self) {
595 unsafe { cugraph_type_erased_host_array_view_free(self.raw) };
599 }
600}
601
602#[cfg(test)]
603mod tests {
604 use static_assertions::assert_not_impl_any;
605
606 use super::*;
607 use crate::error::ErrorKind;
608
609 #[test]
610 fn test_array_types_not_send_sync() {
611 assert_not_impl_any!(DeviceArray<i32>: Send, Sync);
612 assert_not_impl_any!(DeviceArrayView<'static, i32>: Send, Sync);
613 assert_not_impl_any!(HostArrayView<'static>: Send, Sync);
614 }
615
616 #[test]
617 fn test_host_array_view_rejects_undersized_slice() {
618 let err = HostArrayView::new(&[0u8; 3], 2, DataType::Int32).expect_err("must reject undersized slice");
619 assert_eq!(err.kind(), ErrorKind::InvalidInput);
620 assert_eq!(err.operation(), "HostArrayView::new");
621 }
622
623 #[test]
624 fn test_host_array_view_rejects_overflowing_size() {
625 let err = HostArrayView::new(&[], usize::MAX, DataType::Int64).expect_err("must reject overflowing size math");
626 assert_eq!(err.kind(), ErrorKind::InvalidInput);
627 assert_eq!(err.operation(), "HostArrayView::new");
628 }
629
630 #[test]
631 fn test_raw_device_array_view_try_from_raw_rejects_overflowing_size() {
632 let err = unsafe {
635 RawDeviceArrayView::try_from_raw(
636 std::ptr::null_mut(),
637 usize::MAX,
638 DataType::Int64,
639 "RawDeviceArrayView::try_from_raw",
640 )
641 }
642 .expect_err("must reject overflowing size math");
643 assert_eq!(err.kind(), ErrorKind::InvalidInput);
644 assert_eq!(err.operation(), "RawDeviceArrayView::try_from_raw");
645 }
646
647 #[test]
648 fn test_raw_device_array_view_from_borrowed_rejects_null() {
649 let err = unsafe { RawDeviceArrayView::from_borrowed(std::ptr::null_mut(), None, "test_op", "view") }
652 .expect_err("null borrowed view should be rejected");
653 assert_eq!(err.kind(), ErrorKind::Unexpected);
654 assert_eq!(err.operation(), "test_op");
655 }
656
657 #[test]
658 fn test_host_array_view_accepts_exact_size() {
659 let view = HostArrayView::new(&[0u8; 8], 2, DataType::Int32).expect("exact size should succeed");
660 drop(view);
661 }
662
663 #[test]
664 fn test_host_array_view_accepts_oversized_slice() {
665 let view = HostArrayView::new(&[0u8; 16], 2, DataType::Int32).expect("oversized slice should succeed");
666 drop(view);
667 }
668}