Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions vortex-cuda/ffi/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -19,3 +19,7 @@ Use `vx_cuda_session_new` to initialize CUDA once and reuse it across exports.

Use `vx_cuda_array_sink_open_file` to open a standard Vortex file sink configured to produce CUDA-readable files.
Push host-resident arrays and close the sink using the standard `vx_array_sink_*` APIs.

Use `vx_cuda_scan_path_arrow_device_stream` to read such a local file through pinned host buffers
and receive an Arrow C Device stream. Reuse the same CUDA session across scans so the pinned buffer
pool and CUDA state are reused as well.
15 changes: 15 additions & 0 deletions vortex-cuda/ffi/cinclude/vortex_cuda.h
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,21 @@ vx_array_sink *vx_cuda_array_sink_open_file(const vx_session *session,
const vx_dtype *dtype,
vx_error **error_out);

/**
* Scan a local CUDA-compatible Vortex file as an Arrow C Device stream.
*
* Files written by `vx_cuda_array_sink_open_file` are compatible with this path. Reusing the same
* CUDA session across calls also reuses the pinned host buffers used to stage file reads.
*
* On success returns 0 and writes an owned `ArrowDeviceArrayStream` to `out_stream`. The caller
* must release the stream and each produced `ArrowDeviceArray` through their embedded Arrow
* release callbacks. On error returns 1 and writes a `vx_error` to `error_out` when non-NULL.
*/
int vx_cuda_scan_path_arrow_device_stream(const vx_session *session,
vx_view path,
struct ArrowDeviceArrayStream *out_stream,
vx_error **error_out);

