faiss 0.6.1 → 0.6.3

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (186) hide show
  1. checksums.yaml +4 -4
  2. data/CHANGELOG.md +8 -0
  3. data/lib/faiss/version.rb +1 -1
  4. data/vendor/faiss/faiss/AutoTune.cpp +3 -1
  5. data/vendor/faiss/faiss/Clustering.cpp +9 -1
  6. data/vendor/faiss/faiss/IVFlib.cpp +14 -3
  7. data/vendor/faiss/faiss/Index.h +2 -2
  8. data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp +9 -10
  9. data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.cpp +2 -3
  10. data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +1 -2
  11. data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +10 -12
  12. data/vendor/faiss/faiss/IndexBinaryHash.cpp +5 -9
  13. data/vendor/faiss/faiss/IndexBinaryIVF.cpp +5 -7
  14. data/vendor/faiss/faiss/IndexEDEN.cpp +273 -0
  15. data/vendor/faiss/faiss/IndexEDEN.h +57 -0
  16. data/vendor/faiss/faiss/IndexFastScan.cpp +15 -4
  17. data/vendor/faiss/faiss/IndexFlat.cpp +13 -50
  18. data/vendor/faiss/faiss/IndexHNSW.cpp +177 -148
  19. data/vendor/faiss/faiss/IndexIDMap.cpp +16 -3
  20. data/vendor/faiss/faiss/IndexIDMap.h +2 -0
  21. data/vendor/faiss/faiss/IndexIVF.cpp +19 -8
  22. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp +3 -3
  23. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.cpp +3 -4
  24. data/vendor/faiss/faiss/IndexIVFEDEN.cpp +302 -0
  25. data/vendor/faiss/faiss/IndexIVFEDEN.h +70 -0
  26. data/vendor/faiss/faiss/IndexIVFFastScan.cpp +5 -6
  27. data/vendor/faiss/faiss/IndexIVFFlat.cpp +6 -5
  28. data/vendor/faiss/faiss/IndexIVFFlatPanorama.cpp +3 -3
  29. data/vendor/faiss/faiss/IndexIVFIndependentQuantizer.cpp +1 -1
  30. data/vendor/faiss/faiss/IndexIVFPQ.cpp +42 -25
  31. data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +0 -1
  32. data/vendor/faiss/faiss/IndexIVFPQR.cpp +2 -3
  33. data/vendor/faiss/faiss/IndexIVFRaBitQ.cpp +23 -62
  34. data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.cpp +180 -76
  35. data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.h +5 -4
  36. data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +8 -6
  37. data/vendor/faiss/faiss/IndexLSH.cpp +2 -3
  38. data/vendor/faiss/faiss/IndexLattice.cpp +5 -0
  39. data/vendor/faiss/faiss/IndexNNDescent.cpp +10 -3
  40. data/vendor/faiss/faiss/IndexNSG.cpp +8 -4
  41. data/vendor/faiss/faiss/IndexPQ.cpp +6 -8
  42. data/vendor/faiss/faiss/IndexPreTransform.cpp +15 -0
  43. data/vendor/faiss/faiss/IndexRaBitQ.cpp +2 -2
  44. data/vendor/faiss/faiss/IndexRaBitQFastScan.cpp +1 -2
  45. data/vendor/faiss/faiss/IndexRaBitQFastScan.h +5 -1
  46. data/vendor/faiss/faiss/IndexRefine.cpp +30 -1
  47. data/vendor/faiss/faiss/IndexReplicas.cpp +1 -2
  48. data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +68 -6
  49. data/vendor/faiss/faiss/IndexScalarQuantizer.h +10 -0
  50. data/vendor/faiss/faiss/IndexShards.cpp +2 -2
  51. data/vendor/faiss/faiss/IndexShardsIVF.cpp +2 -2
  52. data/vendor/faiss/faiss/MetaIndexes.cpp +2 -4
  53. data/vendor/faiss/faiss/SuperKMeans.cpp +256 -240
  54. data/vendor/faiss/faiss/SuperKMeans.h +30 -0
  55. data/vendor/faiss/faiss/VectorTransform.cpp +33 -2
  56. data/vendor/faiss/faiss/clone_index.cpp +5 -0
  57. data/vendor/faiss/faiss/cppcontrib/SaDecodeKernels.h +1 -1
  58. data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-neon-inl.h +902 -12
  59. data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-neon-inl.h +702 -10
  60. data/vendor/faiss/faiss/factory_tools.cpp +51 -4
  61. data/vendor/faiss/faiss/gpu/GpuCloner.cpp +11 -11
  62. data/vendor/faiss/faiss/gpu/GpuIndex.h +34 -11
  63. data/vendor/faiss/faiss/gpu/GpuIndexCagra.h +47 -0
  64. data/vendor/faiss/faiss/gpu/GpuIndexIVF.h +17 -0
  65. data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h +16 -0
  66. data/vendor/faiss/faiss/gpu/GpuResources.h +3 -2
  67. data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp +11 -12
  68. data/vendor/faiss/faiss/gpu/StandardGpuResources.h +3 -3
  69. data/vendor/faiss/faiss/gpu/perf/PerfClustering.cpp +1 -1
  70. data/vendor/faiss/faiss/gpu/perf/PerfIVFPQAdd.cpp +2 -2
  71. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFScalarQuantizer.cpp +180 -0
  72. data/vendor/faiss/faiss/gpu_metal/MetalDistance.h +87 -0
  73. data/vendor/faiss/faiss/gpu_metal/MetalIndex.h +7 -0
  74. data/vendor/faiss/faiss/gpu_metal/MetalIndexIVFFlat.h +177 -0
  75. data/vendor/faiss/faiss/gpu_metal/MetalIndexIVFPQ.h +88 -0
  76. data/vendor/faiss/faiss/gpu_metal/MetalKernels.h +48 -3
  77. data/vendor/faiss/faiss/gpu_metal/MetalPythonBridge.h +45 -0
  78. data/vendor/faiss/faiss/gpu_metal/impl/MetalIVFFlat.h +193 -0
  79. data/vendor/faiss/faiss/gpu_metal/impl/MetalIVFPQ.h +134 -0
  80. data/vendor/faiss/faiss/impl/ClusteringInitialization.cpp +2 -2
  81. data/vendor/faiss/faiss/impl/DistanceComputer.h +34 -0
  82. data/vendor/faiss/faiss/impl/EDENQuantizer.h +119 -0
  83. data/vendor/faiss/faiss/impl/HNSW.cpp +658 -344
  84. data/vendor/faiss/faiss/impl/HNSW.h +51 -13
  85. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +2 -2
  86. data/vendor/faiss/faiss/impl/NSG.cpp +18 -12
  87. data/vendor/faiss/faiss/impl/Panorama.h +20 -7
  88. data/vendor/faiss/faiss/impl/PolysemousTraining.cpp +152 -84
  89. data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +59 -24
  90. data/vendor/faiss/faiss/impl/RaBitQUtils.cpp +45 -37
  91. data/vendor/faiss/faiss/impl/RaBitQUtils.h +35 -0
  92. data/vendor/faiss/faiss/impl/RaBitQuantizer.cpp +175 -68
  93. data/vendor/faiss/faiss/impl/RaBitQuantizer.h +19 -0
  94. data/vendor/faiss/faiss/impl/RaBitQuantizerMultiBit.cpp +2 -11
  95. data/vendor/faiss/faiss/impl/ResultHandler.h +26 -31
  96. data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +522 -58
  97. data/vendor/faiss/faiss/impl/ScalarQuantizer.h +70 -0
  98. data/vendor/faiss/faiss/impl/ThreadedIndex-inl.h +2 -2
  99. data/vendor/faiss/faiss/impl/VisitedTable.cpp +33 -13
  100. data/vendor/faiss/faiss/impl/VisitedTable.h +88 -33
  101. data/vendor/faiss/faiss/impl/binary_hamming/IndexBinaryIVF_impl.h +1 -1
  102. data/vendor/faiss/faiss/impl/binary_hamming/avx2.cpp +4 -4
  103. data/vendor/faiss/faiss/impl/fast_scan/dispatching.h +38 -3
  104. data/vendor/faiss/faiss/impl/hnsw/LockVector.cpp +1 -1
  105. data/vendor/faiss/faiss/impl/hnsw/MinimaxHeap.cpp +35 -43
  106. data/vendor/faiss/faiss/impl/hnsw/MinimaxHeap.h +64 -15
  107. data/vendor/faiss/faiss/impl/hnsw/avx2.cpp +86 -40
  108. data/vendor/faiss/faiss/impl/hnsw/avx512.cpp +81 -50
  109. data/vendor/faiss/faiss/impl/index_read.cpp +476 -75
  110. data/vendor/faiss/faiss/impl/index_write.cpp +56 -4
  111. data/vendor/faiss/faiss/impl/io_macros.h +25 -0
  112. data/vendor/faiss/faiss/impl/lattice_Zn.cpp +8 -9
  113. data/vendor/faiss/faiss/impl/platform_macros.h +15 -9
  114. data/vendor/faiss/faiss/impl/polysemous_training/avx512.cpp +284 -0
  115. data/vendor/faiss/faiss/impl/polysemous_training/dispatch.h +115 -0
  116. data/vendor/faiss/faiss/impl/pq_code_distance/IVFPQ_QueryTables.cpp +0 -1
  117. data/vendor/faiss/faiss/impl/pq_code_distance/PQDistanceComputer_impl.h +26 -15
  118. data/vendor/faiss/faiss/impl/pq_code_distance/avx2.cpp +6 -4
  119. data/vendor/faiss/faiss/impl/pq_code_distance/avx512.cpp +2 -0
  120. data/vendor/faiss/faiss/impl/pq_code_distance/neon.cpp +2 -0
  121. data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-generic.cpp +20 -0
  122. data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-inl.h +36 -0
  123. data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-sve.cpp +5 -0
  124. data/vendor/faiss/faiss/impl/pq_code_distance/pq_scan_impl.h +105 -0
  125. data/vendor/faiss/faiss/impl/pq_code_distance/rvv.cpp +2 -0
  126. data/vendor/faiss/faiss/impl/result_handler/ResultHandler.cpp +195 -0
  127. data/vendor/faiss/faiss/impl/result_handler/avx2.cpp +133 -0
  128. data/vendor/faiss/faiss/impl/result_handler/avx512.cpp +281 -0
  129. data/vendor/faiss/faiss/impl/scalar_quantizer/EDENQuantizer-avx2.cpp +72 -0
  130. data/vendor/faiss/faiss/impl/scalar_quantizer/EDENQuantizer-avx512.cpp +228 -0
  131. data/vendor/faiss/faiss/impl/scalar_quantizer/EDENQuantizer.cpp +882 -0
  132. data/vendor/faiss/faiss/impl/scalar_quantizer/distance_computers.h +6 -0
  133. data/vendor/faiss/faiss/impl/scalar_quantizer/quantizers.h +336 -26
  134. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx2.cpp +331 -32
  135. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx512-impl.h +553 -0
  136. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx512-spr.cpp +558 -0
  137. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx512.cpp +284 -45
  138. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-dispatch.h +502 -3
  139. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-neon.cpp +157 -32
  140. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-rvv.cpp +26 -0
  141. data/vendor/faiss/faiss/impl/simd_dispatch.h +86 -8
  142. data/vendor/faiss/faiss/index_factory.cpp +37 -7
  143. data/vendor/faiss/faiss/index_io.h +16 -0
  144. data/vendor/faiss/faiss/invlists/DirectMap.cpp +5 -2
  145. data/vendor/faiss/faiss/invlists/InvertedLists.cpp +15 -15
  146. data/vendor/faiss/faiss/invlists/InvertedLists.h +2 -2
  147. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +19 -4
  148. data/vendor/faiss/faiss/python/python_callbacks.cpp +3 -1
  149. data/vendor/faiss/faiss/svs/IndexSVSFaissUtils.h +60 -0
  150. data/vendor/faiss/faiss/svs/IndexSVSFlat.cpp +26 -1
  151. data/vendor/faiss/faiss/svs/IndexSVSFlat.h +13 -0
  152. data/vendor/faiss/faiss/svs/IndexSVSIVF.cpp +1 -1
  153. data/vendor/faiss/faiss/svs/IndexSVSIVFLeanVec.cpp +1 -1
  154. data/vendor/faiss/faiss/svs/IndexSVSVamana.cpp +150 -23
  155. data/vendor/faiss/faiss/svs/IndexSVSVamana.h +30 -7
  156. data/vendor/faiss/faiss/svs/IndexSVSVamanaLVQ.cpp +3 -2
  157. data/vendor/faiss/faiss/svs/IndexSVSVamanaLVQ.h +2 -1
  158. data/vendor/faiss/faiss/svs/IndexSVSVamanaLeanVec.cpp +65 -25
  159. data/vendor/faiss/faiss/svs/IndexSVSVamanaLeanVec.h +3 -2
  160. data/vendor/faiss/faiss/utils/approx_topk_hamming/approx_topk_hamming.h +1 -1
  161. data/vendor/faiss/faiss/utils/bf16.h +34 -0
  162. data/vendor/faiss/faiss/utils/distances.cpp +14 -2
  163. data/vendor/faiss/faiss/utils/distances_simd.cpp +4 -4
  164. data/vendor/faiss/faiss/utils/extra_distances.cpp +4 -14
  165. data/vendor/faiss/faiss/utils/extra_distances.h +1 -2
  166. data/vendor/faiss/faiss/utils/hamming.cpp +9 -9
  167. data/vendor/faiss/faiss/utils/hamming_distance/hamming_avx2.cpp +2 -1
  168. data/vendor/faiss/faiss/utils/hamming_distance/hamming_avx512_spr.cpp +15 -0
  169. data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-avx512.h +6 -30
  170. data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-avx512_spr.h +171 -0
  171. data/vendor/faiss/faiss/utils/partitioning.cpp +0 -2
  172. data/vendor/faiss/faiss/utils/quantize_lut.cpp +29 -8
  173. data/vendor/faiss/faiss/utils/rabitq_simd.h +202 -0
  174. data/vendor/faiss/faiss/utils/simd_impl/distances_avx2.cpp +0 -1
  175. data/vendor/faiss/faiss/utils/simd_impl/distances_avx512.cpp +263 -15
  176. data/vendor/faiss/faiss/utils/simd_impl/distances_rvv.cpp +160 -18
  177. data/vendor/faiss/faiss/utils/simd_impl/partitioning_simdlib256.h +14 -68
  178. data/vendor/faiss/faiss/utils/simd_impl/rabitq_avx2.cpp +245 -0
  179. data/vendor/faiss/faiss/utils/simd_impl/rabitq_avx512.cpp +273 -0
  180. data/vendor/faiss/faiss/utils/simd_impl/rabitq_avx512_spr.cpp +435 -0
  181. data/vendor/faiss/faiss/utils/simd_impl/rabitq_neon.cpp +11 -0
  182. data/vendor/faiss/faiss/utils/simd_impl/rabitq_rvv.cpp +143 -6
  183. data/vendor/faiss/faiss/utils/simd_levels.cpp +56 -2
  184. data/vendor/faiss/faiss/utils/simd_levels.h +14 -0
  185. data/vendor/faiss/faiss/utils/utils.cpp +9 -27
  186. metadata +27 -2
