use dispatcher function to call HammingComputer

Summary:
The HammingComputer class is optimized for several vector sizes. So far it's been the caller's responsiblity to instanciate the relevant optimized version.

This diff introduces a `dispatch_HammingComputer` function that can be called with a template class that is instanciated for all existing optimized HammingComputer's.

Differential Revision: D46858553

fbshipit-source-id: 6ebe0b784755642b416ce4d34fbdc45ef9ea4c49

Author: mdouze

faiss/IndexBinaryHNSW.cpp

@@ -281,31 +281,21 @@ struct FlatHammingDis : DistanceComputer {
     }
 };
 
+struct BuildDistanceComputer {
+    using T = DistanceComputer*;
+    template <class HammingComputer>
+    DistanceComputer* f(IndexBinaryFlat* flat_storage) {
+        return new FlatHammingDis<HammingComputer>(*flat_storage);
+    }
+};
+
 } // namespace
 
 DistanceComputer* IndexBinaryHNSW::get_distance_computer() const {
     IndexBinaryFlat* flat_storage = dynamic_cast<IndexBinaryFlat*>(storage);
-
     FAISS_ASSERT(flat_storage != nullptr);
-
-    switch (code_size) {
-        case 4:
-            return new FlatHammingDis<HammingComputer4>(*flat_storage);
-        case 8:
-            return new FlatHammingDis<HammingComputer8>(*flat_storage);
-        case 16:
-            return new FlatHammingDis<HammingComputer16>(*flat_storage);
-        case 20:
-            return new FlatHammingDis<HammingComputer20>(*flat_storage);
-        case 32:
-            return new FlatHammingDis<HammingComputer32>(*flat_storage);
-        case 64:
-            return new FlatHammingDis<HammingComputer64>(*flat_storage);
-        default:
-            break;
-    }
-
-    return new FlatHammingDis<HammingComputerDefault>(*flat_storage);
+    BuildDistanceComputer bd;
+    return dispatch_HammingComputer(code_size, bd, flat_storage);
 }
 
 } // namespace faiss

faiss/IndexBinaryHash.cpp

@@ -176,6 +176,14 @@ void search_single_query_template(
     } while (fe.next());
 }
 
+struct Run_search_single_query {
+    using T = void;
+    template <class HammingComputer, class... Types>
+    T f(Types... args) {
+        search_single_query_template<HammingComputer>(args...);
+    }
+};
+
 template <class SearchResults>
 void search_single_query(
         const IndexBinaryHash& index,
@@ -184,29 +192,9 @@ void search_single_query(
         size_t& n0,
         size_t& nlist,
         size_t& ndis) {
-#define HC(name) \
-    search_single_query_template<name>(index, q, res, n0, nlist, ndis);
-    switch (index.code_size) {
-        case 4:
-            HC(HammingComputer4);
-            break;
-        case 8:
-            HC(HammingComputer8);
-            break;
-        case 16:
-            HC(HammingComputer16);
-            break;
-        case 20:
-            HC(HammingComputer20);
-            break;
-        case 32:
-            HC(HammingComputer32);
-            break;
-        default:
-            HC(HammingComputerDefault);
-            break;
-    }
-#undef HC
+    Run_search_single_query r;
+    dispatch_HammingComputer(
+            index.code_size, r, index, q, res, n0, nlist, ndis);
 }
 
 } // anonymous namespace
@@ -365,6 +353,14 @@ static void verify_shortlist(
     }
 }
 
+struct Run_verify_shortlist {
+    using T = void;
+    template <class HammingComputer, class... Types>
+    void f(Types... args) {
+        verify_shortlist<HammingComputer>(args...);
+    }
+};
+
 template <class SearchResults>
 void search_1_query_multihash(
         const IndexBinaryMultiHash& index,
@@ -405,29 +401,9 @@ void search_1_query_multihash(
     ndis += shortlist.size();
 
     // verify shortlist
-
-#define HC(name) verify_shortlist<name>(*index.storage, xi, shortlist, res)
-    switch (index.code_size) {
-        case 4:
-            HC(HammingComputer4);
-            break;
-        case 8:
-            HC(HammingComputer8);
-            break;
-        case 16:
-            HC(HammingComputer16);
-            break;
-        case 20:
-            HC(HammingComputer20);
-            break;
-        case 32:
-            HC(HammingComputer32);
-            break;
-        default:
-            HC(HammingComputerDefault);
-            break;
-    }
-#undef HC
+    Run_verify_shortlist r;
+    dispatch_HammingComputer(
+            index.code_size, r, *index.storage, xi, shortlist, res);
 }
 
 } // anonymous namespace

