improve PQ/IVFPQ code_distance() for avx2 for 4 and 8 subquantizers (#…

…2831) Summary: Pull Request resolved: #2831 Reviewed By: mdouze Differential Revision: D45329803 fbshipit-source-id: 627be46036b313f4be05e9a54b13562b414d2789
facebookresearch · May 4, 2023 · 29318b3 · 29318b3
1 parent bc12015
commit 29318b3
Show file tree

Hide file tree

Showing 7 changed files with 585 additions and 74 deletions.
diff --git a/faiss/IndexIVFPQ.cpp b/faiss/IndexIVFPQ.cpp
@@ -937,7 +937,8 @@ struct IVFPQScannerT : QueryTables {
                 float distance_2 = 0;
                 float distance_3 = 0;
                 distance_four_codes<PQDecoder>(
-                        pq,
+                        pq.M,
+                        pq.nbits,
                         sim_table,
                         codes + saved_j[0] * pq.code_size,
                         codes + saved_j[1] * pq.code_size,
@@ -957,24 +958,30 @@ struct IVFPQScannerT : QueryTables {
         }
 
         if (counter >= 1) {
-            float dis =
-                    dis0 +
+            float dis = dis0 +
                     distance_single_code<PQDecoder>(
-                            pq, sim_table, codes + saved_j[0] * pq.code_size);
+                                pq.M,
+                                pq.nbits,
+                                sim_table,
+                                codes + saved_j[0] * pq.code_size);
             res.add(saved_j[0], dis);
         }
         if (counter >= 2) {
-            float dis =
-                    dis0 +
+            float dis = dis0 +
                     distance_single_code<PQDecoder>(
-                            pq, sim_table, codes + saved_j[1] * pq.code_size);
+                                pq.M,
+                                pq.nbits,
+                                sim_table,
+                                codes + saved_j[1] * pq.code_size);
             res.add(saved_j[1], dis);
         }
         if (counter >= 3) {
-            float dis =
-                    dis0 +
+            float dis = dis0 +
                     distance_single_code<PQDecoder>(
-                            pq, sim_table, codes + saved_j[2] * pq.code_size);
+                                pq.M,
+                                pq.nbits,
+                                sim_table,
+                                codes + saved_j[2] * pq.code_size);
             res.add(saved_j[2], dis);
         }
     }
@@ -1137,7 +1144,8 @@ struct IVFPQScannerT : QueryTables {
                 float distance_2 = dis0;
                 float distance_3 = dis0;
                 distance_four_codes<PQDecoder>(
-                        pq,
+                        pq.M,
+                        pq.nbits,
                         sim_table,
                         codes + saved_j[0] * pq.code_size,
                         codes + saved_j[1] * pq.code_size,
@@ -1165,10 +1173,12 @@ struct IVFPQScannerT : QueryTables {
         for (size_t kk = 0; kk < counter; kk++) {
             n_hamming_pass++;
 
-            float dis =
-                    dis0 +
+            float dis = dis0 +
                     distance_single_code<PQDecoder>(
-                            pq, sim_table, codes + saved_j[kk] * pq.code_size);
+                                pq.M,
+                                pq.nbits,
+                                sim_table,
+                                codes + saved_j[kk] * pq.code_size);
 
             res.add(saved_j[kk], dis);
         }
@@ -1185,7 +1195,10 @@ struct IVFPQScannerT : QueryTables {
 
                 float dis = dis0 +
                         distance_single_code<PQDecoder>(
-                                    pq, sim_table, codes + j * code_size);
+                                    pq.M,
+                                    pq.nbits,
+                                    sim_table,
+                                    codes + j * code_size);
 
                 res.add(j, dis);
             }
@@ -1263,7 +1276,7 @@ struct IVFPQScanner : IVFPQScannerT<idx_t, METRIC_TYPE, PQDecoder>,
         assert(precompute_mode == 2);
         float dis = this->dis0 +
                 distance_single_code<PQDecoder>(
-                            this->pq, this->sim_table, code);
+                            this->pq.M, this->pq.nbits, this->sim_table, code);
         return dis;
     }
 

diff --git a/faiss/IndexPQ.cpp b/faiss/IndexPQ.cpp
@@ -86,7 +86,7 @@ struct PQDistanceComputer : FlatCodesDistanceComputer {
         ndis++;
 
         float dis = distance_single_code<PQDecoder>(
-                pq, precomputed_table.data(), code);
+                pq.M, pq.nbits, precomputed_table.data(), code);
         return dis;
     }