bitcount kernel

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

Author: 0ax1

vortex-cuda/benches/arrow_validity_cuda.rs

@@ -179,10 +179,56 @@ fn benchmark_arrow_validity_repack(c: &mut Criterion) {
     group.finish();
 }
 
+fn benchmark_arrow_validity_count_nulls(c: &mut Criterion) {
+    let mut group = c.benchmark_group("cuda");
+
+    for &(len, len_label) in bench_config::BENCH_SIZES {
+        group.throughput(Throughput::Bytes(
+            validity_bitmap_byte_len(len, ARROW_OFFSET) as u64,
+        ));
+        group.bench_with_input(
+            BenchmarkId::new("cuda/arrow_validity/count_nulls", len_label),
+            &len,
+            |b, &len| {
+                b.iter_custom(|iters| {
+                    let timed = TimedLaunchStrategy::default();
+                    let timer = timed.timer();
+
+                    let mut cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())
+                        .vortex_expect("failed to create execution context")
+                        .with_launch_strategy(Arc::new(timed));
+                    let source = BitBuffer::collect_bool(len + ARROW_OFFSET, |idx| {
+                        idx >= ARROW_OFFSET && idx % 3 != 0
+                    });
+                    let (_, _, input_buffer) = source.into_inner();
+                    let input_buffer =
+                        block_on(cuda_ctx.ensure_on_device(BufferHandle::new_host(input_buffer)))
+                            .vortex_expect("failed to copy validity input to device");
+
+                    for _ in 0..iters {
+                        let null_count = test_harness::count_arrow_validity_nulls(
+                            &input_buffer,
+                            len,
+                            ARROW_OFFSET,
+                            &mut cuda_ctx,
+                        )
+                        .vortex_expect("failed to count Arrow validity nulls");
+                        std::hint::black_box(null_count);
+                    }
+
+                    Duration::from_nanos(timer.load(Ordering::Relaxed))
+                });
+            },
+        );
+    }
+
+    group.finish();
+}
+
 criterion::criterion_group! {
     name = benches;
     config = bench_config::cuda_bench_config();
-    targets = benchmark_arrow_validity_repack, benchmark_arrow_validity_export
+    targets = benchmark_arrow_validity_repack, benchmark_arrow_validity_count_nulls, benchmark_arrow_validity_export
 }
 
 #[cuda_available]

vortex-cuda/kernels/src/arrow_validity.cu

@@ -97,6 +97,58 @@ __device__ void arrow_validity_repack_device(const uint8_t *const input,
     }
 }
 
+__device__ uint64_t warp_sum(uint64_t value) {
+    for (int offset = 16; offset > 0; offset >>= 1) {
+        value += __shfl_down_sync(0xffffffff, value, offset);
+    }
+    return value;
+}
+
+__device__ void arrow_validity_count_valid_device(const uint8_t *const input,
+                                                 uint64_t *const output,
+                                                 uint64_t len,
+                                                 uint64_t arrow_offset) {
+    __shared__ uint64_t warp_counts[32];
+
+    const uint32_t thread = threadIdx.x;
+    const uint64_t worker = blockIdx.x * blockDim.x + thread;
+    const uint64_t validity_bits = len + arrow_offset;
+    const uint64_t input_bytes = (validity_bits + 7) / 8;
+    const uint64_t stride = static_cast<uint64_t>(gridDim.x) * blockDim.x;
+
+    uint64_t valid_count = 0;
+    for (uint64_t byte_idx = worker; byte_idx < input_bytes; byte_idx += stride) {
+        const uint64_t byte_start = byte_idx * 8;
+        uint32_t mask = 0xff;
+        if (byte_start < arrow_offset) {
+            const uint64_t lead = arrow_offset - byte_start;
+            mask = lead >= 8 ? 0 : mask << lead;
+        }
+        const uint64_t remaining = validity_bits - byte_start;
+        if (remaining < 8) {
+            mask &= (uint32_t {1} << remaining) - 1;
+        }
+        valid_count += __popc(static_cast<uint32_t>(input[byte_idx]) & mask);
+    }
+
+    const uint32_t lane = thread & 31;
+    const uint32_t warp = thread >> 5;
+    valid_count = warp_sum(valid_count);
+    if (lane == 0) {
+        warp_counts[warp] = valid_count;
+    }
+    __syncthreads();
+
+    valid_count = thread < (blockDim.x + 31) / 32 ? warp_counts[lane] : 0;
+    if (warp == 0) {
+        valid_count = warp_sum(valid_count);
+        if (lane == 0) {
+            atomicAdd(reinterpret_cast<unsigned long long *>(output),
+                      static_cast<unsigned long long>(valid_count));
+        }
+    }
+}
+
 } // namespace
 
 // CUDA entry point for validity bitmap repacking used by Arrow Device export.