/**
* Export a borrowed Vortex array for cuDF's Arrow Device import path.
*
Expand Down
58 changes: 58 additions & 0 deletions vortex-cuda/ffi/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,13 +13,19 @@ use std::ptr;
use std::sync::Arc;

use arrow_schema::ffi::FFI_ArrowSchema;
use vortex::array::stream::ArrayStreamExt;
use vortex::compressor::BtrBlocksCompressorBuilder;
use vortex::error::VortexResult;
use vortex::error::vortex_ensure;
use vortex::file::OpenOptionsSessionExt;
use vortex::file::WriteStrategyBuilder;
use vortex::io::VortexReadAt;
use vortex::io::runtime::BlockingRuntime;
use vortex::io::session::RuntimeSessionExt;
use vortex::session::SessionExt;
use vortex::session::VortexSession;
use vortex_cuda::CudaSession;
use vortex_cuda::PooledFileReadAt;
use vortex_cuda::arrow::ArrowDeviceArray;
use vortex_cuda::arrow::ArrowDeviceArrayStream;
use vortex_cuda::arrow::DeviceArrayExt;
Expand Down Expand Up @@ -104,6 +110,58 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_sink_open_file(
})
}

/// Scan a local Vortex file through pinned host buffers and export an Arrow C Device stream.
///
/// The file must use encodings and layouts supported by the CUDA execution path, such as files
/// written by [`vx_cuda_array_sink_open_file`]. Pinned staging buffers are reused across scans made
/// with the same CUDA session.
///
/// On success returns `0` and writes an owned [`ArrowDeviceArrayStream`] to `out_stream`. The
/// caller must release the stream and each array produced by it through their embedded Arrow
/// release callbacks.
///
/// On error returns `1` and, when `error_out` is non-null, writes a `vx_error` (free with
/// `vx_error_free`).
///
/// # Safety
///
/// `session` must be a valid borrowed handle created by `vortex-ffi`. `path` must be valid for the
/// duration of this call and contain UTF-8. `out_stream` must be a valid writable pointer. 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_scan_path_arrow_device_stream(
session: *const vx_session,
path: vx_view,
out_stream: *mut ArrowDeviceArrayStream,
error_out: *mut *mut vx_error,
) -> c_int {
try_or(error_out, VX_CUDA_ERR, || {
vortex_ensure!(!out_stream.is_null(), "null ArrowDeviceArrayStream output");

let path = unsafe { path.as_str() }?.to_owned();
let session = session_with_cuda(unsafe { vx_session_ref(session) }?)?;
let cuda_session = session.get::<CudaSession>();
let stream = cuda_session.stream()?;
let pool = Arc::clone(cuda_session.pinned_buffer_pool());
drop(cuda_session);

let reader: Arc<dyn VortexReadAt> = Arc::new(PooledFileReadAt::open(
path,
session.handle(),
pool,
stream,
)?);
let array_stream = ffi_runtime().block_on(async {
let file = session.open_options().open(reader).await?;
Ok::<_, vortex::error::VortexError>(file.scan()?.into_array_stream()?.boxed())
})?;
let device_stream = array_stream.export_device_array_stream(&session, ffi_runtime())?;

unsafe { ptr::write(out_stream, device_stream) };
Ok(VX_CUDA_OK)
})
}

/// 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
Expand Down
12 changes: 5 additions & 7 deletions vortex-cuda/gpu-scan-cli/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -29,13 +29,12 @@ use vortex::file::OpenOptionsSessionExt;
use vortex::file::WriteOptionsSessionExt;
use vortex::file::WriteStrategyBuilder;
use vortex::io::session::RuntimeSessionExt;
use vortex::session::SessionExt;
use vortex::session::VortexSession;
use vortex_cuda::CudaSession;
use vortex_cuda::PinnedByteBufferPool;
use vortex_cuda::PooledByteBufferReadAt;
use vortex_cuda::PooledFileReadAt;
use vortex_cuda::TracingLaunchStrategy;
use vortex_cuda::VortexCudaStreamPool;
use vortex_cuda::executor::CudaArrayExt;
use vortex_cuda::layout::CudaFlatLayoutStrategy;
use vortex_cuda::layout::register_cuda_layout;
Expand Down Expand Up @@ -155,11 +154,10 @@ async fn cmd_scan(path: PathBuf, gpu_file: bool, json_output: bool) -> VortexRes
let mut cuda_ctx = CudaSession::create_execution_ctx(&session)?
.with_launch_strategy(Arc::new(TracingLaunchStrategy));

let pool = Arc::new(PinnedByteBufferPool::new(Arc::clone(
cuda_ctx.stream().context(),
)));
let cuda_stream =
VortexCudaStreamPool::new(Arc::clone(cuda_ctx.stream().context()), 1).stream()?;
let cuda_session = session.get::<CudaSession>();
let pool = Arc::clone(cuda_session.pinned_buffer_pool());
let cuda_stream = cuda_session.stream()?;
drop(cuda_session);
let handle = session.handle();

let gpu_file_handle = if gpu_file {
Expand Down
9 changes: 9 additions & 0 deletions vortex-cuda/src/pinned.rs
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,15 @@ pub struct PinnedByteBufferPool {
puts: AtomicU64,
}

impl std::fmt::Debug for PinnedByteBufferPool {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PinnedByteBufferPool")
.field("max_keep_per_size", &self.max_keep_per_size)
.field("stats", &self.stats())
.finish_non_exhaustive()
}
}

struct InflightPinnedBuffer {
event: Arc<CudaEvent>,
buffer: PinnedByteBuffer,
Expand Down
9 changes: 9 additions & 0 deletions vortex-cuda/src/session.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ use crate::executor::CudaExecute;
pub use crate::executor::CudaExecutionCtx;
use crate::initialize_cuda;
use crate::kernel::KernelLoader;
use crate::pinned::PinnedByteBufferPool;
use crate::stream::VortexCudaStream;
use crate::stream_pool::VortexCudaStreamPool;

Expand All @@ -40,6 +41,7 @@ pub struct CudaSession {
export_device_array: Arc<dyn ExportDeviceArray>,
kernel_loader: Arc<KernelLoader>,
stream_pool: Arc<VortexCudaStreamPool>,
pinned_buffer_pool: Arc<PinnedByteBufferPool>,
}

impl CudaSession {
Expand All @@ -57,12 +59,14 @@ impl CudaSession {
Arc::clone(&context),
stream_pool_capacity,
));
let pinned_buffer_pool = Arc::new(PinnedByteBufferPool::new(Arc::clone(&context)));
Self {
context,
kernels: Arc::new(DashMap::default()),
kernel_loader: Arc::new(KernelLoader::new()),
export_device_array: Arc::new(CanonicalDeviceArrayExport),
stream_pool,
pinned_buffer_pool,
}
}

Expand Down Expand Up @@ -105,6 +109,11 @@ impl CudaSession {
self.stream_pool.stream()
}

/// Returns the session-scoped pool used for staging file reads in pinned host memory.
pub fn pinned_buffer_pool(&self) -> &Arc<PinnedByteBufferPool> {
&self.pinned_buffer_pool
}

/// Registers CUDA support for an array encoding.
///
/// # Arguments
Expand Down
4 changes: 2 additions & 2 deletions vortex-ffi/src/string.rs
Original file line number Diff line number Diff line change
Expand Up @@ -61,8 +61,8 @@ impl vx_view {
///
/// # Safety
///
/// Same requirements as in as_bytes
pub(crate) unsafe fn as_str<'a>(&self) -> VortexResult<&'a str> {
/// `self.ptr` must be valid for `self.len` reads, or null when `self.len` is zero.
pub unsafe fn as_str<'a>(&self) -> VortexResult<&'a str> {
str::from_utf8(unsafe { self.as_bytes() }?).map_err(|e| vortex_err!("invalid utf-8: {e}"))
}
}
Expand Down
Loading