Add sa_decode() to IndexIVFAdditiveQuantizer (facebookresearch#2362)

Summary: Pull Request resolved: facebookresearch#2362 Reviewed By: mdouze Differential Revision: D37280450 fbshipit-source-id: 610b553e8219df9a9f52442ffc3942036f47284a
alcidesmorales · Jun 20, 2022 · fb8193d · fb8193d
1 parent 3986ebf
commit fb8193d
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 0 deletions.
diff --git a/faiss/IndexIVFAdditiveQuantizer.cpp b/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() {}
 
 /*********************************************

diff --git a/faiss/IndexIVFAdditiveQuantizer.h b/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;
 };
 

diff --git a/tests/test_standalone_codec.py b/tests/test_standalone_codec.py
@@ -73,6 +73,9 @@ def test_IVFPQ6x8np(self):
     def test_LSH(self):
         self.do_encode_twice('LSHrt')
 
+    def test_RQ6x8(self):
+        self.do_encode_twice('RQ6x8')
+
 
 class TestIndexEquiv(unittest.TestCase):