tests

Author: 0ax1

vortex-cuda/src/kernel/filter.rs

@@ -18,6 +18,7 @@ use vortex_array::arrays::PrimitiveArray;
 use vortex_array::buffer::BufferHandle;
 use vortex_cub::filter::CubFilterable;
 use vortex_cub::filter::cudaStream_t;
+use vortex_cuda_macros::cuda_tests;
 use vortex_dtype::NativePType;
 use vortex_dtype::match_each_native_simd_ptype;
 use vortex_error::VortexResult;
@@ -167,93 +168,96 @@ where
     )))
 }
 
-#[cfg(test)]
+#[cuda_tests]
 mod tests {
-
-    use vortex_cuda_macros::cuda_tests;
-
-    #[cuda_tests]
-    mod cuda {
-        use super::*;
-
-        #[rstest]
-        #[case::i32_sparse(
+    use rstest::rstest;
+    use vortex_array::IntoArray;
+    use vortex_array::arrays::FilterArray;
+    use vortex_array::assert_arrays_eq;
+    use vortex_error::VortexExpect;
+    use vortex_session::VortexSession;
+
+    use super::*;
+    use crate::CanonicalCudaExt;
+    use crate::session::CudaSession;
+
+    #[rstest]
+    #[case::i32_sparse(
             PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6, 7, 8]),
             Mask::from_iter([true, false, true, false, true, false, true, false])
         )]
-        #[case::i32_dense(
+    #[case::i32_dense(
             PrimitiveArray::from_iter([10i32, 20, 30, 40, 50]),
             Mask::from_iter([true, true, true, false, true])
         )]
-        #[case::i64_large(
+    #[case::i64_large(
             PrimitiveArray::from_iter((0..1000i64).collect::<Vec<_>>()),
             Mask::from_iter((0..1000).map(|i| i % 3 == 0))
         )]
-        #[case::f64_values(
+    #[case::f64_values(
             PrimitiveArray::from_iter([1.1f64, 2.2, 3.3, 4.4, 5.5]),
             Mask::from_iter([false, true, false, true, false])
         )]
-        #[case::u8_all_true(
+    #[case::u8_all_true(
             PrimitiveArray::from_iter([1u8, 2, 3, 4, 5]),
             Mask::from_iter([true, true, true, true, true])
         )]
-        #[case::u32_all_false(
+    #[case::u32_all_false(
             PrimitiveArray::from_iter([1u32, 2, 3, 4, 5]),
             Mask::from_iter([false, false, false, false, false])
         )]
-        #[tokio::test]
-        async fn test_gpu_filter(
-            #[case] input: PrimitiveArray,
-            #[case] mask: Mask,
-        ) -> VortexResult<()> {
-            let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())
-                .vortex_expect("failed to create CUDA execution context");
+    #[tokio::test]
+    async fn test_gpu_filter(
+        #[case] input: PrimitiveArray,
+        #[case] mask: Mask,
+    ) -> VortexResult<()> {
+        let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())
+            .vortex_expect("failed to create CUDA execution context");
 
-            let filter_array = FilterArray::try_new(input.clone().into_array(), mask.clone())?;
+        let filter_array = FilterArray::try_new(input.clone().into_array(), mask.clone())?;
 
-            let cpu_result = filter_array.to_canonical()?.into_array();
+        let cpu_result = filter_array.to_canonical()?.into_array();
 
-            let gpu_result = FilterExecutor
-                .execute(filter_array.into_array(), &mut cuda_ctx)
-                .await
-                .vortex_expect("GPU filter failed")
-                .into_host()
-                .await?
-                .into_array();
+        let gpu_result = FilterExecutor
+            .execute(filter_array.into_array(), &mut cuda_ctx)
+            .await
+            .vortex_expect("GPU filter failed")
+            .into_host()
+            .await?
+            .into_array();
 
-            assert_arrays_eq!(cpu_result, gpu_result);
+        assert_arrays_eq!(cpu_result, gpu_result);
 
-            Ok(())
-        }
+        Ok(())
+    }
 
-        #[tokio::test]
-        async fn test_gpu_filter_large_array() -> VortexResult<()> {
-            let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())
-                .vortex_expect("failed to create CUDA execution context");
+    #[tokio::test]
+    async fn test_gpu_filter_large_array() -> VortexResult<()> {
+        let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())
+            .vortex_expect("failed to create CUDA execution context");
 
-            // Create a large array to test multi-block execution
-            let data: Vec<i32> = (0..100_000).collect();
-            let input = PrimitiveArray::from_iter(data);
+        // Create a large array to test multi-block execution
+        let data: Vec<i32> = (0..100_000).collect();
+        let input = PrimitiveArray::from_iter(data);
 
-            // Select every 7th element
-            let mask = Mask::from_iter((0..100_000).map(|i| i % 7 == 0));
+        // Select every 7th element
+        let mask = Mask::from_iter((0..100_000).map(|i| i % 7 == 0));
 
-            let filter_array = FilterArray::try_new(input.into_array(), mask)?;
+        let filter_array = FilterArray::try_new(input.into_array(), mask)?;
 
-            let cpu_result = filter_array.to_canonical()?.into_array();
+        let cpu_result = filter_array.to_canonical()?.into_array();
 
-            let gpu_result = FilterExecutor
-                .execute(filter_array.into_array(), &mut cuda_ctx)
-                .await
-                .vortex_expect("GPU filter failed")
-                .into_host()
-                .await?
-                .into_array();
+        let gpu_result = FilterExecutor
+            .execute(filter_array.into_array(), &mut cuda_ctx)
+            .await
+            .vortex_expect("GPU filter failed")
+            .into_host()
+            .await?
+            .into_array();
 
-            assert_eq!(cpu_result.len(), gpu_result.len());
-            assert_arrays_eq!(cpu_result, gpu_result);
+        assert_eq!(cpu_result.len(), gpu_result.len());
+        assert_arrays_eq!(cpu_result, gpu_result);
 
-            Ok(())
-        }
+        Ok(())
     }
 }