Skip to content

Commit e410d23

Browse files
mdouzefacebook-github-bot
authored andcommitted
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
1 parent f69b1db commit e410d23

11 files changed

+250
-436
lines changed

faiss/IndexBinaryHNSW.cpp

+10-20
Original file line numberDiff line numberDiff line change
@@ -281,31 +281,21 @@ struct FlatHammingDis : DistanceComputer {
281281
}
282282
};
283283

284+
struct BuildDistanceComputer {
285+
using T = DistanceComputer*;
286+
template <class HammingComputer>
287+
DistanceComputer* f(IndexBinaryFlat* flat_storage) {
288+
return new FlatHammingDis<HammingComputer>(*flat_storage);
289+
}
290+
};
291+
284292
} // namespace
285293

286294
DistanceComputer* IndexBinaryHNSW::get_distance_computer() const {
287295
IndexBinaryFlat* flat_storage = dynamic_cast<IndexBinaryFlat*>(storage);
288-
289296
FAISS_ASSERT(flat_storage != nullptr);
290-
291-
switch (code_size) {
292-
case 4:
293-
return new FlatHammingDis<HammingComputer4>(*flat_storage);
294-
case 8:
295-
return new FlatHammingDis<HammingComputer8>(*flat_storage);
296-
case 16:
297-
return new FlatHammingDis<HammingComputer16>(*flat_storage);
298-
case 20:
299-
return new FlatHammingDis<HammingComputer20>(*flat_storage);
300-
case 32:
301-
return new FlatHammingDis<HammingComputer32>(*flat_storage);
302-
case 64:
303-
return new FlatHammingDis<HammingComputer64>(*flat_storage);
304-
default:
305-
break;
306-
}
307-
308-
return new FlatHammingDis<HammingComputerDefault>(*flat_storage);
297+
BuildDistanceComputer bd;
298+
return dispatch_HammingComputer(code_size, bd, flat_storage);
309299
}
310300

311301
} // namespace faiss

faiss/IndexBinaryHash.cpp

+22-46
Original file line numberDiff line numberDiff line change
@@ -176,6 +176,14 @@ void search_single_query_template(
176176
} while (fe.next());
177177
}
178178

179+
struct Run_search_single_query {
180+
using T = void;
181+
template <class HammingComputer, class... Types>
182+
T f(Types... args) {
183+
search_single_query_template<HammingComputer>(args...);
184+
}
185+
};
186+
179187
template <class SearchResults>
180188
void search_single_query(
181189
const IndexBinaryHash& index,
@@ -184,29 +192,9 @@ void search_single_query(
184192
size_t& n0,
185193
size_t& nlist,
186194
size_t& ndis) {
187-
#define HC(name) \
188-
search_single_query_template<name>(index, q, res, n0, nlist, ndis);
189-
switch (index.code_size) {
190-
case 4:
191-
HC(HammingComputer4);
192-
break;
193-
case 8:
194-
HC(HammingComputer8);
195-
break;
196-
case 16:
197-
HC(HammingComputer16);
198-
break;
199-
case 20:
200-
HC(HammingComputer20);
201-
break;
202-
case 32:
203-
HC(HammingComputer32);
204-
break;
205-
default:
206-
HC(HammingComputerDefault);
207-
break;
208-
}
209-
#undef HC
195+
Run_search_single_query r;
196+
dispatch_HammingComputer(
197+
index.code_size, r, index, q, res, n0, nlist, ndis);
210198
}
211199

212200
} // anonymous namespace
@@ -365,6 +353,14 @@ static void verify_shortlist(
365353
}
366354
}
367355

356+
struct Run_verify_shortlist {
357+
using T = void;
358+
template <class HammingComputer, class... Types>
359+
void f(Types... args) {
360+
verify_shortlist<HammingComputer>(args...);
361+
}
362+
};
363+
368364
template <class SearchResults>
369365
void search_1_query_multihash(
370366
const IndexBinaryMultiHash& index,
@@ -405,29 +401,9 @@ void search_1_query_multihash(
405401
ndis += shortlist.size();
406402

407403
// verify shortlist
408-
409-
#define HC(name) verify_shortlist<name>(*index.storage, xi, shortlist, res)
410-
switch (index.code_size) {
411-
case 4:
412-
HC(HammingComputer4);
413-
break;
414-
case 8:
415-
HC(HammingComputer8);
416-
break;
417-
case 16:
418-
HC(HammingComputer16);
419-
break;
420-
case 20:
421-
HC(HammingComputer20);
422-
break;
423-
case 32:
424-
HC(HammingComputer32);
425-
break;
426-
default:
427-
HC(HammingComputerDefault);
428-
break;
429-
}
430-
#undef HC
404+
Run_verify_shortlist r;
405+
dispatch_HammingComputer(
406+
index.code_size, r, *index.storage, xi, shortlist, res);
431407
}
432408

433409
} // anonymous namespace

