From e2fa356adfc41a13ea18da839b53353b774739cc Mon Sep 17 00:00:00 2001
From: Nathan Bossart <nathandbossart@gmail.com>
Date: Sat, 20 Aug 2022 21:44:21 -0700
Subject: [PATCH 2/2] ARM SIMD support

---
 src/include/port/pg_lfind.h |  6 +++++-
 src/include/port/simd.h     | 28 ++++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/src/include/port/pg_lfind.h b/src/include/port/pg_lfind.h
index def858cbe1..04f09200b4 100644
--- a/src/include/port/pg_lfind.h
+++ b/src/include/port/pg_lfind.h
@@ -92,7 +92,7 @@ pg_lfind32(uint32 key, uint32 *base, uint32 nelem)
 	size_t		nelem_for_asserts = nelem;
 #endif
 
-#ifdef USE_SSE2
+#if defined(USE_SSE2) || defined(__ARM_NEON)
 	/*
 	 * A 16-byte register only has four 4-byte lanes. For better
 	 * instruction-level parallelism, each loop iteration operates on a block
@@ -126,7 +126,11 @@ pg_lfind32(uint32 key, uint32 *base, uint32 nelem)
 		result = vector32_vor(tmp1, tmp2);
 
 		/* see if there was a match */
+#ifdef USE_SSE2
 		if (_mm_movemask_epi8(result) != 0)
+#elif defined(__ARM_NEON)
+		if (vmaxvq_u32(result) != 0)
+#endif
 		{
 			Assert(lfind(&key, base, &nelem_for_asserts, sizeof(uint32),
 						 uint32_cmp_eq));
diff --git a/src/include/port/simd.h b/src/include/port/simd.h
index 4dda87f3dd..4f40b31e2a 100644
--- a/src/include/port/simd.h
+++ b/src/include/port/simd.h
@@ -30,6 +30,15 @@
 typedef __m128i Vector8;
 typedef __m128i Vector32;
 
+/*
+ * Include arm_neon.h if the compiler is targeting an architecture that
+ * supports ARM Advanced SIMD (Neon) intrinsics.
+ */
+#elif defined(__ARM_NEON)
+#include <arm_neon.h>
+typedef uint8x16_t Vector8;
+typedef uint32x4_t Vector32;
+
 /*
  * If no SIMD instructions are available, we emulate specialized vector
  * operations using uint64.
@@ -114,6 +123,8 @@ vector8_load(Vector8 *v, const uint8 *s)
 {
 #ifdef USE_SSE2
 	*v = _mm_loadu_si128((const __m128i *) s);
+#elif defined(__ARM_NEON)
+	*v = vld1q_u8(s);
 #else
 	memcpy(v, s, sizeof(Vector8));
 #endif
@@ -124,6 +135,8 @@ vector32_load(Vector32 *v, const uint32 *s)
 {
 #ifdef USE_SSE2
 	*v = _mm_loadu_si128((const __m128i *) s);
+#elif defined(__ARM_NEON)
+	*v = vld1q_u32(s);
 #else
 	elog(ERROR, "vector32() without SIMD not implemented");
 	pg_unreachable();
@@ -140,6 +153,8 @@ vector8_broadcast(const uint8 c)
 {
 #ifdef USE_SSE2
 	return _mm_set1_epi8(c);
+#elif defined(__ARM_NEON)
+	return vdupq_n_u8(c);
 #else
 	return ~UINT64CONST(0) / 0xFF * c;
 #endif
@@ -150,6 +165,8 @@ vector32_broadcast(const uint32 c)
 {
 #ifdef USE_SSE2
 	return _mm_set1_epi32(c);
+#elif defined(__ARM_NEON)
+	return vdupq_n_u32(c);
 #else
 	elog(ERROR, "vector32_broadcast() without SIMD not implemented");
 	pg_unreachable();
@@ -166,6 +183,8 @@ vector8_has_zero(const Vector8 v)
 {
 #ifdef USE_SSE2
 	return _mm_movemask_epi8(_mm_cmpeq_epi8(v, _mm_setzero_si128()));
+#elif defined(__ARM_NEON)
+	return vmaxvq_u8(vceqzq_u8(v));
 #else
 	return vector8_le(v, 0);
 #endif
@@ -178,6 +197,8 @@ vector8_eq(const Vector8 v, const uint8 c)
 
 #ifdef USE_SSE2
 	result = _mm_movemask_epi8(_mm_cmpeq_epi8(v, vector8_broadcast(c)));
+#elif defined(__ARM_NEON)
+	result = vmaxvq_u8(vceqq_u8(v, vector8_broadcast(c)));
 #else
 	/* any bytes in v equal to c will evaluate to zero via XOR */
 	result = vector8_has_zero(v ^ vector8_broadcast(c));
@@ -193,6 +214,8 @@ vector32_veq(const Vector32 v1, const Vector32 v2)
 {
 #ifdef USE_SSE2
 	return _mm_cmpeq_epi32(v1, v2);
+#elif defined(__ARM_NEON)
+	return vceqq_u32(v1, v2);
 #else
 	elog(ERROR, "vector32_veq() without SIMD not implemented");
 	pg_unreachable();
@@ -211,6 +234,9 @@ vector8_le(const Vector8 v, const uint8 c)
 	 */
 	__m128i sub = _mm_subs_epu8(v, vector8_broadcast(c));
 	result = vector8_has_zero(sub);
+#elif __ARM_NEON
+	uint8x16_t sub = vqsubq_u8(v, vector8_broadcast(c));
+	result = vector8_has_zero(sub);
 #else
 	/*
 	 * To find bytes <= c, we can use bitwise operations to find bytes < c + 1,
@@ -248,6 +274,8 @@ vector32_vor(const Vector32 v1, const Vector32 v2)
 {
 #ifdef USE_SSE2
 	return _mm_or_si128(v1, v2);
+#elif defined(__ARM_NEON)
+	return vorrq_u32(v1, v2);
 #else
 	elog(ERROR, "vector32_vor() without SIMD not implemented");
 	pg_unreachable();
-- 
2.25.1