@@ -108,3 +160,11 @@ extern "C" __global__ void arrow_validity_repack(const uint8_t *const input,
                                                  uint64_t input_bytes) {
     arrow_validity_repack_device(input, output, len, input_offset, arrow_offset, input_bytes);
 }
+
+// Kernel entry point for counting valid bits in an Arrow validity bitmap.
+extern "C" __global__ void arrow_validity_count_valid(const uint8_t *const input,
+                                                      uint64_t *const output,
+                                                      uint64_t len,
+                                                      uint64_t arrow_offset) {
+    arrow_validity_count_valid_device(input, output, len, arrow_offset);
+}

vortex-cuda/src/arrow/canonical.rs

@@ -804,7 +804,9 @@ pub(super) async fn export_arrow_validity_buffer(
             } else {
                 repack_arrow_validity_buffer(&bitmap, meta.offset(), len, arrow_offset, ctx)?
             };
-            Ok((Some(bitmap), ARROW_UNKNOWN_NULL_COUNT))
+            // Keep nullable exports self-describing for consumers that require exact null counts.
+            let null_count = count_arrow_validity_nulls(&bitmap, len, arrow_offset, ctx)?;
+            Ok((Some(bitmap), null_count))
         }
     }
 }
@@ -817,9 +819,6 @@ fn validity_bitmap_byte_len(len: usize, arrow_offset: usize) -> VortexResult<usi
         .div_ceil(8))
 }
 
-/// Arrow C Data uses -1 when the null count has not been computed.
-const ARROW_UNKNOWN_NULL_COUNT: i64 = -1;
-
 /// Allocate a zeroed device buffer with cuDF-safe padding for Arrow validity masks.
 fn device_zeroed_byte_buffer(
     byte_len: usize,
@@ -833,6 +832,62 @@ fn device_zeroed_byte_buffer(
     Ok(BufferHandle::new_device(Arc::new(CudaDeviceBuffer::new(buffer))).slice(0..byte_len))
 }
 
+pub(super) fn count_arrow_validity_nulls(
+    bitmap: &BufferHandle,
+    len: usize,
+    arrow_offset: usize,
+    ctx: &mut CudaExecutionCtx,
+) -> VortexResult<i64> {
+    if len == 0 {
+        return Ok(0);
+    }
+
+    let expected_bytes = validity_bitmap_byte_len(len, arrow_offset)?;
+    vortex_ensure!(
+        bitmap.len() >= expected_bytes,
+        "Arrow validity bitmap has {} bytes, expected at least {expected_bytes}",
+        bitmap.len()
+    );
+
+    let mut count = ctx.device_alloc::<u64>(1)?;
+    ctx.stream()
+        .memset_zeros(&mut count)
+        .map_err(|err| vortex_err!("Failed to zero Arrow validity count buffer: {err}"))?;
+    let count = CudaDeviceBuffer::new(count);
+
+    let input_view = bitmap.cuda_view::<u8>()?;
+    let output_view = count.as_view::<u64>();
+    let len = u64::try_from(len)?;
+    let arrow_offset = u64::try_from(arrow_offset)?;
+
+    let kernel = ctx.load_function_with_suffixes("arrow_validity", &["count_valid"])?;
+    const COUNT_THREADS_PER_BLOCK: u32 = 256;
+    const MAX_COUNT_BLOCKS: u32 = 4096;
+    let num_blocks = u32::try_from(expected_bytes.div_ceil(COUNT_THREADS_PER_BLOCK as usize))?
+        .clamp(1, MAX_COUNT_BLOCKS);
+    let config = LaunchConfig {
+        grid_dim: (num_blocks, 1, 1),
+        block_dim: (COUNT_THREADS_PER_BLOCK, 1, 1),
+        shared_mem_bytes: 0,
+    };
+    ctx.launch_kernel_config(&kernel, config, expected_bytes, |args| {
+        args.arg(&input_view)
+            .arg(&output_view)
+            .arg(&len)
+            .arg(&arrow_offset);
+    })?;
+
+    let valid_count = ctx
+        .stream()
+        .clone_dtoh(&output_view)
+        .map_err(|err| vortex_err!("Failed to copy Arrow validity count to host: {err}"))?
+        .into_iter()
+        .next()
+        .ok_or_else(|| vortex_err!("Arrow validity count kernel returned no output"))?;
+
+    Ok(i64::try_from(len - valid_count)?)
+}
+
 /// Repack a validity bitmap into Arrow layout without copying bitmap bits back to the CPU.
 ///
 /// Vortex bitmaps may start at any bit offset. Arrow exposes only a byte-addressed validity buffer
@@ -1665,10 +1720,7 @@ mod tests {
             )
         );
         assert_eq!(exported.array.array.length, 4);