faiss/IndexBinaryIVF.cpp

+41-124
Original file line numberDiff line numberDiff line change
@@ -735,151 +735,68 @@ void search_knn_hamming_per_invlist(
735735
}
736736
}
737737

738+
struct Run_search_knn_hamming_per_invlist {
739+
using T = void;
740+
741+
template <class HammingComputer, class... Types>
742+
void f(Types... args) {
743+
search_knn_hamming_per_invlist<HammingComputer>(args...);
744+
}
745+
};
746+
738747
template <bool store_pairs>
739-
void search_knn_hamming_count_1(
740-
const IndexBinaryIVF& ivf,
741-
size_t nx,
742-
const uint8_t* x,
743-
const idx_t* keys,
744-
int k,
745-
int32_t* distances,
746-
idx_t* labels,
747-
const IVFSearchParameters* params) {
748-
switch (ivf.code_size) {
749-
#define HANDLE_CS(cs) \
750-
case cs: \
751-
search_knn_hamming_count<HammingComputer##cs, store_pairs>( \
752-
ivf, nx, x, keys, k, distances, labels, params); \
753-
break;
754-
HANDLE_CS(4);
755-
HANDLE_CS(8);
756-
HANDLE_CS(16);
757-
HANDLE_CS(20);
758-
HANDLE_CS(32);
759-
HANDLE_CS(64);
760-
#undef HANDLE_CS
761-
default:
762-
search_knn_hamming_count<HammingComputerDefault, store_pairs>(
763-
ivf, nx, x, keys, k, distances, labels, params);
764-
break;
748+
struct Run_search_knn_hamming_count {
749+
using T = void;
750+
751+
template <class HammingComputer, class... Types>
752+
void f(Types... args) {
753+
search_knn_hamming_count<HammingComputer, store_pairs>(args...);
765754
}
766-
}
755+
};
767756

768-
void search_knn_hamming_per_invlist_1(
769-
const IndexBinaryIVF& ivf,
770-
size_t n,
771-
const uint8_t* x,
772-
idx_t k,
773-
const idx_t* keys,
774-
const int32_t* coarse_dis,
775-
int32_t* distances,
776-
idx_t* labels,
777-
bool store_pairs,
778-
const IVFSearchParameters* params) {
779-
switch (ivf.code_size) {
780-
#define HANDLE_CS(cs) \
781-
case cs: \
782-
search_knn_hamming_per_invlist<HammingComputer##cs>( \
783-
ivf, \
784-
n, \
785-
x, \
786-
k, \
787-
keys, \
788-
coarse_dis, \
789-
distances, \
790-
labels, \
791-
store_pairs, \
792-
params); \
793-
break;
794-
HANDLE_CS(4);
795-
HANDLE_CS(8);
796-
HANDLE_CS(16);
797-
HANDLE_CS(20);
798-
HANDLE_CS(32);
799-
HANDLE_CS(64);
800-
#undef HANDLE_CS
801-
default:
802-
search_knn_hamming_per_invlist<HammingComputerDefault>(
803-
ivf,
804-
n,
805-
x,
806-
k,
807-
keys,
808-
coarse_dis,
809-
distances,
810-
labels,
811-
store_pairs,
812-
params);
813-
break;
757+
struct BuildScanner {
758+
using T = BinaryInvertedListScanner*;
759+
760+
template <class HammingComputer>
761+
T f(size_t code_size, bool store_pairs) {
762+
return new IVFBinaryScannerL2<HammingComputer>(code_size, store_pairs);
814763
}
815-
}
764+
};
816765

817766
} // anonymous namespace
818767

819768
BinaryInvertedListScanner* IndexBinaryIVF::get_InvertedListScanner(
820769
bool store_pairs) const {
821-
#define HC(name) return new IVFBinaryScannerL2<name>(code_size, store_pairs)
822-
switch (code_size) {
823-
case 4:
824-
HC(HammingComputer4);
825-
case 8:
826-
HC(HammingComputer8);
827-
case 16:
828-
HC(HammingComputer16);
829-
case 20:
830-
HC(HammingComputer20);
831-
case 32:
832-
HC(HammingComputer32);
833-
case 64:
834-
HC(HammingComputer64);
835-
default:
836-
HC(HammingComputerDefault);
837-
}
838-
#undef HC
770+
BuildScanner bs;
771+
return dispatch_HammingComputer(code_size, bs, code_size, store_pairs);
839772
}
840773

