Repository: valkey-io/valkey (C)

Issue:
[NEW] Implement SIMD instructions for NEON on hashtable
We have support for x86 SIMD [for the main hashtable](https://github.com/valkey-io/valkey/pull/2030). We ideally should implement similar support for arm NEON instructions. 


Discussion: Hi, I’d also be interested in contributing here. Happy to chip in with testing or implementation work alongside @SoftlyRaining 
I appreciate it! I'm quite new to SIMD - I understand the instruction and architecture concepts but I haven't programmed with SIMD before. I currently have a first draft of a NEON version that is only slightly slower than the scalar version. My uninformed suspicion is that there's too much overhead moving to/from NEON registers for this small batch size to be faster. 🤷🏻‍♀️
We've got a pretty good change now which combines ideas from both @ahmadbelb and me. 😁 I'm running tests and benchmarks and getting the PR ready.





Fix (patch):
diff --git a/src/bitops.c b/src/bitops.c
index a5ecdf9821..9d3b4f7945 100644
--- a/src/bitops.c
+++ b/src/bitops.c
@@ -35,7 +35,7 @@
 #define __MM_MALLOC_H
 #include <immintrin.h>
 #endif
-#if defined(__aarch64__)
+#if HAVE_ARM_NEON
 #include <arm_neon.h>
 #endif
 /* -----------------------------------------------------------------------------
@@ -190,7 +190,7 @@ long long popcountScalar(void *s, long count) {
     return bits;
 }
 
-#if defined(__aarch64__)
+#if HAVE_ARM_NEON
 #include <arm_neon.h>
 
 /*  SIMD version of popcount for ARM NEON.
diff --git a/src/config.h b/src/config.h
index 964d18d51a..de0c24dcc2 100644
--- a/src/config.h
+++ b/src/config.h
@@ -374,7 +374,7 @@ void setcpuaffinity(const char *cpulist);
 #define valkey_prefetch(addr) ((void)(addr))
 #endif
 
-/* Check if we can compile SIMD code */
+/* Check if we can compile x86 SIMD code */
 #if defined(__x86_64__) && ((defined(__GNUC__) && __GNUC__ >= 5) || (defined(__clang__) && __clang_major__ >= 4)) && defined(__has_attribute) && __has_attribute(target)
 #define HAVE_X86_SIMD 1
 #else
@@ -391,6 +391,13 @@ void setcpuaffinity(const char *cpulist);
 #define ATTRIBUTE_TARGET_AVX512
 #endif
 
+/* Check if we can compile ARM SIMD code */
+#if defined(__ARM_NEON) || defined(__ARM_NEON__)
+#define HAVE_ARM_NEON 1
+#else
+#define HAVE_ARM_NEON 0
+#endif
+
 #if defined(__linux__) && defined(__GLIBC__) && (defined(__GNUC__) && (__GNUC__ > 4) || defined(__clang__) && (__clang_major__) > 5)
 #define HAVE_IFUNC 1
 #else
diff --git a/src/hashtable.c b/src/hashtable.c
index bd3b54661a..f5f896823b 100644
--- a/src/hashtable.c
+++ b/src/hashtable.c
@@ -60,6 +60,9 @@
 #if HAVE_X86_SIMD
 #include <immintrin.h>
 #endif
+#if HAVE_ARM_NEON
+#include <arm_neon.h>
+#endif
 
 /* The default hashing function uses the SipHash implementation in siphash.c. */
 
@@ -719,6 +722,38 @@ static int findKeyInBucketSSE2(hashtable *ht, bucket *b, uint8_t h2, const void
 }
 #endif
 
+#if HAVE_ARM_NEON && ENTRIES_PER_BUCKET <= 8
+/* 32 bit architectures would require a different implementation
+ * consider that even if they had Neon SIMD support,
+ * they have 12 entries per bucket and only 32-bit scalar registers */
+
+static inline int popMatchBitmask(uint64_t *mask) {
+    /* one byte (8 bits) per item - either 0x80 or 0x00 */
+    int pos = __builtin_ctzll(*mask) >> 3;
+    *mask &= (*mask - 1); /* clear lowest set bit */
+    return pos;
+}
+
+static int findKeyInBucketNeon(hashtable *ht, bucket *b, uint8_t h2, const void *key, int table, int *pos_in_bucket, int *table_index) {
+    const uint8x8_t hash_vector = vld1_u8(b->hashes);             /* simple load */
+    const uint8x8_t h2_vector = vdup_n_u8(h2);                    /* duplicated into every byte */
+    const uint8x8_t equal_mask = vceq_u8(hash_vector, h2_vector); /* compare: 8 bits per item, 0xFF or 0x00 */
+    uint64_t matches = vget_lane_u64(vreinterpret_u64_u8(equal_mask), 0);
+
+    /* reduce each match to one bit and zero out invalid positions */
+    const uint64_t valid_entry_mask = (1ul << (ENTRIES_PER_BUCKET << 3)) - 1ul;
+    matches = matches & 0x8080808080808080ul & valid_entry_mask;
+
+    while (matches) {
+        int pos = popMatchBitmask(&matches);
+        if ((b->presence & (1 << pos)) &&
+            checkCandidateInBucket(ht, b, pos, key, table, pos_in_bucket, table_index))
+            return 1;
+    }
+    return 0; /* No match */
+}
+#endif
+
 /* Finds an entry matching the key. If a match is found, returns a pointer to
  * the bucket containing the matching entry and points 'pos_in_bucket' to the
  * index within the bucket. Returns NULL if no matching entry was found.
@@ -746,6 +781,8 @@ static bucket *findBucket(hashtable *ht, uint64_t hash, const void *key, int *po
 #if HAVE_X86_SIMD
             /* All x86-64 CPUs have SSE2. */
             if (findKeyInBucketSSE2(ht, b, h2, key, table, pos_in_bucket, table_index)) return b;
+#elif HAVE_ARM_NEON && ENTRIES_PER_BUCKET <= 8
+            if (findKeyInBucketNeon(ht, b, h2, key, table, pos_in_bucket, table_index)) return b;
 #else
             /* Find candidate entries with presence flag set and matching h2 hash. */
             for (int pos = 0; pos < numBucketPositions(b); pos++) {
diff --git a/src/hyperloglog.c b/src/hyperloglog.c
index 28c0153dc3..cd3e122cd9 100644
--- a/src/hyperloglog.c
+++ b/src/hyperloglog.c
@@ -47,7 +47,7 @@
 #include <immintrin.h>
 #endif
 
-#ifdef __aarch64__
+#if HAVE_ARM_NEON
 #include <arm_neon.h>
 #endif
 
@@ -224,7 +224,7 @@ struct hllhdr {
 
 static char *invalid_hll_err = "-INVALIDOBJ Corrupted HLL object detected";
 
-#if HAVE_X86_SIMD || defined(__aarch64__)
+#if HAVE_X86_SIMD || HAVE_ARM_NEON
 #define SIMD_SUPPORTED 1
 static int simd_enabled = 1;
 #else
@@ -237,7 +237,7 @@ static int simd_enabled = 1;
 #define HLL_USE_AVX2 0
 #endif
 
-#ifdef __aarch64__
+#if defined(__aarch64__) && HAVE_ARM_NEON
 #define HLL_USE_NEON (simd_enabled)
 #else
 #define HLL_USE_NEON 0
@@ -1235,7 +1235,7 @@ void hllMergeDenseAVX2(uint8_t *reg_raw, const uint8_t *reg_dense) {
 }
 #endif
 
-#if defined(__aarch64__)
+#if HAVE_ARM_NEON
 /*
  * hllMergeDenseNEON is an ARM optimized version of hllMergeDense using NEON
  *
@@ -1322,7 +1322,7 @@ void hllMergeDenseNEON(uint8_t *reg_raw, const uint8_t *reg_dense) {
         }
     }
 }
-#endif // __aarch64__
+#endif /* HAVE_ARM_NEON */
 
 /* Merge dense-encoded registers to raw registers array. */
 void hllMergeDense(uint8_t *reg_raw, const uint8_t *reg_dense) {
@@ -1334,12 +1334,10 @@ void hllMergeDense(uint8_t *reg_raw, const uint8_t *reg_dense) {
         }
     }
 #endif
-#ifdef __aarch64__
-    if (HLL_REGISTERS == 16384 && HLL_BITS == 6) {
-        if (HLL_USE_NEON) {
-            hllMergeDenseNEON(reg_raw, reg_dense);
-            return;
-        }
+#if defined(__aarch64__) && HAVE_ARM_NEON && HLL_REGISTERS == 16384 && HLL_BITS == 6
+    if (HLL_USE_NEON) {
+        hllMergeDenseNEON(reg_raw, reg_dense);
+        return;
     }
 #endif
 
@@ -1508,7 +1506,7 @@ void hllDenseCompressAVX2(uint8_t *reg_dense, const uint8_t *reg_raw) {
 }
 #endif
 
-#if defined(__aarch64__)
+#if HAVE_ARM_NEON
 /*
  * hllDenseCompressNEON is ARM optimized version of hllDenseCompress using NEON.
  *
@@ -1574,26 +1572,22 @@ void hllDenseCompressNEON(uint8_t *reg_dense, const uint8_t *reg_raw) {
         HLL_DENSE_SET_REGISTER(reg_dense, i, reg_raw[i]);
     }
 }
-#endif // __aarch64__
+#endif /* HAVE_ARM_NEON */
 
 /* Compress raw registers to dense representation. */
 void hllDenseCompress(uint8_t *reg_dense, const uint8_t *reg_raw) {
-#if HAVE_X86_SIMD
-    if (HLL_REGISTERS == 16384 && HLL_BITS == 6) {
-        if (HLL_USE_AVX2) {
-            hllDenseCompressAVX2(reg_dense, reg_raw);
-            return;
-        }
+#if HAVE_X86_SIMD && HLL_REGISTERS == 16384 && HLL_BITS == 6
+    if (HLL_USE_AVX2) {
+        hllDenseCompressAVX2(reg_dense, reg_raw);
+        return;
     }
 
 #endif
 
-#ifdef __ARM_NEON
-    if (HLL_REGISTERS == 16384 && HLL_BITS == 6) {
-        if (HLL_USE_NEON) {
-            hllDenseCompressNEON(reg_dense, reg_raw);
-            return;
-        }
+#if HAVE_ARM_NEON && HLL_REGISTERS == 16384 && HLL_BITS == 6
+    if (HLL_USE_NEON) {
+        hllDenseCompressNEON(reg_dense, reg_raw);
+        return;
     }
 #endif