-        assert_eq!(
-            exported.array.array.null_count,
-            super::ARROW_UNKNOWN_NULL_COUNT
-        );
+        assert_eq!(exported.array.array.null_count, 1);
         assert_eq!(exported.array.array.n_buffers, 2);
         assert_eq!(exported.array.array.n_children, 0);
         assert!(!exported.array.array.dictionary.is_null());
@@ -1726,7 +1778,7 @@ mod tests {
         let children = unsafe { std::slice::from_raw_parts(exported.array.array.children, 1) };
         let dict_child = unsafe { &*children[0] };
         assert!(!dict_child.dictionary.is_null());
-        assert_eq!(dict_child.null_count, super::ARROW_UNKNOWN_NULL_COUNT);
+        assert_eq!(dict_child.null_count, 1);
 
         unsafe { release_exported_array(&raw mut exported.array.array) };
         Ok(())
@@ -1759,7 +1811,7 @@ mod tests {
             [0, 1, 0]
         );
         let dictionary = unsafe { &*exported.array.array.dictionary };
-        assert_eq!(dictionary.null_count, super::ARROW_UNKNOWN_NULL_COUNT);
+        assert_eq!(dictionary.null_count, 1);
         assert_eq!(private_data_buffer_i32_values(dictionary, 1)?.len(), 2);
 
         unsafe { release_exported_array(&raw mut exported.array.array) };
@@ -2036,7 +2088,7 @@ mod tests {
         assert_binary_layout(
             &exported.array.array,
             5,
-            super::ARROW_UNKNOWN_NULL_COUNT,
+            1,
             &[0, 0, 3, 3, 8, i32::try_from(8 + out_of_line.len())?],
             &[b"\x00\xff\xfe".as_slice(), b"short", out_of_line].concat(),
         )?;
@@ -2550,7 +2602,7 @@ mod tests {
             [0, 1, 0, 1, 2]
         );
         let dictionary = unsafe { &*elements.dictionary };
-        assert_eq!(dictionary.null_count, super::ARROW_UNKNOWN_NULL_COUNT);
+        assert_eq!(dictionary.null_count, 1);
 
         unsafe { release_exported_array(&raw mut exported.array.array) };
         Ok(())
@@ -2796,7 +2848,7 @@ mod tests {
         let mut exported = utf8.export_device_array_with_schema(&mut ctx).await?;
         let field = Field::try_from(&exported.schema)?;
         assert_eq!(field, Field::new("", DataType::Utf8View, true));
-        assert_varbinview_shape(&exported.array.array, 3, super::ARROW_UNKNOWN_NULL_COUNT)?;
+        assert_varbinview_shape(&exported.array.array, 3, 1)?;
         assert_eq!(exported.array.device_type, ARROW_DEVICE_CUDA);
 
         let private_data = unsafe { &*exported.array.array.private_data.cast::<PrivateData>() };
@@ -2824,7 +2876,7 @@ mod tests {
         assert_binary_layout(
             &exported.array.array,
             3,
-            super::ARROW_UNKNOWN_NULL_COUNT,
+            1,
             &[0, 0, 2, i32::try_from(2 + sliced_out_of_line.len())?],
             &[b"\x00\xff".as_slice(), sliced_out_of_line].concat(),
         )?;
@@ -2942,7 +2994,7 @@ mod tests {
         let mut primitive = assert_nullable_export(
             PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(),
             2,
-            super::ARROW_UNKNOWN_NULL_COUNT,
+            1,
             &mut ctx,
         )
         .await?;
@@ -2971,7 +3023,7 @@ mod tests {
                 .into_array()
                 .slice(1..4)?,
             2,
-            super::ARROW_UNKNOWN_NULL_COUNT,
+            1,
             &mut ctx,
         )
         .await?;
@@ -3020,7 +3072,7 @@ mod tests {
             )
             .into_array(),
             2,
-            super::ARROW_UNKNOWN_NULL_COUNT,
+            1,
             &mut ctx,
         )
         .await?;
@@ -3050,7 +3102,7 @@ mod tests {
             )
             .into_array(),
             2,
-            super::ARROW_UNKNOWN_NULL_COUNT,
+            1,
             &mut ctx,
         )
         .await?;
@@ -3073,7 +3125,7 @@ mod tests {
             ])
             .into_array(),
             4,
-            super::ARROW_UNKNOWN_NULL_COUNT,
+            1,
             &mut ctx,
         )
         .await?;
@@ -3097,7 +3149,7 @@ mod tests {
             )?
             .into_array(),
             1,
-            super::ARROW_UNKNOWN_NULL_COUNT,
+            1,
             &mut ctx,
         )
         .await?;
