Skip to content

perf: enable POPCNT in the BMI2, AVX2 and AVX-512 kernels #10456

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
robert3005 merged 1 commit into from
Oct 12, 2026
Merged

perf: enable POPCNT in the BMI2, AVX2 and AVX-512 kernels #10456

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions vortex-array/src/arrays/filter/execute/simd_compress/tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ fn engages_on_supported_cpus() {
let mask = mask_values(&mask).expect("alternating mask is mixed");

#[cfg(target_arch = "x86_64")]
let expected = is_x86_feature_detected!("avx2");
let expected = is_x86_feature_detected!("avx2") && is_x86_feature_detected!("popcnt");
#[cfg(target_arch = "aarch64")]
let expected = true;

Expand All @@ -109,7 +109,7 @@ fn declines_sparse_and_short_masks() {
#[cfg(all(target_arch = "x86_64", not(miri)))]
#[test]
fn avx2_kernels_match_scalar() {
if !is_x86_feature_detected!("avx2") {
if !(is_x86_feature_detected!("avx2") && is_x86_feature_detected!("popcnt")) {
return;
}

Expand All @@ -126,15 +126,15 @@ fn avx2_kernels_match_scalar() {
);

let mut out = vec![T::default(); mask.true_count() + SLACK_BYTES / size_of::<T>()];
// SAFETY: AVX2 was detected above and the output has a vector of slack.
// SAFETY: AVX2 and POPCNT were detected above and the output has a vector of slack.
let written =
unsafe { kernel_out_of_place(values.as_ptr().cast(), out.as_mut_ptr().cast(), mask) };
assert_eq!(written, mask.true_count());
assert_eq!(&out[..written], expected.as_slice());

let mut compacted = values.to_vec();
let ptr = compacted.as_mut_ptr().cast::<u8>();
// SAFETY: AVX2 was detected above; in-place compaction stays within the slice.
// SAFETY: AVX2 and POPCNT were detected above; in-place compaction stays within the slice.
let written = unsafe { kernel_in_place(ptr.cast_const(), ptr, mask) };
assert_eq!(written, mask.true_count());
assert_eq!(&compacted[..written], expected.as_slice());
Expand Down
17 changes: 9 additions & 8 deletions vortex-array/src/arrays/filter/execute/simd_compress/x86.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,17 +65,18 @@ pub(super) fn select_kernel<T, const IN_PLACE: bool>(mask: &MaskValues) -> Optio
}

fn avx512f() -> bool {
is_x86_feature_detected!("avx512f")
is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("popcnt")
}

fn avx512_vbmi2() -> bool {
is_x86_feature_detected!("avx512f")
&& is_x86_feature_detected!("avx512bw")
&& is_x86_feature_detected!("avx512vbmi2")
&& is_x86_feature_detected!("popcnt")
}

fn avx2() -> bool {
is_x86_feature_detected!("avx2")
is_x86_feature_detected!("avx2") && is_x86_feature_detected!("popcnt")
}

/// Generate an AVX-512 kernel for one element width.
Expand Down Expand Up @@ -187,7 +188,7 @@ macro_rules! avx512_compress_kernel {

avx512_compress_kernel!(
compress_word_avx512_epi8, compress_avx512_epi8,
features: "avx512f,avx512bw,avx512vbmi2",
features: "avx512f,avx512bw,avx512vbmi2,popcnt",
elem: i8,
lanes: 64,
kmask: u64,
Expand All @@ -200,7 +201,7 @@ avx512_compress_kernel!(

avx512_compress_kernel!(
compress_word_avx512_epi16, compress_avx512_epi16,
features: "avx512f,avx512bw,avx512vbmi2",
features: "avx512f,avx512bw,avx512vbmi2,popcnt",
elem: i16,
lanes: 32,
kmask: u32,
Expand All @@ -213,7 +214,7 @@ avx512_compress_kernel!(

avx512_compress_kernel!(
compress_word_avx512_epi32, compress_avx512_epi32,
features: "avx512f",
features: "avx512f,popcnt",
elem: i32,
lanes: 16,
kmask: u16,
Expand All @@ -226,7 +227,7 @@ avx512_compress_kernel!(

avx512_compress_kernel!(
compress_word_avx512_epi64, compress_avx512_epi64,
features: "avx512f",
features: "avx512f,popcnt",
elem: i64,
lanes: 8,
kmask: u8,
Expand Down Expand Up @@ -329,7 +330,7 @@ macro_rules! avx2_compress_kernel {
clippy::cast_possible_truncation,
reason = "deliberate submask narrowing"
)]
#[target_feature(enable = "avx2")]
#[target_feature(enable = "avx2,popcnt")]
#[inline]
unsafe fn $word_fn<const IN_PLACE: bool>(
src: *const u8,
Expand Down Expand Up @@ -409,7 +410,7 @@ macro_rules! avx2_compress_kernel {
/// The CPU must support AVX2 and the pointer contract of
/// [`filter_slice_by_bitmap`](super::filter_slice_by_bitmap) /
/// [`filter_slice_mut_by_bitmap`](super::filter_slice_mut_by_bitmap) must hold.
#[target_feature(enable = "avx2")]
#[target_feature(enable = "avx2,popcnt")]
pub(super) unsafe fn $walk_fn<const IN_PLACE: bool>(
src: *const u8,
dst: *mut u8,
Expand Down
12 changes: 6 additions & 6 deletions vortex-mask/src/intersect_by_rank/bmi2.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,8 +45,8 @@ unsafe fn pdep_bmi2(source: u64, mask: u64) -> u64 {
/// kernel without the feature calls an out-of-line `pdep_bmi2` once per word.
///
/// # Safety
/// Requires BMI2.
#[target_feature(enable = "bmi2")]
/// Requires BMI2 and POPCNT.
#[target_feature(enable = "bmi2,popcnt")]
pub(super) unsafe fn intersect_bit_buffers_bmi2(
self_buffer: &BitBuffer,
mask_buffer: &BitBuffer,
Expand All @@ -58,8 +58,8 @@ pub(super) unsafe fn intersect_bit_buffers_bmi2(
/// See [`intersect_bit_buffers_bmi2`].
///
/// # Safety
/// Requires BMI2.
#[target_feature(enable = "bmi2")]
/// Requires BMI2 and POPCNT.
#[target_feature(enable = "bmi2,popcnt")]
pub(super) unsafe fn intersect_bit_buffer_by_rank_indices_bmi2(
self_buffer: &BitBuffer,
mask_indices: &[usize],
Expand All @@ -70,8 +70,8 @@ pub(super) unsafe fn intersect_bit_buffer_by_rank_indices_bmi2(
/// See [`intersect_bit_buffers_bmi2`].
///
/// # Safety
/// Requires BMI2.
#[target_feature(enable = "bmi2")]
/// Requires BMI2 and POPCNT.
#[target_feature(enable = "bmi2,popcnt")]
pub(super) unsafe fn intersect_mask_driven_bmi2<I>(
self_buffer: &BitBuffer,
mask_indices: I,
Expand Down
13 changes: 9 additions & 4 deletions vortex-mask/src/intersect_by_rank/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -443,7 +443,9 @@ type IntersectRankIndices = unsafe fn(&BitBuffer, &[usize]) -> Mask;
fn select_bit_buffers<P: DepositBits>() -> IntersectBuffers {
#[cfg(target_arch = "x86_64")]
{
if std::arch::is_x86_feature_detected!("bmi2") {
if std::arch::is_x86_feature_detected!("bmi2")
&& std::arch::is_x86_feature_detected!("popcnt")
{
return bmi2::intersect_bit_buffers_bmi2;
}
}
Expand All @@ -460,7 +462,9 @@ fn select_bit_buffers<P: DepositBits>() -> IntersectBuffers {
fn select_rank_indices<P: DepositBits>() -> IntersectRankIndices {
#[cfg(target_arch = "x86_64")]
{
if std::arch::is_x86_feature_detected!("bmi2") {
if std::arch::is_x86_feature_detected!("bmi2")
&& std::arch::is_x86_feature_detected!("popcnt")
{
return bmi2::intersect_bit_buffer_by_rank_indices_bmi2;
}
}
Expand Down Expand Up @@ -523,8 +527,9 @@ where
I: Iterator<Item = usize>,
{
#[cfg(target_arch = "x86_64")]
if std::arch::is_x86_feature_detected!("bmi2") {
// SAFETY: BMI2 was just detected.
if std::arch::is_x86_feature_detected!("bmi2") && std::arch::is_x86_feature_detected!("popcnt")
{
// SAFETY: BMI2 and POPCNT were just detected.
return unsafe { bmi2::intersect_mask_driven_bmi2(self_buffer, mask_indices, true_count) };
}

Expand Down
Loading