diff --git a/vortex-cuda/ffi/README.md b/vortex-cuda/ffi/README.md index 6ccb5f738f3..3c030095374 100644 --- a/vortex-cuda/ffi/README.md +++ b/vortex-cuda/ffi/README.md @@ -14,3 +14,5 @@ Use this crate as the CUDA-enabled FFI artifact. Include both headers: ``` and link the CUDA FFI library (`vortex_cuda_ffi`). Do not pass Vortex handles between independently linked Rust FFI libraries. + +Use `vx_cuda_session_new` to initialize CUDA once and reuse it across exports. diff --git a/vortex-cuda/ffi/cinclude/vortex_cuda.h b/vortex-cuda/ffi/cinclude/vortex_cuda.h index 33e151e807d..ee0f0fb3ca3 100644 --- a/vortex-cuda/ffi/cinclude/vortex_cuda.h +++ b/vortex-cuda/ffi/cinclude/vortex_cuda.h @@ -13,7 +13,11 @@ extern "C" { #endif -#ifndef USE_OWN_ARROW_DEVICE +/* Definitions from the Arrow C Device data interface. Define USE_OWN_ARROW_DEVICE to skip them. + * https://arrow.apache.org/docs/format/CDeviceDataInterface.html */ +#if !defined(ARROW_C_DEVICE_DATA_INTERFACE) && !defined(USE_OWN_ARROW_DEVICE) +#define ARROW_C_DEVICE_DATA_INTERFACE + typedef int32_t ArrowDeviceType; #define ARROW_DEVICE_CPU 1 #define ARROW_DEVICE_CUDA 2 @@ -39,6 +43,14 @@ struct ArrowDeviceArray { }; #endif +/** + * Create a CUDA Vortex session. + * + * Repeated `vx_cuda_array_export_arrow_device` calls reuse this CUDA state. Returns an owned + * session handle, or NULL and an optional `vx_error` on failure. + */ +vx_session *vx_cuda_session_new(vx_error **error_out); + /** * Export a borrowed Vortex array for cuDF's Arrow Device import path. * @@ -48,6 +60,8 @@ struct ArrowDeviceArray { * * `out_array` is exported on `ARROW_DEVICE_CUDA`; struct arrays become table-shaped schemas, * non-struct arrays a single column field. + * + * Export is stream-ordered; `out_array->sync_event` is valid until `out_array` is released. */ int vx_cuda_array_export_arrow_device(const vx_session *session, const vx_array *array, diff --git a/vortex-cuda/ffi/src/lib.rs b/vortex-cuda/ffi/src/lib.rs index 83fae0768d2..2a74cb2666c 100644 --- a/vortex-cuda/ffi/src/lib.rs +++ b/vortex-cuda/ffi/src/lib.rs @@ -24,6 +24,7 @@ use vortex_ffi::vx_array; use vortex_ffi::vx_array_ref; use vortex_ffi::vx_error; use vortex_ffi::vx_session; +use vortex_ffi::vx_session_new_with; use vortex_ffi::vx_session_ref; const VX_CUDA_OK: c_int = 0; @@ -37,6 +38,26 @@ fn session_with_cuda(session: &VortexSession) -> VortexResult { Ok(session.clone().with_some(CudaSession::try_default()?)) } +/// Create a CUDA Vortex session. +/// +/// Repeated [`vx_cuda_array_export_arrow_device`] calls reuse this CUDA state. Returns an owned +/// session handle, or null and an optional `vx_error` on failure. +/// +/// # Safety +/// +/// If `error_out` is non-null, it must be valid for writing one error pointer. +#[unsafe(no_mangle)] +pub unsafe extern "C-unwind" fn vx_cuda_session_new( + error_out: *mut *mut vx_error, +) -> *mut vx_session { + try_or(error_out, ptr::null_mut(), || { + let cuda_session = CudaSession::try_default()?; + Ok(vx_session_new_with(|session| { + session.with_some(cuda_session) + })) + }) +} + /// Export a borrowed Vortex array for cuDF's Arrow Device import path. /// /// On success returns `0` and writes independently releasable `out_schema` and `out_array`; the @@ -47,6 +68,8 @@ fn session_with_cuda(session: &VortexSession) -> VortexResult { /// `out_array` is exported on `ARROW_DEVICE_CUDA`; struct arrays become table-shaped schemas, /// non-struct arrays a single column field. /// +/// Export is stream-ordered; `out_array->sync_event` is valid until `out_array` is released. +/// /// # Safety /// /// `session` and `array` must be valid borrowed handles created by `vortex-ffi`. `out_schema` @@ -235,6 +258,39 @@ mod tests { Ok(()) } + #[cuda_test] + fn test_cuda_session_new_export() { + let mut error = ptr::null_mut(); + let session = unsafe { vx_cuda_session_new(&raw mut error) }; + assert!(error.is_null()); + assert!(!session.is_null()); + + let array = test_array(PrimitiveArray::from_iter(0u32..5)); + let mut schema = FFI_ArrowSchema::empty(); + let mut device_array = empty_device_array(); + + let status = unsafe { + vx_cuda_array_export_arrow_device( + session, + array, + &raw mut schema, + &raw mut device_array, + &raw mut error, + ) + }; + assert_eq!(status, VX_CUDA_OK); + assert!(error.is_null()); + assert_eq!(device_array.array.length, 5); + assert_eq!(device_array.device_type, ARROW_DEVICE_CUDA); + + unsafe { + release_device_array(&mut device_array); + release_schema(&mut schema); + free_test_array(array); + vortex_ffi::vx_session_free(session); + } + } + #[cuda_not_available] #[test] fn test_export_reports_cuda_initialization_error() { diff --git a/vortex-ffi/src/lib.rs b/vortex-ffi/src/lib.rs index d0830b5ca90..f4ce314c073 100644 --- a/vortex-ffi/src/lib.rs +++ b/vortex-ffi/src/lib.rs @@ -36,6 +36,8 @@ pub use error::vx_error; pub use error::vx_error_free; pub use log::vx_log_level; pub use session::vx_session; +pub use session::vx_session_free; +pub use session::vx_session_new_with; pub use session::vx_session_ref; use vortex::dtype::FieldName; use vortex::error::VortexExpect; diff --git a/vortex-ffi/src/session.rs b/vortex-ffi/src/session.rs index 7e42b4e8c70..2af3d997444 100644 --- a/vortex-ffi/src/session.rs +++ b/vortex-ffi/src/session.rs @@ -17,12 +17,21 @@ box_wrapper!( vx_session ); +/// Create an FFI session from a configured default session. +pub fn vx_session_new_with( + configure: impl FnOnce(VortexSession) -> VortexSession, +) -> *mut vx_session { + vx_session::new(configure( + VortexSession::default().with_handle(RUNTIME.handle()), + )) +} + /// Create a new Vortex session. /// /// The caller is responsible for freeing the session with [`vx_session_free`]. #[unsafe(no_mangle)] pub unsafe extern "C-unwind" fn vx_session_new() -> *mut vx_session { - vx_session::new(VortexSession::default().with_handle(RUNTIME.handle())) + vx_session_new_with(|session| session) } /// Clone a Vortex session, returning an owned copy.