diff --git a/crates/core_simd/src/simd/num/int.rs b/crates/core_simd/src/simd/num/int.rs index 5561fe20e61..3a51235ff95 100644 --- a/crates/core_simd/src/simd/num/int.rs +++ b/crates/core_simd/src/simd/num/int.rs @@ -219,6 +219,12 @@ pub trait SimdInt: Copy + Sealed { /// The least significant bit becomes the most significant bit, second least-significant bit becomes second most-significant bit, etc. fn reverse_bits(self) -> Self; + /// Returns the number of ones in the binary representation of each element. + fn count_ones(self) -> Self::Unsigned; + + /// Returns the number of zeros in the binary representation of each element. + fn count_zeros(self) -> Self::Unsigned; + /// Returns the number of leading zeros in the binary representation of each element. fn leading_zeros(self) -> Self::Unsigned; @@ -367,6 +373,16 @@ macro_rules! impl_trait { unsafe { core::intrinsics::simd::simd_bitreverse(self) } } + #[inline] + fn count_ones(self) -> Self::Unsigned { + self.cast::<$unsigned>().count_ones() + } + + #[inline] + fn count_zeros(self) -> Self::Unsigned { + self.cast::<$unsigned>().count_zeros() + } + #[inline] fn leading_zeros(self) -> Self::Unsigned { self.cast::<$unsigned>().leading_zeros() diff --git a/crates/core_simd/src/simd/num/uint.rs b/crates/core_simd/src/simd/num/uint.rs index 7cc1b5bf8b6..1ab2d8c7b73 100644 --- a/crates/core_simd/src/simd/num/uint.rs +++ b/crates/core_simd/src/simd/num/uint.rs @@ -101,6 +101,12 @@ pub trait SimdUint: Copy + Sealed { /// The least significant bit becomes the most significant bit, second least-significant bit becomes second most-significant bit, etc. fn reverse_bits(self) -> Self; + /// Returns the number of ones in the binary representation of each element. + fn count_ones(self) -> Self; + + /// Returns the number of zeros in the binary representation of each element. + fn count_zeros(self) -> Self; + /// Returns the number of leading zeros in the binary representation of each element. fn leading_zeros(self) -> Self; @@ -215,6 +221,17 @@ macro_rules! impl_trait { unsafe { core::intrinsics::simd::simd_bitreverse(self) } } + #[inline] + fn count_ones(self) -> Self { + // Safety: `self` is an integer vector + unsafe { core::intrinsics::simd::simd_ctpop(self) } + } + + #[inline] + fn count_zeros(self) -> Self { + (!self).count_ones() + } + #[inline] fn leading_zeros(self) -> Self { // Safety: `self` is an integer vector diff --git a/crates/core_simd/tests/ops_macros.rs b/crates/core_simd/tests/ops_macros.rs index 6503cc00053..6e64bfcb424 100644 --- a/crates/core_simd/tests/ops_macros.rs +++ b/crates/core_simd/tests/ops_macros.rs @@ -216,6 +216,22 @@ macro_rules! impl_common_integer_tests { ) } + fn count_ones() { + test_helpers::test_unary_elementwise( + &$vector::::count_ones, + &|x| x.count_ones() as _, + &|_| true, + ) + } + + fn count_zeros() { + test_helpers::test_unary_elementwise( + &$vector::::count_zeros, + &|x| x.count_zeros() as _, + &|_| true, + ) + } + fn leading_zeros() { test_helpers::test_unary_elementwise( &$vector::::leading_zeros,