| // Copyright 2026 Google LLC |
| // SPDX-License-Identifier: Apache-2.0 |
| // |
| // Licensed under the Apache License, Version 2.0 (the "License"); |
| // you may not use this file except in compliance with the License. |
| // You may obtain a copy of the License at |
| // |
| // http://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, software |
| // distributed under the License is distributed on an "AS IS" BASIS, |
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| // See the License for the specific language governing permissions and |
| // limitations under the License. |
| |
| #include <stddef.h> |
| #include <stdint.h> |
| |
| // Per-target include guard |
| #if defined(HIGHWAY_HWY_CONTRIB_ALGO_COUNT_INL_H_) == \ |
| defined(HWY_TARGET_TOGGLE) // NOLINT |
| #ifdef HIGHWAY_HWY_CONTRIB_ALGO_COUNT_INL_H_ |
| #undef HIGHWAY_HWY_CONTRIB_ALGO_COUNT_INL_H_ |
| #else |
| #define HIGHWAY_HWY_CONTRIB_ALGO_COUNT_INL_H_ |
| #endif |
| |
| #include "third_party/highway/hwy/highway.h" |
| |
| HWY_BEFORE_NAMESPACE(); |
| namespace hwy { |
| namespace HWY_NAMESPACE { |
| |
| // Returns the number of elements in `in[0, count)` equal to `value`. |
| template <class D, typename T = TFromD<D>> |
| size_t Count(D d, T value, const T* HWY_RESTRICT in, size_t count) { |
| const size_t N = Lanes(d); |
| using V = Vec<D>; |
| const V broadcasted = Set(d, value); |
| const RebindToSigned<D> di; |
| using TI = TFromD<decltype(di)>; |
| using VI = Vec<decltype(di)>; |
| |
| size_t total = 0; |
| size_t i = 0; |
| |
| // Min 4 lanes needed for two pairwise widenings, 8->16->32 |
| if constexpr (sizeof(T) == 1 && HWY_MAX_LANES_D(D) >= 4) { |
| const VI k1 = Set(di, TI{1}); |
| const RebindToUnsigned<decltype(di)> du; |
| const RepartitionToWide<decltype(di)> di16; |
| const Repartition<int32_t, D> di32; |
| auto wide_sum = Zero(di32); |
| |
| if (count >= 4 * N && N >= 4) { |
| while (i <= count - 4 * N) { |
| VI acc0 = Zero(di); |
| VI acc1 = Zero(di); |
| VI acc2 = Zero(di); |
| VI acc3 = Zero(di); |
| const size_t cap = HWY_MIN(i + 128 * 4 * N, count); |
| |
| if constexpr (HWY_NATIVE_MASK) { |
| for (; i <= cap - 4 * N; i += 4 * N) { |
| const auto m0 = RebindMask(di, Eq(broadcasted, LoadU(d, in + i))); |
| const auto m1 = |
| RebindMask(di, Eq(broadcasted, LoadU(d, in + i + N))); |
| const auto m2 = |
| RebindMask(di, Eq(broadcasted, LoadU(d, in + i + 2 * N))); |
| const auto m3 = |
| RebindMask(di, Eq(broadcasted, LoadU(d, in + i + 3 * N))); |
| acc0 = MaskedAddOr(acc0, m0, acc0, k1); |
| acc1 = MaskedAddOr(acc1, m1, acc1, k1); |
| acc2 = MaskedAddOr(acc2, m2, acc2, k1); |
| acc3 = MaskedAddOr(acc3, m3, acc3, k1); |
| } |
| } else { |
| for (; i <= cap - 4 * N; i += 4 * N) { |
| const auto v0 = VecFromMask(d, Eq(broadcasted, LoadU(d, in + i))); |
| const auto v1 = |
| VecFromMask(d, Eq(broadcasted, LoadU(d, in + i + N))); |
| const auto v2 = |
| VecFromMask(d, Eq(broadcasted, LoadU(d, in + i + 2 * N))); |
| const auto v3 = |
| VecFromMask(d, Eq(broadcasted, LoadU(d, in + i + 3 * N))); |
| acc0 = Add(acc0, BitCast(di, v0)); |
| acc1 = Add(acc1, BitCast(di, v1)); |
| acc2 = Add(acc2, BitCast(di, v2)); |
| acc3 = Add(acc3, BitCast(di, v3)); |
| } |
| |
| acc0 = Neg(acc0); |
| acc1 = Neg(acc1); |
| acc2 = Neg(acc2); |
| acc3 = Neg(acc3); |
| } |
| |
| const auto w0 = SatWidenMulPairwiseAdd(di16, BitCast(du, acc0), k1); |
| const auto w1 = SatWidenMulPairwiseAdd(di16, BitCast(du, acc1), k1); |
| const auto w2 = SatWidenMulPairwiseAdd(di16, BitCast(du, acc2), k1); |
| const auto w3 = SatWidenMulPairwiseAdd(di16, BitCast(du, acc3), k1); |
| const auto sum16 = Add(Add(w0, w1), Add(w2, w3)); |
| wide_sum = SatWidenMulPairwiseAccumulate( |
| di32, sum16, Set(di16, int16_t{1}), wide_sum); |
| } |
| } |
| total += static_cast<size_t>(ReduceSum(di32, wide_sum)); |
| } else if constexpr (sizeof(T) == 2 && HWY_MAX_LANES_D(D) >= 2) { |
| // Min 2 lanes needed for pairwise widening, 16->32 |
| const Repartition<int32_t, D> di32; |
| auto wide_sum = Zero(di32); |
| |
| if (count >= 4 * N && N >= 2) { |
| while (i <= count - 4 * N) { |
| VI acc0 = Zero(di); |
| VI acc1 = Zero(di); |
| VI acc2 = Zero(di); |
| VI acc3 = Zero(di); |
| const size_t cap = HWY_MIN(i + 32768 * 4 * N, count); |
| |
| if constexpr (HWY_NATIVE_MASK) { |
| const auto k1 = Set(di, TI{1}); |
| for (; i <= cap - 4 * N; i += 4 * N) { |
| const auto m0 = RebindMask(di, Eq(broadcasted, LoadU(d, in + i))); |
| const auto m1 = |
| RebindMask(di, Eq(broadcasted, LoadU(d, in + i + N))); |
| const auto m2 = |
| RebindMask(di, Eq(broadcasted, LoadU(d, in + i + 2 * N))); |
| const auto m3 = |
| RebindMask(di, Eq(broadcasted, LoadU(d, in + i + 3 * N))); |
| acc0 = MaskedAddOr(acc0, m0, acc0, k1); |
| acc1 = MaskedAddOr(acc1, m1, acc1, k1); |
| acc2 = MaskedAddOr(acc2, m2, acc2, k1); |
| acc3 = MaskedAddOr(acc3, m3, acc3, k1); |
| } |
| } else { |
| for (; i <= cap - 4 * N; i += 4 * N) { |
| const auto v0 = VecFromMask(d, Eq(broadcasted, LoadU(d, in + i))); |
| const auto v1 = |
| VecFromMask(d, Eq(broadcasted, LoadU(d, in + i + N))); |
| const auto v2 = |
| VecFromMask(d, Eq(broadcasted, LoadU(d, in + i + 2 * N))); |
| const auto v3 = |
| VecFromMask(d, Eq(broadcasted, LoadU(d, in + i + 3 * N))); |
| acc0 = Add(acc0, BitCast(di, v0)); |
| acc1 = Add(acc1, BitCast(di, v1)); |
| acc2 = Add(acc2, BitCast(di, v2)); |
| acc3 = Add(acc3, BitCast(di, v3)); |
| } |
| } |
| const auto mul = Set(di, HWY_NATIVE_MASK ? TI{1} : TI{-1}); |
| wide_sum = SatWidenMulPairwiseAccumulate(di32, acc0, mul, wide_sum); |
| wide_sum = SatWidenMulPairwiseAccumulate(di32, acc1, mul, wide_sum); |
| wide_sum = SatWidenMulPairwiseAccumulate(di32, acc2, mul, wide_sum); |
| wide_sum = SatWidenMulPairwiseAccumulate(di32, acc3, mul, wide_sum); |
| } |
| } |
| total += static_cast<size_t>(ReduceSum(di32, wide_sum)); |
| } else { |
| // Lane type wide enough to accumulate directly |
| if (count >= 4 * N) { |
| VI acc0 = Zero(di); |
| VI acc1 = Zero(di); |
| VI acc2 = Zero(di); |
| VI acc3 = Zero(di); |
| |
| if constexpr (HWY_NATIVE_MASK) { |
| const auto k1 = Set(di, TI{1}); |
| for (; i <= count - 4 * N; i += 4 * N) { |
| const auto m0 = RebindMask(di, Eq(broadcasted, LoadU(d, in + i))); |
| const auto m1 = RebindMask(di, Eq(broadcasted, LoadU(d, in + i + N))); |
| const auto m2 = |
| RebindMask(di, Eq(broadcasted, LoadU(d, in + i + 2 * N))); |
| const auto m3 = |
| RebindMask(di, Eq(broadcasted, LoadU(d, in + i + 3 * N))); |
| acc0 = MaskedAddOr(acc0, m0, acc0, k1); |
| acc1 = MaskedAddOr(acc1, m1, acc1, k1); |
| acc2 = MaskedAddOr(acc2, m2, acc2, k1); |
| acc3 = MaskedAddOr(acc3, m3, acc3, k1); |
| } |
| acc0 = Add(Add(acc0, acc1), Add(acc2, acc3)); |
| } else { |
| for (; i <= count - 4 * N; i += 4 * N) { |
| const auto v0 = VecFromMask(d, Eq(broadcasted, LoadU(d, in + i))); |
| const auto v1 = VecFromMask(d, Eq(broadcasted, LoadU(d, in + i + N))); |
| const auto v2 = |
| VecFromMask(d, Eq(broadcasted, LoadU(d, in + i + 2 * N))); |
| const auto v3 = |
| VecFromMask(d, Eq(broadcasted, LoadU(d, in + i + 3 * N))); |
| acc0 = Add(acc0, BitCast(di, v0)); |
| acc1 = Add(acc1, BitCast(di, v1)); |
| acc2 = Add(acc2, BitCast(di, v2)); |
| acc3 = Add(acc3, BitCast(di, v3)); |
| } |
| acc0 = Neg(Add(Add(acc0, acc1), Add(acc2, acc3))); |
| } |
| total += static_cast<size_t>(ReduceSum(di, acc0)); |
| } |
| } |
| |
| if (count >= N) { |
| for (; i <= count - N; i += N) { |
| total += CountTrue(d, Eq(broadcasted, LoadU(d, in + i))); |
| } |
| } |
| |
| if (i != count) { |
| const size_t remaining = count - i; |
| HWY_DASSERT(0 != remaining && remaining < N); |
| const V v = LoadN(d, in + i, remaining); |
| total += CountTrue(d, And(Eq(broadcasted, v), FirstN(d, remaining))); |
| } |
| |
| return total; |
| } |
| |
| // Returns the number of elements in `in[0, count)` for which `func(d, vec)` |
| // returns true. |
| template <class D, class Func, typename T = TFromD<D>> |
| size_t CountIf(D d, const T* HWY_RESTRICT in, size_t count, const Func& func) { |
| const size_t N = Lanes(d); |
| using V = Vec<D>; |
| const RebindToSigned<D> di; |
| using TI = TFromD<decltype(di)>; |
| using VI = Vec<decltype(di)>; |
| using MI = Mask<decltype(di)>; |
| |
| size_t total = 0; |
| size_t i = 0; |
| |
| // Min 4 lanes needed for two pairwise widenings, 8->16->32 |
| if constexpr (sizeof(T) == 1 && HWY_MAX_LANES_D(D) >= 4) { |
| const VI k1 = Set(di, TI{1}); |
| const RebindToUnsigned<decltype(di)> du; |
| const RepartitionToWide<decltype(di)> di16; |
| const Repartition<int32_t, D> di32; |
| using VI16 = Vec<decltype(di16)>; |
| using VI32 = Vec<decltype(di32)>; |
| VI32 wide_sum = Zero(di32); |
| |
| if (count >= 4 * N && N >= 4) { |
| while (i <= count - 4 * N) { |
| VI acc0 = Zero(di); |
| VI acc1 = Zero(di); |
| VI acc2 = Zero(di); |
| VI acc3 = Zero(di); |
| const size_t cap = HWY_MIN(i + 128 * 4 * N, count); |
| |
| if constexpr (HWY_NATIVE_MASK) { |
| for (; i <= cap - 4 * N; i += 4 * N) { |
| const MI m0 = RebindMask(di, func(d, LoadU(d, in + i))); |
| const MI m1 = RebindMask(di, func(d, LoadU(d, in + i + N))); |
| const MI m2 = RebindMask(di, func(d, LoadU(d, in + i + 2 * N))); |
| const MI m3 = RebindMask(di, func(d, LoadU(d, in + i + 3 * N))); |
| acc0 = MaskedAddOr(acc0, m0, acc0, k1); |
| acc1 = MaskedAddOr(acc1, m1, acc1, k1); |
| acc2 = MaskedAddOr(acc2, m2, acc2, k1); |
| acc3 = MaskedAddOr(acc3, m3, acc3, k1); |
| } |
| } else { |
| for (; i <= cap - 4 * N; i += 4 * N) { |
| const V v0 = VecFromMask(d, func(d, LoadU(d, in + i))); |
| const V v1 = VecFromMask(d, func(d, LoadU(d, in + i + N))); |
| const V v2 = VecFromMask(d, func(d, LoadU(d, in + i + 2 * N))); |
| const V v3 = VecFromMask(d, func(d, LoadU(d, in + i + 3 * N))); |
| acc0 = Add(acc0, BitCast(di, v0)); |
| acc1 = Add(acc1, BitCast(di, v1)); |
| acc2 = Add(acc2, BitCast(di, v2)); |
| acc3 = Add(acc3, BitCast(di, v3)); |
| } |
| |
| acc0 = Neg(acc0); |
| acc1 = Neg(acc1); |
| acc2 = Neg(acc2); |
| acc3 = Neg(acc3); |
| } |
| |
| const VI16 w0 = SatWidenMulPairwiseAdd(di16, BitCast(du, acc0), k1); |
| const VI16 w1 = SatWidenMulPairwiseAdd(di16, BitCast(du, acc1), k1); |
| const VI16 w2 = SatWidenMulPairwiseAdd(di16, BitCast(du, acc2), k1); |
| const VI16 w3 = SatWidenMulPairwiseAdd(di16, BitCast(du, acc3), k1); |
| const VI16 sum16 = Add(Add(w0, w1), Add(w2, w3)); |
| wide_sum = SatWidenMulPairwiseAccumulate( |
| di32, sum16, Set(di16, int16_t{1}), wide_sum); |
| } |
| } |
| total += static_cast<size_t>(ReduceSum(di32, wide_sum)); |
| } else if constexpr (sizeof(T) == 2 && HWY_MAX_LANES_D(D) >= 2) { |
| // Min 2 lanes needed for pairwise widening, 16->32 |
| const Repartition<int32_t, D> di32; |
| using VI32 = Vec<decltype(di32)>; |
| VI32 wide_sum = Zero(di32); |
| |
| if (count >= 4 * N && N >= 2) { |
| while (i <= count - 4 * N) { |
| VI acc0 = Zero(di); |
| VI acc1 = Zero(di); |
| VI acc2 = Zero(di); |
| VI acc3 = Zero(di); |
| const size_t cap = HWY_MIN(i + 32768 * 4 * N, count); |
| |
| if constexpr (HWY_NATIVE_MASK) { |
| const VI k1 = Set(di, TI{1}); |
| for (; i <= cap - 4 * N; i += 4 * N) { |
| const MI m0 = RebindMask(di, func(d, LoadU(d, in + i))); |
| const MI m1 = RebindMask(di, func(d, LoadU(d, in + i + N))); |
| const MI m2 = RebindMask(di, func(d, LoadU(d, in + i + 2 * N))); |
| const MI m3 = RebindMask(di, func(d, LoadU(d, in + i + 3 * N))); |
| acc0 = MaskedAddOr(acc0, m0, acc0, k1); |
| acc1 = MaskedAddOr(acc1, m1, acc1, k1); |
| acc2 = MaskedAddOr(acc2, m2, acc2, k1); |
| acc3 = MaskedAddOr(acc3, m3, acc3, k1); |
| } |
| } else { |
| for (; i <= cap - 4 * N; i += 4 * N) { |
| const V v0 = VecFromMask(d, func(d, LoadU(d, in + i))); |
| const V v1 = VecFromMask(d, func(d, LoadU(d, in + i + N))); |
| const V v2 = VecFromMask(d, func(d, LoadU(d, in + i + 2 * N))); |
| const V v3 = VecFromMask(d, func(d, LoadU(d, in + i + 3 * N))); |
| acc0 = Add(acc0, BitCast(di, v0)); |
| acc1 = Add(acc1, BitCast(di, v1)); |
| acc2 = Add(acc2, BitCast(di, v2)); |
| acc3 = Add(acc3, BitCast(di, v3)); |
| } |
| } |
| const VI mul = Set(di, HWY_NATIVE_MASK ? TI{1} : TI{-1}); |
| wide_sum = SatWidenMulPairwiseAccumulate(di32, acc0, mul, wide_sum); |
| wide_sum = SatWidenMulPairwiseAccumulate(di32, acc1, mul, wide_sum); |
| wide_sum = SatWidenMulPairwiseAccumulate(di32, acc2, mul, wide_sum); |
| wide_sum = SatWidenMulPairwiseAccumulate(di32, acc3, mul, wide_sum); |
| } |
| } |
| total += static_cast<size_t>(ReduceSum(di32, wide_sum)); |
| } else { |
| // Lane type wide enough to accumulate directly |
| if (count >= 4 * N) { |
| VI acc0 = Zero(di); |
| VI acc1 = Zero(di); |
| VI acc2 = Zero(di); |
| VI acc3 = Zero(di); |
| |
| if constexpr (HWY_NATIVE_MASK) { |
| const VI k1 = Set(di, TI{1}); |
| for (; i <= count - 4 * N; i += 4 * N) { |
| const MI m0 = RebindMask(di, func(d, LoadU(d, in + i))); |
| const MI m1 = RebindMask(di, func(d, LoadU(d, in + i + N))); |
| const MI m2 = RebindMask(di, func(d, LoadU(d, in + i + 2 * N))); |
| const MI m3 = RebindMask(di, func(d, LoadU(d, in + i + 3 * N))); |
| acc0 = MaskedAddOr(acc0, m0, acc0, k1); |
| acc1 = MaskedAddOr(acc1, m1, acc1, k1); |
| acc2 = MaskedAddOr(acc2, m2, acc2, k1); |
| acc3 = MaskedAddOr(acc3, m3, acc3, k1); |
| } |
| acc0 = Add(Add(acc0, acc1), Add(acc2, acc3)); |
| } else { |
| for (; i <= count - 4 * N; i += 4 * N) { |
| const V v0 = VecFromMask(d, func(d, LoadU(d, in + i))); |
| const V v1 = VecFromMask(d, func(d, LoadU(d, in + i + N))); |
| const V v2 = VecFromMask(d, func(d, LoadU(d, in + i + 2 * N))); |
| const V v3 = VecFromMask(d, func(d, LoadU(d, in + i + 3 * N))); |
| acc0 = Add(acc0, BitCast(di, v0)); |
| acc1 = Add(acc1, BitCast(di, v1)); |
| acc2 = Add(acc2, BitCast(di, v2)); |
| acc3 = Add(acc3, BitCast(di, v3)); |
| } |
| acc0 = Neg(Add(Add(acc0, acc1), Add(acc2, acc3))); |
| } |
| total += static_cast<size_t>(ReduceSum(di, acc0)); |
| } |
| } |
| |
| if (count >= N) { |
| for (; i <= count - N; i += N) { |
| total += CountTrue(d, func(d, LoadU(d, in + i))); |
| } |
| } |
| |
| if (i != count) { |
| const size_t remaining = count - i; |
| HWY_DASSERT(0 != remaining && remaining < N); |
| const V v = LoadN(d, in + i, remaining); |
| total += CountTrue(d, And(func(d, v), FirstN(d, remaining))); |
| } |
| |
| return total; |
| } |
| |
| // NOLINTNEXTLINE(google-readability-namespace-comments) |
| } // namespace HWY_NAMESPACE |
| } // namespace hwy |
| HWY_AFTER_NAMESPACE(); |
| |
| #endif // HIGHWAY_HWY_CONTRIB_ALGO_COUNT_INL_H_ |