Add sa_decode() to IndexIVFAdditiveQuantizer (#2362)

Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/2362

Reviewed By: mdouze

Differential Revision: D37280450

fbshipit-source-id: 610b553e8219df9a9f52442ffc3942036f47284a

faiss/IndexIVFAdditiveQuantizer.cpp

@@ -100,6 +100,32 @@ void IndexIVFAdditiveQuantizer::encode_vectors(
     }
 }
 
+void IndexIVFAdditiveQuantizer::sa_decode(
+        idx_t n,
+        const uint8_t* codes,
+        float* x) const {
+    const size_t coarse_size = coarse_code_size();
+
+#pragma omp parallel if (n > 1000)
+    {
+        std::vector<float> residual(d);
+
+#pragma omp for
+        for (idx_t i = 0; i < n; i++) {
+            const uint8_t* code = codes + i * (code_size + coarse_size);
+            int64_t list_no = decode_listno(code);
+            float* xi = x + i * d;
+            aq->decode(code + coarse_size, xi, 1);
+            if (by_residual) {
+                quantizer->reconstruct(list_no, residual.data());
+                for (size_t j = 0; j < d; j++) {
+                    xi[j] += residual[j];
+                }
+            }
+        }
+    }
+}
+
 IndexIVFAdditiveQuantizer::~IndexIVFAdditiveQuantizer() {}
 
 /*********************************************

faiss/IndexIVFAdditiveQuantizer.h

@@ -51,6 +51,8 @@ struct IndexIVFAdditiveQuantizer : IndexIVF {
     InvertedListScanner* get_InvertedListScanner(
             bool store_pairs) const override;
 
+    void sa_decode(idx_t n, const uint8_t* codes, float* x) const override;
+
     ~IndexIVFAdditiveQuantizer() override;
 };
 

tests/test_standalone_codec.py

@@ -73,6 +73,9 @@ class TestEncodeDecode(unittest.TestCase):
     def test_LSH(self):
         self.do_encode_twice('LSHrt')
 
+    def test_RQ6x8(self):
+        self.do_encode_twice('RQ6x8')
+
 
 class TestIndexEquiv(unittest.TestCase):