forked from weaselab/weaseldb
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:
+39
-3
@@ -24,6 +24,8 @@
|
|||||||
|
|
||||||
#if defined(__x86_64__) || defined(__amd64__) || defined(_M_X64)
|
#if defined(__x86_64__) || defined(__amd64__) || defined(_M_X64)
|
||||||
#include <immintrin.h>
|
#include <immintrin.h>
|
||||||
|
#elif defined(__aarch64__)
|
||||||
|
#include <arm_neon.h>
|
||||||
#endif
|
#endif
|
||||||
#include <simdutf.h>
|
#include <simdutf.h>
|
||||||
|
|
||||||
@@ -1400,8 +1402,8 @@ void Gauge::set(double x) {
|
|||||||
Histogram::Histogram() = default;
|
Histogram::Histogram() = default;
|
||||||
|
|
||||||
// Vectorized histogram bucket updates with mutex protection for consistency
|
// Vectorized histogram bucket updates with mutex protection for consistency
|
||||||
// AVX-optimized implementation for high performance on x86-64, scalar fallback
|
// AVX-optimized implementation for high performance on x86-64, NEON-optimized
|
||||||
// on other architectures (e.g., ARM64).
|
// implementation on ARM64, and a scalar fallback for other architectures.
|
||||||
|
|
||||||
#if defined(__x86_64__) || defined(__amd64__) || defined(_M_X64)
|
#if defined(__x86_64__) || defined(__amd64__) || defined(_M_X64)
|
||||||
__attribute__((target("avx"))) static void
|
__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
|
#endif
|
||||||
|
|
||||||
static void update_histogram_buckets(std::span<const double> thresholds,
|
static void update_histogram_buckets(std::span<const double> thresholds,
|
||||||
std::span<uint64_t> counts, double x,
|
std::span<uint64_t> counts, double x,
|
||||||
size_t start_idx) {
|
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);
|
update_histogram_buckets_simd(thresholds, counts, x, start_idx);
|
||||||
#else
|
#else
|
||||||
const size_t size = thresholds.size();
|
const size_t size = thresholds.size();
|
||||||
|
|||||||
Reference in New Issue
Block a user