841774
void IndexBinaryIVF::search_preassigned(
842775
idx_t n,
843776
const uint8_t* x,
844777
idx_t k,
845-
const idx_t* idx,
846-
const int32_t* coarse_dis,
847-
int32_t* distances,
848-
idx_t* labels,
778+
const idx_t* cidx,
779+
const int32_t* cdis,
780+
int32_t* dis,
781+
idx_t* idx,
849782
bool store_pairs,
850783
const IVFSearchParameters* params) const {
851784
if (per_invlist_search) {
852-
search_knn_hamming_per_invlist_1(
853-
*this,
854-
n,
855-
x,
856-
k,
857-
idx,
858-
coarse_dis,
859-
distances,
860-
labels,
861-
store_pairs,
862-
params);
785+
Run_search_knn_hamming_per_invlist r;
786+
dispatch_HammingComputer(
787+
code_size, r, *this, n, x, k,
788+
cidx, cdis, dis, idx, store_pairs, params);
863789
} else if (use_heap) {
864790
search_knn_hamming_heap(
865-
*this,
866-
n,
867-
x,
868-
k,
869-
idx,
870-
coarse_dis,
871-
distances,
872-
labels,
873-
store_pairs,
874-
params);
791+
*this, n, x, k, cidx, cdis, dis, idx, store_pairs, params);
792+
} else if (store_pairs) {
793+
Run_search_knn_hamming_count<true> r;
794+
dispatch_HammingComputer(
795+
code_size, r, *this, n, x, idx, k, dis, idx, params);
875796
} else {
876-
if (store_pairs) {
877-
search_knn_hamming_count_1<true>(
878-
*this, n, x, idx, k, distances, labels, params);
879-
} else {
880-
search_knn_hamming_count_1<false>(
881-
*this, n, x, idx, k, distances, labels, params);
882-
}
797+
Run_search_knn_hamming_count<false> r;
798+
dispatch_HammingComputer(
799+
code_size, r, *this, n, x, idx, k, dis, idx, params);
883800
}
884801
}
885802

faiss/IndexIVFPQ.cpp

+12-19
Original file line numberDiff line numberDiff line change
@@ -1154,30 +1154,23 @@ struct IVFPQScannerT : QueryTables {
11541154
{ indexIVFPQ_stats.n_hamming_pass += n_hamming_pass; }
11551155
}
11561156

1157+
template <class SearchResultType>
1158+
struct Run_scan_list_polysemous_hc {
1159+
using T = void;
1160+
template <class HammingComputer, class... Types>
1161+
void f(const IVFPQScannerT* scanner, Types... args) {
1162+
scanner->scan_list_polysemous_hc<HammingComputer, SearchResultType>(
1163+
args...);
1164+
}
1165+
};
1166+
11571167
template <class SearchResultType>
11581168
void scan_list_polysemous(
11591169
size_t ncode,
11601170
const uint8_t* codes,
11611171
SearchResultType& res) const {
1162-
switch (pq.code_size) {
1163-
#define HANDLE_CODE_SIZE(cs) \
1164-
case cs: \
1165-
scan_list_polysemous_hc<HammingComputer##cs, SearchResultType>( \
1166-
ncode, codes, res); \
1167-
break
1168-
HANDLE_CODE_SIZE(4);
1169-
HANDLE_CODE_SIZE(8);
1170-
HANDLE_CODE_SIZE(16);
1171-
HANDLE_CODE_SIZE(20);
1172-
HANDLE_CODE_SIZE(32);
1173-
HANDLE_CODE_SIZE(64);
1174-
#undef HANDLE_CODE_SIZE
1175-
default:
1176-
scan_list_polysemous_hc<
1177-
HammingComputerDefault,
1178-
SearchResultType>(ncode, codes, res);
1179-
break;
1180-
}
1172+
Run_scan_list_polysemous_hc<SearchResultType> r;
1173+
dispatch_HammingComputer(pq.code_size, r, this, ncode, codes, res);
11811174
}
11821175
};
11831176

0 commit comments

Comments
 (0)