faiss/IndexBinaryIVF.cpp

@@ -735,151 +735,68 @@ void search_knn_hamming_per_invlist(
     }
 }
 
+struct Run_search_knn_hamming_per_invlist {
+    using T = void;
+
+    template <class HammingComputer, class... Types>
+    void f(Types... args) {
+        search_knn_hamming_per_invlist<HammingComputer>(args...);
+    }
+};
+
 template <bool store_pairs>
-void search_knn_hamming_count_1(
-        const IndexBinaryIVF& ivf,
-        size_t nx,
-        const uint8_t* x,
-        const idx_t* keys,
-        int k,
-        int32_t* distances,
-        idx_t* labels,
-        const IVFSearchParameters* params) {
-    switch (ivf.code_size) {
-#define HANDLE_CS(cs)                                               \
-    case cs:                                                        \
-        search_knn_hamming_count<HammingComputer##cs, store_pairs>( \
-                ivf, nx, x, keys, k, distances, labels, params);    \
-        break;
-        HANDLE_CS(4);
-        HANDLE_CS(8);
-        HANDLE_CS(16);
-        HANDLE_CS(20);
-        HANDLE_CS(32);
-        HANDLE_CS(64);
-#undef HANDLE_CS
-        default:
-            search_knn_hamming_count<HammingComputerDefault, store_pairs>(
-                    ivf, nx, x, keys, k, distances, labels, params);
-            break;
+struct Run_search_knn_hamming_count {
+    using T = void;
+
+    template <class HammingComputer, class... Types>
+    void f(Types... args) {
+        search_knn_hamming_count<HammingComputer, store_pairs>(args...);
     }
-}
+};
 
