Introduce simd_any()/simd_eq_all(), optimize simd_hmax*() for Armv8 Neon

nskyav · milot-mirdita · commit 103fe7967a0a · 2025-06-09T14:17:30.000+09:00
diff --git a/lib/simd/simd.h b/lib/simd/simd.h
@@ -166,6 +166,7 @@ typedef __m512i simd_int;
 uint32_t simd_hmax32_sse(const __m128i buffer);
 uint16_t simd_hmax16_sse(const __m128i buffer);
 uint8_t simd_hmax8_sse(const __m128i buffer);
+bool simd_any_sse(const __m128i buffer);
 
 inline uint32_t simd_hmax32_avx(const __m256i buffer) {
     const __m128i abcd = _mm256_castsi256_si128(buffer);
@@ -191,6 +192,34 @@ inline uint8_t simd_hmax8_avx(const __m256i buffer) {
     return std::max(first, second);
 }
 
+inline bool simd_any_avx(const __m256i buffer) {
+#if defined(SIMDE_ARM_NEON_A64V8_NATIVE)
+    const __m128i lo = _mm256_castsi256_si128(buffer);
+    const __m128i hi = _mm256_extracti128_si256(buffer, 1);
+    const __m128i lo_or_hi = _mm_or_si128(lo, hi);
+    return simd_any_sse(lo_or_hi);
+#else
+    const __m256i vZero = _mm256_set1_epi32(0);
+    const __m256i vTemp = _mm256_cmpeq_epi8(buffer, vZero);
+    const uint32_t mask = _mm256_movemask_epi8(vTemp);
+    return (mask != 0xffffffff);
+#endif
+}
+
+inline bool simd_eq_all_avx(const __m256i a, const __m256i b) {
+    const __m256i vector_mask = _mm256_cmpeq_epi8(a, b);
+#if defined(SIMDE_ARM_NEON_A64V8_NATIVE)
+    const __m128i lo = _mm256_castsi256_si128(vector_mask);
+    const __m128i hi = _mm256_extracti128_si256(vector_mask, 1);
+    const __m128i lo_and_hi = _mm_and_si128(lo, hi);
+    const uint32_t min_dword = vminvq_u32(vreinterpretq_u32_s64(lo_and_hi));
+    return (min_dword == 0xffffffff);
+#else
+    const uint32_t bit_mask = _mm256_movemask_epi8(vector_mask);
+    return (bit_mask == 0xffffffff);
+#endif
+}
+
 inline float simdf32_hadd(const __m256 x) {
     /* ( x3+x7, x2+x6, x1+x5, x0+x4 ) */
     const __m128 x128 = _mm_add_ps(_mm256_extractf128_ps(x, 1), _mm256_castps256_ps128(x));
@@ -250,6 +279,8 @@ typedef __m256i simd_int;
 #define simdi16_hmax(x)     simd_hmax16_avx(x)
 #define simdui8_max(x,y)    _mm256_max_epu8(x,y)
 #define simdi8_hmax(x)      simd_hmax8_avx(x)
+#define simd_any(x)         simd_any_avx(x)
+#define simd_eq_all(x,y)    simd_eq_all_avx(x,y)
 #define simdi_load(x)       _mm256_load_si256(x)
 #define simdf_load(x)       _mm256_load_ps(x)
 #define simdui8_avg(x,y)    _mm256_avg_epu8(x,y)
@@ -352,24 +383,59 @@ typedef __m256 simd_float;
 #include <simde/x86/sse4.1.h>
 // see https://stackoverflow.com/questions/6996764/fastest-way-to-do-horizontal-sse-vector-sum-or-other-reduction
 inline uint32_t simd_hmax32_sse(const __m128i buffer) {
+#if defined(SIMDE_ARM_NEON_A64V8_NATIVE)
+    return vmaxvq_u32(vreinterpretq_u32_s64(buffer));
+#else
     __m128i hi64  = _mm_shuffle_epi32(buffer, _MM_SHUFFLE(1, 0, 3, 2));
     __m128i max64 = _mm_max_epi32(hi64, buffer);
     __m128i hi32  = _mm_shufflelo_epi16(max64, _MM_SHUFFLE(1, 0, 3, 2)); // Swap the low two elements
     __m128i max32 = _mm_max_epi32(max64, hi32);
     return _mm_cvtsi128_si32(max32); // SSE2 movd
+#endif
 }
 
 inline uint16_t simd_hmax16_sse(const __m128i buffer) {
+#if defined(SIMDE_ARM_NEON_A64V8_NATIVE)
+    return vmaxvq_u16(vreinterpretq_u16_s64(buffer));
+#else
     __m128i tmp1 = _mm_subs_epu16(_mm_set1_epi16((short)65535), buffer);
     __m128i tmp3 = _mm_minpos_epu16(tmp1);
     return (65535 - _mm_cvtsi128_si32(tmp3));
+#endif
 }
 
 inline uint8_t simd_hmax8_sse(const __m128i buffer) {
+#if defined(SIMDE_ARM_NEON_A64V8_NATIVE)
+    return vmaxvq_u8(vreinterpretq_u8_s64(buffer));
+#else
     __m128i tmp1 = _mm_subs_epu8(_mm_set1_epi8((char)255), buffer);
     __m128i tmp2 = _mm_min_epu8(tmp1, _mm_srli_epi16(tmp1, 8));
     __m128i tmp3 = _mm_minpos_epu16(tmp2);
     return (int8_t)(255 -(int8_t) _mm_cvtsi128_si32(tmp3));
+#endif
+}
+
+inline bool simd_any_sse(const __m128i buffer) {
+#if defined(SIMDE_ARM_NEON_A64V8_NATIVE)
+    const uint32_t non_zero = vmaxvq_u32(vreinterpretq_u32_s64(buffer));
+    return static_cast<bool>(non_zero);
+#else
+    const __m128i vZero = _mm_set1_epi32(0);
+    const __m128i vTemp = _mm_cmpeq_epi8(buffer, vZero);
+    const uint32_t mask = _mm_movemask_epi8(vTemp);
+    return (mask != 0xffff);
+#endif
+}
+
+inline bool simd_eq_all_sse(const __m128i a, const __m128i b) {
+    const __m128i vector_mask = _mm_cmpeq_epi8(a, b);
+#if defined(SIMDE_ARM_NEON_A64V8_NATIVE)
+    const uint32_t min_dword = vminvq_u32(vreinterpretq_u32_s64(vector_mask));
+    return (min_dword == 0xffffffff);
+#else
+    const uint32_t bit_mask = _mm_movemask_epi8(vector_mask);
+    return (bit_mask == 0xffff);
+#endif
 }
 
 // see https://stackoverflow.com/questions/6996764/fastest-way-to-do-horizontal-sse-vector-sum-or-other-reduction
@@ -471,6 +537,8 @@ typedef __m128i simd_int;
 #define simdi16_hmax(x)     simd_hmax16_sse(x)
 #define simdui8_max(x,y)    _mm_max_epu8(x,y)
 #define simdi8_hmax(x)      simd_hmax8_sse(x)
+#define simd_any(x)         simd_any_sse(x)
+#define simd_eq_all(x,y)    simd_eq_all_sse(x,y)
 #define simdui8_avg(x,y)    _mm_avg_epu8(x,y)
 #define simdui16_avg(x,y)   _mm_avg_epu16(x,y)
 #define simdi_load(x)       _mm_load_si128(x)
diff --git a/src/alignment/StripedSmithWaterman.cpp b/src/alignment/StripedSmithWaterman.cpp
@@ -657,7 +657,6 @@ std::pair<SmithWaterman::alignment_end, SmithWaterman::alignment_end> SmithWater
                                                          is set to 0, it will not be used */
 														   uint8_t bias,  /* Shift 0 point to a positive value. */
 														   int32_t maskLen) {
-#define max16(m, vm) ((m) = simdi8_hmax((vm)));
 
 	uint8_t max = 0;		                     /* the max alignment score */
 	int32_t end_query = query_length - 1;
@@ -835,9 +834,7 @@ std::pair<SmithWaterman::alignment_end, SmithWaterman::alignment_end> SmithWater
         }
 #endif
 		vTemp = simdui8_subs (vF, vTemp);
-		vTemp = simdi8_eq (vTemp, vZero);
-		uint32_t cmp = simdi8_movemask (vTemp);
-		while (cmp != SIMD_MOVEMASK_MAX) {
+		while (simd_any(vTemp)) {
 #ifdef GAP_POS_SCORING
             if (posSpecificGaps) {
                 vH = simdui8_max (vH, simdui8_subs(vF, simdi_load(gap_close_del + j)));
@@ -871,19 +868,13 @@ std::pair<SmithWaterman::alignment_end, SmithWaterman::alignment_end> SmithWater
             }
 #endif
 			vTemp = simdui8_subs (vF, vTemp);
-			vTemp = simdi8_eq (vTemp, vZero);
-			cmp  = simdi8_movemask (vTemp);
 		}
 
 		vMaxScore = simdui8_max(vMaxScore, vMaxColumn);
-		vTemp = simdi8_eq(vMaxMark, vMaxScore);
-		cmp = simdi8_movemask(vTemp);
-		if (cmp != SIMD_MOVEMASK_MAX) {
+		if (!simd_eq_all(vMaxMark, vMaxScore)) {
 			uint8_t temp;
 			vMaxMark = vMaxScore;
-			max16(temp, vMaxScore);
-			vMaxScore = vMaxMark;
-
+			temp = simdi8_hmax(vMaxScore);
 
 			if (LIKELY(temp > max)) {
 			    max = temp;
@@ -904,7 +895,7 @@ std::pair<SmithWaterman::alignment_end, SmithWaterman::alignment_end> SmithWater
         //fprintf(stderr, "\n");
 
 		/* Record the max score of current column. */
-		max16(maxColumn[i], vMaxColumn);
+		maxColumn[i] = simdi8_hmax(vMaxColumn);
 		//		fprintf(stderr, "maxColumn[%d]: %d\n", i, maxColumn[i]);
 		if (maxColumn[i] == terminate) break;
 	}
@@ -946,7 +937,6 @@ std::pair<SmithWaterman::alignment_end, SmithWaterman::alignment_end> SmithWater
 		}
 	}
 	return std::make_pair(best0, best1);
-#undef max16
 }
 
 template <const unsigned int type, const bool posSpecificGaps>
@@ -968,8 +958,6 @@ std::pair<SmithWaterman::alignment_end, SmithWaterman::alignment_end> SmithWater
                                                            const uint16_t bias,
                                                            int32_t maskLen) {
 
-#define max8(m, vm) ((m) = simdi16_hmax((vm)));
-
 	uint16_t max = 0;		                     /* the max alignment score */
 	int32_t end_read = query_length - 1;
 	int32_t end_ref = 0; /* 1_based best alignment ending point; Initialized as isn't aligned - 0. */
@@ -1003,7 +991,6 @@ std::pair<SmithWaterman::alignment_end, SmithWaterman::alignment_end> SmithWater
 	//simd_int vBias = simdi16_set(-bias);    // set as a negative value for simd use
 	simd_int vMaxScore = vZero; /* Trace the highest score of the whole SW matrix. */
 	simd_int vMaxMark = vZero; /* Trace the highest score till the previous column. */
-	simd_int vTemp;
 	int32_t edge, begin = 0, end = db_length, step = 1;
 
     //fprintf(stderr, "start alignment of length %d [%d]\n", query_length, segLen * SIMD_SIZE);
@@ -1098,6 +1085,7 @@ std::pair<SmithWaterman::alignment_end, SmithWaterman::alignment_end> SmithWater
 
 			/* Update vE value. */
 #ifdef GAP_POS_SCORING
+            simd_int vTemp;
             if (posSpecificGaps) {
                 // copy vH for update of vF
                 vTemp = vH;
@@ -1156,19 +1144,16 @@ std::pair<SmithWaterman::alignment_end, SmithWaterman::alignment_end> SmithWater
                 }
 #endif
 				vF = simdui16_subs(vF, vGapE);
-				if (UNLIKELY(! simdi8_movemask(simdi16_gt(vF, vH)))) goto end;
+				if (UNLIKELY(!simd_any(simdi16_gt(vF, vH)))) goto end;
 			}
 		}
 
 		end:
 		vMaxScore = simdi16_max(vMaxScore, vMaxColumn);
-		vTemp = simdi16_eq(vMaxMark, vMaxScore);
-		uint32_t cmp = simdi8_movemask(vTemp);
-		if (cmp != SIMD_MOVEMASK_MAX) {
+		if (!simd_eq_all(vMaxMark, vMaxScore)) {
 			uint16_t temp;
 			vMaxMark = vMaxScore;
-			max8(temp, vMaxScore);
-			vMaxScore = vMaxMark;
+			temp = simdi16_hmax(vMaxScore);
 
 			if (LIKELY(temp > max)) {
 				max = temp;
@@ -1184,7 +1169,7 @@ std::pair<SmithWaterman::alignment_end, SmithWaterman::alignment_end> SmithWater
         //fprintf(stderr, "\n");
 
 		/* Record the max score of current column. */
-		max8(maxColumn[i], vMaxColumn);
+		maxColumn[i] = simdi16_hmax(vMaxColumn);
 		if (maxColumn[i] == terminate) break;
 	}
 
@@ -1226,7 +1211,6 @@ std::pair<SmithWaterman::alignment_end, SmithWaterman::alignment_end> SmithWater
 	}
 
 	return std::make_pair(best0, best1);
-#undef max8
 }
 
 void SmithWaterman::ssw_init(const Sequence* q,
