faiss 0.6.0 → 0.6.1

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 (361) hide show
  1. checksums.yaml +4 -4
  2. data/CHANGELOG.md +4 -0
  3. data/ext/faiss/extconf.rb +2 -1
  4. data/ext/faiss/{index_rb.cpp → index.cpp} +1 -1
  5. data/ext/faiss/index_binary.cpp +1 -1
  6. data/ext/faiss/kmeans.cpp +1 -1
  7. data/ext/faiss/pca_matrix.cpp +1 -1
  8. data/ext/faiss/product_quantizer.cpp +1 -1
  9. data/ext/faiss/{utils_rb.cpp → utils.cpp} +1 -1
  10. data/lib/faiss/version.rb +1 -1
  11. data/vendor/faiss/faiss/AutoTune.cpp +93 -80
  12. data/vendor/faiss/faiss/Clustering.cpp +39 -240
  13. data/vendor/faiss/faiss/Clustering.h +6 -0
  14. data/vendor/faiss/faiss/IVFlib.cpp +41 -21
  15. data/vendor/faiss/faiss/Index.cpp +6 -5
  16. data/vendor/faiss/faiss/Index.h +5 -5
  17. data/vendor/faiss/faiss/Index2Layer.cpp +37 -53
  18. data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp +49 -37
  19. data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.cpp +36 -34
  20. data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.h +4 -1
  21. data/vendor/faiss/faiss/IndexBinary.cpp +5 -3
  22. data/vendor/faiss/faiss/IndexBinary.h +4 -4
  23. data/vendor/faiss/faiss/IndexBinaryFlat.cpp +1 -1
  24. data/vendor/faiss/faiss/IndexBinaryFlat.h +1 -1
  25. data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +4 -4
  26. data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +84 -92
  27. data/vendor/faiss/faiss/IndexBinaryHNSW.h +9 -3
  28. data/vendor/faiss/faiss/IndexBinaryHash.cpp +45 -236
  29. data/vendor/faiss/faiss/IndexBinaryHash.h +6 -6
  30. data/vendor/faiss/faiss/IndexBinaryIVF.cpp +87 -415
  31. data/vendor/faiss/faiss/IndexFastScan.cpp +72 -109
  32. data/vendor/faiss/faiss/IndexFastScan.h +25 -23
  33. data/vendor/faiss/faiss/IndexFlat.cpp +27 -20
  34. data/vendor/faiss/faiss/IndexFlat.h +21 -18
  35. data/vendor/faiss/faiss/IndexFlatCodes.cpp +42 -19
  36. data/vendor/faiss/faiss/IndexHNSW.cpp +283 -145
  37. data/vendor/faiss/faiss/IndexHNSW.h +16 -2
  38. data/vendor/faiss/faiss/IndexIDMap.cpp +25 -21
  39. data/vendor/faiss/faiss/IndexIDMap.h +9 -7
  40. data/vendor/faiss/faiss/IndexIVF.cpp +465 -362
  41. data/vendor/faiss/faiss/IndexIVF.h +33 -12
  42. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp +77 -74
  43. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.cpp +96 -93
  44. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.h +4 -1
  45. data/vendor/faiss/faiss/IndexIVFFastScan.cpp +357 -238
  46. data/vendor/faiss/faiss/IndexIVFFastScan.h +42 -41
  47. data/vendor/faiss/faiss/IndexIVFFlat.cpp +36 -68
  48. data/vendor/faiss/faiss/IndexIVFFlat.h +32 -0
  49. data/vendor/faiss/faiss/IndexIVFFlatPanorama.cpp +53 -30
  50. data/vendor/faiss/faiss/IndexIVFFlatPanorama.h +3 -1
  51. data/vendor/faiss/faiss/IndexIVFIndependentQuantizer.cpp +18 -15
  52. data/vendor/faiss/faiss/IndexIVFPQ.cpp +71 -843
  53. data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +151 -121
  54. data/vendor/faiss/faiss/IndexIVFPQFastScan.h +3 -0
  55. data/vendor/faiss/faiss/IndexIVFPQR.cpp +21 -17
  56. data/vendor/faiss/faiss/IndexIVFRaBitQ.cpp +26 -39
  57. data/vendor/faiss/faiss/IndexIVFRaBitQ.h +2 -1
  58. data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.cpp +475 -476
  59. data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.h +248 -93
  60. data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +41 -127
  61. data/vendor/faiss/faiss/IndexIVFSpectralHash.h +1 -1
  62. data/vendor/faiss/faiss/IndexLSH.cpp +36 -19
  63. data/vendor/faiss/faiss/IndexLattice.cpp +13 -13
  64. data/vendor/faiss/faiss/IndexNNDescent.cpp +36 -21
  65. data/vendor/faiss/faiss/IndexNNDescent.h +2 -2
  66. data/vendor/faiss/faiss/IndexNSG.cpp +39 -23
  67. data/vendor/faiss/faiss/IndexNeuralNetCodec.cpp +31 -11
  68. data/vendor/faiss/faiss/IndexPQ.cpp +128 -221
  69. data/vendor/faiss/faiss/IndexPQ.h +3 -2
  70. data/vendor/faiss/faiss/IndexPQFastScan.cpp +20 -14
  71. data/vendor/faiss/faiss/IndexPQFastScan.h +3 -0
  72. data/vendor/faiss/faiss/IndexPreTransform.cpp +25 -18
  73. data/vendor/faiss/faiss/IndexPreTransform.h +1 -1
  74. data/vendor/faiss/faiss/IndexRaBitQ.cpp +11 -36
  75. data/vendor/faiss/faiss/IndexRaBitQ.h +2 -1
  76. data/vendor/faiss/faiss/IndexRaBitQFastScan.cpp +41 -277
  77. data/vendor/faiss/faiss/IndexRaBitQFastScan.h +183 -27
  78. data/vendor/faiss/faiss/IndexRefine.cpp +30 -25
  79. data/vendor/faiss/faiss/IndexRefine.h +4 -4
  80. data/vendor/faiss/faiss/IndexReplicas.cpp +6 -6
  81. data/vendor/faiss/faiss/IndexRowwiseMinMax.cpp +15 -14
  82. data/vendor/faiss/faiss/IndexRowwiseMinMax.h +1 -1
  83. data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +82 -14
  84. data/vendor/faiss/faiss/IndexShards.cpp +10 -9
  85. data/vendor/faiss/faiss/IndexShardsIVF.cpp +21 -15
  86. data/vendor/faiss/faiss/MatrixStats.cpp +5 -4
  87. data/vendor/faiss/faiss/MetaIndexes.cpp +19 -17
  88. data/vendor/faiss/faiss/MetaIndexes.h +1 -1
  89. data/vendor/faiss/faiss/MetricType.h +14 -7
  90. data/vendor/faiss/faiss/SuperKMeans.cpp +656 -0
  91. data/vendor/faiss/faiss/SuperKMeans.h +97 -0
  92. data/vendor/faiss/faiss/VectorTransform.cpp +237 -149
  93. data/vendor/faiss/faiss/VectorTransform.h +16 -16
  94. data/vendor/faiss/faiss/build.cpp +23 -0
  95. data/vendor/faiss/faiss/build.h +15 -0
  96. data/vendor/faiss/faiss/clone_index.cpp +48 -47
  97. data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-avx2-inl.h +47 -47
  98. data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-inl.h +11 -0
  99. data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-avx2-inl.h +38 -38
  100. data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-inl.h +11 -0
  101. data/vendor/faiss/faiss/factory_tools.cpp +5 -0
  102. data/vendor/faiss/faiss/gpu/GpuIndexCagra.h +6 -5
  103. data/vendor/faiss/faiss/gpu/GpuResources.h +1 -1
  104. data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp +9 -9
  105. data/vendor/faiss/faiss/gpu/StandardGpuResources.h +4 -3
  106. data/vendor/faiss/faiss/gpu/test/TestGpuIndexFlat.cpp +46 -0
  107. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFFlat.cpp +56 -0
  108. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFPQ.cpp +78 -1
  109. data/vendor/faiss/faiss/gpu/test/TestUtils.cpp +72 -0
  110. data/vendor/faiss/faiss/gpu/test/TestUtils.h +23 -0
  111. data/vendor/faiss/faiss/gpu/utils/CuvsFilterConvert.h +1 -1
  112. data/vendor/faiss/faiss/gpu/utils/CuvsUtils.h +21 -10
  113. data/vendor/faiss/faiss/gpu_metal/GpuIndexFlat.h +22 -0
  114. data/vendor/faiss/faiss/gpu_metal/MetalCloner.h +35 -0
  115. data/vendor/faiss/faiss/gpu_metal/MetalFlatKernels.h +40 -0
  116. data/vendor/faiss/faiss/gpu_metal/MetalIndex.h +51 -0
  117. data/vendor/faiss/faiss/gpu_metal/MetalIndexFlat.h +65 -0
  118. data/vendor/faiss/faiss/gpu_metal/MetalKernels.h +66 -0
  119. data/vendor/faiss/faiss/gpu_metal/MetalResources.h +79 -0
  120. data/vendor/faiss/faiss/gpu_metal/StandardMetalResources.h +35 -0
  121. data/vendor/faiss/faiss/impl/AdSampling.cpp +103 -0
  122. data/vendor/faiss/faiss/impl/AdSampling.h +35 -0
  123. data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp +29 -25
  124. data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +1 -0
  125. data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp +10 -9
  126. data/vendor/faiss/faiss/impl/AuxIndexStructures.h +3 -0
  127. data/vendor/faiss/faiss/impl/ClusteringHelpers.cpp +244 -0
  128. data/vendor/faiss/faiss/impl/ClusteringHelpers.h +94 -0
  129. data/vendor/faiss/faiss/impl/ClusteringInitialization.cpp +16 -16
  130. data/vendor/faiss/faiss/impl/CodePacker.cpp +3 -3
  131. data/vendor/faiss/faiss/impl/CodePackerRaBitQ.cpp +1 -1
  132. data/vendor/faiss/faiss/impl/DistanceComputer.h +8 -8
  133. data/vendor/faiss/faiss/impl/FaissAssert.h +6 -3
  134. data/vendor/faiss/faiss/impl/FaissException.h +50 -3
  135. data/vendor/faiss/faiss/impl/HNSW.cpp +92 -317
  136. data/vendor/faiss/faiss/impl/HNSW.h +13 -34
  137. data/vendor/faiss/faiss/impl/IDSelector.cpp +15 -11
  138. data/vendor/faiss/faiss/impl/IDSelector.h +8 -8
  139. data/vendor/faiss/faiss/impl/InvertedListScannerStats.h +26 -0
  140. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +82 -77
  141. data/vendor/faiss/faiss/impl/NNDescent.cpp +62 -25
  142. data/vendor/faiss/faiss/impl/NNDescent.h +6 -2
  143. data/vendor/faiss/faiss/impl/NSG.cpp +38 -21
  144. data/vendor/faiss/faiss/impl/NSG.h +4 -4
  145. data/vendor/faiss/faiss/impl/Panorama.cpp +23 -6
  146. data/vendor/faiss/faiss/impl/Panorama.h +258 -87
  147. data/vendor/faiss/faiss/impl/PdxLayout.cpp +93 -0
  148. data/vendor/faiss/faiss/impl/PdxLayout.h +41 -0
  149. data/vendor/faiss/faiss/impl/PolysemousTraining.cpp +46 -32
  150. data/vendor/faiss/faiss/impl/PolysemousTraining.h +3 -3
  151. data/vendor/faiss/faiss/impl/ProductAdditiveQuantizer.cpp +35 -35
  152. data/vendor/faiss/faiss/impl/ProductQuantizer-inl.h +21 -16
  153. data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +30 -23
  154. data/vendor/faiss/faiss/impl/Quantizer.h +2 -2
  155. data/vendor/faiss/faiss/impl/RaBitQUtils.cpp +55 -49
  156. data/vendor/faiss/faiss/impl/RaBitQUtils.h +65 -0
  157. data/vendor/faiss/faiss/impl/RaBitQuantizer.cpp +296 -283
  158. data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp +26 -23
  159. data/vendor/faiss/faiss/impl/ResidualQuantizer.h +1 -1
  160. data/vendor/faiss/faiss/impl/ResultHandler.h +99 -75
  161. data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +52 -4
  162. data/vendor/faiss/faiss/impl/ScalarQuantizer.h +27 -1
  163. data/vendor/faiss/faiss/impl/ThreadedIndex-inl.h +14 -11
  164. data/vendor/faiss/faiss/impl/VisitedTable.h +7 -0
  165. data/vendor/faiss/faiss/impl/approx_topk/approx_topk.h +276 -0
  166. data/vendor/faiss/faiss/impl/approx_topk/avx2.cpp +68 -0
  167. data/vendor/faiss/faiss/{utils → impl}/approx_topk/generic.h +15 -8
  168. data/vendor/faiss/faiss/impl/approx_topk/neon.cpp +68 -0
  169. data/vendor/faiss/faiss/impl/approx_topk/rq_beam_search_tab-inl.h +169 -0
  170. data/vendor/faiss/faiss/impl/approx_topk/rq_beam_search_tab.h +117 -0
  171. data/vendor/faiss/faiss/impl/approx_topk/simdlib256-inl.h +146 -0
  172. data/vendor/faiss/faiss/impl/binary_hamming/IndexBinaryHNSW_impl.h +73 -0
  173. data/vendor/faiss/faiss/impl/binary_hamming/IndexBinaryHash_impl.h +270 -0
  174. data/vendor/faiss/faiss/impl/binary_hamming/IndexBinaryIVF_impl.h +460 -0
  175. data/vendor/faiss/faiss/impl/binary_hamming/IndexIVFSpectralHash_impl.h +159 -0
  176. data/vendor/faiss/faiss/impl/binary_hamming/IndexPQ_impl.h +92 -0
  177. data/vendor/faiss/faiss/impl/binary_hamming/avx2.cpp +26 -0
  178. data/vendor/faiss/faiss/impl/binary_hamming/avx512.cpp +26 -0
  179. data/vendor/faiss/faiss/impl/binary_hamming/dispatch.h +143 -0
  180. data/vendor/faiss/faiss/impl/binary_hamming/neon.cpp +26 -0
  181. data/vendor/faiss/faiss/impl/binary_hamming/rvv.cpp +26 -0
  182. data/vendor/faiss/faiss/impl/expanded_scanners.h +8 -3
  183. data/vendor/faiss/faiss/impl/{FastScanDistancePostProcessing.h → fast_scan/FastScanDistancePostProcessing.h} +13 -6
  184. data/vendor/faiss/faiss/impl/{LookupTableScaler.h → fast_scan/LookupTableScaler.h} +16 -5
  185. data/vendor/faiss/faiss/impl/fast_scan/accumulate_loops.h +237 -0
  186. data/vendor/faiss/faiss/impl/fast_scan/accumulate_loops_512.h +185 -0
  187. data/vendor/faiss/faiss/impl/fast_scan/decompose_qbs.h +229 -0
  188. data/vendor/faiss/faiss/impl/fast_scan/dispatching.h +268 -0
  189. data/vendor/faiss/faiss/impl/{pq4_fast_scan.cpp → fast_scan/fast_scan.cpp} +169 -2
  190. data/vendor/faiss/faiss/impl/fast_scan/fast_scan.h +341 -0
  191. data/vendor/faiss/faiss/impl/fast_scan/impl-avx2.cpp +36 -0
  192. data/vendor/faiss/faiss/impl/fast_scan/impl-avx512.cpp +40 -0
  193. data/vendor/faiss/faiss/impl/fast_scan/impl-neon.cpp +120 -0
  194. data/vendor/faiss/faiss/impl/fast_scan/impl-riscv.cpp +104 -0
  195. data/vendor/faiss/faiss/impl/fast_scan/kernels_simd256.h +213 -0
  196. data/vendor/faiss/faiss/impl/{pq4_fast_scan_search_qbs.cpp → fast_scan/kernels_simd512.h} +26 -356
  197. data/vendor/faiss/faiss/impl/fast_scan/rabitq_dispatching.h +90 -0
  198. data/vendor/faiss/faiss/impl/fast_scan/rabitq_result_handler.h +108 -0
  199. data/vendor/faiss/faiss/impl/{simd_result_handlers.h → fast_scan/simd_result_handlers.h} +282 -134
  200. data/vendor/faiss/faiss/impl/hnsw/LockVector.cpp +54 -0
  201. data/vendor/faiss/faiss/impl/hnsw/LockVector.h +64 -0
  202. data/vendor/faiss/faiss/impl/hnsw/MinimaxHeap.cpp +91 -0
  203. data/vendor/faiss/faiss/impl/hnsw/MinimaxHeap.h +64 -0
  204. data/vendor/faiss/faiss/impl/hnsw/avx2.cpp +104 -0
  205. data/vendor/faiss/faiss/impl/hnsw/avx512.cpp +111 -0
  206. data/vendor/faiss/faiss/impl/index_read.cpp +1132 -45
  207. data/vendor/faiss/faiss/impl/index_read_utils.h +1 -1
  208. data/vendor/faiss/faiss/impl/index_write.cpp +95 -13
  209. data/vendor/faiss/faiss/impl/io.cpp +6 -6
  210. data/vendor/faiss/faiss/impl/io_macros.h +33 -16
  211. data/vendor/faiss/faiss/impl/kmeans1d.cpp +10 -10
  212. data/vendor/faiss/faiss/impl/lattice_Zn.cpp +37 -23
  213. data/vendor/faiss/faiss/impl/lattice_Zn.h +6 -6
  214. data/vendor/faiss/faiss/impl/mapped_io.cpp +6 -6
  215. data/vendor/faiss/faiss/impl/platform_macros.h +11 -4
  216. data/vendor/faiss/faiss/impl/pq_code_distance/IVFPQScanner_impl.h +549 -0
  217. data/vendor/faiss/faiss/impl/pq_code_distance/IVFPQ_QueryTables.cpp +245 -0
  218. data/vendor/faiss/faiss/impl/pq_code_distance/IVFPQ_QueryTables.h +105 -0
  219. data/vendor/faiss/faiss/impl/pq_code_distance/PQDistanceComputer_impl.h +106 -0
  220. data/vendor/faiss/faiss/impl/pq_code_distance/avx2.cpp +21 -0
  221. data/vendor/faiss/faiss/impl/pq_code_distance/avx512.cpp +21 -0
  222. data/vendor/faiss/faiss/impl/pq_code_distance/neon.cpp +21 -0
  223. data/vendor/faiss/faiss/impl/pq_code_distance/{pq_code_distance-avx2.cpp → pq_code_distance-avx2.h} +9 -13
  224. data/vendor/faiss/faiss/impl/pq_code_distance/{pq_code_distance-avx512.cpp → pq_code_distance-avx512.h} +9 -57
  225. data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-generic.cpp +29 -111
  226. data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-generic.h +96 -0
  227. data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-inl.h +238 -5
  228. data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-sve.cpp +5 -7
  229. data/vendor/faiss/faiss/impl/pq_code_distance/rvv.cpp +68 -0
  230. data/vendor/faiss/faiss/impl/residual_quantizer_encode_steps.cpp +311 -477
  231. data/vendor/faiss/faiss/impl/residual_quantizer_encode_steps.h +1 -1
  232. data/vendor/faiss/faiss/impl/scalar_quantizer/codecs.h +1 -1
  233. data/vendor/faiss/faiss/impl/scalar_quantizer/distance_computers.h +3 -2
  234. data/vendor/faiss/faiss/impl/scalar_quantizer/quantizers.h +102 -11
  235. data/vendor/faiss/faiss/impl/scalar_quantizer/scanners.h +27 -1
  236. data/vendor/faiss/faiss/impl/scalar_quantizer/similarities.h +3 -3
  237. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx2.cpp +148 -0
  238. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx512.cpp +167 -0
  239. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-dispatch.h +59 -0
  240. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-neon.cpp +163 -0
  241. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-rvv.cpp +311 -0
  242. data/vendor/faiss/faiss/impl/scalar_quantizer/training.cpp +192 -8
  243. data/vendor/faiss/faiss/impl/scalar_quantizer/training.h +12 -0
  244. data/vendor/faiss/faiss/impl/simd_dispatch.h +100 -66
  245. data/vendor/faiss/faiss/impl/simdlib/simdlib.h +57 -0
  246. data/vendor/faiss/faiss/{utils → impl/simdlib}/simdlib_avx2.h +264 -172
  247. data/vendor/faiss/faiss/impl/simdlib/simdlib_avx512.h +414 -0
  248. data/vendor/faiss/faiss/impl/simdlib/simdlib_dispatch.h +44 -0
  249. data/vendor/faiss/faiss/{utils → impl/simdlib}/simdlib_emulated.h +231 -166
  250. data/vendor/faiss/faiss/{utils → impl/simdlib}/simdlib_neon.h +270 -218
  251. data/vendor/faiss/faiss/{utils → impl/simdlib}/simdlib_ppc64.h +201 -160
  252. data/vendor/faiss/faiss/impl/svs_io.cpp +12 -3
  253. data/vendor/faiss/faiss/impl/svs_io.h +8 -2
  254. data/vendor/faiss/faiss/index_factory.cpp +86 -18
  255. data/vendor/faiss/faiss/index_io.h +24 -0
  256. data/vendor/faiss/faiss/invlists/BlockInvertedLists.cpp +66 -16
  257. data/vendor/faiss/faiss/invlists/DirectMap.cpp +24 -14
  258. data/vendor/faiss/faiss/invlists/DirectMap.h +4 -3
  259. data/vendor/faiss/faiss/invlists/InvertedLists.cpp +157 -73
  260. data/vendor/faiss/faiss/invlists/InvertedLists.h +86 -23
  261. data/vendor/faiss/faiss/invlists/InvertedListsIOHook.cpp +4 -4
  262. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +13 -13
  263. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.h +1 -1
  264. data/vendor/faiss/faiss/svs/IndexSVSFaissUtils.h +1 -1
  265. data/vendor/faiss/faiss/svs/IndexSVSFlat.cpp +2 -2
  266. data/vendor/faiss/faiss/svs/IndexSVSIVF.cpp +350 -0
  267. data/vendor/faiss/faiss/svs/IndexSVSIVF.h +128 -0
  268. data/vendor/faiss/faiss/svs/IndexSVSIVFLVQ.cpp +40 -0
  269. data/vendor/faiss/faiss/svs/IndexSVSIVFLVQ.h +43 -0
  270. data/vendor/faiss/faiss/svs/IndexSVSIVFLeanVec.cpp +225 -0
  271. data/vendor/faiss/faiss/svs/IndexSVSIVFLeanVec.h +71 -0
  272. data/vendor/faiss/faiss/svs/IndexSVSVamana.cpp +25 -1
  273. data/vendor/faiss/faiss/svs/IndexSVSVamana.h +18 -2
  274. data/vendor/faiss/faiss/svs/IndexSVSVamanaLVQ.h +1 -1
  275. data/vendor/faiss/faiss/svs/IndexSVSVamanaLeanVec.cpp +12 -3
  276. data/vendor/faiss/faiss/svs/IndexSVSVamanaLeanVec.h +7 -2
  277. data/vendor/faiss/faiss/utils/Heap.cpp +10 -10
  278. data/vendor/faiss/faiss/utils/NeuralNet.cpp +47 -36
  279. data/vendor/faiss/faiss/utils/NeuralNet.h +1 -1
  280. data/vendor/faiss/faiss/utils/approx_topk_hamming/approx_topk_hamming.h +10 -4
  281. data/vendor/faiss/faiss/utils/distances.cpp +390 -560
  282. data/vendor/faiss/faiss/utils/distances.h +20 -1
  283. data/vendor/faiss/faiss/utils/distances_dispatch.h +117 -37
  284. data/vendor/faiss/faiss/utils/distances_fused/avx512.cpp +8 -7
  285. data/vendor/faiss/faiss/utils/distances_fused/distances_fused.cpp +33 -14
  286. data/vendor/faiss/faiss/utils/distances_fused/distances_fused.h +12 -1
  287. data/vendor/faiss/faiss/utils/distances_fused/simdlib_based.cpp +16 -293
  288. data/vendor/faiss/faiss/utils/distances_fused/simdlib_based_neon.cpp +57 -0
  289. data/vendor/faiss/faiss/utils/distances_fused/simdlib_kernel-inl.h +290 -0
  290. data/vendor/faiss/faiss/utils/distances_simd.cpp +5 -177
  291. data/vendor/faiss/faiss/utils/extra_distances.cpp +9 -8
  292. data/vendor/faiss/faiss/utils/extra_distances.h +32 -6
  293. data/vendor/faiss/faiss/utils/hamming-inl.h +13 -11
  294. data/vendor/faiss/faiss/utils/hamming.cpp +66 -517
  295. data/vendor/faiss/faiss/utils/hamming.h +92 -2
  296. data/vendor/faiss/faiss/utils/hamming_distance/common.h +287 -10
  297. data/vendor/faiss/faiss/utils/hamming_distance/hamming_avx2.cpp +15 -0
  298. data/vendor/faiss/faiss/utils/hamming_distance/hamming_avx512.cpp +15 -0
  299. data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-avx2.h +142 -0
  300. data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-avx512.h +234 -0
  301. data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-generic.h +368 -0
  302. data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-neon.h +322 -0
  303. data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-rvv.h +39 -0
  304. data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer.h +146 -0
  305. data/vendor/faiss/faiss/utils/hamming_distance/hamming_impl.h +481 -0
  306. data/vendor/faiss/faiss/utils/hamming_distance/hamming_neon.cpp +15 -0
  307. data/vendor/faiss/faiss/utils/hamming_distance/hamming_rvv.cpp +15 -0
  308. data/vendor/faiss/faiss/utils/partitioning.cpp +66 -987
  309. data/vendor/faiss/faiss/utils/partitioning.h +31 -0
  310. data/vendor/faiss/faiss/utils/popcount.h +29 -0
  311. data/vendor/faiss/faiss/utils/pq_code_distance.h +2 -2
  312. data/vendor/faiss/faiss/utils/prefetch.h +2 -2
  313. data/vendor/faiss/faiss/utils/quantize_lut.cpp +30 -30
  314. data/vendor/faiss/faiss/utils/quantize_lut.h +1 -1
  315. data/vendor/faiss/faiss/utils/rabitq_simd.h +57 -536
  316. data/vendor/faiss/faiss/utils/random.cpp +6 -6
  317. data/vendor/faiss/faiss/utils/simd_impl/IVFFlatScanner-inl.h +51 -0
  318. data/vendor/faiss/faiss/utils/simd_impl/distances_aarch64.cpp +5 -1
  319. data/vendor/faiss/faiss/utils/simd_impl/distances_arm_sve.cpp +213 -4
  320. data/vendor/faiss/faiss/utils/simd_impl/distances_autovec-inl.h +163 -10
  321. data/vendor/faiss/faiss/utils/simd_impl/distances_avx2.cpp +250 -4
  322. data/vendor/faiss/faiss/utils/simd_impl/distances_avx512.cpp +7 -4
  323. data/vendor/faiss/faiss/utils/simd_impl/distances_rvv.cpp +189 -0
  324. data/vendor/faiss/faiss/utils/simd_impl/distances_simdlib256.h +195 -0
  325. data/vendor/faiss/faiss/utils/simd_impl/distances_sse-inl.h +2 -1
  326. data/vendor/faiss/faiss/utils/{distances_fused/simdlib_based.h → simd_impl/exhaustive_L2sqr_blas_cmax.h} +5 -10
  327. data/vendor/faiss/faiss/utils/simd_impl/hamming_impl.h +481 -0
  328. data/vendor/faiss/faiss/utils/simd_impl/partitioning_avx2.cpp +14 -0
  329. data/vendor/faiss/faiss/utils/simd_impl/partitioning_neon.cpp +14 -0
  330. data/vendor/faiss/faiss/utils/simd_impl/partitioning_simdlib256.h +1085 -0
  331. data/vendor/faiss/faiss/utils/simd_impl/rabitq_avx2.cpp +355 -0
  332. data/vendor/faiss/faiss/utils/simd_impl/rabitq_avx512.cpp +477 -0
  333. data/vendor/faiss/faiss/utils/simd_impl/rabitq_neon.cpp +55 -0
  334. data/vendor/faiss/faiss/utils/simd_impl/rabitq_rvv.cpp +55 -0
  335. data/vendor/faiss/faiss/utils/simd_impl/super_kmeans_dispatch.h +32 -0
  336. data/vendor/faiss/faiss/utils/simd_impl/super_kmeans_kernels.h +43 -0
  337. data/vendor/faiss/faiss/utils/simd_impl/super_kmeans_kernels_avx2.cpp +57 -0
  338. data/vendor/faiss/faiss/utils/simd_impl/super_kmeans_kernels_avx512.cpp +45 -0
  339. data/vendor/faiss/faiss/utils/simd_levels.cpp +17 -5
  340. data/vendor/faiss/faiss/utils/simd_levels.h +93 -1
  341. data/vendor/faiss/faiss/utils/sorting.cpp +48 -36
  342. data/vendor/faiss/faiss/utils/utils.cpp +5 -5
  343. data/vendor/faiss/faiss/utils/utils.h +3 -3
  344. metadata +119 -34
  345. data/vendor/faiss/faiss/impl/RaBitQStats.cpp +0 -29
  346. data/vendor/faiss/faiss/impl/RaBitQStats.h +0 -56
  347. data/vendor/faiss/faiss/impl/pq4_fast_scan.h +0 -224
  348. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp +0 -230
  349. data/vendor/faiss/faiss/utils/approx_topk/approx_topk.h +0 -84
  350. data/vendor/faiss/faiss/utils/approx_topk/avx2-inl.h +0 -196
  351. data/vendor/faiss/faiss/utils/approx_topk/mode.h +0 -34
  352. data/vendor/faiss/faiss/utils/distances_fused/avx512.h +0 -36
  353. data/vendor/faiss/faiss/utils/extra_distances-inl.h +0 -235
  354. data/vendor/faiss/faiss/utils/hamming_distance/avx2-inl.h +0 -462
  355. data/vendor/faiss/faiss/utils/hamming_distance/avx512-inl.h +0 -490
  356. data/vendor/faiss/faiss/utils/hamming_distance/generic-inl.h +0 -449
  357. data/vendor/faiss/faiss/utils/hamming_distance/hamdis-inl.h +0 -87
  358. data/vendor/faiss/faiss/utils/hamming_distance/neon-inl.h +0 -524
  359. data/vendor/faiss/faiss/utils/simdlib.h +0 -42
  360. data/vendor/faiss/faiss/utils/simdlib_avx512.h +0 -365
  361. /data/ext/faiss/{utils_rb.h → utils.h} +0 -0
