Use AArch64 NEON intrinsics for histogram bucket updates

Replace the scalar ARM fallback in update_histogram_buckets with a NEON
implementation that processes two buckets per iteration, matching the
existing AVX path.  The wrapper now dispatches to the SIMD path on both
x86-64 and AArch64 and falls back to scalar code on other architectures.
This commit is contained in:
2026-06-26 11:08:03 -04:00
parent d7de96ef94
commit a377772e63
+39 -3
View File
@@ -24,6 +24,8 @@
#if defined(__x86_64__) || defined(__amd64__) || defined(_M_X64)
#include <immintrin.h>
#elif defined(__aarch64__)
#include <arm_neon.h>
#endif
#include <simdutf.h>
@@ -1400,8 +1402,8 @@ void Gauge::set(double x) {
Histogram::Histogram() = default;
// Vectorized histogram bucket updates with mutex protection for consistency
// AVX-optimized implementation for high performance on x86-64, scalar fallback
// on other architectures (e.g., ARM64).
// AVX-optimized implementation for high performance on x86-64, NEON-optimized
// implementation on ARM64, and a scalar fallback for other architectures.
#if defined(__x86_64__) || defined(__amd64__) || defined(_M_X64)
__attribute__((target("avx"))) static void
@@ -1443,12 +1445,46 @@ update_histogram_buckets_simd(std::span<const double> thresholds,
}
}
}
#elif defined(__aarch64__)
static void
update_histogram_buckets_simd(std::span<const double> thresholds,
std::span<uint64_t> counts, double x,
size_t start_idx) {
const size_t size = thresholds.size();
size_t i = start_idx;
// Process 2 buckets at a time with 128-bit NEON vectors
const float64x2_t x_vec = vdupq_n_f64(x);
const uint64x2_t one = vdupq_n_u64(1);
for (; i + 2 <= size; i += 2) {
// Compare x <= thresholds per lane; true lanes are all ones.
float64x2_t thresholds_vec = vld1q_f64(&thresholds[i]);
uint64x2_t cmp_result = vcleq_f64(x_vec, thresholds_vec);
// Convert all-ones/all-zeros masks to per-lane 1/0 increments.
uint64x2_t increments = vandq_u64(cmp_result, one);
// Load current counts, add increments, and store back.
uint64x2_t current_counts = vld1q_u64(&counts[i]);
uint64x2_t updated_counts = vaddq_u64(current_counts, increments);
vst1q_u64(&counts[i], updated_counts);
}
// Handle remainder with scalar operations
for (; i < size; ++i) {
if (x <= thresholds[i]) {
counts[i]++;
}
}
}
#endif
static void update_histogram_buckets(std::span<const double> thresholds,
std::span<uint64_t> counts, double x,
size_t start_idx) {
#if defined(__x86_64__) || defined(__amd64__) || defined(_M_X64)
#if defined(__x86_64__) || defined(__amd64__) || defined(_M_X64) || \
defined(__aarch64__)
update_histogram_buckets_simd(thresholds, counts, x, start_idx);
#else
const size_t size = thresholds.size();