faiss 0.2.3 → 0.2.5

data/vendor/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp

@@ -8,6 +8,7 @@
 #include <faiss/impl/pq4_fast_scan.h>
 
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/LookupTableScaler.h>
 #include <faiss/impl/simd_result_handlers.h>
 #include <faiss/utils/simdlib.h>
 
@@ -27,15 +28,17 @@ namespace {
  * writes results in a ResultHandler
  */
 
-template <int NQ, class ResultHandler>
+template <int NQ, class ResultHandler, class Scaler>
 void kernel_accumulate_block(
         int nsq,
         const uint8_t* codes,
         const uint8_t* LUT,
-        ResultHandler& res) {
+        ResultHandler& res,
+        const Scaler& scaler) {
     // dummy alloc to keep the windows compiler happy
     constexpr int NQA = NQ > 0 ? NQ : 1;
     // distance accumulators
+    // layout: accu[q][b]: distance accumulator for vectors 8*b..8*b+7
     simd16uint16 accu[NQA][4];
 
     for (int q = 0; q < NQ; q++) {
@@ -45,7 +48,7 @@ void kernel_accumulate_block(
     }
 
     // _mm_prefetch(codes + 768, 0);
-    for (int sq = 0; sq < nsq; sq += 2) {
+    for (int sq = 0; sq < nsq - scaler.nscale; sq += 2) {
         // prefetch
         simd32uint8 c(codes);
         codes += 32;
@@ -71,6 +74,31 @@ void kernel_accumulate_block(
         }
     }
 
+    for (int sq = 0; sq < scaler.nscale; sq += 2) {
+        // prefetch
+        simd32uint8 c(codes);
+        codes += 32;
+
+        simd32uint8 mask(0xf);
+        // shift op does not exist for int8...
+        simd32uint8 chi = simd32uint8(simd16uint16(c) >> 4) & mask;
+        simd32uint8 clo = c & mask;
+
+        for (int q = 0; q < NQ; q++) {
+            // load LUTs for 2 quantizers
+            simd32uint8 lut(LUT);
+            LUT += 32;
+
+            simd32uint8 res0 = scaler.lookup(lut, clo);
+            accu[q][0] += scaler.scale_lo(res0); // handle vectors 0..7
+            accu[q][1] += scaler.scale_hi(res0); // handle vectors 8..15
+
+            simd32uint8 res1 = scaler.lookup(lut, chi);
+            accu[q][2] += scaler.scale_lo(res1); // handle vectors 16..23
+            accu[q][3] += scaler.scale_hi(res1); //  handle vectors 24..31
+        }
+    }
+
     for (int q = 0; q < NQ; q++) {
         accu[q][0] -= accu[q][1] << 8;
         simd16uint16 dis0 = combine2x2(accu[q][0], accu[q][1]);
@@ -81,13 +109,14 @@ void kernel_accumulate_block(
 }
 
 // handle at most 4 blocks of queries
-template <int QBS, class ResultHandler>
+template <int QBS, class ResultHandler, class Scaler>
 void accumulate_q_4step(
         size_t ntotal2,
         int nsq,
         const uint8_t* codes,
         const uint8_t* LUT0,
-        ResultHandler& res) {
+        ResultHandler& res,
+        const Scaler& scaler) {
     constexpr int Q1 = QBS & 15;
     constexpr int Q2 = (QBS >> 4) & 15;
     constexpr int Q3 = (QBS >> 8) & 15;
@@ -97,21 +126,21 @@ void accumulate_q_4step(
     for (int64_t j0 = 0; j0 < ntotal2; j0 += 32) {
         FixedStorageHandler<SQ, 2> res2;
         const uint8_t* LUT = LUT0;
-        kernel_accumulate_block<Q1>(nsq, codes, LUT, res2);
+        kernel_accumulate_block<Q1>(nsq, codes, LUT, res2, scaler);
         LUT += Q1 * nsq * 16;
         if (Q2 > 0) {
             res2.set_block_origin(Q1, 0);
-            kernel_accumulate_block<Q2>(nsq, codes, LUT, res2);
+            kernel_accumulate_block<Q2>(nsq, codes, LUT, res2, scaler);
             LUT += Q2 * nsq * 16;
         }
         if (Q3 > 0) {
             res2.set_block_origin(Q1 + Q2, 0);
-            kernel_accumulate_block<Q3>(nsq, codes, LUT, res2);
+            kernel_accumulate_block<Q3>(nsq, codes, LUT, res2, scaler);
             LUT += Q3 * nsq * 16;
         }
         if (Q4 > 0) {
             res2.set_block_origin(Q1 + Q2 + Q3, 0);
-            kernel_accumulate_block<Q4>(nsq, codes, LUT, res2);
+            kernel_accumulate_block<Q4>(nsq, codes, LUT, res2, scaler);
         }
         res.set_block_origin(0, j0);
         res2.to_other_handler(res);
@@ -119,29 +148,31 @@ void accumulate_q_4step(
     }
 }
 
-template <int NQ, class ResultHandler>
+template <int NQ, class ResultHandler, class Scaler>
 void kernel_accumulate_block_loop(
         size_t ntotal2,
         int nsq,
         const uint8_t* codes,
         const uint8_t* LUT,
-        ResultHandler& res) {
+        ResultHandler& res,
+        const Scaler& scaler) {
     for (int64_t j0 = 0; j0 < ntotal2; j0 += 32) {
         res.set_block_origin(0, j0);
         kernel_accumulate_block<NQ, ResultHandler>(
-                nsq, codes + j0 * nsq / 2, LUT, res);
+                nsq, codes + j0 * nsq / 2, LUT, res, scaler);
     }
 }
 
 // non-template version of accumulate kernel -- dispatches dynamically
-template <class ResultHandler>
+template <class ResultHandler, class Scaler>
 void accumulate(
         int nq,
         size_t ntotal2,
         int nsq,
         const uint8_t* codes,
         const uint8_t* LUT,
-        ResultHandler& res) {
+        ResultHandler& res,
+        const Scaler& scaler) {
     assert(nsq % 2 == 0);
     assert(is_aligned_pointer(codes));
     assert(is_aligned_pointer(LUT));
@@ -149,7 +180,7 @@ void accumulate(
 #define DISPATCH(NQ)                                     \
     case NQ:                                             \
         kernel_accumulate_block_loop<NQ, ResultHandler>( \
-                ntotal2, nsq, codes, LUT, res);          \
+                ntotal2, nsq, codes, LUT, res, scaler);  \
         return
 
     switch (nq) {
@@ -165,23 +196,24 @@ void accumulate(
 
 } // namespace
 
-template <class ResultHandler>
+template <class ResultHandler, class Scaler>
 void pq4_accumulate_loop_qbs(
         int qbs,
         size_t ntotal2,
         int nsq,
         const uint8_t* codes,
         const uint8_t* LUT0,
-        ResultHandler& res) {
+        ResultHandler& res,
+        const Scaler& scaler) {
     assert(nsq % 2 == 0);
     assert(is_aligned_pointer(codes));
     assert(is_aligned_pointer(LUT0));
 
     // try out optimized versions
     switch (qbs) {
-#define DISPATCH(QBS)                                            \
-    case QBS:                                                    \
-        accumulate_q_4step<QBS>(ntotal2, nsq, codes, LUT0, res); \
+#define DISPATCH(QBS)                                                    \
+    case QBS:                                                            \
+        accumulate_q_4step<QBS>(ntotal2, nsq, codes, LUT0, res, scaler); \
         return;
         DISPATCH(0x3333); // 12
         DISPATCH(0x2333); // 11
@@ -219,9 +251,10 @@ void pq4_accumulate_loop_qbs(
             int nq = qi & 15;
             qi >>= 4;
             res.set_block_origin(i0, j0);
-#define DISPATCH(NQ)                                                      \
-    case NQ:                                                              \
-        kernel_accumulate_block<NQ, ResultHandler>(nsq, codes, LUT, res); \
+#define DISPATCH(NQ)                                \
+    case NQ:                                        \
+        kernel_accumulate_block<NQ, ResultHandler>( \
+                nsq, codes, LUT, res, scaler);      \
         break
             switch (nq) {
                 DISPATCH(1);
@@ -241,9 +274,23 @@ void pq4_accumulate_loop_qbs(
 
 // explicit template instantiations
 
-#define INSTANTIATE_ACCUMULATE_Q(RH)           \
-    template void pq4_accumulate_loop_qbs<RH>( \
-            int, size_t, int, const uint8_t*, const uint8_t*, RH&);
+#define INSTANTIATE_ACCUMULATE_Q(RH)                            \
+    template void pq4_accumulate_loop_qbs<RH, DummyScaler>(     \
+            int,                                                \
+            size_t,                                             \
+            int,                                                \
+            const uint8_t*,                                     \
+            const uint8_t*,                                     \
+            RH&,                                                \
+            const DummyScaler&);                                \
+    template void pq4_accumulate_loop_qbs<RH, NormTableScaler>( \
+            int,                                                \
+            size_t,                                             \
+            int,                                                \
+            const uint8_t*,                                     \
+            const uint8_t*,                                     \
+            RH&,                                                \
+            const NormTableScaler&);
 
 using Csi = CMax<uint16_t, int>;
 INSTANTIATE_ACCUMULATE_Q(SingleResultHandler<Csi>)
@@ -293,7 +340,8 @@ void accumulate_to_mem(
         uint16_t* accu) {
     FAISS_THROW_IF_NOT(ntotal2 % 32 == 0);
     StoreResultHandler handler(accu, ntotal2);
-    accumulate(nq, ntotal2, nsq, codes, LUT, handler);
+    DummyScaler scaler;
+    accumulate(nq, ntotal2, nsq, codes, LUT, handler, scaler);
 }
 
 int pq4_preferred_qbs(int n) {