1use std::{
41 cell::{Cell, RefCell},
42 marker::PhantomData,
43 mem::ManuallyDrop,
44 ptr,
45};
46
47use cudf::{
48 interop::contract::{InteropError, LeaseCompletion},
49 runtime::{CudaEvent, ExecutionBindingId},
50};
51use cugraph_sys::{cugraph_decompress_to_edgelist, cugraph_edgelist_t, cugraph_error_t, cugraph_graph_t};
52
53use crate::{
54 array::{DeviceArray, DeviceArrayView},
55 cuda::{DeviceGuard, DeviceResourceKind, drop_device_resource},
56 error::{Error, ErrorKind, Result, check_error},
57 handle::{EnteredContext, ResourceHandleLease},
58 random::RngState,
59 sampling::select_random_vertices,
60 structure::{ResolvedVertexSelection, VertexSelection},
61 types::{CugraphElement, DataType, GraphProperties, to_bool_t},
62};
63
64mod algorithms;
65mod builder;
66mod edgelist;
67mod parked;
68
69pub use builder::GraphBuilder;
70pub(crate) use edgelist::DecompressedEdgelist;
71pub use parked::ParkedGraph;
72
73mod sealed {
74 pub trait Sealed {}
75 impl Sealed for super::Forward {}
76 impl Sealed for super::Transposed {}
77}
78
79pub trait Transposition: sealed::Sealed {
85 #[doc(hidden)]
87 fn store_transposed() -> bool;
88}
89
90pub struct Forward;
94impl Transposition for Forward {
95 fn store_transposed() -> bool {
96 false
97 }
98}
99
100pub struct Transposed;
104impl Transposition for Transposed {
105 fn store_transposed() -> bool {
106 true
107 }
108}
109
110pub struct Graph<T = Forward> {
134 raw: *mut cugraph_graph_t,
135 device: i32,
136 properties: GraphProperties,
137 vertex_dtype: DataType,
138 weight_dtype: DataType,
139 has_weights: bool,
140 edge_id_dtype: Option<DataType>,
141 pending_input_completion: Option<DeferredInputCompletion>,
142 resource_handle_leases: RefCell<Vec<ResourceHandleLease>>,
148 ordered_binding: Cell<Option<ExecutionBindingId>>,
154 untracked_resource_handle_use: Cell<bool>,
158 _nosend_nosync: PhantomData<*mut ()>,
159 _transposition: PhantomData<T>,
160}
161
162pub(crate) struct DeferredInputCompletion {
163 completions: Vec<LeaseCompletion>,
164}
165
166impl DeferredInputCompletion {
167 pub(crate) fn new(completions: Vec<LeaseCompletion>) -> Option<Self> {
168 if completions.is_empty() {
169 return None;
170 }
171
172 Some(Self { completions })
173 }
174
175 pub(crate) fn synchronize(&mut self, operation: &'static str) -> Result<()> {
176 for completion in &mut self.completions {
177 completion.synchronize().map_err(|error| map_completion_witness_error(error, operation))?;
178 }
179 Ok(())
180 }
181}
182
183fn map_completion_witness_error(error: InteropError, operation: &'static str) -> Error {
184 let mut mapped = Error::new(ErrorKind::InvalidState, operation)
185 .with_message("failed to synchronize graph input completion witness")
186 .with_context("interop_kind", format!("{:?}", error.kind()))
187 .with_context("interop_message", error.message().to_string());
188
189 for (key, value) in error.context() {
190 mapped = mapped.with_context(key, value.clone());
191 }
192
193 mapped.set_source(error)
194}
195
196fn map_cudf_runtime_error(error: cudf::Error, operation: &'static str, message: &'static str) -> Error {
197 let kind = match error.code() {
198 cudf::ErrorCode::OutOfMemory | cudf::ErrorCode::ForeignResourceExhausted => ErrorKind::OutOfMemory,
199 cudf::ErrorCode::CudaRuntime
200 | cudf::ErrorCode::FatalCuda
201 | cudf::ErrorCode::ReadinessUnavailable
202 | cudf::ErrorCode::ForeignTransientDependencyFailure
203 | cudf::ErrorCode::ForeignPersistentDependencyFailure => ErrorKind::Unexpected,
204 cudf::ErrorCode::RuntimeBindingMismatch
205 | cudf::ErrorCode::ExecutionBindingUnavailable
206 | cudf::ErrorCode::InvariantViolation
207 | cudf::ErrorCode::ForeignInvariantViolation => ErrorKind::InvalidState,
208 cudf::ErrorCode::CrossDeviceInputs | cudf::ErrorCode::CrossProvenanceInputs => ErrorKind::InvalidInput,
209 cudf::ErrorCode::UnsupportedDtype
210 | cudf::ErrorCode::UnsupportedLayout
211 | cudf::ErrorCode::UnsupportedOperation => ErrorKind::UnsupportedType,
212 _ => ErrorKind::Unexpected,
213 };
214
215 let source_status = error.status();
216 let mut mapped = Error::new(kind, operation)
217 .with_message(message)
218 .with_context("cudf_code", error.code().to_string())
219 .with_context("cudf_kind", error.kind().to_string())
220 .with_context("cudf_status", source_status.to_string())
221 .with_context("cudf_operation", error.operation().to_string());
222
223 mapped = match source_status {
224 cudf::ErrorStatus::Permanent => mapped.set_permanent(),
225 cudf::ErrorStatus::Temporary => mapped.set_temporary(),
226 cudf::ErrorStatus::Persistent => mapped.set_persistent(),
227 };
228
229 mapped.set_source(error)
230}
231
232impl<T> Graph<T> {
233 pub fn vertex_dtype(&self) -> DataType {
235 self.vertex_dtype
236 }
237
238 pub fn weight_dtype(&self) -> DataType {
242 self.weight_dtype
243 }
244
245 pub(crate) fn has_edge_ids(&self) -> bool {
246 self.edge_id_dtype.is_some()
247 }
248
249 pub(crate) fn edge_id_dtype(&self) -> Option<DataType> {
250 self.edge_id_dtype
251 }
252
253 pub(crate) fn has_weights(&self) -> bool {
254 self.has_weights
255 }
256
257 pub(crate) fn with_pending_input_completion(mut self, pending: DeferredInputCompletion) -> Self {
258 self.pending_input_completion = Some(pending);
259 self
260 }
261
262 pub fn synchronize_input_reads(&mut self) -> Result<()> {
270 self.finish_input_reads("Graph::synchronize_input_reads")
271 }
272
273 fn finish_input_reads(&mut self, operation: &'static str) -> Result<()> {
274 if let Some(pending) = self.pending_input_completion.as_mut() {
275 pending.synchronize(operation)?;
276 }
277 self.pending_input_completion = None;
278 Ok(())
279 }
280
281 pub fn park(mut self, handle: &EnteredContext<'_>) -> Result<ParkedGraph<T>>
296 where T: Send {
297 let operation = "Graph::park";
298 self.validate_runtime_bound_transfer(handle, operation)?;
299
300 self.finish_input_reads(operation)?;
301 self.park_recorded(handle, operation)
302 }
303
304 pub fn park_under(mut self, handle: &EnteredContext<'_>) -> Result<ParkedGraph<T>>
315 where T: Send {
316 let operation = "Graph::park_under";
317 if handle.device() != self.device {
318 return Err(Error::new(ErrorKind::InvalidInput, operation)
319 .with_message("graph and resource handle must be on the same CUDA device")
320 .with_context("graph_device", self.device.to_string())
321 .with_context("handle_device", handle.device().to_string()));
322 }
323 let Some(handle_lease) = handle.runtime_bound_lease() else {
324 return Err(Error::new(ErrorKind::InvalidState, operation)
325 .with_message("park_under requires a runtime-bound resource handle"));
326 };
327 if self.untracked_resource_handle_use.get() {
328 return Err(Error::new(ErrorKind::InvalidState, operation)
329 .with_message("graph was used with a non-runtime-bound resource handle and cannot be parked"));
330 }
331 if self.resource_handle_leases.get_mut().is_empty() {
332 return Err(Error::new(ErrorKind::InvalidState, operation)
333 .with_message("graph has no runtime-bound handle leases to prove park stream ownership"));
334 }
335 let entering = handle_lease.runtime().binding_id();
336 if self.ordered_binding.get() != Some(entering) {
337 return Err(Error::new(ErrorKind::InvalidState, operation)
338 .with_message("park handle binding does not match the ordered graph execution binding")
339 .with_context(
340 "graph_binding_id",
341 self.ordered_binding.get().map(|id| id.to_string()).as_deref().unwrap_or("none"),
342 )
343 .with_context("handle_binding_id", entering.to_string()));
344 }
345 self.finish_input_reads(operation)?;
346 self.park_recorded(handle, operation)
347 }
348
349 fn park_recorded(self, handle: &EnteredContext<'_>, operation: &'static str) -> Result<ParkedGraph<T>> {
350 let _guard = DeviceGuard::enter_owner_device(self.device, operation)?;
351 let event = CudaEvent::new()
352 .map_err(|error| map_cudf_runtime_error(error, operation, "failed to create graph park CUDA event"))?;
353 let runtime = handle.runtime().expect("validated runtime-bound handle");
354 event
355 .record(runtime.stream())
356 .map_err(|error| map_cudf_runtime_error(error, operation, "failed to record graph park CUDA event"))?;
357 let park_binding = runtime.binding_id();
358
359 let this = ManuallyDrop::new(self);
360 let resource_handle_leases = unsafe { ptr::read(&this.resource_handle_leases).into_inner() };
363
364 Ok(ParkedGraph {
365 raw: this.raw,
366 device: this.device,
367 properties: this.properties,
368 vertex_dtype: this.vertex_dtype,
369 weight_dtype: this.weight_dtype,
370 has_weights: this.has_weights,
371 edge_id_dtype: this.edge_id_dtype,
372 park_event: event,
373 park_binding,
374 resource_handle_leases,
375 _not_sync: PhantomData,
376 _transposition: PhantomData,
377 })
378 }
379
380 fn validate_runtime_bound_transfer(&mut self, handle: &EnteredContext<'_>, operation: &'static str) -> Result<()> {
381 if handle.device() != self.device {
382 return Err(Error::new(ErrorKind::InvalidInput, operation)
383 .with_message("graph and resource handle must be on the same CUDA device")
384 .with_context("graph_device", self.device.to_string())
385 .with_context("handle_device", handle.device().to_string()));
386 }
387
388 let Some(handle_lease) = handle.runtime_bound_lease() else {
389 return Err(Error::new(ErrorKind::InvalidState, operation)
390 .with_message("parked graph transfer requires a runtime-bound resource handle"));
391 };
392 if self.untracked_resource_handle_use.get() {
393 return Err(Error::new(ErrorKind::InvalidState, operation)
394 .with_message("graph was used with a non-runtime-bound resource handle and cannot be parked"));
395 }
396
397 let leases = self.resource_handle_leases.get_mut();
398 if leases.is_empty() {
399 return Err(Error::new(ErrorKind::InvalidState, operation)
400 .with_message("graph has no runtime-bound handle leases to prove park stream ownership"));
401 }
402 if !leases.iter().all(|lease| lease.same_binding(&handle_lease)) {
403 return Err(Error::new(ErrorKind::InvalidState, operation)
404 .with_message("graph execution binding does not match the park handle")
405 .with_context("lease_count", leases.len().to_string()));
406 }
407 let entering = handle_lease.runtime().binding_id();
411 if self.ordered_binding.get().is_some_and(|ordered| ordered != entering) {
412 return Err(Error::new(ErrorKind::InvalidState, operation)
413 .with_message("park handle binding does not match the ordered graph execution binding")
414 .with_context(
415 "graph_binding_id",
416 self.ordered_binding.get().map(|id| id.to_string()).as_deref().unwrap_or("none"),
417 )
418 .with_context("handle_binding_id", entering.to_string()));
419 }
420
421 Ok(())
422 }
423
424 fn retain_runtime_bound_handle(&self, handle: &EnteredContext<'_>, operation: &'static str) -> Result<()> {
425 if handle.device() != self.device {
426 return Err(Error::new(ErrorKind::InvalidInput, operation)
427 .with_message("graph and resource handle must be on the same CUDA device")
428 .with_context("graph_device", self.device.to_string())
429 .with_context("handle_device", handle.device().to_string()));
430 }
431
432 let Some(lease) = handle.runtime_bound_lease() else {
433 self.untracked_resource_handle_use.set(true);
434 return Ok(());
435 };
436
437 let mut leases = self.resource_handle_leases.borrow_mut();
438
439 let entering = lease.runtime().binding_id();
444 match self.ordered_binding.get() {
445 Some(ordered) if ordered != entering => {
446 return Err(Error::new(ErrorKind::InvalidState, operation)
447 .with_message("graph execution binding does not match the resource handle")
448 .with_context("graph_binding_id", ordered.to_string())
449 .with_context("handle_binding_id", entering.to_string()));
450 },
451 None => self.ordered_binding.set(Some(entering)),
452 Some(_) => {},
453 }
454
455 let owns_teardown = leases.first().is_none_or(|existing| existing.same_binding(&lease));
462 if owns_teardown && !leases.iter().any(|existing| existing.same_native_owner(&lease)) {
463 leases.push(lease);
464 }
465
466 Ok(())
467 }
468
469 pub fn properties(&self) -> GraphProperties {
477 self.properties
478 }
479
480 pub fn edge_count(&mut self, handle: &EnteredContext<'_>) -> Result<usize> {
493 self.edge_count_impl(handle, false, "Graph::edge_count")
494 }
495
496 pub fn edge_count_checked(&mut self, handle: &EnteredContext<'_>, do_expensive_check: bool) -> Result<usize> {
501 self.edge_count_impl(handle, do_expensive_check, "Graph::edge_count_checked")
502 }
503
504 pub fn select_random_vertices<V: CugraphElement>(
509 &self,
510 handle: &EnteredContext<'_>,
511 rng_state: &mut RngState,
512 num_vertices: usize,
513 ) -> Result<DeviceArray<V>> {
514 select_random_vertices(handle, self, rng_state, num_vertices)
515 }
516
517 pub(crate) fn as_ptr(&self, handle: &EnteredContext<'_>, operation: &'static str) -> Result<*mut cugraph_graph_t> {
518 self.retain_runtime_bound_handle(handle, operation)?;
519 Ok(self.raw)
520 }
521
522 fn edge_count_impl(
523 &mut self,
524 handle: &EnteredContext<'_>,
525 do_expensive_check: bool,
526 operation: &'static str,
527 ) -> Result<usize> {
528 let mut edgelist: *mut cugraph_edgelist_t = ptr::null_mut();
529 let mut error: *mut cugraph_error_t = ptr::null_mut();
530
531 unsafe {
534 let code = cugraph_decompress_to_edgelist(
535 handle.handle_ptr(),
536 self.as_ptr(handle, operation)?,
537 to_bool_t(do_expensive_check),
538 &mut edgelist,
539 &mut error,
540 );
541 check_error(code, error, handle.device(), operation)?;
542 }
543
544 let edgelist =
547 unsafe { DecompressedEdgelist::from_raw(edgelist, handle.device(), handle.runtime_bound_lease()) }?;
548 edgelist.edge_count(handle, self.properties, operation)
549 }
550
551 pub(crate) fn validate_vertex_input_dtype(
552 &self,
553 dtype: DataType,
554 op: &'static str,
555 input_name: &'static str,
556 ) -> Result<()> {
557 if !dtype.is_supported_index_type() {
558 return Err(Error::new(ErrorKind::UnsupportedType, op)
559 .with_message(format!("{input_name} must use Int32 or Int64 vertex ids, got {dtype}")));
560 }
561
562 if dtype != self.vertex_dtype() {
563 return Err(Error::new(ErrorKind::UnsupportedType, op).with_message(format!(
564 "{input_name} dtype {dtype} does not match graph vertex dtype {}",
565 self.vertex_dtype()
566 )));
567 }
568
569 Ok(())
570 }
571
572 pub(crate) fn validate_vertex_selection<V: CugraphElement>(
573 &self,
574 vertices: &VertexSelection<'_, V>,
575 op: &'static str,
576 input_name: &'static str,
577 allow_all: bool,
578 ) -> Result<()> {
579 match vertices {
580 VertexSelection::All if allow_all => {},
581 VertexSelection::All => {
582 return Err(Error::new(ErrorKind::InvalidInput, op).with_message(format!(
583 "{input_name} must specify an explicit vertex subset; VertexSelection::All is not supported"
584 )));
585 },
586 VertexSelection::Device(_) | VertexSelection::Host(_) => {
587 self.validate_vertex_input_dtype(V::data_type(), op, input_name)?;
588 },
589 }
590 Ok(())
591 }
592
593 pub(crate) fn resolved_vertex_view<'a, V: CugraphElement>(
594 &self,
595 resolved: &'a ResolvedVertexSelection<'_, V>,
596 op: &'static str,
597 input_name: &'static str,
598 ) -> Result<DeviceArrayView<'a, V>> {
599 resolved.view()?.ok_or_else(|| {
600 Error::new(ErrorKind::InvalidInput, op).with_message(format!(
601 "{input_name} must specify an explicit vertex subset; VertexSelection::All is not supported"
602 ))
603 })
604 }
605}
606
607impl<T> std::fmt::Debug for Graph<T> {
608 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
609 f.debug_struct("Graph")
610 .field("properties", &self.properties)
611 .field("vertex_dtype", &self.vertex_dtype)
612 .field("weight_dtype", &self.weight_dtype)
613 .field("has_weights", &self.has_weights)
614 .finish_non_exhaustive()
615 }
616}
617
618fn free_graph_raw(raw: &mut *mut cugraph_graph_t, device: i32, operation: &'static str) {
619 drop_device_resource(raw, device, operation, DeviceResourceKind::Graph);
620}
621
622fn free_graph_raw_with_leases(
623 raw: &mut *mut cugraph_graph_t,
624 device: i32,
625 operation: &'static str,
626 leases: &[ResourceHandleLease],
627) {
628 if let Some(lease) = leases.first() {
629 lease.with_teardown_gate(|| free_graph_raw(raw, device, operation));
632 } else {
633 free_graph_raw(raw, device, operation);
634 }
635}
636
637impl<T> Drop for Graph<T> {
638 fn drop(&mut self) {
639 drop(self.pending_input_completion.take());
642 free_graph_raw_with_leases(&mut self.raw, self.device, "Graph::drop", self.resource_handle_leases.get_mut());
643 }
644}
645
646#[cfg(test)]
647mod tests;