diff --git a/vortex-cuda/benches/alp_cuda.rs b/vortex-cuda/benches/alp_cuda.rs index 11e9b894dd6..0f006018a58 100644 --- a/vortex-cuda/benches/alp_cuda.rs +++ b/vortex-cuda/benches/alp_cuda.rs @@ -31,6 +31,7 @@ use vortex::encodings::alp::ALPFloat; use vortex::encodings::alp::alp_encode; use vortex::error::VortexExpect; use vortex::session::VortexSession; +use vortex_cuda::CudaDispatchMode; use vortex_cuda::CudaSession; use vortex_cuda::executor::CudaArrayExt; use vortex_cuda_macros::cuda_available; @@ -109,6 +110,7 @@ where let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) .vortex_expect("failed to create execution context") + .with_dispatch_mode(CudaDispatchMode::StandaloneOnly) .with_launch_strategy(Arc::new(timed)); for _ in 0..iters { diff --git a/vortex-cuda/benches/bitpacked_cuda.rs b/vortex-cuda/benches/bitpacked_cuda.rs index 449061e6978..7568db82ec5 100644 --- a/vortex-cuda/benches/bitpacked_cuda.rs +++ b/vortex-cuda/benches/bitpacked_cuda.rs @@ -31,6 +31,7 @@ use vortex::encodings::fastlanes::BitPackedData; use vortex::encodings::fastlanes::unpack_iter::BitPacked; use vortex::error::VortexExpect; use vortex::session::VortexSession; +use vortex_cuda::CudaDispatchMode; use vortex_cuda::CudaSession; use vortex_cuda::executor::CudaArrayExt; use vortex_cuda_macros::cuda_available; @@ -127,6 +128,7 @@ where let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) .vortex_expect("failed to create execution context") + .with_dispatch_mode(CudaDispatchMode::StandaloneOnly) .with_launch_strategy(Arc::new(timed)); for _ in 0..iters { @@ -172,6 +174,7 @@ where let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) .vortex_expect("failed to create execution context") + .with_dispatch_mode(CudaDispatchMode::StandaloneOnly) .with_launch_strategy(Arc::new(timed)); for _ in 0..iters { diff --git a/vortex-cuda/benches/date_time_parts_cuda.rs b/vortex-cuda/benches/date_time_parts_cuda.rs index da696681418..59a8c8aa9df 100644 --- a/vortex-cuda/benches/date_time_parts_cuda.rs +++ b/vortex-cuda/benches/date_time_parts_cuda.rs @@ -30,6 +30,7 @@ use vortex::error::VortexExpect; use vortex::extension::datetime::TimeUnit; use vortex::extension::datetime::Timestamp; use vortex::session::VortexSession; +use vortex_cuda::CudaDispatchMode; use vortex_cuda::CudaSession; use vortex_cuda::executor::CudaArrayExt; use vortex_cuda_macros::cuda_available; @@ -68,6 +69,7 @@ fn benchmark_datetimeparts(c: &mut Criterion) { let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) .vortex_expect("failed to create execution context") + .with_dispatch_mode(CudaDispatchMode::StandaloneOnly) .with_launch_strategy(Arc::new(timed)); for _ in 0..iters { diff --git a/vortex-cuda/benches/dict_cuda.rs b/vortex-cuda/benches/dict_cuda.rs index 1e1a51d3a8a..4f1b72a49e3 100644 --- a/vortex-cuda/benches/dict_cuda.rs +++ b/vortex-cuda/benches/dict_cuda.rs @@ -26,6 +26,7 @@ use vortex::buffer::Buffer; use vortex::dtype::NativePType; use vortex::error::VortexExpect; use vortex::session::VortexSession; +use vortex_cuda::CudaDispatchMode; use vortex_cuda::CudaSession; use vortex_cuda::executor::CudaArrayExt; use vortex_cuda_macros::cuda_available; @@ -96,6 +97,7 @@ where let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) .vortex_expect("failed to create execution context") + .with_dispatch_mode(CudaDispatchMode::StandaloneOnly) .with_launch_strategy(Arc::new(timed)); for _ in 0..iters { diff --git a/vortex-cuda/benches/for_cuda.rs b/vortex-cuda/benches/for_cuda.rs index 29402737a52..c1c098f1f4b 100644 --- a/vortex-cuda/benches/for_cuda.rs +++ b/vortex-cuda/benches/for_cuda.rs @@ -33,6 +33,7 @@ use vortex::encodings::fastlanes::FoRArray; use vortex::error::VortexExpect; use vortex::scalar::Scalar; use vortex::session::VortexSession; +use vortex_cuda::CudaDispatchMode; use vortex_cuda::CudaSession; use vortex_cuda::executor::CudaArrayExt; use vortex_cuda_macros::cuda_available; @@ -91,6 +92,7 @@ where let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) .vortex_expect("failed to create execution context") + .with_dispatch_mode(CudaDispatchMode::StandaloneOnly) .with_launch_strategy(Arc::new(timed)); for _ in 0..iters { @@ -129,6 +131,7 @@ where let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty()) .vortex_expect("failed to create execution context") + .with_dispatch_mode(CudaDispatchMode::StandaloneOnly) .with_launch_strategy(Arc::new(timed)); for _ in 0..iters { diff --git a/vortex-cuda/src/executor.rs b/vortex-cuda/src/executor.rs index 8234bfb8efe..d490d3f03fe 100644 --- a/vortex-cuda/src/executor.rs +++ b/vortex-cuda/src/executor.rs @@ -57,6 +57,25 @@ impl CudaKernelEvents { } } +/// Controls which GPU dispatch strategy is used when executing arrays. +/// +/// By default, `execute_cuda` tries fused dynamic dispatch first, +/// then falls back to standalone per-encoding kernels. Benchmarks and tests +/// can force a specific strategy to get stable, isolated measurements. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum CudaDispatchMode { + /// Automatically choose the best strategy (current default behavior). + /// Tries fused → partially-fused → standalone → CPU fallback. + #[default] + Auto, + /// Only use standalone per-encoding `CudaExecute` kernels. + /// Skips the fused/partially-fused dynamic dispatch planner entirely. + StandaloneOnly, + /// Only use fused or partially-fused dynamic dispatch. + /// Returns an error if the array is not dyn-dispatch-compatible. + DynDispatchOnly, +} + /// CUDA execution context. /// /// Provides access to the CUDA context and stream for kernel execution. @@ -66,6 +85,7 @@ pub struct CudaExecutionCtx { ctx: ExecutionCtx, cuda_session: CudaSession, strategy: Arc, + dispatch_mode: CudaDispatchMode, } impl CudaExecutionCtx { @@ -77,6 +97,7 @@ impl CudaExecutionCtx { ctx, cuda_session, strategy: Arc::new(DefaultLaunchStrategy), + dispatch_mode: CudaDispatchMode::Auto, } } @@ -94,6 +115,15 @@ impl CudaExecutionCtx { self } + /// Set the dispatch mode for the execution context. + /// + /// This controls whether `execute_cuda` uses fused dynamic dispatch, + /// standalone per-encoding kernels, or the automatic (default) strategy. + pub fn with_dispatch_mode(mut self, dispatch_mode: CudaDispatchMode) -> Self { + self.dispatch_mode = dispatch_mode; + self + } + /// Perform an external kernel launch, with events created and logged via the configured /// [`LaunchStrategy`]. /// @@ -276,6 +306,11 @@ impl CudaExecutionCtx { self.ctx.session() } + /// Returns the current dispatch mode. + pub fn dispatch_mode(&self) -> CudaDispatchMode { + self.dispatch_mode + } + /// Returns a reference to the CUDA session. pub(crate) fn cuda_session(&self) -> &CudaSession { &self.cuda_session diff --git a/vortex-cuda/src/hybrid_dispatch/mod.rs b/vortex-cuda/src/hybrid_dispatch/mod.rs index 36636076da7..59ec18aa938 100644 --- a/vortex-cuda/src/hybrid_dispatch/mod.rs +++ b/vortex-cuda/src/hybrid_dispatch/mod.rs @@ -70,8 +70,21 @@ pub async fn try_gpu_dispatch( array: &ArrayRef, ctx: &mut CudaExecutionCtx, ) -> VortexResult { + use crate::executor::CudaDispatchMode; + trace!(encoding = %array.encoding_id(), dtype = ?array.dtype(), len = array.len(), "try GPU dispatch"); + // Short-circuit: standalone-only skips the plan builder entirely. + if ctx.dispatch_mode() == CudaDispatchMode::StandaloneOnly { + trace!(encoding = %array.encoding_id(), "standalone-only dispatch"); + return ctx + .cuda_session() + .kernel(&array.encoding_id()) + .ok_or_else(|| vortex_err!("No CUDA kernel for encoding {:?}", array.encoding_id()))? + .execute(array.clone(), ctx) + .await; + } + match DispatchPlan::new(array)? { DispatchPlan::Fused(plan) => { let materialized = plan.materialize(ctx).await?; @@ -99,6 +112,14 @@ pub async fn try_gpu_dispatch( materialized.execute(array.len(), ctx) } DispatchPlan::Unfused => { + // In dyn-dispatch-only mode, don't fall back to standalone kernels. + if ctx.dispatch_mode() == CudaDispatchMode::DynDispatchOnly { + return Err(vortex_err!( + "Array with encoding {:?} is not dyn-dispatch-compatible", + array.encoding_id() + )); + } + // Unfused kernel dispatch fallback. ctx.cuda_session() .kernel(&array.encoding_id()) diff --git a/vortex-cuda/src/lib.rs b/vortex-cuda/src/lib.rs index 7d9caf6aa99..bbebef9430d 100644 --- a/vortex-cuda/src/lib.rs +++ b/vortex-cuda/src/lib.rs @@ -27,6 +27,7 @@ pub use canonical::CanonicalCudaExt; pub use device_buffer::CudaBufferExt; pub use device_buffer::CudaDeviceBuffer; pub use device_read_at::CopyDeviceReadAt; +pub use executor::CudaDispatchMode; pub use executor::CudaExecutionCtx; pub use executor::CudaKernelEvents; use kernel::ALPExecutor;