@@ -8,6 +8,7 @@
8
8
  #include <faiss/IndexHNSW.h>
9
9
 
10
10
  #include <omp.h>
11
+ #include <atomic>
11
12
  #include <cinttypes>
12
13
  #include <cstdio>
13
14
  #include <cstdlib>
@@ -26,14 +27,15 @@
26
27
  #include <faiss/IndexIVFPQ.h>
27
28
  #include <faiss/impl/AuxIndexStructures.h>
28
29
  #include <faiss/impl/FaissAssert.h>
30
+ #include <faiss/impl/FaissException.h>
29
31
  #include <faiss/impl/ResultHandler.h>
30
32
  #include <faiss/impl/VisitedTable.h>
33
+ #include <faiss/impl/hnsw/MinimaxHeap.h>
31
34
  #include <faiss/utils/random.h>
32
35
  #include <faiss/utils/sorting.h>
33
36
 
34
37
  namespace faiss {
35
38
 
36
- using MinimaxHeap = HNSW::MinimaxHeap;
37
39
  using storage_idx_t = HNSW::storage_idx_t;
38
40
  using NodeDistFarther = HNSW::NodeDistFarther;
39
41
 
@@ -82,10 +84,8 @@ void hnsw_add_vertices(
82
84
  printf(" max_level = %d\n", max_level);
83
85
  }
84
86
 
85
- std::vector<omp_lock_t> locks(ntotal);
86
- for (int i = 0; i < ntotal; i++) {
87
- omp_init_lock(&locks[i]);
88
- }
87
+ auto& locks = index_hnsw.locks;
88
+ locks.prepare(ntotal);
89
89
 
90
90
  // add vectors from highest to lowest level
91
91
  std::vector<int> hist;
@@ -94,10 +94,10 @@ void hnsw_add_vertices(
94
94
  { // make buckets with vectors of the same level
95
95
 
96
96
  // build histogram
97
- for (int i = 0; i < n; i++) {
98
- storage_idx_t pt_id = i + n0;
97
+ for (size_t i = 0; i < n; i++) {
98
+ storage_idx_t pt_id = static_cast<storage_idx_t>(i + n0);
99
99
  int pt_level = hnsw.levels[pt_id] - 1;
100
- while (pt_level >= hist.size()) {
100
+ while (pt_level >= static_cast<int>(hist.size())) {
101
101
  hist.push_back(0);
102
102
  }
103
103
  hist[pt_level]++;
@@ -105,13 +105,13 @@ void hnsw_add_vertices(
105
105
 
106
106
  // accumulate
107
107
  std::vector<int> offsets(hist.size() + 1, 0);
108
- for (int i = 0; i < hist.size() - 1; i++) {
108
+ for (size_t i = 0; i < hist.size() - 1; i++) {
109
109
  offsets[i + 1] = offsets[i] + hist[i];
110
110
  }
111
111
 
112
112
  // bucket sort
113
- for (int i = 0; i < n; i++) {
114
- storage_idx_t pt_id = i + n0;
113
+ for (size_t i = 0; i < n; i++) {
114
+ storage_idx_t pt_id = static_cast<storage_idx_t>(i + n0);
115
115
  int pt_level = hnsw.levels[pt_id] - 1;
116
116
  order[offsets[pt_level]++] = pt_id;
117
117
  }
@@ -123,20 +123,22 @@ void hnsw_add_vertices(
123
123
  { // perform add
124
124
  RandomGenerator rng2(789);
125
125
 
126
- int i1 = n;
126
+ size_t i1 = static_cast<int>(n);
127
127
 
128
- for (int pt_level = hist.size() - 1;
128
+ for (int pt_level = static_cast<int>(hist.size()) - 1;
129
129
  pt_level >= int(!index_hnsw.init_level0);
130
130
  pt_level--) {
131
- int i0 = i1 - hist[pt_level];
131
+ size_t i0 = i1 - hist[pt_level];
132
132
 
133
133
  if (verbose) {
134
- printf("Adding %d elements at level %d\n", i1 - i0, pt_level);
134
+ printf("Adding %zu elements at level %d\n", i1 - i0, pt_level);
135
135
  }
136
136
 
137
137
  // random permutation to get rid of dataset order bias
138
- for (int j = i0; j < i1; j++) {
139
- std::swap(order[j], order[j + rng2.rand_int(i1 - j)]);
138
+ for (size_t j = i0; j < i1; j++) {
139
+ std::swap(
140
+ order[j],
141
+ order[j + rng2.rand_int(static_cast<int>(i1 - j))]);
140
142
  }
141
143
 
142
144
  bool interrupt = false;
@@ -147,15 +149,15 @@ void hnsw_add_vertices(
147
149
 
148
150
  std::unique_ptr<DistanceComputer> dis(
149
151
  storage_distance_computer(index_hnsw.storage));
150
- int prev_display =
151
- verbose && omp_get_thread_num() == 0 ? 0 : -1;
152
+ bool do_display = verbose && omp_get_thread_num() == 0;
153
+ size_t prev_display = 0;
152
154
  size_t counter = 0;
153
155
 
154
156
  // here we should do schedule(dynamic) but this segfaults for
155
157
  // some versions of LLVM. The performance impact should not be
156
158
  // too large when (i1 - i0) / num_threads >> 1
157
159
  #pragma omp for schedule(static)
158
- for (int i = i0; i < i1; i++) {
160
+ for (int64_t i = i0; i < i1; i++) {
159
161
  storage_idx_t pt_id = order[i];
160
162
  dis->set_query(x + (pt_id - n0) * d);
161
163
 
@@ -172,9 +174,9 @@ void hnsw_add_vertices(
172
174
  vt,
173
175
  index_hnsw.keep_max_size_level0 && (pt_level == 0));
174
176
 
175
- if (prev_display >= 0 && i - i0 > prev_display + 10000) {
177
+ if (do_display && i - i0 > prev_display + 10000) {
176
178
  prev_display = i - i0;
177
- printf(" %d / %d\r", i - i0, i1 - i0);
179
+ printf(" %zu / %zu\r", i - i0, i1 - i0);
178
180
  fflush(stdout);
179
181
  }
180
182
  if (counter % check_period == 0) {
@@ -199,9 +201,8 @@ void hnsw_add_vertices(
199
201
  if (verbose) {
200
202
  printf("Done in %.3f ms\n", getmillisecs() - t0);
201
203
  }
202
-
203
- for (int i = 0; i < ntotal; i++) {
204
- omp_destroy_lock(&locks[i]);
204
+ if (!index_hnsw.retain_locks) {
205
+ locks.clear();
205
206
  }
206
207
  }
207
208
 
@@ -211,11 +212,13 @@ void hnsw_add_vertices(
211
212
  * IndexHNSW implementation
212
213
  **************************************************************/
213
214
 
214
- IndexHNSW::IndexHNSW(int d, int M, MetricType metric)
215
- : Index(d, metric), hnsw(M) {}
215
+ IndexHNSW::IndexHNSW(int d_in, int M, MetricType metric)
216
+ : Index(d_in, metric), hnsw(M) {}
216
217
 
217
- IndexHNSW::IndexHNSW(Index* storage, int M)
218
- : Index(storage->d, storage->metric_type), hnsw(M), storage(storage) {
218
+ IndexHNSW::IndexHNSW(Index* storage_in, int M)
219
+ : Index(storage_in->d, storage_in->metric_type),
220
+ hnsw(M),
221
+ storage(storage_in) {
219
222
  metric_arg = storage->metric_arg;
220
223
  }
221
224
 
@@ -263,28 +266,48 @@ void hnsw_search(
263
266
 
264
267
  for (idx_t i0 = 0; i0 < n; i0 += check_period) {
265
268
  idx_t i1 = std::min(i0 + check_period, n);
269
+ std::exception_ptr ex;
270
+ std::atomic<bool> interrupt{false};
266
271
 
267
272
  #pragma omp parallel if (i1 - i0 > 1)
268
273
  {
269
- VisitedTable vt(index->ntotal, hnsw.use_visited_hashset);
270
- typename BlockResultHandler::SingleResultHandler res(bres);
271
-
272
- std::unique_ptr<DistanceComputer> dis(
273
- storage_distance_computer(index->storage));
274
+ std::unique_ptr<VisitedTable> vt;
275
+ std::unique_ptr<typename BlockResultHandler::SingleResultHandler>
276
+ res;
277
+ std::unique_ptr<DistanceComputer> dis;
278
+ try {
279
+ vt = std::make_unique<VisitedTable>(
280
+ index->ntotal, hnsw.use_visited_hashset);
281
+ res = std::make_unique<
282
+ typename BlockResultHandler::SingleResultHandler>(bres);
283
+ dis.reset(storage_distance_computer(index->storage));
284
+ } catch (...) {
285
+ omp_capture_exception(ex, [&] { interrupt = true; });
286
+ }
274
287
 
275
288
  #pragma omp for reduction(+ : n1, n2, ndis, nhops) schedule(guided)
276
289
  for (idx_t i = i0; i < i1; i++) {
277
- res.begin(i);
278
- dis->set_query(x + i * index->d);
279
-
280
- HNSWStats stats = hnsw.search(*dis, index, res, vt, params);
281
- n1 += stats.n1;
282
- n2 += stats.n2;
283
- ndis += stats.ndis;
284
- nhops += stats.nhops;
285
- res.end();
290
+ if (interrupt.load(std::memory_order_relaxed)) {
291
+ continue;
292
+ }
293
+ try {
294
+ res->begin(i);
295
+ dis->set_query(x + i * index->d);
296
+
297
+ HNSWStats stats =
298
+ hnsw.search(*dis, index, *res, *vt, params);
299
+ n1 += stats.n1;
300
+ n2 += stats.n2;
301
+ ndis += stats.ndis;
302
+ nhops += stats.nhops;
303
+ res->end();
304
+ vt->advance();
305
+ } catch (...) {
306
+ omp_capture_exception(ex, [&] { interrupt = true; });
307
+ }
286
308
  }
287
309
  }
310
+ omp_rethrow_if_exception(ex);
288
311
  InterruptCallback::check();
289
312
  }
290
313
 
@@ -309,7 +332,7 @@ void IndexHNSW::search(
309
332
 
310
333
  if (is_similarity_metric(this->metric_type)) {
311
334
  // we need to revert the negated distances
312
- for (size_t i = 0; i < k * n; i++) {
335
+ for (idx_t i = 0; i < k * n; i++) {
313
336
  distances[i] = -distances[i];
314
337
  }
315
338
  }
@@ -347,15 +370,22 @@ void IndexHNSW::add(idx_t n, const float* x) {
347
370
  storage,
348
371
  "Please use IndexHNSWFlat (or variants) instead of IndexHNSW directly");
349
372
  FAISS_THROW_IF_NOT(is_trained);
350
- int n0 = ntotal;
373
+ size_t n0 = ntotal;
351
374
  storage->add(n, x);
352
375
  ntotal = storage->ntotal;
353
376
 
354
- hnsw_add_vertices(*this, n0, n, x, verbose, hnsw.levels.size() == ntotal);
377
+ hnsw_add_vertices(
378
+ *this,
379
+ n0,
380
+ n,
381
+ x,
382
+ verbose,
383
+ hnsw.levels.size() == static_cast<size_t>(ntotal));
355
384
  }
356
385
 
357
386
  void IndexHNSW::reset() {
358
387
  hnsw.reset();
388
+ locks.clear();
359
389
  storage->reset();
360
390
  ntotal = 0;
361
391
  }
@@ -427,42 +457,59 @@ void IndexHNSW::search_level_0(
427
457
  FAISS_THROW_IF_NOT(k > 0);
428
458
  FAISS_THROW_IF_NOT(nprobe > 0);
429
459
 
430
- storage_idx_t ntotal = hnsw.levels.size();
460
+ size_t hnsw_ntotal = hnsw.levels.size();
431
461
 
432
462
  using RH = HeapBlockResultHandler<HNSW::C>;
433
463
  RH bres(n, distances, labels, k);
434
464
 
465
+ std::exception_ptr ex;
466
+ std::atomic<bool> interrupt{false};
435
467
  #pragma omp parallel
436
468
  {
437
- std::unique_ptr<DistanceComputer> qdis(
438
- storage_distance_computer(storage));
469
+ std::unique_ptr<DistanceComputer> qdis;
439
470
  HNSWStats search_stats;
440
- VisitedTable vt(ntotal, hnsw.use_visited_hashset);
441
- RH::SingleResultHandler res(bres);
471
+ std::unique_ptr<VisitedTable> vt;
472
+ std::unique_ptr<RH::SingleResultHandler> res;
473
+ try {
474
+ qdis.reset(storage_distance_computer(storage));
475
+ vt = std::make_unique<VisitedTable>(
476
+ hnsw_ntotal, hnsw.use_visited_hashset);
477
+ res = std::make_unique<RH::SingleResultHandler>(bres);
478
+ } catch (...) {
479
+ omp_capture_exception(ex, [&] { interrupt = true; });
480
+ }
442
481
 
443
482
  #pragma omp for
444
483
  for (idx_t i = 0; i < n; i++) {
445
- res.begin(i);
446
- qdis->set_query(x + i * d);
447
-
448
- hnsw.search_level_0(
449
- *qdis.get(),
450
- res,
451
- nprobe,
452
- nearest + i * nprobe,
453
- nearest_d + i * nprobe,
454
- search_type,
455
- search_stats,
456
- vt,
457
- params);
458
- res.end();
459
- vt.advance();
484
+ if (interrupt.load(std::memory_order_relaxed)) {
485
+ continue;
486
+ }
487
+ try {
488
+ res->begin(i);
489
+ qdis->set_query(x + i * d);
490
+
491
+ hnsw.search_level_0(
492
+ *qdis.get(),
493
+ *res,
494
+ nprobe,
495
+ nearest + i * nprobe,
496
+ nearest_d + i * nprobe,
497
+ search_type,
498
+ search_stats,
499
+ *vt,
500
+ params);
501
+ res->end();
502
+ vt->advance();
503
+ } catch (...) {
504
+ omp_capture_exception(ex, [&] { interrupt = true; });
505
+ }
460
506
  }
461
507
  #pragma omp critical
462
508
  {
463
509
  hnsw_stats.combine(search_stats);
464
510
  }
465
511
  }
512
+ omp_rethrow_if_exception(ex);
466
513
  if (is_similarity_metric(this->metric_type)) {
467
514
  // we need to revert the negated distances
468
515
  #pragma omp parallel for
@@ -487,8 +534,8 @@ void IndexHNSW::init_level_0_from_knngraph(
487
534
 
488
535
  std::priority_queue<NodeDistFarther> initial_list;
489
536
 
490
- for (size_t j = 0; j < k; j++) {
491
- int v1 = I[i * k + j];
537
+ for (int j = 0; j < k; j++) {
538
+ int v1 = static_cast<int>(I[i * k + j]);
492
539
  if (v1 == i) {
493
540
  continue;
494
541
  }
@@ -518,10 +565,7 @@ void IndexHNSW::init_level_0_from_entry_points(
518
565
  int n,
519
566
  const storage_idx_t* points,
520
567
  const storage_idx_t* nearests) {
521
- std::vector<omp_lock_t> locks(ntotal);
522
- for (int i = 0; i < ntotal; i++) {
523
- omp_init_lock(&locks[i]);
524
- }
568
+ locks.prepare(ntotal);
525
569
 
526
570
  #pragma omp parallel
527
571
  {
@@ -539,7 +583,7 @@ void IndexHNSW::init_level_0_from_entry_points(
539
583
  dis->set_query(vec.data());
540
584
 
541
585
  hnsw.add_links_starting_from(
542
- *dis, pt_id, nearest, (*dis)(nearest), 0, locks.data(), vt);
586
+ *dis, pt_id, nearest, (*dis)(nearest), 0, locks, vt);
543
587
 
544
588
  if (verbose && i % 10000 == 0) {
545
589
  printf(" %d / %d\r", i, n);
@@ -551,8 +595,8 @@ void IndexHNSW::init_level_0_from_entry_points(
551
595
  printf("\n");
552
596
  }
553
597
 
554
- for (int i = 0; i < ntotal; i++) {
555
- omp_destroy_lock(&locks[i]);
598
+ if (!retain_locks) {
599
+ locks.clear();
556
600
  }
557
601
  }
558
602
 
@@ -595,7 +639,7 @@ void IndexHNSW::link_singletons() {
595
639
 
596
640
  std::vector<bool> seen(ntotal);
597
641
 
598
- for (size_t i = 0; i < ntotal; i++) {
642
+ for (idx_t i = 0; i < ntotal; i++) {
599
643
  size_t begin, end;
600
644
  hnsw.neighbor_range(i, 0, &begin, &end);
601
645
  for (size_t j = begin; j < end; j++) {
@@ -624,7 +668,7 @@ void IndexHNSW::link_singletons() {
624
668
  n_sing_l1);
625
669
 
626
670
  std::vector<float> recons(singletons.size() * d);
627
- for (int i = 0; i < singletons.size(); i++) {
671
+ for (size_t i = 0; i < singletons.size(); i++) {
628
672
  FAISS_ASSERT(false); // not implemented
629
673
  }
630
674
  }
@@ -649,10 +693,10 @@ IndexHNSWFlat::IndexHNSWFlat() {
649
693
  is_trained = true;
650
694
  }
651
695
 
652
- IndexHNSWFlat::IndexHNSWFlat(int d, int M, MetricType metric)
696
+ IndexHNSWFlat::IndexHNSWFlat(int d_in, int M, MetricType metric)
653
697
  : IndexHNSW(
654
- (metric == METRIC_L2) ? new IndexFlatL2(d)
655
- : new IndexFlat(d, metric),
698
+ (metric == METRIC_L2) ? new IndexFlatL2(d_in)
699
+ : new IndexFlat(d_in, metric),
656
700
  M) {
657
701
  own_fields = true;
658
702
  is_trained = true;
@@ -663,17 +707,20 @@ IndexHNSWFlat::IndexHNSWFlat(int d, int M, MetricType metric)
663
707
  **************************************************************/
664
708
 
665
709
  IndexHNSWFlatPanorama::IndexHNSWFlatPanorama()
666
- : IndexHNSWFlat(), cum_sums(), pano(0, 1, 1), num_panorama_levels(0) {}
710
+ : IndexHNSWFlat(),
711
+ cum_sums(),
712
+ pano(sizeof(float), 1, 1),
713
+ num_panorama_levels(0) {}
667
714
 
668
715
  IndexHNSWFlatPanorama::IndexHNSWFlatPanorama(
669
- int d,
716
+ int d_in,
670
717
  int M,
671
- int num_panorama_levels,
718
+ int num_panorama_levels_in,
672
719
  MetricType metric)
673
- : IndexHNSWFlat(d, M, metric),
720
+ : IndexHNSWFlat(d_in, M, metric),
674
721
  cum_sums(),
675
- pano(d * sizeof(float), num_panorama_levels, 1),
676
- num_panorama_levels(num_panorama_levels) {
722
+ pano(d_in * sizeof(float), num_panorama_levels_in, 1),
723
+ num_panorama_levels(num_panorama_levels_in) {
677
724
  // For now, we only support L2 distance.
678
725
  // Supporting dot product and cosine distance is a trivial addition
679
726
  // left for future work.
@@ -718,12 +765,12 @@ void IndexHNSWFlatPanorama::permute_entries(const idx_t* perm) {
718
765
  IndexHNSWPQ::IndexHNSWPQ() = default;
719
766
 
720
767
  IndexHNSWPQ::IndexHNSWPQ(
721
- int d,
768
+ int d_in,
722
769
  int pq_m,
723
770
  int M,
724
771
  int pq_nbits,
725
772
  MetricType metric)
726
- : IndexHNSW(new IndexPQ(d, pq_m, pq_nbits, metric), M) {
773
+ : IndexHNSW(new IndexPQ(d_in, pq_m, pq_nbits, metric), M) {
727
774
  own_fields = true;
728
775
  is_trained = false;
729
776
  }
@@ -738,11 +785,11 @@ void IndexHNSWPQ::train(idx_t n, const float* x) {
738
785
  **************************************************************/
739
786
 
740
787
  IndexHNSWSQ::IndexHNSWSQ(
741
- int d,
788
+ int d_in,
742
789
  ScalarQuantizer::QuantizerType qtype,
743
790
  int M,
744
791
  MetricType metric)
745
- : IndexHNSW(new IndexScalarQuantizer(d, qtype, metric), M) {
792
+ : IndexHNSW(new IndexScalarQuantizer(d_in, qtype, metric), M) {
746
793
  is_trained = this->storage->is_trained;
747
794
  own_fields = true;
748
795
  }
@@ -855,8 +902,11 @@ void IndexHNSW2Level::search(
855
902
 
856
903
  const IndexIVFPQ* index_ivfpq =
857
904
  dynamic_cast<const IndexIVFPQ*>(storage);
905
+ FAISS_THROW_IF_NOT_MSG(
906
+ index_ivfpq,
907
+ "IndexHNSW2Level mixed search requires IndexIVFPQ storage");
858
908
 
859
- int nprobe = index_ivfpq->nprobe;
909
+ size_t nprobe = index_ivfpq->nprobe;
860
910
 
861
911
  std::unique_ptr<idx_t[]> coarse_assign(new idx_t[n * nprobe]);
862
912
  std::unique_ptr<float[]> coarse_dis(new float[n * nprobe]);
@@ -874,72 +924,88 @@ void IndexHNSW2Level::search(
874
924
  labels,
875
925
  false);
876
926
 
927
+ std::exception_ptr ex;
928
+ std::atomic<bool> interrupt{false};
877
929
  #pragma omp parallel
878
930
  {
879
931
  // visited table (not hash set) for tri-state flags.
880
- VisitedTable vt(ntotal, /*use_hashset=*/false);
881
- std::unique_ptr<DistanceComputer> dis(
882
- storage_distance_computer(storage));
883
-
932
+ std::unique_ptr<VisitedTable> vt;
933
+ std::unique_ptr<DistanceComputer> dis;
884
934
  constexpr int candidates_size = 1;
885
- MinimaxHeap candidates(candidates_size);
935
+ std::unique_ptr<MinimaxHeap> candidates;
936
+ try {
937
+ vt = std::make_unique<VisitedTable>(
938
+ ntotal, /*use_hashset=*/false);
939
+ dis.reset(storage_distance_computer(storage));
940
+ candidates = std::make_unique<MinimaxHeap>(candidates_size);
941
+ } catch (...) {
942
+ omp_capture_exception(ex, [&] { interrupt = true; });
943
+ }
886
944
 
887
945
  #pragma omp for reduction(+ : n1, n2, ndis, nhops)
888
946
  for (idx_t i = 0; i < n; i++) {
889
- idx_t* idxi = labels + i * k;
890
- float* simi = distances + i * k;
891
- dis->set_query(x + i * d);
892
-
893
- // mark all inverted list elements as visited
894
-
895
- for (int j = 0; j < nprobe; j++) {
896
- idx_t key = coarse_assign[j + i * nprobe];
897
- if (key < 0) {
898
- break;
899
- }
900
- size_t list_length = index_ivfpq->get_list_size(key);
901
- const idx_t* ids = index_ivfpq->invlists->get_ids(key);
947
+ if (interrupt.load(std::memory_order_relaxed)) {
948
+ continue;
949
+ }
950
+ try {
951
+ idx_t* idxi = labels + i * k;
952
+ float* simi = distances + i * k;
953
+ dis->set_query(x + i * d);
954
+
955
+ // mark all inverted list elements as visited
956
+ for (size_t j = 0; j < nprobe; j++) {
957
+ idx_t key = coarse_assign[j + i * nprobe];
958
+ if (key < 0) {
959
+ break;
960
+ }
961
+ size_t list_length = index_ivfpq->get_list_size(key);
962
+ const idx_t* ids = index_ivfpq->invlists->get_ids(key);
902
963
 
903
- for (int jj = 0; jj < list_length; jj++) {
904
- vt.set(ids[jj]);
964
+ for (size_t jj = 0; jj < list_length; jj++) {
965
+ vt->set(ids[jj]);
966
+ }
905
967
  }
906
- }
907
968
 
908
- candidates.clear();
969
+ candidates->clear();
909
970
 
910
- for (int j = 0; j < k; j++) {
911
- if (idxi[j] < 0) {
912
- break;
971
+ for (int j = 0; j < k; j++) {
972
+ if (idxi[j] < 0) {
973
+ break;
974
+ }
975
+ candidates->push(
976
+ static_cast<storage_idx_t>(idxi[j]), simi[j]);
913
977
  }
914
- candidates.push(idxi[j], simi[j]);
915
- }
916
978
 
917
- // reorder from sorted to heap
918
- maxheap_heapify(k, simi, idxi, simi, idxi, k);
919
-
920
- HNSWStats search_stats;
921
- search_from_candidates_2(
922
- hnsw,
923
- *dis,
924
- k,
925
- idxi,
926
- simi,
927
- candidates,
928
- vt,
929
- search_stats,
930
- 0,
931
- k);
932
- n1 += search_stats.n1;
933
- n2 += search_stats.n2;
934
- ndis += search_stats.ndis;
935
- nhops += search_stats.nhops;
936
-
937
- vt.advance();
938
- vt.advance();
979
+ // reorder from sorted to heap
980
+ maxheap_heapify(k, simi, idxi, simi, idxi, k);
939
981
 
940
- maxheap_reorder(k, simi, idxi);
982
+ HNSWStats search_stats;
983
+ search_from_candidates_2(
984
+ hnsw,
985
+ *dis,
986
+ k,
987
+ idxi,
988
+ simi,
989
+ *candidates,
990
+ *vt,
991
+ search_stats,
992
+ 0,
993
+ k);
994
+ n1 += search_stats.n1;
995
+ n2 += search_stats.n2;
996
+ ndis += search_stats.ndis;
997
+ nhops += search_stats.nhops;
998
+
999
+ vt->advance();
1000
+ vt->advance();
1001
+
1002
+ maxheap_reorder(k, simi, idxi);
1003
+ } catch (...) {
1004
+ omp_capture_exception(ex, [&] { interrupt = true; });
1005
+ }
941
1006
  }
942
1007
  }
1008
+ omp_rethrow_if_exception(ex);
943
1009
 
944
1010
  hnsw_stats.combine({n1, n2, ndis, nhops});
945
1011
  }
@@ -976,11 +1042,11 @@ IndexHNSWCagra::IndexHNSWCagra() {
976
1042
  }
977
1043
 
978
1044
  IndexHNSWCagra::IndexHNSWCagra(
979
- int d,
1045
+ int d_in,
980
1046
  int M,
981
1047
  MetricType metric,
982
1048
  NumericType numeric_type)
983
- : IndexHNSW(d, M, metric) {
1049
+ : IndexHNSW(d_in, M, metric) {
984
1050
  FAISS_THROW_IF_NOT_MSG(
985
1051
  ((metric == METRIC_L2) || (metric == METRIC_INNER_PRODUCT)),
986
1052
  "unsupported metric type for IndexHNSWCagra");
@@ -1024,10 +1090,18 @@ void IndexHNSWCagra::search(
1024
1090
  if (!base_level_only) {
1025
1091
  IndexHNSW::search(n, x, k, distances, labels, params);
1026
1092
  } else {
1093
+ if (ntotal == 0) {
1094
+ std::fill(
1095
+ distances,
1096
+ distances + n * k,
1097
+ std::numeric_limits<float>::max());
1098
+ std::fill(labels, labels + n * k, -1);
1099
+ return;
1100
+ }
1027
1101
  std::vector<storage_idx_t> nearest(n);
1028
1102
  std::vector<float> nearest_d(n);
1029
1103
 
1030
- #pragma omp for
1104
+ #pragma omp parallel for
1031
1105
  for (idx_t i = 0; i < n; i++) {
1032
1106
  std::unique_ptr<DistanceComputer> dis(
1033
1107
  storage_distance_computer(this->storage));
@@ -1043,7 +1117,7 @@ void IndexHNSWCagra::search(
1043
1117
  auto idx = distrib(gen);
1044
1118
  auto distance = (*dis)(idx);
1045
1119
  if (distance < nearest_d[i]) {
1046
- nearest[i] = idx;
1120
+ nearest[i] = static_cast<storage_idx_t>(idx);
1047
1121
  nearest_d[i] = distance;
1048
1122
  }
1049
1123
  }
@@ -1065,6 +1139,70 @@ void IndexHNSWCagra::search(
1065
1139
  }
1066
1140
  }
1067
1141
 
1142
+ void IndexHNSWCagra::range_search(
1143
+ idx_t n,
1144
+ const float* x,
1145
+ float radius,
1146
+ RangeSearchResult* result,
1147
+ const SearchParameters* params) const {
1148
+ if (!base_level_only) {
1149
+ IndexHNSW::range_search(n, x, radius, result, params);
1150
+ return;
1151
+ }
1152
+
1153
+ const HNSW& hnsw = this->hnsw;
1154
+ size_t n1 = 0, n2 = 0, ndis = 0, nhops = 0;
1155
+ float threshold = is_similarity_metric(metric_type) ? -radius : radius;
1156
+ RangeSearchPartialResult pres(result);
1157
+
1158
+ for (idx_t i = 0; i < n; i++) {
1159
+ std::unique_ptr<DistanceComputer> dis(
1160
+ storage_distance_computer(storage));
1161
+ dis->set_query(x + i * d);
1162
+
1163
+ storage_idx_t nearest = -1;
1164
+ float nearest_d = std::numeric_limits<float>::max();
1165
+
1166
+ std::random_device rd;
1167
+ std::mt19937 gen(rd());
1168
+ std::uniform_int_distribution<idx_t> distrib(0, ntotal - 1);
1169
+
1170
+ for (idx_t j = 0; j < num_base_level_search_entrypoints; j++) {
1171
+ auto idx = distrib(gen);
1172
+ auto distance = (*dis)(idx);
1173
+ if (distance < nearest_d) {
1174
+ nearest = idx;
1175
+ nearest_d = distance;
1176
+ }
1177
+ }
1178
+ FAISS_THROW_IF_NOT_MSG(
1179
+ nearest >= 0, "Could not find a valid entrypoint.");
1180
+
1181
+ RangeQueryResult& qres = pres.new_result(i);
1182
+ RangeResultHandler<HNSW::C> res(&qres, threshold);
1183
+ VisitedTable vt(ntotal, hnsw.use_visited_hashset);
1184
+ HNSWStats stats;
1185
+ hnsw.search_level_0(
1186
+ *dis, res, 1, &nearest, &nearest_d, 1, stats, vt, params);
1187
+ n1 += stats.n1;
1188
+ n2 += stats.n2;
1189
+ ndis += stats.ndis;
1190
+ nhops += stats.nhops;
1191
+ }
1192
+
1193
+ pres.set_lims();
1194
+ result->do_allocation();
1195
+ pres.copy_result();
1196
+
1197
+ hnsw_stats.combine({n1, n2, ndis, nhops});
1198
+
1199
+ if (is_similarity_metric(metric_type)) {
1200
+ for (size_t i = 0; i < result->lims[result->nq]; i++) {
1201
+ result->distances[i] = -result->distances[i];
1202
+ }
1203
+ }
1204
+ }
1205
+
1068
1206
  faiss::NumericType IndexHNSWCagra::get_numeric_type() const {
1069
1207
  return numeric_type_;
1070
1208
  }