mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
Accept ragged code sizes in faiss::hammings(), and fix the vectorization axis for wide codes (#5520)
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5520 Two changes to `faiss::hammings()`: it accepts ragged code sizes, and the kernel for wide codes is told which axis to vectorize. Throughout, **ncodes** is number of bytes per vector (so, 32 for a 256-bit vector), a **word** is one 64-bit (8-byte) unit. A code size is **word-multiple** when `ncodes % 8 == 0` and **ragged** otherwise, so a code is `nwords = ncodes / 8` whole words followed by `rem = ncodes % 8` trailing bytes. ## 1. "Ragged" code sizes Ragged here doesn't mean variable vector lengths, it means non-divisible by 8 ncodes. `hammings()` throws unless the code size is word-multiple, even though the other Hamming entry points (`hammings_knn_hc`, `hamming_range_search`) have always accepted ragged sizes. Callers who need them have to reach past `hammings()` into `with_simd_level_a0_spr` and `with_HammingComputer` and hand-roll the sweep; Laser does exactly that in `KnnFaissBinaryIndex::measureAllCentroids` (D115840549) and can drop it after this. Ragged sizes go to `hammings_ragged_dispatch()`, behind its own `hammings_ragged_fixSL<SL>` entry point so the word-multiple kernels keep the code they generate today. ## 2. Vectorization axis for wide codes All-pairs Hamming is a doubly nested loop: over candidates, and over the words of one code. Only one of them gets vectorized, and clang picks the candidate loop. That means each vector lane wants word `k` of a different candidate, and since candidates are `ncodes` bytes apart the load becomes a strided gather: ``` across candidates (today) within one code (this diff) vpgatherqq (%r8,%zmm19,1),%zmm20 vpxorq (%rbx,%r12,1),%zmm4,%zmm5 ``` Vectorizing the other loop instead reads one whole code contiguously, which is a plain vector load. `#pragma clang loop vectorize(disable)` on the candidate loop is what expresses that: it does not disable vectorization, it takes the candidate loop out of the running so the word loop is what gets vectorized. The arithmetic is identical either way. It is applied only when `nwords >= 4`, i.e. once a code is at least a full 256-bit register wide. That is the widest threshold that costs nothing anywhere: | threshold | ncodes 8 | ncodes 16 | ncodes 32 | ncodes 64 | | --- | --- | --- | --- | --- | | not applied | +0.0 / -0.1 | +0.0 / -0.6 | +0.1 / -0.2 | -0.1 / +0.2 | | `nwords >= 1` | **+45.3 / +63.5** | -10.4 / **+12.2** | -12.6 / -22.7 | -40.7 / -63.1 | | `nwords >= 2` | +0.1 / +0.0 | -10.5 / **+11.7** | -12.6 / -23.7 | -40.7 / -63.1 | | `nwords >= 4` | +0.2 / +0.1 | -0.1 / +0.1 | -12.7 / -23.7 | -40.7 / -63.1 | (AVX2 / AVX512, percent against the parent commit.) At one word the candidate loop is already contiguous, so taking it out of the running costs 45-63%. At two words it helps AVX2 at ncodes 16 and hurts AVX512 by the same amount, which would make the behaviour SIMD-level dependent. Four words is where both levels agree. The pragma is clang-only and guarded with `#if defined(__clang__)`, since MSVC warns C4068 on an unknown pragma. Other compilers therefore never see it and keep today's behaviour: correct, just without the gain. ## 3. Why the ragged kernel is templated A ragged code's loads must be `memcpy` because it is not 8-byte aligned. Both `nwords` and `rem` are template parameters, which is worth an order of magnitude: | ncodes | `<NW, REM>` | `<NW>`, rem runtime | neither templated | | --- | --- | --- | --- | | 2 | **0.487** | 8.720 | 8.724 | | 4 | **0.433** | 9.419 | 9.164 | | 10 | **0.868** | 9.344 | 10.160 | | 12 | **0.692** | 10.387 | 11.472 | | 20 | **1.066** | 9.879 | 10.730 | | 33 | **2.646** | 10.757 | 10.769 | Fig 1: ragged table comparing templates NW and REM (AVX2, ns per pair.) `rem` is the critical part: at a constant size `memcpy` folds into a single instruction, whereas at a runtime size it becomes a real `memcpy` call per candidate, costing 10-20x. That is why the dispatch switches on `ncodes % 8` rather than passing the remainder as an argument. `NW` is worth a further ~2x, but only for 1 to 3 whole words, where a compile-time trip count lets the word loop unroll to straight-line popcounts: | ncodes | NW templated | NW runtime | | --- | --- | --- | | 2, 4 (0 words) | 0.489 / 0.433 | 0.489 / 0.413 | | 10 (1 word) | **0.873** | 1.730 | | 12 (1 word) | **0.696** | 1.725 | | 20 (2 words) | **1.065** | 2.086 | | 33 (4 words) | 2.644 | 2.589 | Fig 2: Ragged table comparing just NW templating vs not having it templated So `NW` is templated for 0-3 words and left runtime from 4 up, which is where it stops paying. ## Results AMD Genoa (EPYC, Zen 4), `mode/opt`, interleaved A/B against the parent commit, one code size per process, 250 samples per cell, min ns per (query, candidate) pair. Word-multiple sizes: | ncodes | AVX2 | AVX512 | | --- | --- | --- | | 32 | **-12.6%** (1.440 -> 1.258) | **-23.3%** (1.352 -> 1.037) | | 64 | **-40.6%** (3.421 -> 2.031) | **-63.1%** (5.680 -> 2.094) | | 40 | +0.0% | **-4.0%** (2.954 -> 2.837) | | 48 | +0.1% | **-4.7%** (3.530 -> 3.365) | | 56 | +0.0% | **+4.0%** (3.944 -> 4.101) | | 8, 16, 24, 128, 160 | within +/-0.3% | within +/-0.8% | Fig 3: result table comparing against baseline (aka before this diff) ncodes 32 and 64 are the sizes whose codes are exactly one vector register wide, which is where the gather cost the most. Absolute cost of every code size with this diff, so the ragged ones can be read against the word-multiple ones: | ncodes | kind | AVX2 ns/pair | AVX2 ns/byte | AVX512 ns/pair | AVX512 ns/byte | | --- | --- | --- | --- | --- | --- | | 2 | ragged | 0.487 | 0.243 | 0.476 | 0.238 | | 4 | ragged | 0.434 | 0.108 | 0.441 | 0.110 | | 8 | word-multiple | 0.240 | 0.030 | 0.215 | 0.027 | | 10 | ragged | 0.879 | 0.088 | 0.789 | 0.079 | | 12 | ragged | 0.693 | 0.058 | 0.693 | 0.058 | | 16 | word-multiple | 0.713 | 0.045 | 0.573 | 0.036 | | 20 | ragged | 1.065 | 0.053 | 1.127 | 0.056 | | 24 | word-multiple | 2.082 | 0.087 | 2.081 | 0.087 | | 32 | word-multiple | 1.257 | 0.039 | 1.036 | 0.032 | | 33 | ragged | 2.644 | 0.080 | 3.160 | 0.096 | | 40 | word-multiple | 2.650 | 0.066 | 2.837 | 0.071 | | 48 | word-multiple | 3.110 | 0.065 | 3.361 | 0.070 | | 56 | word-multiple | 3.452 | 0.062 | 4.096 | 0.073 | | 64 | word-multiple | 2.039 | 0.032 | 2.102 | 0.033 | | 128 | word-multiple | 6.033 | 0.047 | 5.348 | 0.042 | | 160 | word-multiple | 8.164 | 0.051 | 5.731 | 0.036 | Fig 4: result table showing the ns/pair and ns/byte performance. ragged has no "before" because it is newly supported with this diff. Read down the ns/byte columns and the sizes with their own kernel (8, 16, 32, 64) sit at 0.027-0.045, the new ragged kernel at 0.053-0.096 once past a single word, and the shared runtime-word-count kernel (24, 40, 48, 56, 128, 160) at 0.036-0.087. That last group is the weakest: ncodes 24 costs 2.082 ns/pair against 1.257 for ncodes 32, and ncodes 56 costs 3.452 against 2.039 for ncodes 64, both more expensive than a larger code. It vectorizes across candidates for the same reason described above, but its word count is not known at compile time so the same fix does not apply. Left alone here. The two ragged sizes below one word (2 and 4) are dominated by fixed per-call work rather than the compare itself, which is why their ns/byte is high. Reviewed By: alibeklfc Differential Revision: D115982541 fbshipit-source-id: 081c600973fbdf30852527fb6beaae4e29cadfa3
This commit is contained in:
Michael Norris
committed by
meta-codesync[bot]
meta-codesync[bot]
parent
2027d3b331
commit
e2cf9e1ebf
@@ -147,7 +147,13 @@ void hammings(
|
||||
size_t ncodes,
|
||||
hamdis_t* __restrict dis) {
|
||||
with_simd_level_a0_spr([&]<SIMDLevel SL>() {
|
||||
hammings_fixSL<SL>(a, b, na, nb, ncodes, dis);
|
||||
// Ragged sizes have their own kernel; keeping it out of
|
||||
// hammings_fixSL() leaves the word-level paths untouched.
|
||||
if (ncodes % 8 != 0) {
|
||||
hammings_ragged_fixSL<SL>(a, b, na, nb, ncodes, dis);
|
||||
} else {
|
||||
hammings_fixSL<SL>(a, b, na, nb, ncodes, dis);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
+10
-1
@@ -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(
|
||||
@@ -327,6 +327,15 @@ void hamming_range_search_fixSL(
|
||||
RangeSearchResult* result,
|
||||
const IDSelector* sel);
|
||||
|
||||
template <SIMDLevel SL>
|
||||
void hammings_ragged_fixSL(
|
||||
const uint8_t* a,
|
||||
const uint8_t* b,
|
||||
size_t na,
|
||||
size_t nb,
|
||||
size_t ncodes,
|
||||
hamdis_t* dis);
|
||||
|
||||
template <SIMDLevel SL>
|
||||
void hammings_fixSL(
|
||||
const uint8_t* a,
|
||||
|
||||
@@ -210,12 +210,23 @@ inline void hammings_impl(
|
||||
size_t n2,
|
||||
hamdis_t* __restrict dis) {
|
||||
size_t i, j;
|
||||
const size_t nwords = nbits / 64;
|
||||
constexpr size_t nwords = nbits / 64;
|
||||
for (i = 0; i < n1; i++) {
|
||||
const uint64_t* __restrict bs1_ = bs1 + i * nwords;
|
||||
hamdis_t* __restrict dis_ = dis + i * n2;
|
||||
for (j = 0; j < n2; j++) {
|
||||
dis_[j] = hamming<nbits>(bs1_, bs2 + j * nwords);
|
||||
if constexpr (nwords >= 4) {
|
||||
// Wide codes are quicker one candidate at a time; otherwise the
|
||||
// compiler batches candidates and must gather scattered bytes.
|
||||
#if defined(__clang__)
|
||||
#pragma clang loop vectorize(disable)
|
||||
#endif
|
||||
for (j = 0; j < n2; j++) {
|
||||
dis_[j] = hamming<nbits>(bs1_, bs2 + j * nwords);
|
||||
}
|
||||
} else {
|
||||
for (j = 0; j < n2; j++) {
|
||||
dis_[j] = hamming<nbits>(bs1_, bs2 + j * nwords);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -36,6 +36,7 @@
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdio>
|
||||
#include <cstring>
|
||||
#include <limits>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
@@ -298,6 +299,132 @@ void generalized_hammings_knn_hc_impl(
|
||||
}
|
||||
}
|
||||
|
||||
/* Both number of words and remainder are constants, so the compiler flattens
|
||||
* the per-chunk loop. The query is copied into qw/qt once, not re-read for
|
||||
* every candidate. */
|
||||
template <int NW, int REM>
|
||||
void hammings_ragged_fixed(
|
||||
const uint8_t* __restrict a,
|
||||
const uint8_t* __restrict b,
|
||||
size_t na,
|
||||
size_t nb,
|
||||
size_t ncodes,
|
||||
hamdis_t* __restrict dis) {
|
||||
for (size_t i = 0; i < na; i++) {
|
||||
const uint8_t* __restrict ai = a + i * ncodes;
|
||||
hamdis_t* __restrict dis_ = dis + i * nb;
|
||||
// Read the query once here, not once per candidate.
|
||||
uint64_t qw[NW == 0 ? 1 : NW];
|
||||
for (int k = 0; k < NW; k++) {
|
||||
memcpy(&qw[k], ai + k * 8, 8);
|
||||
}
|
||||
uint64_t qt = 0;
|
||||
if constexpr (REM > 0) {
|
||||
memcpy(&qt, ai + NW * 8, REM);
|
||||
}
|
||||
for (size_t j = 0; j < nb; j++) {
|
||||
const uint8_t* __restrict bj = b + j * ncodes;
|
||||
hamdis_t h = 0;
|
||||
uint64_t y;
|
||||
for (int k = 0; k < NW; k++) {
|
||||
memcpy(&y, bj + k * 8, 8);
|
||||
h += popcount64(qw[k] ^ y);
|
||||
}
|
||||
if constexpr (REM > 0) {
|
||||
y = 0;
|
||||
memcpy(&y, bj + NW * 8, REM);
|
||||
h += popcount64(qt ^ y);
|
||||
}
|
||||
dis_[j] = h;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/* Same loop with the word count read at run time, for codes long enough that
|
||||
* a constant count gives no performance improvement. */
|
||||
template <int REM>
|
||||
void hammings_ragged_var(
|
||||
const uint8_t* __restrict a,
|
||||
const uint8_t* __restrict b,
|
||||
size_t na,
|
||||
size_t nb,
|
||||
size_t ncodes,
|
||||
hamdis_t* __restrict dis) {
|
||||
const size_t nwords = ncodes / 8;
|
||||
for (size_t i = 0; i < na; i++) {
|
||||
const uint8_t* __restrict ai = a + i * ncodes;
|
||||
hamdis_t* __restrict dis_ = dis + i * nb;
|
||||
for (size_t j = 0; j < nb; j++) {
|
||||
const uint8_t* __restrict bj = b + j * ncodes;
|
||||
hamdis_t h = 0;
|
||||
uint64_t x, y;
|
||||
for (size_t k = 0; k < nwords; k++) {
|
||||
memcpy(&x, ai + k * 8, 8);
|
||||
memcpy(&y, bj + k * 8, 8);
|
||||
h += popcount64(x ^ y);
|
||||
}
|
||||
if constexpr (REM > 0) {
|
||||
x = 0;
|
||||
y = 0;
|
||||
memcpy(&x, ai + nwords * 8, REM);
|
||||
memcpy(&y, bj + nwords * 8, REM);
|
||||
h += popcount64(x ^ y);
|
||||
}
|
||||
dis_[j] = h;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/* Makes the word count a constant for short codes when calling
|
||||
* hammings_ragged_fixed, worth roughly 2x perf improvement. Longer codes
|
||||
* gain nothing from it, so they share hammings_ragged_var. */
|
||||
template <int REM>
|
||||
void hammings_ragged_by_nwords(
|
||||
const uint8_t* __restrict a,
|
||||
const uint8_t* __restrict b,
|
||||
size_t na,
|
||||
size_t nb,
|
||||
size_t ncodes,
|
||||
hamdis_t* __restrict dis) {
|
||||
switch (ncodes / 8) {
|
||||
case 0:
|
||||
return hammings_ragged_fixed<0, REM>(a, b, na, nb, ncodes, dis);
|
||||
case 1:
|
||||
return hammings_ragged_fixed<1, REM>(a, b, na, nb, ncodes, dis);
|
||||
case 2:
|
||||
return hammings_ragged_fixed<2, REM>(a, b, na, nb, ncodes, dis);
|
||||
case 3:
|
||||
return hammings_ragged_fixed<3, REM>(a, b, na, nb, ncodes, dis);
|
||||
default:
|
||||
return hammings_ragged_var<REM>(a, b, na, nb, ncodes, dis);
|
||||
}
|
||||
}
|
||||
|
||||
void hammings_ragged_dispatch(
|
||||
const uint8_t* __restrict a,
|
||||
const uint8_t* __restrict b,
|
||||
size_t na,
|
||||
size_t nb,
|
||||
size_t ncodes,
|
||||
hamdis_t* __restrict dis) {
|
||||
switch (ncodes % 8) {
|
||||
case 1:
|
||||
return hammings_ragged_by_nwords<1>(a, b, na, nb, ncodes, dis);
|
||||
case 2:
|
||||
return hammings_ragged_by_nwords<2>(a, b, na, nb, ncodes, dis);
|
||||
case 3:
|
||||
return hammings_ragged_by_nwords<3>(a, b, na, nb, ncodes, dis);
|
||||
case 4:
|
||||
return hammings_ragged_by_nwords<4>(a, b, na, nb, ncodes, dis);
|
||||
case 5:
|
||||
return hammings_ragged_by_nwords<5>(a, b, na, nb, ncodes, dis);
|
||||
case 6:
|
||||
return hammings_ragged_by_nwords<6>(a, b, na, nb, ncodes, dis);
|
||||
default:
|
||||
return hammings_ragged_by_nwords<7>(a, b, na, nb, ncodes, dis);
|
||||
}
|
||||
}
|
||||
|
||||
} // anonymous namespace
|
||||
|
||||
/******************************************************************
|
||||
@@ -356,6 +483,19 @@ void hamming_range_search_fixSL<THE_SIMD_LEVEL>(
|
||||
});
|
||||
}
|
||||
|
||||
/* Its own entry point, so the word-multiple kernels keep the code the
|
||||
* compiler already generates for them. */
|
||||
template <>
|
||||
void hammings_ragged_fixSL<THE_SIMD_LEVEL>(
|
||||
const uint8_t* a,
|
||||
const uint8_t* b,
|
||||
size_t na,
|
||||
size_t nb,
|
||||
size_t ncodes,
|
||||
hamdis_t* dis) {
|
||||
hammings_ragged_dispatch(a, b, na, nb, ncodes, dis);
|
||||
}
|
||||
|
||||
template <>
|
||||
void hammings_fixSL<THE_SIMD_LEVEL>(
|
||||
const uint8_t* a,
|
||||
|
||||
+40
-1
@@ -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,42 @@ TEST(TestHamming, test_hamming_knn) {
|
||||
EXPECT_EQ(dist_gen, *true_bit_distances) << assert_str.str();
|
||||
}
|
||||
}
|
||||
|
||||
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) {
|
||||
uint8_t x = a[i * code_size + c] ^ b[j * code_size + c];
|
||||
while (x) {
|
||||
d += x & 1;
|
||||
x >>= 1;
|
||||
}
|
||||
}
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user