Route Hamming code sizes to the fastest kernel, and accept ragged sizes

faiss/utils/hamming.h

@@ -111,7 +111,7 @@ FAISS_API extern size_t hamming_batch_size;
  *
  * @param  a             size na * nbytespercode
  * @param  b             size nb * nbytespercode
- * @param  nbytespercode should be multiple of 8
+ * @param  nbytespercode any size; multiples of 8 take a faster kernel
  * @param  dis           output distances, size na * nb
  */
 void hammings(

faiss/utils/hamming_distance/hamming_impl.h

@@ -298,6 +298,41 @@ void generalized_hammings_knn_hc_impl(
     }
 }
 
+void hammings_with_computer(
+        const uint8_t* __restrict a,
+        const uint8_t* __restrict b,
+        size_t na,
+        size_t nb,
+        size_t ncodes,
+        hamdis_t* __restrict dis) {
+    with_HammingComputer<THE_SIMD_LEVEL>(ncodes, [&]<class HammingComputer>() {
+        for (size_t i = 0; i < na; i++) {
+            HammingComputer hc(a + i * ncodes, ncodes);
+            const uint8_t* bi = b;
+            for (size_t j = 0; j < nb; j++) {
+                dis[i * nb + j] = hc.hamming(bi);
+                bi += ncodes;
+            }
+        }
+    });
+}
+
+// Word-multiple sizes whose HammingComputer beats the kernel below; measured
+// per SIMD level on x86, unmeasured levels keep the kernels. See diff summary.
+constexpr bool prefer_hamming_computer(size_t ncodes) {
+    constexpr bool is_avx512 = THE_SIMD_LEVEL == SIMDLevel::AVX512 ||
+            THE_SIMD_LEVEL == SIMDLevel::AVX512_SPR;
+    constexpr bool is_avx2 = THE_SIMD_LEVEL == SIMDLevel::AVX2;
+
+    if (ncodes == 64) {
+        return is_avx2 || is_avx512;
+    }
+    if (ncodes == 16 || ncodes == 32) {
+        return is_avx2;
+    }
+    return false;
+}
+
 } // anonymous namespace
 
 /******************************************************************
@@ -364,7 +399,10 @@ void hammings_fixSL<THE_SIMD_LEVEL>(
         size_t nb,
         size_t ncodes,
         hamdis_t* dis) {
-    FAISS_THROW_IF_NOT(ncodes % 8 == 0);
+    if (ncodes % 8 != 0 || prefer_hamming_computer(ncodes)) {
+        hammings_with_computer(a, b, na, nb, ncodes, dis);
+        return;
+    }
     switch (ncodes) {
         case 8:
             hammings_impl<64>(C64(a), C64(b), na, nb, dis);

tests/test_hamming.cpp

@@ -316,7 +316,7 @@ TEST(TestHamming, test_hamming_knn) {
         ASSERT_EQ(dist_ham_knn, *true_bit_distances) << assert_str.str();
     }
 
-    for (auto code_size : {8, 16, 24, 32}) {
+    for (auto code_size : {8, 16, 24, 32, 64}) {
         std::stringstream assert_str = get_correct_hamming_example(
                 na,
                 nb,
@@ -333,3 +333,44 @@ TEST(TestHamming, test_hamming_knn) {
         EXPECT_EQ(dist_gen, *true_bit_distances) << assert_str.str();
     }
 }
+
+// Code sizes that are not a whole number of 64-bit words take the byte-tail
+// kernel rather than the word-level one. The knn entry points have always
+// accepted them; hammings() used to throw instead.
+TEST(TestHamming, test_hammings_ragged_code_size) {
+    std::default_random_engine rng(123);
+    std::uniform_int_distribution<int32_t> uniform(0, 255);
+
+    const size_t na = 3;
+    const size_t nb = 7;
+
+    for (auto code_size : {1, 2, 5, 12, 20, 33}) {
+        std::vector<uint8_t> a(na * code_size);
+        std::vector<uint8_t> b(nb * code_size);
+        for (auto& v : a) {
+            v = uniform(rng);
+        }
+        for (auto& v : b) {
+            v = uniform(rng);
+        }
+
+        // Reference: byte-wise popcount of the XOR, independent of faiss.
+        std::vector<hamdis_t> expected(na * nb);
+        for (size_t i = 0; i < na; ++i) {
+            for (size_t j = 0; j < nb; ++j) {
+                int d = 0;
+                for (int c = 0; c < code_size; ++c) {
+                    d += __builtin_popcount(
+                            static_cast<unsigned>(
+                                    a[i * code_size + c] ^
+                                    b[j * code_size + c]));
+                }
+                expected[i * nb + j] = d;
+            }
+        }
+
+        std::vector<hamdis_t> dis(na * nb);
+        faiss::hammings(a.data(), b.data(), na, nb, code_size, dis.data());
+        EXPECT_EQ(dis, expected) << "code_size = " << code_size;
+    }
+}