-void search_knn_hamming_per_invlist_1(
-        const IndexBinaryIVF& ivf,
-        size_t n,
-        const uint8_t* x,
-        idx_t k,
-        const idx_t* keys,
-        const int32_t* coarse_dis,
-        int32_t* distances,
-        idx_t* labels,
-        bool store_pairs,
-        const IVFSearchParameters* params) {
-    switch (ivf.code_size) {
-#define HANDLE_CS(cs)                                        \
-    case cs:                                                 \
-        search_knn_hamming_per_invlist<HammingComputer##cs>( \
-                ivf,                                         \
-                n,                                           \
-                x,                                           \
-                k,                                           \
-                keys,                                        \
-                coarse_dis,                                  \
-                distances,                                   \
-                labels,                                      \
-                store_pairs,                                 \
-                params);                                     \
-        break;
-        HANDLE_CS(4);
-        HANDLE_CS(8);
-        HANDLE_CS(16);
-        HANDLE_CS(20);
-        HANDLE_CS(32);
-        HANDLE_CS(64);
-#undef HANDLE_CS
-        default:
-            search_knn_hamming_per_invlist<HammingComputerDefault>(
-                    ivf,
-                    n,
-                    x,
-                    k,
-                    keys,
-                    coarse_dis,
-                    distances,
-                    labels,
-                    store_pairs,
-                    params);
-            break;
+struct BuildScanner {
+    using T = BinaryInvertedListScanner*;
+
+    template <class HammingComputer>
+    T f(size_t code_size, bool store_pairs) {
+        return new IVFBinaryScannerL2<HammingComputer>(code_size, store_pairs);
     }
-}
+};
 
 } // anonymous namespace
 
 BinaryInvertedListScanner* IndexBinaryIVF::get_InvertedListScanner(
         bool store_pairs) const {
-#define HC(name) return new IVFBinaryScannerL2<name>(code_size, store_pairs)
-    switch (code_size) {
-        case 4:
-            HC(HammingComputer4);
-        case 8:
-            HC(HammingComputer8);
-        case 16:
-            HC(HammingComputer16);
-        case 20:
-            HC(HammingComputer20);
-        case 32:
-            HC(HammingComputer32);
-        case 64:
-            HC(HammingComputer64);
-        default:
-            HC(HammingComputerDefault);
-    }
-#undef HC
+    BuildScanner bs;
+    return dispatch_HammingComputer(code_size, bs, code_size, store_pairs);
 }
 
 void IndexBinaryIVF::search_preassigned(
         idx_t n,
         const uint8_t* x,
         idx_t k,
-        const idx_t* idx,
-        const int32_t* coarse_dis,
-        int32_t* distances,
-        idx_t* labels,
+        const idx_t* cidx,
+        const int32_t* cdis,
+        int32_t* dis,
+        idx_t* idx,
         bool store_pairs,
         const IVFSearchParameters* params) const {
     if (per_invlist_search) {
-        search_knn_hamming_per_invlist_1(
-                *this,
-                n,
-                x,
-                k,
-                idx,
-                coarse_dis,
-                distances,
-                labels,
-                store_pairs,
-                params);
+        Run_search_knn_hamming_per_invlist r;
+        dispatch_HammingComputer(
+                code_size, r, *this, n, x, k,
+                cidx, cdis, dis, idx, store_pairs, params);
     } else if (use_heap) {
         search_knn_hamming_heap(
-                *this,
-                n,
-                x,
-                k,
-                idx,
-                coarse_dis,
-                distances,
-                labels,
-                store_pairs,
-                params);
+                *this, n, x, k, cidx, cdis, dis, idx, store_pairs, params);
+    } else if (store_pairs) {
+        Run_search_knn_hamming_count<true> r;
+        dispatch_HammingComputer(
+                code_size, r, *this, n, x, idx, k, dis, idx, params);
     } else {
-        if (store_pairs) {
-            search_knn_hamming_count_1<true>(
-                    *this, n, x, idx, k, distances, labels, params);
-        } else {
-            search_knn_hamming_count_1<false>(
-                    *this, n, x, idx, k, distances, labels, params);
-        }
+        Run_search_knn_hamming_count<false> r;
+        dispatch_HammingComputer(
+                code_size, r, *this, n, x, idx, k, dis, idx, params);
     }
 }
 

faiss/IndexIVFPQ.cpp

@@ -1154,30 +1154,23 @@ struct IVFPQScannerT : QueryTables {
         { indexIVFPQ_stats.n_hamming_pass += n_hamming_pass; }
     }
 
+    template <class SearchResultType>
+    struct Run_scan_list_polysemous_hc {
+        using T = void;
+        template <class HammingComputer, class... Types>
+        void f(const IVFPQScannerT* scanner, Types... args) {
+            scanner->scan_list_polysemous_hc<HammingComputer, SearchResultType>(
+                    args...);
+        }
+    };
+
     template <class SearchResultType>
     void scan_list_polysemous(
             size_t ncode,
             const uint8_t* codes,
             SearchResultType& res) const {
-        switch (pq.code_size) {
-#define HANDLE_CODE_SIZE(cs)                                            \
-    case cs:                                                            \
-        scan_list_polysemous_hc<HammingComputer##cs, SearchResultType>( \
-                ncode, codes, res);                                     \
-        break
-            HANDLE_CODE_SIZE(4);
-            HANDLE_CODE_SIZE(8);
-            HANDLE_CODE_SIZE(16);
-            HANDLE_CODE_SIZE(20);
-            HANDLE_CODE_SIZE(32);
-            HANDLE_CODE_SIZE(64);
-#undef HANDLE_CODE_SIZE
-            default:
-                scan_list_polysemous_hc<
-                        HammingComputerDefault,
-                        SearchResultType>(ncode, codes, res);
-                break;
-        }
+        Run_scan_list_polysemous_hc<SearchResultType> r;
+        dispatch_HammingComputer(pq.code_size, r, this, ncode, codes, res);
     }
 };