@@ -3106,10 +3158,7 @@ mod tests {
         Ok(())
     }
 
-    // A container's row validity may be a non-canonical bool array (e.g. dict-encoded, or
-    // produced by take/scan). The container export paths bypass execute_cuda, so
-    // export_arrow_validity_buffer must canonicalize the validity on the GPU itself instead of
-    // bailing.
+    // Non-canonical row validity should export as a device-resident bitmap.
     #[crate::test]
     async fn test_export_struct_non_canonical_validity() -> VortexResult<()> {
         let mut ctx = CudaSession::create_execution_ctx(&VortexSession::empty())
@@ -3131,7 +3180,7 @@ mod tests {
             )?
             .into_array(),
             1,
-            super::ARROW_UNKNOWN_NULL_COUNT,
+            1,
             &mut ctx,
         )
         .await?;

vortex-cuda/src/arrow/mod.rs

@@ -70,8 +70,19 @@ pub mod test_harness {
     use vortex::error::VortexResult;
 
     use crate::CudaExecutionCtx;
+    use crate::arrow::canonical::count_arrow_validity_nulls as count_arrow_validity_nulls_impl;
     use crate::arrow::canonical::repack_arrow_validity_buffer as repack_arrow_validity_buffer_impl;
 
+    /// Count null bits in an Arrow validity bitmap.
+    pub fn count_arrow_validity_nulls(
+        bitmap: &BufferHandle,
+        len: usize,
+        arrow_offset: usize,
+        ctx: &mut CudaExecutionCtx,
+    ) -> VortexResult<i64> {
+        count_arrow_validity_nulls_impl(bitmap, len, arrow_offset, ctx)
+    }
+
     /// Repack a validity bitmap into Arrow's byte-addressed bitmap layout on the active stream.
     pub fn repack_arrow_validity_buffer(
         input_buffer: &BufferHandle,