fix: dyn dispatch requires same bit width types

Signed-off-by: Alexander Droste <alexander.droste@protonmail.com>

Author: 0ax1

vortex-cuda/src/dynamic_dispatch/mod.rs

@@ -803,6 +803,54 @@ mod tests {
         Ok(())
     }
 
+    /// Dict with mismatched code/value ptypes (u8 codes, u32 values) must be
+    /// rejected by the plan builder so the single-kernel `DictExecutor` handles
+    /// it instead. The dynamic dispatch kernel operates on a single type T, so
+    /// loading narrower codes as wider elements would read out of bounds.
+    #[crate::test]
+    fn test_dict_mismatched_ptypes_rejected() -> VortexResult<()> {
+        let dict_values: Vec<u32> = vec![100, 200, 300, 400];
+        let len = 3000;
+        let codes: Vec<u8> = (0..len).map(|i| (i % dict_values.len()) as u8).collect();
+
+        let codes_prim = PrimitiveArray::new(Buffer::from(codes), NonNullable);
+        let values_prim = PrimitiveArray::new(Buffer::from(dict_values), NonNullable);
+        let dict = DictArray::try_new(codes_prim.into_array(), values_prim.into_array())?;
+
+        let cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?;
+        // build_plan should fail because u8 codes != u32 values in byte width.
+        assert!(
+            build_plan(&dict.into_array(), &cuda_ctx).is_err(),
+            "Dict with mismatched code/value ptypes should not be fused"
+        );
+
+        Ok(())
+    }
+
+    /// RunEnd with mismatched ends/values ptypes (u64 ends, i32 values) must be
+    /// rejected by the plan builder so the single-kernel `RunEndExecutor`
+    /// handles it instead. The dynamic dispatch kernel loads both ends and
+    /// values into shared memory as the single uniform type T; mismatched byte
+    /// widths would misinterpret the raw data.
+    #[crate::test]
+    fn test_runend_mismatched_ptypes_rejected() -> VortexResult<()> {
+        let ends: Vec<u64> = vec![1000, 2000, 3000];
+        let values: Vec<i32> = vec![10, 20, 30];
+
+        let ends_arr = PrimitiveArray::new(Buffer::from(ends), NonNullable).into_array();
+        let values_arr = PrimitiveArray::new(Buffer::from(values), NonNullable).into_array();
+        let re = RunEndArray::new(ends_arr, values_arr);
+
+        let cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?;
+        // build_plan should fail because u64 ends != i32 values in byte width.
+        assert!(
+            build_plan(&re.into_array(), &cuda_ctx).is_err(),
+            "RunEnd with mismatched ends/values ptypes should not be fused"
+        );
+
+        Ok(())
+    }
+
     #[rstest]
     #[case(0, 1024)]
     #[case(0, 3000)]

vortex-cuda/src/dynamic_dispatch/plan_builder.rs

@@ -178,10 +178,36 @@ fn is_dyn_dispatch_compatible(array: &ArrayRef) -> bool {
         }
         return false;
     }
+    if id == Dict::ID {
+        if let Ok(a) = array.clone().try_into::<Dict>() {
+            // As of now the dict dyn dispatch kernel requires
+            // codes and values to have the same byte width.
+            return match (
+                PType::try_from(a.values().dtype()),
+                PType::try_from(a.codes().dtype()),
+            ) {
+                (Ok(values), Ok(codes)) => values.byte_width() == codes.byte_width(),
+                _ => false,
+            };
+        }
+        return false;
+    }
+    if id == RunEnd::ID {
+        if let Ok(a) = array.clone().try_into::<RunEnd>() {
+            // As of now the run-end dyn dispatch kernel requires
+            // ends and values to have the same byte width.
+            return match (
+                PType::try_from(a.ends().dtype()),
+                PType::try_from(a.values().dtype()),
+            ) {
+                (Ok(e), Ok(v)) => e.byte_width() == v.byte_width(),
+                _ => false,
+            };
+        }
+        return false;
+    }
     id == FoR::ID
         || id == ZigZag::ID
-        || id == Dict::ID
-        || id == RunEnd::ID
         || id == Primitive::ID
         || id == Slice::ID
         || id == Sequence::ID