refactor(simd): share quantizer, reduction, rescale-search and dispatch boilerplate

include/rabitqlib/fastscan/detail/simd.hpp

@@ -58,7 +58,6 @@ void accumulate_avx2(
     size_t dim
 );
 void transfer_lut_hacc_generic(const uint16_t* lut, size_t dim, uint8_t* hc_lut);
-void transfer_lut_hacc_avx2(const uint16_t* lut, size_t dim, uint8_t* hc_lut);
 void accumulate_hacc_generic(
     const uint8_t* __restrict__ codes,
     const uint8_t* __restrict__ lut,

src/fastscan/fastscan_avx2.cpp

@@ -195,33 +195,6 @@ void accumulate_avx2(
     }
 }
 
-void transfer_lut_hacc_avx2(const uint16_t* lut, size_t dim, uint8_t* hc_lut) {
-    size_t num_codebook = dim >> 2;
-
-    for (size_t i = 0; i < num_codebook; i++) {
-        constexpr size_t kRegBits = 256;
-        constexpr size_t kLaneBits = 128;
-        constexpr size_t kByteBits = 8;
-
-        constexpr size_t kLutPerIter = kRegBits / kLaneBits;
-        constexpr size_t kCodePerIter = 2 * kRegBits / kByteBits;
-        constexpr size_t kCodePerLine = kLaneBits / kByteBits;
-
-        uint8_t* fill_lo =
-            hc_lut + (i / kLutPerIter * kCodePerIter) + ((i % kLutPerIter) * kCodePerLine);
-        uint8_t* fill_hi = fill_lo + (kRegBits / kByteBits);
-
-        for (size_t j = 0; j < 16; ++j) {
-            int tmp = lut[j];
-            uint8_t lo = static_cast<uint8_t>(tmp);
-            uint8_t hi = static_cast<uint8_t>(tmp >> 8);
-            fill_lo[j] = lo;
-            fill_hi[j] = hi;
-        }
-        lut += 16;
-    }
-}
-
 namespace {
 
 void accumulate_hacc_chunk_avx2(

src/fastscan/fastscan_dispatch.cpp

@@ -37,7 +37,8 @@ constexpr Backend kAvx2Backend{
     simd::pack_lut_avx2,
     simd::pack_lut_range_avx2,
     rabitqlib::simd::scalar_quantize_nonnegative_uint8_avx2,
-    simd::transfer_lut_hacc_avx2,
+    // AVX2 consumes the Pair32 layout the generic producer writes.
+    simd::transfer_lut_hacc_generic,
     simd::accumulate_avx2,
     simd::accumulate_hacc_avx2,
     rabitqlib::simd::qg_batch_estdist_avx2,

src/simd/dispatch.cpp

@@ -90,51 +90,36 @@ void pairwise_distances_lower(
 }
 
 float euclidean_sqr(const float* a, const float* b, size_t dim) {
-#if defined(__aarch64__)
-    static const auto kEuclideanSqrFn =
-        resolve_optional_kernel(euclidean_sqr_neon, euclidean_sqr_generic, cpu::has_neon());
-#else
-    static const auto kEuclideanSqrFn =
-        RABITQ_RESOLVE(euclidean_sqr_avx512, euclidean_sqr_avx2, euclidean_sqr_generic);
-#endif
+    static const auto kEuclideanSqrFn = RABITQ_RESOLVE_NEON(
+        euclidean_sqr_neon, euclidean_sqr_avx512, euclidean_sqr_avx2, euclidean_sqr_generic
+    );
 
     return kEuclideanSqrFn(a, b, dim);
 }
 
 float dot_product(const float* a, const float* b, size_t dim) {
-#if defined(__aarch64__)
-    static const auto kDotProductFn =
-        resolve_optional_kernel(dot_product_neon, dot_product_generic, cpu::has_neon());
-#else
-    static const auto kDotProductFn =
-        RABITQ_RESOLVE(dot_product_avx512, dot_product_avx2, dot_product_generic);
-#endif
+    static const auto kDotProductFn = RABITQ_RESOLVE_NEON(
+        dot_product_neon, dot_product_avx512, dot_product_avx2, dot_product_generic
+    );
 
     return kDotProductFn(a, b, dim);
 }
 
 float dot_product_dis(const float* a, const float* b, size_t dim) {
-#if defined(__aarch64__)
-    static const auto kDotProductDisFn = resolve_optional_kernel(
-        dot_product_dis_neon, dot_product_dis_generic, cpu::has_neon()
+    static const auto kDotProductDisFn = RABITQ_RESOLVE_NEON(
+        dot_product_dis_neon,
+        dot_product_dis_avx512,
+        dot_product_dis_avx2,
+        dot_product_dis_generic
     );
-#else
-    static const auto kDotProductDisFn = RABITQ_RESOLVE(
-        dot_product_dis_avx512, dot_product_dis_avx2, dot_product_dis_generic
-    );
-#endif
 
     return kDotProductDisFn(a, b, dim);
 }
 
 float l2norm_sqr(const float* a, size_t dim) {
-#if defined(__aarch64__)
-    static const auto kL2normSqrFn =
-        resolve_optional_kernel(l2norm_sqr_neon, l2norm_sqr_generic, cpu::has_neon());
-#else
-    static const auto kL2normSqrFn =
-        RABITQ_RESOLVE(l2norm_sqr_avx512, l2norm_sqr_avx2, l2norm_sqr_generic);
-#endif
+    static const auto kL2normSqrFn = RABITQ_RESOLVE_NEON(
+        l2norm_sqr_neon, l2norm_sqr_avx512, l2norm_sqr_avx2, l2norm_sqr_generic
+    );
     return kL2normSqrFn(a, dim);
 }
 
@@ -186,13 +171,9 @@ void fht_rotate(
     float scale_factor,
     const uint8_t* flip
 ) {
-#if defined(__aarch64__) || defined(_M_ARM64)
-    static const auto kFhtRotateFn =
-        resolve_optional_kernel(fht_rotate_neon, fht_rotate_generic, cpu::has_neon());
-#else
-    static const auto kFhtRotateFn =
-        RABITQ_RESOLVE(fht_rotate_avx512, fht_rotate_avx2, fht_rotate_generic);
-#endif
+    static const auto kFhtRotateFn = RABITQ_RESOLVE_NEON(
+        fht_rotate_neon, fht_rotate_avx512, fht_rotate_avx2, fht_rotate_generic
+    );
 
     kFhtRotateFn(data, rotated_data, dim, padded_dim, trunc_dim, scale_factor, flip);
 }
@@ -256,60 +237,42 @@ static ExcodeIpTable resolve_excode_ip_table() {
 }
 
 void flip_sign(const uint8_t* flip, float* data, size_t dim) {
-#if defined(__aarch64__) || defined(_M_ARM64)
-    static const auto kFlipSignFn =
-        resolve_optional_kernel(flip_sign_neon, flip_sign_generic, cpu::has_neon());
-#else
-    static const auto kFlipSignFn =
-        RABITQ_RESOLVE(flip_sign_avx512, flip_sign_avx2, flip_sign_generic);
-#endif
+    static const auto kFlipSignFn = RABITQ_RESOLVE_NEON(
+        flip_sign_neon, flip_sign_avx512, flip_sign_avx2, flip_sign_generic
+    );
 
     kFlipSignFn(flip, data, dim);
 }
 
 void kacs_walk(float* data, size_t len) {
-#if defined(__aarch64__) || defined(_M_ARM64)
-    static const auto kKacsWalkFn =
-        resolve_optional_kernel(kacs_walk_neon, kacs_walk_generic, cpu::has_neon());
-#else
-    static const auto kKacsWalkFn =
-        RABITQ_RESOLVE(kacs_walk_avx512, kacs_walk_avx2, kacs_walk_generic);
-#endif
+    static const auto kKacsWalkFn = RABITQ_RESOLVE_NEON(
+        kacs_walk_neon, kacs_walk_avx512, kacs_walk_avx2, kacs_walk_generic
+    );
     kKacsWalkFn(data, len);
 }
 
 void scalar_quantize_uint8(
     uint8_t* result, const float* vec0, size_t dim, float lo, float delta
 ) {
-#if defined(__aarch64__) || defined(_M_ARM64)
-    static const auto kScalarQuantizeUint8Fn = resolve_optional_kernel(
-        scalar_quantize_uint8_neon, scalar_quantize_uint8_generic, cpu::has_neon()
-    );
-#else
-    static const auto kScalarQuantizeUint8Fn = RABITQ_RESOLVE(
+    static const auto kScalarQuantizeUint8Fn = RABITQ_RESOLVE_NEON(
+        scalar_quantize_uint8_neon,
         scalar_quantize_uint8_avx512,
         scalar_quantize_uint8_avx2,
         scalar_quantize_uint8_generic
     );
-#endif
 
     kScalarQuantizeUint8Fn(result, vec0, dim, lo, delta);
 }
 
 void scalar_quantize_uint16(
     uint16_t* result, const float* vec0, size_t dim, float lo, float delta
 ) {
-#if defined(__aarch64__) || defined(_M_ARM64)
-    static const auto kScalarQuantizeUint16Fn = resolve_optional_kernel(
-        scalar_quantize_uint16_neon, scalar_quantize_uint16_generic, cpu::has_neon()
-    );
-#else
-    static const auto kScalarQuantizeUint16Fn = RABITQ_RESOLVE(
+    static const auto kScalarQuantizeUint16Fn = RABITQ_RESOLVE_NEON(
+        scalar_quantize_uint16_neon,
         scalar_quantize_uint16_avx512,
         scalar_quantize_uint16_avx2,
         scalar_quantize_uint16_generic
     );
-#endif
 
     kScalarQuantizeUint16Fn(result, vec0, dim, lo, delta);
 }
@@ -424,55 +387,36 @@ float excode_ipimpl::ip64_fxu7_avx(
 }
 
 void new_transpose_bin(const uint16_t* q, uint64_t* tq, size_t padded_dim, size_t b_query) {
-#if defined(__aarch64__) || defined(_M_ARM64)
-    static const auto kNewTransposeBinFn = simd::resolve_optional_kernel(
-        simd::new_transpose_bin_neon, simd::new_transpose_bin_generic, cpu::has_neon()
-    );
-#else
-    static const auto kNewTransposeBinFn = RABITQ_RESOLVE(
+    static const auto kNewTransposeBinFn = RABITQ_RESOLVE_NEON(
+        simd::new_transpose_bin_neon,
         simd::new_transpose_bin_avx512,
         simd::new_transpose_bin_avx2,
         simd::new_transpose_bin_generic
     );
-#endif
 
     kNewTransposeBinFn(q, tq, padded_dim, b_query);
 }
 
 void new_transpose_bin_512(
     const uint8_t* q, uint64_t* tq, size_t padded_dim, size_t b_query
 ) {
-#if defined(__aarch64__) || defined(_M_ARM64)
-    static const auto kNewTransposeBin512Fn = simd::resolve_optional_kernel(
+    static const auto kNewTransposeBin512Fn = RABITQ_RESOLVE_NEON(
         simd::new_transpose_bin_512_neon,
-        simd::new_transpose_bin_512_generic,
-        cpu::has_neon()
-    );
-#else
-    static const auto kNewTransposeBin512Fn = RABITQ_RESOLVE(
         simd::new_transpose_bin_512_avx512,
         simd::new_transpose_bin_512_avx2,
         simd::new_transpose_bin_512_generic
     );
-#endif
 
     kNewTransposeBin512Fn(q, tq, padded_dim, b_query);
 }
 
 float mask_ip_x0_q(const float* query, const uint8_t* data, size_t padded_dim) {
-#if defined(__aarch64__) || defined(_M_ARM64)
-    static const MaskIpX0QFn kMaskIpX0QFn = simd::resolve_optional_kernel(
+    static const MaskIpX0QFn kMaskIpX0QFn = RABITQ_RESOLVE_NEON(
         static_cast<MaskIpX0QFn>(simd::mask_ip_x0_q_neon),
-        static_cast<MaskIpX0QFn>(simd::mask_ip_x0_q_generic),
-        cpu::has_neon()
-    );
-#else
-    static const MaskIpX0QFn kMaskIpX0QFn = RABITQ_RESOLVE(
         static_cast<MaskIpX0QFn>(simd::mask_ip_x0_q_avx512),
         static_cast<MaskIpX0QFn>(simd::mask_ip_x0_q_avx2),
         static_cast<MaskIpX0QFn>(simd::mask_ip_x0_q_generic)
     );
-#endif
 
     return kMaskIpX0QFn(query, data, padded_dim);
 }
@@ -496,20 +440,13 @@ float warmup_ip_x0_q_512(
     size_t padded_dim,
     size_t b_query
 ) {
-#if defined(__aarch64__) || defined(_M_ARM64)
-    static const WarmupIpX0Q512Fn kWarmupIpX0Q512Fn = simd::resolve_optional_kernel(
+    static const WarmupIpX0Q512Fn kWarmupIpX0Q512Fn = RABITQ_RESOLVE_NEON(
         static_cast<WarmupIpX0Q512Fn>(simd::warmup_ip_x0_q_512_neon),
-        static_cast<WarmupIpX0Q512Fn>(simd::warmup_ip_x0_q_512_generic),
-        cpu::has_neon()
-    );
-#else
-    static const WarmupIpX0Q512Fn kWarmupIpX0Q512Fn = RABITQ_RESOLVE(
         static_cast<WarmupIpX0Q512Fn>(rabitqlib::simd::warmup_ip_x0_q_512_avx512),
         static_cast<WarmupIpX0Q512Fn>(rabitqlib::simd::warmup_ip_x0_q_512_avx2),
         static_cast<WarmupIpX0Q512Fn>(rabitqlib::simd::warmup_ip_x0_q_512_generic),
         cpu::has_avx512_popcnt()
     );
-#endif
 
     return kWarmupIpX0Q512Fn(data, query, delta, vl, padded_dim, b_query);
 }

src/simd/dispatch_resolver.hpp

@@ -48,3 +48,12 @@ Function resolve_optional_kernel(Function preferred, Function fallback, bool sup
 #else
 #define RABITQ_RESOLVE(avx512, avx2, fallback, ...) fallback
 #endif
+
+// NEON when available on ARM64, otherwise RABITQ_RESOLVE over the AVX kernels and an
+// optional AVX-512 predicate. Each target references only its own ISA's symbols.
+#if defined(__aarch64__) || defined(_M_ARM64)
+#define RABITQ_RESOLVE_NEON(neon, avx512, avx2, fallback, ...) \
+    rabitqlib::simd::resolve_optional_kernel(neon, fallback, rabitqlib::cpu::has_neon())
+#else
+#define RABITQ_RESOLVE_NEON(neon, ...) RABITQ_RESOLVE(__VA_ARGS__)
+#endif

src/simd/packed_tail_avx2.hpp

@@ -83,6 +83,13 @@ inline void accumulate_excode_tail_avx2(
     }
 }
 
+// Sum the eight lanes as a tree rather than serially.
+static inline float reduce_add_avx2(__m256 v) {
+    __m128 lanes = _mm_add_ps(_mm256_castps256_ps128(v), _mm256_extractf128_ps(v, 1));
+    lanes = _mm_add_ps(lanes, _mm_movehl_ps(lanes, lanes));
+    return _mm_cvtss_f32(_mm_add_ss(lanes, _mm_movehdup_ps(lanes)));
+}
+
 static inline float mask_ip_avx2(
     const float* query, const uint8_t* data, size_t padded_dim
 ) {
@@ -139,9 +146,8 @@ static inline float mask_ip_avx2(
         sum2 = accumulate(sum2, it_query + 16, bits, shifts2);
         sum3 = accumulate(sum3, it_query + 24, bits, shifts3);
     }
-    const __m256 sum = _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3));
-    __m128 lanes = _mm_add_ps(_mm256_castps256_ps128(sum), _mm256_extractf128_ps(sum, 1));
-    lanes = _mm_add_ps(lanes, _mm_movehl_ps(lanes, lanes));
-    return _mm_cvtss_f32(_mm_add_ss(lanes, _mm_movehdup_ps(lanes)));
+    return reduce_add_avx2(
+        _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3))
+    );
 }
 }  // namespace rabitqlib::simd::detail