Merge pull request 'Support building on ARM with NEON histogram intrinsics' (#4) from weaselbot/weaseldb:weaselbot/issue-3 into main

Reviewed-on: weaselab/weaseldb#4
Reviewed-by: andrew <andrew@weaselab.dev>

src/cpu_work.cpp

@@ -55,7 +55,7 @@ asm(".text\n"
     "    b.ne .L_loop\n"    // Branch back if not zero
     ".L_end:\n"             // End
     "    ret\n"             // Return
-    ".size spend_cpu_cycles, spend_cpu_cycles\n");
+    ".size spend_cpu_cycles, .-spend_cpu_cycles\n");
 #endif
 
 #endif

src/metric.cpp

@@ -22,7 +22,11 @@
 #include <unordered_set>
 #include <vector>
 
+#if defined(__x86_64__) || defined(__amd64__) || defined(_M_X64)
 #include <immintrin.h>
+#elif defined(__aarch64__)
+#include <arm_neon.h>
+#endif
 #include <simdutf.h>
 
 #include "arena.hpp"
@@ -1398,8 +1402,10 @@ void Gauge::set(double x) {
 Histogram::Histogram() = default;
 
 // Vectorized histogram bucket updates with mutex protection for consistency
-// AVX-optimized implementation for high performance
+// 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
 update_histogram_buckets_simd(std::span<const double> thresholds,
                               std::span<uint64_t> counts, double x,
@@ -1439,6 +1445,55 @@ 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) ||            \
+    defined(__aarch64__)
+  update_histogram_buckets_simd(thresholds, counts, x, start_idx);
+#else
+  const size_t size = thresholds.size();
+  for (size_t i = start_idx; i < size; ++i) {
+    if (x <= thresholds[i]) {
+      counts[i]++;
+    }
+  }
+#endif
+}
 
 void Histogram::observe(double x) {
   assert(p->thresholds.size() == p->shared.bucket_counts.size());
@@ -1459,15 +1514,14 @@ void Histogram::observe(double x) {
     }
 
     // Update shared directly
-    update_histogram_buckets_simd(p->thresholds, p->shared.bucket_counts, x, 0);
+    update_histogram_buckets(p->thresholds, p->shared.bucket_counts, x, 0);
     p->shared.sum += x;
     p->shared.observations++;
 
     p->mutex.unlock();
   } else {
     // Slow path: accumulate in pending (lock-free)
-    update_histogram_buckets_simd(p->thresholds, p->pending.bucket_counts, x,
-                                  0);
+    update_histogram_buckets(p->thresholds, p->pending.bucket_counts, x, 0);
     p->pending.sum += x;
     p->pending.observations++;
   }