diff --git a/vortex-cuda/kernels/src/dynamic_dispatch.cu b/vortex-cuda/kernels/src/dynamic_dispatch.cu index 3a7f28e1576..99050c0a022 100644 --- a/vortex-cuda/kernels/src/dynamic_dispatch.cu +++ b/vortex-cuda/kernels/src/dynamic_dispatch.cu @@ -170,7 +170,8 @@ __device__ inline void scalar_op(T *values, const struct ScalarOp &op, char *__r const T *dict = reinterpret_cast(smem + op.params.dict.values_smem_byte_offset); #pragma unroll for (uint32_t i = 0; i < N; ++i) { - values[i] = dict[static_cast(values[i])]; + const uint32_t code = static_cast(values[i]); + values[i] = code < op.params.dict.values_len ? dict[code] : T {}; } break; } diff --git a/vortex-cuda/kernels/src/dynamic_dispatch.h b/vortex-cuda/kernels/src/dynamic_dispatch.h index e3b1d0f8df3..4403921e2b7 100644 --- a/vortex-cuda/kernels/src/dynamic_dispatch.h +++ b/vortex-cuda/kernels/src/dynamic_dispatch.h @@ -197,6 +197,7 @@ union ScalarParams { /// element type. struct DictParams { uint32_t values_smem_byte_offset; // byte offset to decoded dict values in smem + uint32_t values_len; } dict; }; diff --git a/vortex-cuda/src/dynamic_dispatch/mod.rs b/vortex-cuda/src/dynamic_dispatch/mod.rs index ad9c02ba6a1..a0005b5565c 100644 --- a/vortex-cuda/src/dynamic_dispatch/mod.rs +++ b/vortex-cuda/src/dynamic_dispatch/mod.rs @@ -415,13 +415,14 @@ impl ScalarOp { /// Dictionary gather: use current value as index into decoded values /// in shared memory (populated by an earlier input stage). - pub fn dict(values_smem_byte_offset: u32, output_ptype: PTypeTag) -> Self { + pub fn dict(values_smem_byte_offset: u32, values_len: u32, output_ptype: PTypeTag) -> Self { Self { op_code: ScalarOp_ScalarOpCode_DICT, output_ptype, params: ScalarParams { dict: ScalarParams_DictParams { values_smem_byte_offset, + values_len, }, }, } @@ -691,7 +692,7 @@ mod tests { SourceOp::bitunpack(6, 0), &[ ScalarOp::frame_of_ref(42, PTypeTag_PTYPE_U32), - ScalarOp::dict(0, PTypeTag_PTYPE_U32), + ScalarOp::dict(0, 256, PTypeTag_PTYPE_U32), ], ), ], @@ -2519,20 +2520,38 @@ mod tests { Ok(()) } - /// Dict with nullable codes must fall back to Unfused (not fused). + /// Dict with nullable codes fuses. Null positions may contain arbitrary + /// physical codes, so exercise out-of-range values as well. #[crate::test] - fn test_dict_nullable_codes_rejected() -> VortexResult<()> { + async fn test_dict_nullable_codes_fuses() -> VortexResult<()> { use vortex::buffer::buffer; - let codes = PrimitiveArray::from_option_iter([Some(0u32), None, Some(1), None, Some(2)]); - let values = PrimitiveArray::new(buffer![10u32, 20, 30], NonNullable); + let mut cpu_ctx = array_session().create_execution_ctx(); + let mut cuda_ctx = CudaSession::create_execution_ctx(&cuda_session())?; + + let codes = PrimitiveArray::new( + buffer![0u32, u32::MAX, 1, 99, 2], + Validity::from_iter([true, false, true, false, true]), + ); + let values = PrimitiveArray::from_option_iter([Some(10u32), None, Some(30)]); let dict = DictArray::try_new(codes.into_array(), values.into_array())?; - let plan = DispatchPlan::new(&dict.into_array(), CudaDispatchMode::Auto)?; + let plan = DispatchPlan::new(&dict.clone().into_array(), CudaDispatchMode::Auto)?; assert!( - matches!(plan, DispatchPlan::Unfused), - "Dict with nullable codes should fall back to Unfused" + matches!(plan, DispatchPlan::Fused(..)), + "Dict with nullable codes should fuse" ); + + let gpu = dict + .clone() + .into_array() + .execute_cuda(&mut cuda_ctx) + .await? + .into_host() + .await? + .into_array(); + + assert_arrays_eq!(dict, gpu, &mut cpu_ctx); Ok(()) } diff --git a/vortex-cuda/src/dynamic_dispatch/plan_builder.rs b/vortex-cuda/src/dynamic_dispatch/plan_builder.rs index 43f6a5a3712..4743ac3c451 100644 --- a/vortex-cuda/src/dynamic_dispatch/plan_builder.rs +++ b/vortex-cuda/src/dynamic_dispatch/plan_builder.rs @@ -93,11 +93,6 @@ fn is_dyn_dispatch_compatible(array: &ArrayRef) -> bool { } if id == Dict.id() { let arr = array.as_::(); - // Nullable codes could hold garbage values at null positions, causing - // out-of-bounds shared memory reads in the DICT gather scalar op. - if arr.codes().dtype().is_nullable() { - return false; - } // Dict codes and values may have different byte widths. // The kernel handles mixed widths via widening input stages, // but only when codes are no wider than values (the output type). @@ -303,9 +298,9 @@ impl DispatchPlan { /// - **F16 primitives** are not supported (no reinterpret path in the kernel). /// - **ALP** is supported for f32 and f64 only (including patches). /// - **BitPacked** with patches is supported. - /// - **Dict** with nullable codes is rejected (garbage at null positions - /// could OOB the DICT gather). Dict with codes wider than values is - /// also rejected (load would truncate code indices). + /// - **Dict** with codes wider than values is rejected (load would truncate + /// code indices). Nullable codes are safe because the gather bounds-checks + /// their unspecified physical values and validity is propagated separately. /// - **RunEnd** with nullable ends is rejected (garbage values break the /// binary search). RunEnd with ends wider than values is also rejected. /// - Validity is propagated from the root array to the output. @@ -757,8 +752,6 @@ impl FusedPlan { /// Cases that require a separate kernel dispatch: /// /// - **F16 primitives** — no reinterpret path in the kernel. - /// - **Dict with nullable codes** — garbage at null positions could OOB - /// the DICT gather in shared memory. /// - **Dict with codes wider than values** — `load_element()` would /// truncate the code indices. /// - **RunEnd with nullable ends** — garbage values break the binary @@ -806,7 +799,11 @@ impl FusedPlan { // DICT scalar op: pass byte offset directly (C ABI uses byte offsets). // output_ptype is the values' ptype — DICT transforms codes → values. pipeline.scalar_ops.push(( - ScalarOp::dict(values_smem_byte_offset, ptype_to_tag(values_ptype)), + ScalarOp::dict( + values_smem_byte_offset, + values_len, + ptype_to_tag(values_ptype), + ), None, )); Ok(pipeline)