Additional C++ templates for fast sa_decode: add IndexPQDecoder (face…

…bookresearch#2444)

Summary:
Pull Request resolved: facebookresearch#2444

Add IndexPQDecoder. The following codecs are supported:
* PQ[1]x8

Additionally, AVX2 and ARM versions support the following codecs:
* PQ[1]x10
* PQ[1]x16

Reviewed By: mdouze

Differential Revision: D39176423

fbshipit-source-id: b002b3d3b0533849f72f3660e8088d8dc44a66d6

faiss/cppcontrib/SaDecodeKernels.h

@@ -6,13 +6,24 @@
 //   function for the following index families:
 //   * IVF256,PQ[1]x8np
 //   * Residual[1]x8,PQ[2]x8
-//   * IVF[9-16 bit],PQ[1]x8 (such as IVF1024,PQ16np)
+//   * IVF[2^9-2^16 bit],PQ[1]x8 (such as IVF1024,PQ16np)
 //   * Residual1x[9-16 bit],PQ[1]x8 (such as Residual1x9,PQ8)
+//   * PQ[1]x8
 // Additionally, AVX2 and ARM versions support
 //   * Residual[1]x8,PQ[2]x10
 //   * Residual[1]x8,PQ[2]x16
+//   * Residual[1]x10,PQ[2]x10
+//   * Residual[1]x10,PQ[2]x16
+//   * Residual[1]x16,PQ[2]x10
+//   * Residual[1]x16,PQ[2]x16
 //   * Residual1x[9-16 bit],PQ[1]x10 (such as Residual1x9,PQ16x10)
+//   * * (use with COARSE_BITS=16)
 //   * Residual1x[9-16 bit],PQ[1]x16 (such as Residual1x9,PQ16x16)
+//   * * (use with COARSE_BITS=16)
+//   * PQ[1]x10
+//   * PQ[1]x16
+// Unfortunately, currently Faiss does not support something like
+//   IVF256,PQ16x10np
 //
 // The goal was to achieve the maximum performance, so the template version it
 // is. The provided index families share the same code for sa_decode.
@@ -25,26 +36,49 @@
 //        intptr_t DIM,
 //        intptr_t COARSE_SIZE,
 //        intptr_t FINE_SIZE,
-//        intptr_t COARSE_BITS = 8>
+//        intptr_t COARSE_BITS = 8
+//        intptr_t FINE_BITS = 8>
 //     struct Index2LevelDecoder { /*...*/ };
 //   }
 // * DIM is the dimensionality of data
 // * COARSE_SIZE is the dimensionality of the coarse quantizer (IVF, Residual)
 // * FINE_SIZE is the dimensionality of the ProductQuantizer dsq
 // * COARSE_BITS is the number of bits that are needed to represent a coarse
 //   quantizer code.
+// * FINE_BITS is the number of bits that are needed to represent a fine
+//   quantizer code.
 // For example, "IVF256,PQ8np" for 160-dim data translates into
 //   Index2LevelDecoder<160,160,20,8>
 // For example, "Residual4x8,PQ16" for 256-dim data translates into
 //   Index2LevelDecoder<256,64,1,8>
 // For example, "IVF1024,PQ16np" for 256-dim data translates into
-//   Index2LevelDecoder<256,256,16,16>
+//   Index2LevelDecoder<256,256,16,10>. But as there are only 1 coarse code
+//   element, Index2LevelDecoder<256,256,16,16> can be used as a faster
+//   decoder.
+// For example, "Residual4x10,PQ16x10np" for 256-dim data translates into
+//   Index2LevelDecoder<256,64,16,10,10>
+//
+// Additional supported values for COARSE_BITS and FINE_BITS may be added later.
 //
-// Additional supported COARSE_BITS values may be added later.
-// FINE_BITS might be added later.
+// Second one:
+//   {
+//     template <
+//        intptr_t DIM,
+//        intptr_t FINE_SIZE,
+//        intptr_t FINE_BITS = 8>
+//     struct IndexPQDecoder { /*...*/ };
+//   }
+// * DIM is the dimensionality of data
+// * FINE_SIZE is the dimensionality of the ProductQuantizer dsq
+// * FINE_BITS is the number of bits that are needed to represent a fine
+//   quantizer code.
+// For example, "PQ8np" for 160-dim data translates into
+//   IndexPQDecoder<160,20>
 //
 // Unlike the general purpose version in faiss::Index::sa_decode(),
-//   this version provides the following functions:
+//   this version provides the following functions (please note that
+//   pqCoarseCentroids params are not available for IndexPQDecoder,
+//   but the functionality is the same as for Index2LevelDecoder):
 // * ::store, which is similar to sa_decode(1, input, output),
 //   The method signature is the following:
 //   {
@@ -101,8 +135,11 @@
 
 #ifdef __AVX2__
 #include <faiss/cppcontrib/sa_decode/Level2-avx2-inl.h>
+#include <faiss/cppcontrib/sa_decode/PQ-avx2-inl.h>
 #elif defined(__ARM_NEON)
 #include <faiss/cppcontrib/sa_decode/Level2-neon-inl.h>
+#include <faiss/cppcontrib/sa_decode/PQ-neon-inl.h>
 #else
 #include <faiss/cppcontrib/sa_decode/Level2-inl.h>
+#include <faiss/cppcontrib/sa_decode/PQ-inl.h>
 #endif