@@ -25,7 +25,9 @@ namespace faiss {
25
25
  // Forward declarations to avoid circular dependency.
26
26
  struct IndexHNSW;
27
27
  struct IndexHNSWFlatPanorama;
28
- struct MinimaxHeap;
28
+ template <class HC_>
29
+ struct MinimaxHeapT;
30
+ using MinimaxHeap = MinimaxHeapT<CMax<float, int32_t>>;
29
31
  class LockVector;
30
32
 
31
33
  /** Implementation of the Hierarchical Navigable Small World
@@ -60,30 +62,52 @@ struct HNSW {
60
62
  /// internal storage of vectors (32 bits: this is expensive)
61
63
  using storage_idx_t = int32_t;
62
64
 
63
- // for now we do only these distances
64
- using C = CMax<float, int64_t>;
65
+ // The two comparator flavors HNSW supports. CMax (smaller-is-better)
66
+ // is the default; CMin (larger-is-better) is used when `is_similarity`
67
+ // is set on the owning index.
68
+ using C_distance = CMax<float, int64_t>;
69
+ using C_similarity = CMin<float, int64_t>;
70
+
71
+ // Back-compat alias: keeps `HNSW::C` resolving to the distance
72
+ // (CMax) comparator everywhere the type is referenced directly.
73
+ using C = C_distance;
65
74
 
66
75
  typedef std::pair<float, storage_idx_t> Node;
67
76
 
68
77
  /// to sort pairs of (id, distance) from nearest to farthest or the reverse
69
- struct NodeDistCloser {
78
+ template <class CT>
79
+ struct NodeDistCloserT {
70
80
  float d;
71
81
  int id;
72
- NodeDistCloser(float d_in, int id_in) : d(d_in), id(id_in) {}
73
- bool operator<(const NodeDistCloser& obj1) const {
74
- return d < obj1.d;
82
+ NodeDistCloserT(float d_in, int id_in) : d(d_in), id(id_in) {}
83
+ bool operator<(const NodeDistCloserT& obj1) const {
84
+ // priority_queue keeps the "worst" element at the top so that
85
+ // when the queue is full we can pop it. For CMax (distance) the
86
+ // worst element is the largest d; for CMin (similarity) it is
87
+ // the smallest d. Equivalent to: obj1.d "better than" d.
88
+ return CT::cmp(obj1.d, d);
75
89
  }
76
90
  };
77
91
 
78
- struct NodeDistFarther {
92
+ template <class CT>
93
+ struct NodeDistFartherT {
79
94
  float d;
80
95
  int id;
81
- NodeDistFarther(float d_in, int id_in) : d(d_in), id(id_in) {}
82
- bool operator<(const NodeDistFarther& obj1) const {
83
- return d > obj1.d;
96
+ NodeDistFartherT(float d_in, int id_in) : d(d_in), id(id_in) {}
97
+ bool operator<(const NodeDistFartherT& obj1) const {
98
+ // priority_queue here keeps the "best" element at the top so we
99
+ // can process the nearest candidate first. For CMax (distance)
100
+ // the best is the smallest d; for CMin (similarity) the best is
101
+ // the largest d. Equivalent to: d "better than" obj1.d.
102
+ return CT::cmp(d, obj1.d);
84
103
  }
85
104
  };
86
105
 
106
+ // Back-compat aliases: default to the distance (CMax) comparator so
107
+ // existing call sites that mention `HNSW::NodeDist*` keep working.
108
+ using NodeDistCloser = NodeDistCloserT<C_distance>;
109
+ using NodeDistFarther = NodeDistFartherT<C_distance>;
110
+
87
111
  /// assignment probability to each layer (sum=1)
88
112
  std::vector<double> assign_probas;
89
113
 
@@ -131,6 +155,12 @@ struct HNSW {
131
155
  /// use Panorama progressive pruning in search
132
156
  bool is_panorama = false;
133
157
 
158
+ /// distance comparison semantics: when true, distances are treated as
159
+ /// similarity scores (larger is better). Default false matches the
160
+ /// historical L2/Hamming behavior (smaller is better).
161
+ /// Not serialized: must be re-set by the owning Index after loading.
162
+ bool is_similarity = false;
163
+
134
164
  // See impl/VisitedTable.h.
135
165
  std::optional<bool> use_visited_hashset;
136
166
 
@@ -216,10 +246,11 @@ struct HNSW {
216
246
 
217
247
  int prepare_level_tab(size_t n, bool preset_levels = false);
218
248
 
249
+ template <class C = C_distance>
219
250
  static void shrink_neighbor_list(
220
251
  DistanceComputer& qdis,
221
- std::priority_queue<NodeDistFarther>& input,
222
- std::vector<NodeDistFarther>& output,
252
+ std::priority_queue<NodeDistFartherT<C>>& input,
253
+ std::vector<NodeDistFartherT<C>>& output,
223
254
  size_t max_size,
224
255
  bool keep_max_size_level0 = false);
225
256
 
@@ -250,6 +281,11 @@ struct HNSWStats {
250
281
  // global var that collects them all
251
282
  FAISS_API extern HNSWStats hnsw_stats;
252
283
 
284
+ /// Internal HNSW algorithm helpers. These are not part of the public API; they
285
+ /// are exposed here only so that unit tests (and a few cross-TU callers such as
286
+ /// the Panorama search variant) can reach them.
287
+ namespace hnsw_detail {
288
+
253
289
  int search_from_candidates(
254
290
  const HNSW& hnsw,
255
291
  DistanceComputer& qdis,
@@ -302,4 +338,6 @@ void search_neighbors_to_add(
302
338
  VisitedTable& vt,
303
339
  bool reference_version = false);
304
340
 
341
+ } // namespace hnsw_detail
342
+
305
343
  } // namespace faiss
@@ -597,8 +597,8 @@ void LocalSearchQuantizer::icm_encode_step(
597
597
  const float* binaries,
598
598
  size_t n,
599
599
  size_t n_iters) const {
600
- FAISS_THROW_IF_NOT(M != 0 && K != 0);
601
- FAISS_THROW_IF_NOT(binaries != nullptr);
600
+ FAISS_THROW_IF_MSG(M == 0 || K == 0, "M and K must be nonzero");
601
+ FAISS_THROW_IF_NOT(binaries);
602
602
 
603
603
  // Resolve SIMD level once, not per iteration of the n × n_iters × M loop.
604
604
  with_simd_level_256bit([&]<SIMDLevel SL>() {
@@ -142,7 +142,9 @@ void NSG::build(
142
142
  idx_t n,
143
143
  const nsg::Graph<idx_t>& knn_graph,
144
144
  bool verbose) {
145
- FAISS_THROW_IF_NOT(!is_built && ntotal == 0);
145
+ FAISS_THROW_IF_MSG(
146
+ is_built || ntotal != 0,
147
+ "NSG graph must be empty and not yet built");
146
148
 
147
149
  if (verbose) {
148
150
  printf("NSG::build R=%d, L=%d, C=%d\n", R, L, C);
@@ -234,10 +236,11 @@ void NSG::init_graph(Index* storage, const nsg::Graph<idx_t>& knn_graph) {
234
236
  std::unique_ptr<DistanceComputer> dis(storage_distance_computer(storage));
235
237
 
236
238
  dis->set_query(center.get());
237
- VisitedTable vt(ntotal, use_visited_hashset);
239
+ std::unique_ptr<VisitedTable> vt =
240
+ VisitedTable::create(ntotal, use_visited_hashset);
238
241
 
239
242
  // Do not collect the visited nodes
240
- search_on_graph<false>(knn_graph, *dis, vt, ep, L, retset, tmpset);
243
+ search_on_graph<false>(knn_graph, *dis, *vt, ep, L, retset, tmpset);
241
244
 
242
245
  // set enterpoint
243
246
  enterpoint = retset[0].id;
@@ -344,7 +347,8 @@ void NSG::link(
344
347
  std::vector<Node> pool;
345
348
  std::vector<Neighbor> tmp;
346
349
 
347
- VisitedTable vt(ntotal, use_visited_hashset);
350
+ std::unique_ptr<VisitedTable> vt =
351
+ VisitedTable::create(ntotal, use_visited_hashset);
348
352
  std::unique_ptr<DistanceComputer> dis(
349
353
  storage_distance_computer(storage));
350
354
 
@@ -355,13 +359,13 @@ void NSG::link(
355
359
 
356
360
  // Collect the visited nodes into pool
357
361
  search_on_graph<true>(
358
- knn_graph, *dis, vt, enterpoint, L, tmp, pool);
362
+ knn_graph, *dis, *vt, enterpoint, L, tmp, pool);
359
363
 
360
- sync_prune(i, pool, *dis, vt, knn_graph, graph);
364
+ sync_prune(i, pool, *dis, *vt, knn_graph, graph);
361
365
 
362
366
  pool.clear();
363
367
  tmp.clear();
364
- vt.advance();
368
+ vt->advance();
365
369
  }
366
370
  } // omp parallel
367
371
 
@@ -531,19 +535,21 @@ void NSG::add_reverse_links(
531
535
 
532
536
  int NSG::tree_grow(Index* storage, std::vector<int>& degrees) {
533
537
  int root = enterpoint;
534
- VisitedTable vt(ntotal, use_visited_hashset);
535
- VisitedTable vt2(ntotal, use_visited_hashset);
538
+ std::unique_ptr<VisitedTable> vt =
539
+ VisitedTable::create(ntotal, use_visited_hashset);
540
+ std::unique_ptr<VisitedTable> vt2 =
541
+ VisitedTable::create(ntotal, use_visited_hashset);
536
542
 
537
543
  int num_attached = 0;
538
544
  int cnt = 0;
539
545
  while (true) {
540
- cnt = dfs(vt, root, cnt);
546
+ cnt = dfs(*vt, root, cnt);
541
547
  if (cnt >= ntotal) {
542
548
  break;
543
549
  }
544
550
 
545
- root = attach_unlinked(storage, vt, vt2, degrees);
546
- vt2.advance();
551
+ root = attach_unlinked(storage, *vt, *vt2, degrees);
552
+ vt2->advance();
547
553
  num_attached += 1;
548
554
  }
549
555
 
@@ -36,7 +36,12 @@ namespace faiss {
36
36
  /// from active_indices (subsequent levels after pruning).
37
37
  /// @tparam LevelWidth Compile-time level width in floats (0 = use runtime
38
38
  /// level_width_dims). Enables full loop unrolling.
39
+ // Skip pragmas under nvcc: its EDG frontend warns on `#pragma GCC optimize`
40
+ // (#1675-D) for every `.cu` that transitively includes this header. These
41
+ // templates are CPU-only, so the hint is irrelevant during nvcc parse.
42
+ #if !defined(__NVCC__)
39
43
  FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
44
+ #endif
40
45
  template <bool AllActive = false, size_t LevelWidth = 0>
41
46
  static inline void compute_level_dot_kernel(
42
47
  const float* FAISS_RESTRICT query_level,
@@ -44,18 +49,20 @@ static inline void compute_level_dot_kernel(
44
49
  const uint32_t* active_indices,
45
50
  const size_t num_active,
46
51
  const size_t level_width_dims,
47
- float* FAISS_RESTRICT dot_products) {
52
+ float* FAISS_RESTRICT dot_products,
53
+ size_t stride = 0) {
48
54
  const size_t width = LevelWidth > 0 ? LevelWidth : level_width_dims;
55
+ const size_t row_stride = stride == 0 ? width : stride;
49
56
  size_t i = 0;
50
57
  for (; i + 4 <= num_active; i += 4) {
51
58
  const float* y0 = level_storage +
52
- (AllActive ? (i + 0) : active_indices[i + 0]) * width;
59
+ (AllActive ? (i + 0) : active_indices[i + 0]) * row_stride;
53
60
  const float* y1 = level_storage +
54
- (AllActive ? (i + 1) : active_indices[i + 1]) * width;
61
+ (AllActive ? (i + 1) : active_indices[i + 1]) * row_stride;
55
62
  const float* y2 = level_storage +
56
- (AllActive ? (i + 2) : active_indices[i + 2]) * width;
63
+ (AllActive ? (i + 2) : active_indices[i + 2]) * row_stride;
57
64
  const float* y3 = level_storage +
58
- (AllActive ? (i + 3) : active_indices[i + 3]) * width;
65
+ (AllActive ? (i + 3) : active_indices[i + 3]) * row_stride;
59
66
 
60
67
  float dp0 = 0, dp1 = 0, dp2 = 0, dp3 = 0;
61
68
  FAISS_PRAGMA_IMPRECISE_LOOP
@@ -73,8 +80,8 @@ static inline void compute_level_dot_kernel(
73
80
  dot_products[i + 3] = dp3;
74
81
  }
75
82
  for (; i < num_active; i++) {
76
- const float* yj =
77
- level_storage + (AllActive ? i : active_indices[i]) * width;
83
+ const float* yj = level_storage +
84
+ (AllActive ? i : active_indices[i]) * row_stride;
78
85
  float dp = 0;
79
86
  FAISS_PRAGMA_IMPRECISE_LOOP
80
87
  for (size_t j = 0; j < width; j++) {
@@ -83,7 +90,9 @@ static inline void compute_level_dot_kernel(
83
90
  dot_products[i] = dp;
84
91
  }
85
92
  }
93
+ #if !defined(__NVCC__)
86
94
  FAISS_PRAGMA_IMPRECISE_FUNCTION_END
95
+ #endif
87
96
 
88
97
  /// Update exact distances with the current level's dot products, then apply
89
98
  /// Panorama pruning: for each active vector, compute a lower bound on
@@ -92,7 +101,9 @@ FAISS_PRAGMA_IMPRECISE_FUNCTION_END
92
101
  ///
93
102
  /// Uses `if constexpr` on C::is_max rather than C::cmp() to ensure the
94
103
  /// comparison autovectorizes (C::cmp generates scalar function calls).
104
+ #if !defined(__NVCC__)
95
105
  FAISS_PRAGMA_IMPRECISE_FUNCTION_BEGIN
106
+ #endif
96
107
  template <bool AllActive, typename C, MetricType M>
97
108
  static inline void prune_kernel(
98
109
  float* FAISS_RESTRICT exact_distances,
@@ -128,7 +139,9 @@ static inline void prune_kernel(
128
139
  }
129
140
  }
130
141
  }
142
+ #if !defined(__NVCC__)
131
143
  FAISS_PRAGMA_IMPRECISE_FUNCTION_END
144
+ #endif
132
145
 
133
146
  /// Compact active_indices in-place, removing entries where active_byteset[i]
134
147
  /// is zero. Returns the new count of active elements. Uses a branchless BMI2 +
@@ -26,6 +26,8 @@
26
26
 
27
27
  #include <faiss/impl/FaissAssert.h>
28
28
 
29
+ #include <faiss/impl/polysemous_training/dispatch.h>
30
+
29
31
  /*****************************************
30
32
  * Mixed PQ / Hamming
31
33
  ******************************************/
@@ -168,6 +170,94 @@ static inline int hamming_dis(uint64_t a, uint64_t b) {
168
170
  return popcount64(a ^ b);
169
171
  }
170
172
 
173
+ static inline double sqr(double x) {
174
+ return x * x;
175
+ }
176
+
177
+ // Scalar (NONE) kernels for the two objectives. The dispatch boundary lives in
178
+ // the objective methods below, which route to these or to the AVX-512
179
+ // specializations (in polysemous_training/avx512.cpp) via
180
+ // with_selected_simd_levels.
181
+ namespace polysemous_training {
182
+
183
+ template <>
184
+ double hamming_compute_cost<SIMDLevel::NONE>(
185
+ int n,
186
+ const int* perm,
187
+ const double* target_dis,
188
+ const double* weights) {
189
+ double cost = 0;
190
+ for (int i = 0; i < n; i++) {
191
+ for (int j = 0; j < n; j++) {
192
+ double wanted = target_dis[i * n + j];
193
+ double w = weights[i * n + j];
194
+ double actual = hamming_dis(perm[i], perm[j]);
195
+ cost += w * sqr(wanted - actual);
196
+ }
197
+ }
198
+ return cost;
199
+ }
200
+
201
+ template <>
202
+ double hamming_cost_update<SIMDLevel::NONE>(
203
+ int n,
204
+ const int* perm,
205
+ int iw,
206
+ int jw,
207
+ const double* target_dis,
208
+ const double* weights) {
209
+ double delta_cost = 0;
210
+
211
+ for (int i = 0; i < n; i++) {
212
+ if (i == iw) {
213
+ for (int j = 0; j < n; j++) {
214
+ double wanted = target_dis[i * n + j], w = weights[i * n + j];
215
+ double actual = hamming_dis(perm[i], perm[j]);
216
+ delta_cost -= w * sqr(wanted - actual);
217
+ double new_actual = hamming_dis(
218
+ perm[jw],
219
+ perm[j == iw ? jw
220
+ : j == jw ? iw
221
+ : j]);
222
+ delta_cost += w * sqr(wanted - new_actual);
223
+ }
224
+ } else if (i == jw) {
225
+ for (int j = 0; j < n; j++) {
226
+ double wanted = target_dis[i * n + j], w = weights[i * n + j];
227
+ double actual = hamming_dis(perm[i], perm[j]);
228
+ delta_cost -= w * sqr(wanted - actual);
229
+ double new_actual = hamming_dis(
230
+ perm[iw],
231
+ perm[j == iw ? jw
232
+ : j == jw ? iw
233
+ : j]);
234
+ delta_cost += w * sqr(wanted - new_actual);
235
+ }
236
+ } else {
237
+ int j = iw;
238
+ {
239
+ double wanted = target_dis[i * n + j], w = weights[i * n + j];
240
+ double actual = hamming_dis(perm[i], perm[j]);
241
+ delta_cost -= w * sqr(wanted - actual);
242
+ double new_actual = hamming_dis(perm[i], perm[jw]);
243
+ delta_cost += w * sqr(wanted - new_actual);
244
+ }
245
+ j = jw;
246
+ {
247
+ double wanted = target_dis[i * n + j], w = weights[i * n + j];
248
+ double actual = hamming_dis(perm[i], perm[j]);
249
+ delta_cost -= w * sqr(wanted - actual);
250
+ double new_actual = hamming_dis(perm[i], perm[iw]);
251
+ delta_cost += w * sqr(wanted - new_actual);
252
+ }
253
+ }
254
+ }
255
+
256
+ return delta_cost;
257
+ }
258
+
259
+ } // namespace polysemous_training
260
+
171
261
  namespace {
172
262
 
173
263
  /// optimize permutation to reproduce a distance table with Hamming distances
@@ -175,10 +265,6 @@ struct ReproduceWithHammingObjective : PermutationObjective {
175
265
  int nbits;
176
266
  double dis_weight_factor;
177
267
 
178
- static double sqr(double x) {
179
- return x * x;
180
- }
181
-
182
268
  // weighting of distances: it is more important to reproduce small
183
269
  // distances well
184
270
  double dis_weight(double x) const {
@@ -190,73 +276,21 @@ struct ReproduceWithHammingObjective : PermutationObjective {
190
276
 
191
277
  // cost = quadratic difference between actual distance and Hamming distance
192
278
  double compute_cost(const int* perm) const override {
193
- double cost = 0;
194
- for (int i = 0; i < n; i++) {
195
- for (int j = 0; j < n; j++) {
196
- double wanted = target_dis[i * n + j];
197
- double w = weights[i * n + j];
198
- double actual = hamming_dis(perm[i], perm[j]);
199
- cost += w * sqr(wanted - actual);
200
- }
201
- }
202
- return cost;
279
+ return with_selected_simd_levels<polysemous_training::SIMD_LEVELS>(
280
+ [&]<SIMDLevel SL>() {
281
+ return polysemous_training::hamming_compute_cost<SL>(
282
+ n, perm, target_dis.data(), weights.data());
283
+ });
203
284
  }
204
285
 
205
286
  // what would the cost update be if iw and jw were swapped?
206
287
  // computed in O(n) instead of O(n^2) for the full re-computation
207
288
  double cost_update(const int* perm, int iw, int jw) const override {
208
- double delta_cost = 0;
209
-
210
- for (int i = 0; i < n; i++) {
211
- if (i == iw) {
212
- for (int j = 0; j < n; j++) {
213
- double wanted = target_dis[i * n + j],
214
- w = weights[i * n + j];
215
- double actual = hamming_dis(perm[i], perm[j]);
216
- delta_cost -= w * sqr(wanted - actual);
217
- double new_actual = hamming_dis(
218
- perm[jw],
219
- perm[j == iw ? jw
220
- : j == jw ? iw
221
- : j]);
222
- delta_cost += w * sqr(wanted - new_actual);
223
- }
224
- } else if (i == jw) {
225
- for (int j = 0; j < n; j++) {
226
- double wanted = target_dis[i * n + j],
227
- w = weights[i * n + j];
228
- double actual = hamming_dis(perm[i], perm[j]);
229
- delta_cost -= w * sqr(wanted - actual);
230
- double new_actual = hamming_dis(
231
- perm[iw],
232
- perm[j == iw ? jw
233
- : j == jw ? iw
234
- : j]);
235
- delta_cost += w * sqr(wanted - new_actual);
236
- }
237
- } else {
238
- int j = iw;
239
- {
240
- double wanted = target_dis[i * n + j],
241
- w = weights[i * n + j];
242
- double actual = hamming_dis(perm[i], perm[j]);
243
- delta_cost -= w * sqr(wanted - actual);
244
- double new_actual = hamming_dis(perm[i], perm[jw]);
245
- delta_cost += w * sqr(wanted - new_actual);
246
- }
247
- j = jw;
248
- {
249
- double wanted = target_dis[i * n + j],
250
- w = weights[i * n + j];
251
- double actual = hamming_dis(perm[i], perm[j]);
252
- delta_cost -= w * sqr(wanted - actual);
253
- double new_actual = hamming_dis(perm[i], perm[iw]);
254
- delta_cost += w * sqr(wanted - new_actual);
255
- }
256
- }
257
- }
258
-
259
- return delta_cost;
289
+ return with_selected_simd_levels<polysemous_training::SIMD_LEVELS>(
290
+ [&]<SIMDLevel SL>() {
291
+ return polysemous_training::hamming_cost_update<SL>(
292
+ n, perm, iw, jw, target_dis.data(), weights.data());
293
+ });
260
294
  }
261
295
 
262
296
  ReproduceWithHammingObjective(
@@ -306,14 +340,20 @@ double ReproduceDistancesObjective::get_source_dis(int i, int j) const {
306
340
  return source_dis[i * n + j];
307
341
  }
308
342
 
343
+ namespace polysemous_training {
344
+
309
345
  // cost = quadratic difference between actual distance and Hamming distance
310
- double ReproduceDistancesObjective::compute_cost(const int* perm) const {
346
+ template <>
347
+ double distances_compute_cost<SIMDLevel::NONE>(
348
+ const ReproduceDistancesObjective& obj,
349
+ const int* perm) {
350
+ const int n = obj.n;
311
351
  double cost = 0;
312
352
  for (int i = 0; i < n; i++) {
313
353
  for (int j = 0; j < n; j++) {
314
- double wanted = target_dis[i * n + j];
315
- double w = weights[i * n + j];
316
- double actual = get_source_dis(perm[i], perm[j]);
354
+ double wanted = obj.target_dis[i * n + j];
355
+ double w = obj.weights[i * n + j];
356
+ double actual = obj.get_source_dis(perm[i], perm[j]);
317
357
  cost += w * sqr(wanted - actual);
318
358
  }
319
359
  }
@@ -322,16 +362,22 @@ double ReproduceDistancesObjective::compute_cost(const int* perm) const {
322
362
 
323
363
  // what would the cost update be if iw and jw were swapped?
324
364
  // computed in O(n) instead of O(n^2) for the full re-computation
325
- double ReproduceDistancesObjective::cost_update(const int* perm, int iw, int jw)
326
- const {
365
+ template <>
366
+ double distances_cost_update<SIMDLevel::NONE>(
367
+ const ReproduceDistancesObjective& obj,
368
+ const int* perm,
369
+ int iw,
370
+ int jw) {
371
+ const int n = obj.n;
327
372
  double delta_cost = 0;
328
373
  for (int i = 0; i < n; i++) {
329
374
  if (i == iw) {
330
375
  for (int j = 0; j < n; j++) {
331
- double wanted = target_dis[i * n + j], w = weights[i * n + j];
332
- double actual = get_source_dis(perm[i], perm[j]);
376
+ double wanted = obj.target_dis[i * n + j],
377
+ w = obj.weights[i * n + j];
378
+ double actual = obj.get_source_dis(perm[i], perm[j]);
333
379
  delta_cost -= w * sqr(wanted - actual);
334
- double new_actual = get_source_dis(
380
+ double new_actual = obj.get_source_dis(
335
381
  perm[jw],
336
382
  perm[j == iw ? jw
337
383
  : j == jw ? iw
@@ -340,10 +386,11 @@ double ReproduceDistancesObjective::cost_update(const int* perm, int iw, int jw)
340
386
  }
341
387
  } else if (i == jw) {
342
388
  for (int j = 0; j < n; j++) {
343
- double wanted = target_dis[i * n + j], w = weights[i * n + j];
344
- double actual = get_source_dis(perm[i], perm[j]);
389
+ double wanted = obj.target_dis[i * n + j],
390
+ w = obj.weights[i * n + j];
391
+ double actual = obj.get_source_dis(perm[i], perm[j]);
345
392
  delta_cost -= w * sqr(wanted - actual);
346
- double new_actual = get_source_dis(
393
+ double new_actual = obj.get_source_dis(
347
394
  perm[iw],
348
395
  perm[j == iw ? jw
349
396
  : j == jw ? iw
@@ -353,18 +400,20 @@ double ReproduceDistancesObjective::cost_update(const int* perm, int iw, int jw)
353
400
  } else {
354
401
  int j = iw;
355
402
  {
356
- double wanted = target_dis[i * n + j], w = weights[i * n + j];
357
- double actual = get_source_dis(perm[i], perm[j]);
403
+ double wanted = obj.target_dis[i * n + j],
404
+ w = obj.weights[i * n + j];
405
+ double actual = obj.get_source_dis(perm[i], perm[j]);
358
406
  delta_cost -= w * sqr(wanted - actual);
359
- double new_actual = get_source_dis(perm[i], perm[jw]);
407
+ double new_actual = obj.get_source_dis(perm[i], perm[jw]);
360
408
  delta_cost += w * sqr(wanted - new_actual);
361
409
  }
362
410
  j = jw;
363
411
  {
364
- double wanted = target_dis[i * n + j], w = weights[i * n + j];
365
- double actual = get_source_dis(perm[i], perm[j]);
412
+ double wanted = obj.target_dis[i * n + j],
413
+ w = obj.weights[i * n + j];
414
+ double actual = obj.get_source_dis(perm[i], perm[j]);
366
415
  delta_cost -= w * sqr(wanted - actual);
367
- double new_actual = get_source_dis(perm[i], perm[iw]);
416
+ double new_actual = obj.get_source_dis(perm[i], perm[iw]);
368
417
  delta_cost += w * sqr(wanted - new_actual);
369
418
  }
370
419
  }
@@ -372,6 +421,25 @@ double ReproduceDistancesObjective::cost_update(const int* perm, int iw, int jw)
372
421
  return delta_cost;
373
422
  }
374
423
 
424
+ } // namespace polysemous_training
425
+
426
+ double ReproduceDistancesObjective::compute_cost(const int* perm) const {
427
+ return with_selected_simd_levels<polysemous_training::SIMD_LEVELS>(
428
+ [&]<SIMDLevel SL>() {
429
+ return polysemous_training::distances_compute_cost<SL>(
430
+ *this, perm);
431
+ });
432
+ }
433
+
434
+ double ReproduceDistancesObjective::cost_update(const int* perm, int iw, int jw)
435
+ const {
436
+ return with_selected_simd_levels<polysemous_training::SIMD_LEVELS>(
437
+ [&]<SIMDLevel SL>() {
438
+ return polysemous_training::distances_cost_update<SL>(
439
+ *this, perm, iw, jw);
440
+ });
441
+ }
442
+
375
443
  ReproduceDistancesObjective::ReproduceDistancesObjective(
376
444
  int n_in,
377
445
  const double* source_dis_in,