faiss 0.5.3 → 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 (379) hide show
  1. checksums.yaml +4 -4
  2. data/CHANGELOG.md +12 -0
  3. data/ext/faiss/ext.cpp +1 -1
  4. data/ext/faiss/extconf.rb +4 -4
  5. data/ext/faiss/index.cpp +63 -45
  6. data/ext/faiss/index_binary.cpp +37 -27
  7. data/ext/faiss/kmeans.cpp +9 -8
  8. data/ext/faiss/pca_matrix.cpp +9 -7
  9. data/ext/faiss/product_quantizer.cpp +13 -11
  10. data/ext/faiss/utils.cpp +4 -2
  11. data/ext/faiss/utils.h +4 -0
  12. data/lib/faiss/version.rb +1 -1
  13. data/lib/faiss.rb +1 -1
  14. data/vendor/faiss/faiss/AutoTune.cpp +214 -82
  15. data/vendor/faiss/faiss/AutoTune.h +14 -1
  16. data/vendor/faiss/faiss/Clustering.cpp +97 -249
  17. data/vendor/faiss/faiss/Clustering.h +18 -0
  18. data/vendor/faiss/faiss/IVFlib.cpp +67 -44
  19. data/vendor/faiss/faiss/Index.cpp +25 -12
  20. data/vendor/faiss/faiss/Index.h +26 -4
  21. data/vendor/faiss/faiss/Index2Layer.cpp +37 -53
  22. data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp +68 -61
  23. data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.cpp +36 -34
  24. data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.h +4 -1
  25. data/vendor/faiss/faiss/IndexBinary.cpp +6 -3
  26. data/vendor/faiss/faiss/IndexBinary.h +4 -4
  27. data/vendor/faiss/faiss/IndexBinaryFlat.cpp +1 -1
  28. data/vendor/faiss/faiss/IndexBinaryFlat.h +1 -1
  29. data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +4 -4
  30. data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +92 -95
  31. data/vendor/faiss/faiss/IndexBinaryHNSW.h +9 -3
  32. data/vendor/faiss/faiss/IndexBinaryHash.cpp +45 -236
  33. data/vendor/faiss/faiss/IndexBinaryHash.h +6 -6
  34. data/vendor/faiss/faiss/IndexBinaryIVF.cpp +120 -414
  35. data/vendor/faiss/faiss/IndexFastScan.cpp +105 -129
  36. data/vendor/faiss/faiss/IndexFastScan.h +35 -24
  37. data/vendor/faiss/faiss/IndexFlat.cpp +216 -152
  38. data/vendor/faiss/faiss/IndexFlat.h +32 -14
  39. data/vendor/faiss/faiss/IndexFlatCodes.cpp +88 -41
  40. data/vendor/faiss/faiss/IndexFlatCodes.h +7 -1
  41. data/vendor/faiss/faiss/IndexHNSW.cpp +299 -187
  42. data/vendor/faiss/faiss/IndexHNSW.h +30 -14
  43. data/vendor/faiss/faiss/IndexIDMap.cpp +26 -22
  44. data/vendor/faiss/faiss/IndexIDMap.h +9 -7
  45. data/vendor/faiss/faiss/IndexIVF.cpp +535 -405
  46. data/vendor/faiss/faiss/IndexIVF.h +47 -16
  47. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp +77 -74
  48. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.cpp +105 -99
  49. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.h +6 -3
  50. data/vendor/faiss/faiss/IndexIVFFastScan.cpp +379 -249
  51. data/vendor/faiss/faiss/IndexIVFFastScan.h +65 -60
  52. data/vendor/faiss/faiss/IndexIVFFlat.cpp +41 -124
  53. data/vendor/faiss/faiss/IndexIVFFlat.h +32 -0
  54. data/vendor/faiss/faiss/IndexIVFFlatPanorama.cpp +89 -138
  55. data/vendor/faiss/faiss/IndexIVFFlatPanorama.h +3 -1
  56. data/vendor/faiss/faiss/IndexIVFIndependentQuantizer.cpp +18 -15
  57. data/vendor/faiss/faiss/IndexIVFPQ.cpp +77 -907
  58. data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +184 -122
  59. data/vendor/faiss/faiss/IndexIVFPQFastScan.h +3 -0
  60. data/vendor/faiss/faiss/IndexIVFPQR.cpp +23 -18
  61. data/vendor/faiss/faiss/IndexIVFRaBitQ.cpp +59 -60
  62. data/vendor/faiss/faiss/IndexIVFRaBitQ.h +4 -3
  63. data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.cpp +564 -416
  64. data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.h +269 -111
  65. data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +41 -127
  66. data/vendor/faiss/faiss/IndexIVFSpectralHash.h +1 -1
  67. data/vendor/faiss/faiss/IndexLSH.cpp +44 -25
  68. data/vendor/faiss/faiss/IndexLattice.cpp +41 -36
  69. data/vendor/faiss/faiss/IndexNNDescent.cpp +37 -21
  70. data/vendor/faiss/faiss/IndexNNDescent.h +2 -2
  71. data/vendor/faiss/faiss/IndexNSG.cpp +40 -23
  72. data/vendor/faiss/faiss/IndexNSG.h +0 -2
  73. data/vendor/faiss/faiss/IndexNeuralNetCodec.cpp +32 -12
  74. data/vendor/faiss/faiss/IndexPQ.cpp +129 -213
  75. data/vendor/faiss/faiss/IndexPQ.h +3 -2
  76. data/vendor/faiss/faiss/IndexPQFastScan.cpp +20 -14
  77. data/vendor/faiss/faiss/IndexPQFastScan.h +3 -0
  78. data/vendor/faiss/faiss/IndexPreTransform.cpp +25 -18
  79. data/vendor/faiss/faiss/IndexPreTransform.h +1 -1
  80. data/vendor/faiss/faiss/IndexRaBitQ.cpp +31 -43
  81. data/vendor/faiss/faiss/IndexRaBitQ.h +4 -3
  82. data/vendor/faiss/faiss/IndexRaBitQFastScan.cpp +135 -317
  83. data/vendor/faiss/faiss/IndexRaBitQFastScan.h +192 -34
  84. data/vendor/faiss/faiss/IndexRefine.cpp +30 -55
  85. data/vendor/faiss/faiss/IndexRefine.h +4 -4
  86. data/vendor/faiss/faiss/IndexReplicas.cpp +6 -6
  87. data/vendor/faiss/faiss/IndexRowwiseMinMax.cpp +15 -14
  88. data/vendor/faiss/faiss/IndexRowwiseMinMax.h +1 -1
  89. data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +82 -14
  90. data/vendor/faiss/faiss/IndexShards.cpp +13 -13
  91. data/vendor/faiss/faiss/IndexShardsIVF.cpp +21 -15
  92. data/vendor/faiss/faiss/MatrixStats.cpp +5 -4
  93. data/vendor/faiss/faiss/MetaIndexes.cpp +19 -17
  94. data/vendor/faiss/faiss/MetaIndexes.h +1 -1
  95. data/vendor/faiss/faiss/MetricType.h +29 -6
  96. data/vendor/faiss/faiss/SuperKMeans.cpp +656 -0
  97. data/vendor/faiss/faiss/SuperKMeans.h +97 -0
  98. data/vendor/faiss/faiss/VectorTransform.cpp +349 -141
  99. data/vendor/faiss/faiss/VectorTransform.h +39 -16
  100. data/vendor/faiss/faiss/build.cpp +23 -0
  101. data/vendor/faiss/faiss/build.h +15 -0
  102. data/vendor/faiss/faiss/clone_index.cpp +55 -51
  103. data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-avx2-inl.h +47 -47
  104. data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-inl.h +11 -0
  105. data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-avx2-inl.h +38 -38
  106. data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-inl.h +11 -0
  107. data/vendor/faiss/faiss/{cppcontrib/factory_tools.cpp → factory_tools.cpp} +6 -1
  108. data/vendor/faiss/faiss/gpu/GpuCloner.cpp +1 -1
  109. data/vendor/faiss/faiss/gpu/GpuIndexCagra.h +6 -5
  110. data/vendor/faiss/faiss/gpu/GpuResources.h +1 -1
  111. data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp +9 -9
  112. data/vendor/faiss/faiss/gpu/StandardGpuResources.h +4 -3
  113. data/vendor/faiss/faiss/gpu/test/TestGpuIndexFlat.cpp +46 -0
  114. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFFlat.cpp +56 -0
  115. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFPQ.cpp +78 -1
  116. data/vendor/faiss/faiss/gpu/test/TestUtils.cpp +72 -0
  117. data/vendor/faiss/faiss/gpu/test/TestUtils.h +23 -0
  118. data/vendor/faiss/faiss/gpu/utils/CuvsFilterConvert.h +1 -1
  119. data/vendor/faiss/faiss/gpu/utils/CuvsUtils.h +21 -10
  120. data/vendor/faiss/faiss/gpu_metal/GpuIndexFlat.h +22 -0
  121. data/vendor/faiss/faiss/gpu_metal/MetalCloner.h +35 -0
  122. data/vendor/faiss/faiss/gpu_metal/MetalFlatKernels.h +40 -0
  123. data/vendor/faiss/faiss/gpu_metal/MetalIndex.h +51 -0
  124. data/vendor/faiss/faiss/gpu_metal/MetalIndexFlat.h +65 -0
  125. data/vendor/faiss/faiss/gpu_metal/MetalKernels.h +66 -0
  126. data/vendor/faiss/faiss/gpu_metal/MetalResources.h +79 -0
  127. data/vendor/faiss/faiss/gpu_metal/StandardMetalResources.h +35 -0
  128. data/vendor/faiss/faiss/impl/AdSampling.cpp +103 -0
  129. data/vendor/faiss/faiss/impl/AdSampling.h +35 -0
  130. data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp +64 -34
  131. data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +1 -0
  132. data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp +10 -9
  133. data/vendor/faiss/faiss/impl/AuxIndexStructures.h +3 -28
  134. data/vendor/faiss/faiss/impl/ClusteringHelpers.cpp +244 -0
  135. data/vendor/faiss/faiss/impl/ClusteringHelpers.h +94 -0
  136. data/vendor/faiss/faiss/impl/ClusteringInitialization.cpp +367 -0
  137. data/vendor/faiss/faiss/impl/ClusteringInitialization.h +107 -0
  138. data/vendor/faiss/faiss/impl/CodePacker.cpp +7 -3
  139. data/vendor/faiss/faiss/impl/CodePacker.h +11 -3
  140. data/vendor/faiss/faiss/impl/CodePackerRaBitQ.cpp +83 -0
  141. data/vendor/faiss/faiss/impl/CodePackerRaBitQ.h +47 -0
  142. data/vendor/faiss/faiss/impl/DistanceComputer.h +8 -8
  143. data/vendor/faiss/faiss/impl/FaissAssert.h +64 -3
  144. data/vendor/faiss/faiss/impl/FaissException.h +50 -3
  145. data/vendor/faiss/faiss/impl/HNSW.cpp +117 -351
  146. data/vendor/faiss/faiss/impl/HNSW.h +21 -40
  147. data/vendor/faiss/faiss/impl/IDSelector.cpp +15 -11
  148. data/vendor/faiss/faiss/impl/IDSelector.h +8 -8
  149. data/vendor/faiss/faiss/impl/InvertedListScannerStats.h +26 -0
  150. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +114 -102
  151. data/vendor/faiss/faiss/impl/NNDescent.cpp +63 -26
  152. data/vendor/faiss/faiss/impl/NNDescent.h +6 -2
  153. data/vendor/faiss/faiss/impl/NSG.cpp +44 -26
  154. data/vendor/faiss/faiss/impl/NSG.h +20 -10
  155. data/vendor/faiss/faiss/impl/Panorama.cpp +76 -52
  156. data/vendor/faiss/faiss/impl/Panorama.h +265 -78
  157. data/vendor/faiss/faiss/impl/PdxLayout.cpp +93 -0
  158. data/vendor/faiss/faiss/impl/PdxLayout.h +41 -0
  159. data/vendor/faiss/faiss/impl/PolysemousTraining.cpp +62 -37
  160. data/vendor/faiss/faiss/impl/PolysemousTraining.h +3 -3
  161. data/vendor/faiss/faiss/impl/ProductAdditiveQuantizer.cpp +35 -35
  162. data/vendor/faiss/faiss/impl/ProductQuantizer-inl.h +21 -16
  163. data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +99 -80
  164. data/vendor/faiss/faiss/impl/Quantizer.h +2 -2
  165. data/vendor/faiss/faiss/impl/RaBitQUtils.cpp +135 -37
  166. data/vendor/faiss/faiss/impl/RaBitQUtils.h +148 -21
  167. data/vendor/faiss/faiss/impl/RaBitQuantizer.cpp +298 -301
  168. data/vendor/faiss/faiss/impl/RaBitQuantizer.h +3 -10
  169. data/vendor/faiss/faiss/impl/RaBitQuantizerMultiBit.cpp +15 -41
  170. data/vendor/faiss/faiss/impl/RaBitQuantizerMultiBit.h +0 -4
  171. data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp +40 -32
  172. data/vendor/faiss/faiss/impl/ResidualQuantizer.h +1 -1
  173. data/vendor/faiss/faiss/impl/ResultHandler.h +218 -113
  174. data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +119 -2362
  175. data/vendor/faiss/faiss/impl/ScalarQuantizer.h +27 -3
  176. data/vendor/faiss/faiss/impl/ThreadedIndex-inl.h +14 -11
  177. data/vendor/faiss/faiss/impl/VisitedTable.cpp +42 -0
  178. data/vendor/faiss/faiss/impl/VisitedTable.h +76 -0
  179. data/vendor/faiss/faiss/impl/approx_topk/approx_topk.h +276 -0
  180. data/vendor/faiss/faiss/impl/approx_topk/avx2.cpp +68 -0
  181. data/vendor/faiss/faiss/{utils → impl}/approx_topk/generic.h +15 -8
  182. data/vendor/faiss/faiss/impl/approx_topk/neon.cpp +68 -0
  183. data/vendor/faiss/faiss/impl/approx_topk/rq_beam_search_tab-inl.h +169 -0
  184. data/vendor/faiss/faiss/impl/approx_topk/rq_beam_search_tab.h +117 -0
  185. data/vendor/faiss/faiss/impl/approx_topk/simdlib256-inl.h +146 -0
  186. data/vendor/faiss/faiss/impl/binary_hamming/IndexBinaryHNSW_impl.h +73 -0
  187. data/vendor/faiss/faiss/impl/binary_hamming/IndexBinaryHash_impl.h +270 -0
  188. data/vendor/faiss/faiss/impl/binary_hamming/IndexBinaryIVF_impl.h +460 -0
  189. data/vendor/faiss/faiss/impl/binary_hamming/IndexIVFSpectralHash_impl.h +159 -0
  190. data/vendor/faiss/faiss/impl/binary_hamming/IndexPQ_impl.h +92 -0
  191. data/vendor/faiss/faiss/impl/binary_hamming/avx2.cpp +26 -0
  192. data/vendor/faiss/faiss/impl/binary_hamming/avx512.cpp +26 -0
  193. data/vendor/faiss/faiss/impl/binary_hamming/dispatch.h +143 -0
  194. data/vendor/faiss/faiss/impl/binary_hamming/neon.cpp +26 -0
  195. data/vendor/faiss/faiss/impl/binary_hamming/rvv.cpp +26 -0
  196. data/vendor/faiss/faiss/impl/expanded_scanners.h +163 -0
  197. data/vendor/faiss/faiss/impl/{FastScanDistancePostProcessing.h → fast_scan/FastScanDistancePostProcessing.h} +13 -6
  198. data/vendor/faiss/faiss/impl/{LookupTableScaler.h → fast_scan/LookupTableScaler.h} +16 -5
  199. data/vendor/faiss/faiss/impl/fast_scan/accumulate_loops.h +237 -0
  200. data/vendor/faiss/faiss/impl/fast_scan/accumulate_loops_512.h +185 -0
  201. data/vendor/faiss/faiss/impl/fast_scan/decompose_qbs.h +229 -0
  202. data/vendor/faiss/faiss/impl/fast_scan/dispatching.h +268 -0
  203. data/vendor/faiss/faiss/impl/{pq4_fast_scan.cpp → fast_scan/fast_scan.cpp} +176 -4
  204. data/vendor/faiss/faiss/impl/fast_scan/fast_scan.h +341 -0
  205. data/vendor/faiss/faiss/impl/fast_scan/impl-avx2.cpp +36 -0
  206. data/vendor/faiss/faiss/impl/fast_scan/impl-avx512.cpp +40 -0
  207. data/vendor/faiss/faiss/impl/fast_scan/impl-neon.cpp +120 -0
  208. data/vendor/faiss/faiss/impl/fast_scan/impl-riscv.cpp +104 -0
  209. data/vendor/faiss/faiss/impl/fast_scan/kernels_simd256.h +213 -0
  210. data/vendor/faiss/faiss/impl/{pq4_fast_scan_search_qbs.cpp → fast_scan/kernels_simd512.h} +26 -348
  211. data/vendor/faiss/faiss/impl/fast_scan/rabitq_dispatching.h +90 -0
  212. data/vendor/faiss/faiss/impl/fast_scan/rabitq_result_handler.h +108 -0
  213. data/vendor/faiss/faiss/impl/{simd_result_handlers.h → fast_scan/simd_result_handlers.h} +290 -142
  214. data/vendor/faiss/faiss/impl/hnsw/LockVector.cpp +54 -0
  215. data/vendor/faiss/faiss/impl/hnsw/LockVector.h +64 -0
  216. data/vendor/faiss/faiss/impl/hnsw/MinimaxHeap.cpp +91 -0
  217. data/vendor/faiss/faiss/impl/hnsw/MinimaxHeap.h +64 -0
  218. data/vendor/faiss/faiss/impl/hnsw/avx2.cpp +104 -0
  219. data/vendor/faiss/faiss/impl/hnsw/avx512.cpp +111 -0
  220. data/vendor/faiss/faiss/impl/index_read.cpp +1950 -505
  221. data/vendor/faiss/faiss/impl/index_read_utils.h +1 -2
  222. data/vendor/faiss/faiss/impl/index_write.cpp +112 -21
  223. data/vendor/faiss/faiss/impl/io.cpp +6 -6
  224. data/vendor/faiss/faiss/impl/io_macros.h +33 -16
  225. data/vendor/faiss/faiss/impl/kmeans1d.cpp +10 -10
  226. data/vendor/faiss/faiss/impl/lattice_Zn.cpp +81 -40
  227. data/vendor/faiss/faiss/impl/lattice_Zn.h +6 -6
  228. data/vendor/faiss/faiss/impl/mapped_io.cpp +15 -8
  229. data/vendor/faiss/faiss/impl/platform_macros.h +11 -4
  230. data/vendor/faiss/faiss/impl/pq_code_distance/IVFPQScanner_impl.h +549 -0
  231. data/vendor/faiss/faiss/impl/pq_code_distance/IVFPQ_QueryTables.cpp +245 -0
  232. data/vendor/faiss/faiss/impl/pq_code_distance/IVFPQ_QueryTables.h +105 -0
  233. data/vendor/faiss/faiss/impl/pq_code_distance/PQDistanceComputer_impl.h +106 -0
  234. data/vendor/faiss/faiss/impl/pq_code_distance/avx2.cpp +21 -0
  235. data/vendor/faiss/faiss/impl/pq_code_distance/avx512.cpp +21 -0
  236. data/vendor/faiss/faiss/impl/pq_code_distance/neon.cpp +21 -0
  237. data/vendor/faiss/faiss/impl/{code_distance/code_distance-avx2.h → pq_code_distance/pq_code_distance-avx2.h} +43 -220
  238. data/vendor/faiss/faiss/impl/{code_distance/code_distance-avx512.h → pq_code_distance/pq_code_distance-avx512.h} +25 -112
  239. data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-generic.cpp +59 -0
  240. data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-generic.h +96 -0
  241. data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-inl.h +256 -0
  242. data/vendor/faiss/faiss/impl/{code_distance/code_distance-sve.h → pq_code_distance/pq_code_distance-sve.cpp} +57 -146
  243. data/vendor/faiss/faiss/impl/pq_code_distance/rvv.cpp +68 -0
  244. data/vendor/faiss/faiss/impl/residual_quantizer_encode_steps.cpp +320 -483
  245. data/vendor/faiss/faiss/impl/residual_quantizer_encode_steps.h +1 -1
  246. data/vendor/faiss/faiss/impl/scalar_quantizer/codecs.h +121 -0
  247. data/vendor/faiss/faiss/impl/scalar_quantizer/distance_computers.h +137 -0
  248. data/vendor/faiss/faiss/impl/scalar_quantizer/quantizers.h +371 -0
  249. data/vendor/faiss/faiss/impl/scalar_quantizer/scanners.h +190 -0
  250. data/vendor/faiss/faiss/impl/scalar_quantizer/similarities.h +94 -0
  251. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx2.cpp +603 -0
  252. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx512.cpp +597 -0
  253. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-dispatch.h +388 -0
  254. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-neon.cpp +630 -0
  255. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-rvv.cpp +311 -0
  256. data/vendor/faiss/faiss/impl/scalar_quantizer/training.cpp +387 -0
  257. data/vendor/faiss/faiss/impl/scalar_quantizer/training.h +54 -0
  258. data/vendor/faiss/faiss/impl/simd_dispatch.h +173 -0
  259. data/vendor/faiss/faiss/impl/simdlib/simdlib.h +57 -0
  260. data/vendor/faiss/faiss/{utils → impl/simdlib}/simdlib_avx2.h +274 -171
  261. data/vendor/faiss/faiss/impl/simdlib/simdlib_avx512.h +414 -0
  262. data/vendor/faiss/faiss/impl/simdlib/simdlib_dispatch.h +44 -0
  263. data/vendor/faiss/faiss/{utils → impl/simdlib}/simdlib_emulated.h +231 -166
  264. data/vendor/faiss/faiss/{utils → impl/simdlib}/simdlib_neon.h +275 -217
  265. data/vendor/faiss/faiss/{utils → impl/simdlib}/simdlib_ppc64.h +201 -160
  266. data/vendor/faiss/faiss/impl/svs_io.cpp +12 -3
  267. data/vendor/faiss/faiss/impl/svs_io.h +8 -2
  268. data/vendor/faiss/faiss/index_factory.cpp +115 -28
  269. data/vendor/faiss/faiss/index_io.h +53 -3
  270. data/vendor/faiss/faiss/invlists/BlockInvertedLists.cpp +73 -20
  271. data/vendor/faiss/faiss/invlists/DirectMap.cpp +24 -14
  272. data/vendor/faiss/faiss/invlists/DirectMap.h +4 -3
  273. data/vendor/faiss/faiss/invlists/InvertedLists.cpp +157 -73
  274. data/vendor/faiss/faiss/invlists/InvertedLists.h +86 -23
  275. data/vendor/faiss/faiss/invlists/InvertedListsIOHook.cpp +4 -4
  276. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +14 -14
  277. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.h +1 -1
  278. data/vendor/faiss/faiss/svs/IndexSVSFaissUtils.h +9 -19
  279. data/vendor/faiss/faiss/svs/IndexSVSFlat.cpp +2 -2
  280. data/vendor/faiss/faiss/svs/IndexSVSFlat.h +2 -0
  281. data/vendor/faiss/faiss/svs/IndexSVSIVF.cpp +350 -0
  282. data/vendor/faiss/faiss/svs/IndexSVSIVF.h +128 -0
  283. data/vendor/faiss/faiss/svs/IndexSVSIVFLVQ.cpp +40 -0
  284. data/vendor/faiss/faiss/svs/IndexSVSIVFLVQ.h +43 -0
  285. data/vendor/faiss/faiss/svs/IndexSVSIVFLeanVec.cpp +225 -0
  286. data/vendor/faiss/faiss/svs/IndexSVSIVFLeanVec.h +71 -0
  287. data/vendor/faiss/faiss/svs/IndexSVSVamana.cpp +25 -1
  288. data/vendor/faiss/faiss/svs/IndexSVSVamana.h +19 -2
  289. data/vendor/faiss/faiss/svs/IndexSVSVamanaLVQ.h +1 -1
  290. data/vendor/faiss/faiss/svs/IndexSVSVamanaLeanVec.cpp +19 -2
  291. data/vendor/faiss/faiss/svs/IndexSVSVamanaLeanVec.h +14 -0
  292. data/vendor/faiss/faiss/utils/Heap.cpp +56 -10
  293. data/vendor/faiss/faiss/utils/Heap.h +21 -0
  294. data/vendor/faiss/faiss/utils/NeuralNet.cpp +54 -40
  295. data/vendor/faiss/faiss/utils/NeuralNet.h +1 -1
  296. data/vendor/faiss/faiss/utils/approx_topk_hamming/approx_topk_hamming.h +10 -4
  297. data/vendor/faiss/faiss/utils/distances.cpp +507 -559
  298. data/vendor/faiss/faiss/utils/distances.h +118 -1
  299. data/vendor/faiss/faiss/utils/distances_dispatch.h +250 -0
  300. data/vendor/faiss/faiss/utils/distances_fused/avx512.cpp +8 -7
  301. data/vendor/faiss/faiss/utils/distances_fused/distances_fused.cpp +33 -14
  302. data/vendor/faiss/faiss/utils/distances_fused/distances_fused.h +12 -1
  303. data/vendor/faiss/faiss/utils/distances_fused/simdlib_based.cpp +16 -293
  304. data/vendor/faiss/faiss/utils/distances_fused/simdlib_based_neon.cpp +57 -0
  305. data/vendor/faiss/faiss/utils/distances_fused/simdlib_kernel-inl.h +290 -0
  306. data/vendor/faiss/faiss/utils/distances_simd.cpp +72 -3681
  307. data/vendor/faiss/faiss/utils/extra_distances.cpp +60 -102
  308. data/vendor/faiss/faiss/utils/extra_distances.h +79 -7
  309. data/vendor/faiss/faiss/utils/hamming-inl.h +13 -11
  310. data/vendor/faiss/faiss/utils/hamming.cpp +66 -517
  311. data/vendor/faiss/faiss/utils/hamming.h +92 -2
  312. data/vendor/faiss/faiss/utils/hamming_distance/common.h +287 -10
  313. data/vendor/faiss/faiss/utils/hamming_distance/hamming_avx2.cpp +15 -0
  314. data/vendor/faiss/faiss/utils/hamming_distance/hamming_avx512.cpp +15 -0
  315. data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-avx2.h +142 -0
  316. data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-avx512.h +234 -0
  317. data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-generic.h +368 -0
  318. data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-neon.h +322 -0
  319. data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-rvv.h +39 -0
  320. data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer.h +146 -0
  321. data/vendor/faiss/faiss/utils/hamming_distance/hamming_impl.h +481 -0
  322. data/vendor/faiss/faiss/utils/hamming_distance/hamming_neon.cpp +15 -0
  323. data/vendor/faiss/faiss/utils/hamming_distance/hamming_rvv.cpp +15 -0
  324. data/vendor/faiss/faiss/utils/partitioning.cpp +66 -987
  325. data/vendor/faiss/faiss/utils/partitioning.h +31 -0
  326. data/vendor/faiss/faiss/utils/popcount.h +29 -0
  327. data/vendor/faiss/faiss/utils/pq_code_distance.h +251 -0
  328. data/vendor/faiss/faiss/utils/prefetch.h +2 -2
  329. data/vendor/faiss/faiss/utils/quantize_lut.cpp +30 -30
  330. data/vendor/faiss/faiss/utils/quantize_lut.h +1 -1
  331. data/vendor/faiss/faiss/utils/rabitq_simd.h +124 -343
  332. data/vendor/faiss/faiss/utils/random.cpp +6 -6
  333. data/vendor/faiss/faiss/utils/simd_impl/IVFFlatScanner-inl.h +51 -0
  334. data/vendor/faiss/faiss/utils/simd_impl/distances_aarch64.cpp +154 -0
  335. data/vendor/faiss/faiss/utils/simd_impl/distances_arm_sve.cpp +777 -0
  336. data/vendor/faiss/faiss/utils/simd_impl/distances_autovec-inl.h +306 -0
  337. data/vendor/faiss/faiss/utils/simd_impl/distances_avx2.cpp +1431 -0
  338. data/vendor/faiss/faiss/utils/simd_impl/distances_avx512.cpp +1095 -0
  339. data/vendor/faiss/faiss/utils/simd_impl/distances_rvv.cpp +189 -0
  340. data/vendor/faiss/faiss/utils/simd_impl/distances_simdlib256.h +195 -0
  341. data/vendor/faiss/faiss/utils/simd_impl/distances_sse-inl.h +392 -0
  342. data/vendor/faiss/faiss/utils/{distances_fused/simdlib_based.h → simd_impl/exhaustive_L2sqr_blas_cmax.h} +5 -10
  343. data/vendor/faiss/faiss/utils/simd_impl/hamming_impl.h +481 -0
  344. data/vendor/faiss/faiss/utils/simd_impl/partitioning_avx2.cpp +14 -0
  345. data/vendor/faiss/faiss/utils/simd_impl/partitioning_neon.cpp +14 -0
  346. data/vendor/faiss/faiss/utils/simd_impl/partitioning_simdlib256.h +1085 -0
  347. data/vendor/faiss/faiss/utils/simd_impl/rabitq_avx2.cpp +355 -0
  348. data/vendor/faiss/faiss/utils/simd_impl/rabitq_avx512.cpp +477 -0
  349. data/vendor/faiss/faiss/utils/simd_impl/rabitq_neon.cpp +55 -0
  350. data/vendor/faiss/faiss/utils/simd_impl/rabitq_rvv.cpp +55 -0
  351. data/vendor/faiss/faiss/utils/simd_impl/super_kmeans_dispatch.h +32 -0
  352. data/vendor/faiss/faiss/utils/simd_impl/super_kmeans_kernels.h +43 -0
  353. data/vendor/faiss/faiss/utils/simd_impl/super_kmeans_kernels_avx2.cpp +57 -0
  354. data/vendor/faiss/faiss/utils/simd_impl/super_kmeans_kernels_avx512.cpp +45 -0
  355. data/vendor/faiss/faiss/utils/simd_levels.cpp +334 -0
  356. data/vendor/faiss/faiss/utils/simd_levels.h +183 -0
  357. data/vendor/faiss/faiss/utils/sorting.cpp +48 -36
  358. data/vendor/faiss/faiss/utils/utils.cpp +21 -14
  359. data/vendor/faiss/faiss/utils/utils.h +3 -3
  360. metadata +156 -42
  361. data/vendor/faiss/faiss/impl/RaBitQStats.cpp +0 -29
  362. data/vendor/faiss/faiss/impl/RaBitQStats.h +0 -56
  363. data/vendor/faiss/faiss/impl/code_distance/code_distance-generic.h +0 -81
  364. data/vendor/faiss/faiss/impl/code_distance/code_distance.h +0 -186
  365. data/vendor/faiss/faiss/impl/pq4_fast_scan.h +0 -216
  366. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp +0 -224
  367. data/vendor/faiss/faiss/utils/approx_topk/approx_topk.h +0 -84
  368. data/vendor/faiss/faiss/utils/approx_topk/avx2-inl.h +0 -196
  369. data/vendor/faiss/faiss/utils/approx_topk/mode.h +0 -34
  370. data/vendor/faiss/faiss/utils/distances_fused/avx512.h +0 -36
  371. data/vendor/faiss/faiss/utils/extra_distances-inl.h +0 -228
  372. data/vendor/faiss/faiss/utils/hamming_distance/avx2-inl.h +0 -462
  373. data/vendor/faiss/faiss/utils/hamming_distance/avx512-inl.h +0 -490
  374. data/vendor/faiss/faiss/utils/hamming_distance/generic-inl.h +0 -450
  375. data/vendor/faiss/faiss/utils/hamming_distance/hamdis-inl.h +0 -87
  376. data/vendor/faiss/faiss/utils/hamming_distance/neon-inl.h +0 -524
  377. data/vendor/faiss/faiss/utils/simdlib.h +0 -42
  378. data/vendor/faiss/faiss/utils/simdlib_avx512.h +0 -296
  379. /data/vendor/faiss/faiss/{cppcontrib/factory_tools.h → factory_tools.h} +0 -0
@@ -28,6 +28,7 @@
28
28
  #include <faiss/gpu/GpuIndexIVFPQ.h>
29
29
 
30
30
  #include <variant>
31
+ #include <vector>
31
32
  #include "faiss/Index.h"
32
33
 
33
34
  namespace faiss {
@@ -193,14 +194,14 @@ struct GpuIndexCagraConfig : public GpuIndexConfig {
193
194
 
194
195
  enum class search_algo {
195
196
  /// For large batch sizes.
196
- SINGLE_CTA,
197
+ SINGLE_CTA = 0,
197
198
  /// For small batch sizes.
198
- MULTI_CTA,
199
- MULTI_KERNEL,
200
- AUTO
199
+ MULTI_CTA = 1,
200
+ MULTI_KERNEL = 2,
201
+ AUTO = 100
201
202
  };
202
203
 
203
- enum class hash_mode { HASH, SMALL, AUTO };
204
+ enum class hash_mode { HASH = 0, SMALL = 1, AUTO = 100 };
204
205
 
205
206
  struct SearchParametersCagra : SearchParameters {
206
207
  /// Maximum number of queries to search at the same time (batch size). Auto
@@ -33,7 +33,7 @@
33
33
 
34
34
  #if defined USE_NVIDIA_CUVS
35
35
  #include <raft/core/device_resources.hpp>
36
- #include <rmm/mr/device/device_memory_resource.hpp>
36
+ #include <rmm/mr/device_memory_resource.hpp>
37
37
  #endif
38
38
 
39
39
  namespace faiss {
@@ -23,9 +23,9 @@
23
23
 
24
24
  #if defined USE_NVIDIA_CUVS
25
25
  #include <raft/core/device_resources.hpp>
26
- #include <rmm/mr/device/managed_memory_resource.hpp>
27
- #include <rmm/mr/device/per_device_resource.hpp>
28
- #include <rmm/mr/host/pinned_memory_resource.hpp>
26
+ #include <rmm/mr/managed_memory_resource.hpp>
27
+ #include <rmm/mr/per_device_resource.hpp>
28
+ #include <rmm/mr/pinned_host_memory_resource.hpp>
29
29
  #include <memory>
30
30
  #endif
31
31
 
@@ -93,7 +93,7 @@ StandardGpuResourcesImpl::StandardGpuResourcesImpl()
93
93
  :
94
94
  #if defined USE_NVIDIA_CUVS
95
95
  mmr_(new rmm::mr::managed_memory_resource),
96
- pmr_(new rmm::mr::pinned_memory_resource),
96
+ pmr_(new rmm::mr::pinned_host_memory_resource),
97
97
  #endif
98
98
  pinnedMemAlloc_(nullptr),
99
99
  pinnedMemAllocSize_(0),
@@ -164,7 +164,7 @@ StandardGpuResourcesImpl::~StandardGpuResourcesImpl() {
164
164
 
165
165
  if (pinnedMemAlloc_) {
166
166
  #if defined USE_NVIDIA_CUVS
167
- pmr_->deallocate(pinnedMemAlloc_, pinnedMemAllocSize_);
167
+ pmr_->deallocate_sync(pinnedMemAlloc_, pinnedMemAllocSize_);
168
168
  #else
169
169
  auto err = cudaFreeHost(pinnedMemAlloc_);
170
170
  FAISS_ASSERT_FMT(
@@ -350,7 +350,7 @@ void StandardGpuResourcesImpl::initializeForDevice(int device) {
350
350
  // pinned memory allocation
351
351
  if (defaultStreams_.empty() && pinnedMemSize_ > 0) {
352
352
  try {
353
- pinnedMemAlloc_ = pmr_->allocate(pinnedMemSize_);
353
+ pinnedMemAlloc_ = pmr_->allocate_sync(pinnedMemSize_);
354
354
  } catch (const std::bad_alloc& rmm_ex) {
355
355
  FAISS_THROW_MSG("CUDA memory allocation error");
356
356
  }
@@ -549,7 +549,7 @@ void* StandardGpuResourcesImpl::allocMemory(const AllocRequest& req) {
549
549
  rmm::mr::device_memory_resource* current_mr =
550
550
  rmm::mr::get_per_device_resource(
551
551
  rmm::cuda_device_id{adjReq.device});
552
- p = current_mr->allocate_async(adjReq.size, adjReq.stream);
552
+ p = current_mr->allocate(adjReq.stream, adjReq.size);
553
553
  adjReq.mr = current_mr;
554
554
  } catch (const std::bad_alloc& rmm_ex) {
555
555
  FAISS_THROW_MSG("CUDA memory allocation error");
@@ -584,7 +584,7 @@ void* StandardGpuResourcesImpl::allocMemory(const AllocRequest& req) {
584
584
  // TODO: change this to use the current device resource once RMM has
585
585
  // a way to retrieve a "guaranteed" managed memory resource for a
586
586
  // device.
587
- p = mmr_->allocate_async(adjReq.size, adjReq.stream);
587
+ p = mmr_->allocate(adjReq.stream, adjReq.size);
588
588
  adjReq.mr = mmr_.get();
589
589
  } catch (const std::bad_alloc& rmm_ex) {
590
590
  FAISS_THROW_MSG("CUDA memory allocation error");
@@ -648,7 +648,7 @@ void StandardGpuResourcesImpl::deallocMemory(int device, void* p) {
648
648
  req.space == MemorySpace::Device ||
649
649
  req.space == MemorySpace::Unified) {
650
650
  #if defined USE_NVIDIA_CUVS
651
- req.mr->deallocate_async(p, req.size, req.stream);
651
+ req.mr->deallocate(req.stream, p, req.size);
652
652
  #else
653
653
  auto err = cudaFree(p);
654
654
  FAISS_ASSERT_FMT(
@@ -25,7 +25,8 @@
25
25
 
26
26
  #if defined USE_NVIDIA_CUVS
27
27
  #include <raft/core/device_resources.hpp>
28
- #include <rmm/mr/host/pinned_memory_resource.hpp>
28
+ #include <rmm/mr/device_memory_resource.hpp>
29
+ #include <rmm/mr/pinned_host_memory_resource.hpp>
29
30
  #endif
30
31
 
31
32
  #include <faiss/gpu/GpuResources.h>
@@ -172,8 +173,8 @@ class StandardGpuResourcesImpl : public GpuResources {
172
173
  // managed_memory_resource
173
174
  std::unique_ptr<rmm::mr::device_memory_resource> mmr_;
174
175
 
175
- // pinned_memory_resource
176
- std::unique_ptr<rmm::mr::host_memory_resource> pmr_;
176
+ // pinned_host_memory_resource
177
+ std::unique_ptr<rmm::mr::pinned_host_memory_resource> pmr_;
177
178
  #endif
178
179
 
179
180
  /// Pinned memory allocation for use with this GPU
@@ -11,6 +11,7 @@
11
11
  #include <faiss/gpu/impl/IndexUtils.h>
12
12
  #include <faiss/gpu/test/TestUtils.h>
13
13
  #include <faiss/gpu/utils/DeviceUtils.h>
14
+ #include <faiss/impl/IDSelector.h>
14
15
  #include <gtest/gtest.h>
15
16
  #include <sstream>
16
17
  #include <unordered_map>
@@ -761,6 +762,51 @@ TEST(TestCuvsGpuIndexFlat, SearchAndReconstruct) {
761
762
  }
762
763
  #endif
763
764
 
765
+ void testIDSelectorFlat(faiss::MetricType metricType) {
766
+ int numAdd = faiss::gpu::randVal(2000, 5000);
767
+ int dim = faiss::gpu::randVal(64, 200);
768
+ int numQuery = faiss::gpu::randVal(32, 100);
769
+ int k = faiss::gpu::randVal(10, 30);
770
+ int device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
771
+
772
+ faiss::gpu::StandardGpuResources res;
773
+ res.noTempMemory();
774
+
775
+ auto queryVecs = faiss::gpu::randVecs(numQuery, dim);
776
+ faiss::gpu::GpuIndexFlatConfig config;
777
+ config.device = device;
778
+ config.use_cuvs = true;
779
+ for (bool useFloat16 : {false, true}) {
780
+ config.useFloat16 = useFloat16;
781
+ faiss::gpu::GpuIndexFlat gpuIndex(&res, dim, metricType, config);
782
+ std::vector<float> addVecs = faiss::gpu::randVecs(numAdd, dim);
783
+ gpuIndex.add(numAdd, addVecs.data());
784
+
785
+ faiss::gpu::TestIDSelectorStruct selector_struct(numAdd);
786
+ faiss::SearchParameters search_params;
787
+ for (auto& [selectorName, selector] : selector_struct.selector_map) {
788
+ search_params.sel = selector.get();
789
+ faiss::gpu::testIDSelectorSearch(
790
+ &gpuIndex,
791
+ &search_params,
792
+ queryVecs,
793
+ numQuery,
794
+ k,
795
+ selectorName);
796
+ }
797
+ }
798
+ }
799
+
800
+ #if defined USE_NVIDIA_CUVS
801
+ TEST(TestCuvsGpuIndexFlat, IDSelector_L2) {
802
+ testIDSelectorFlat(faiss::METRIC_L2);
803
+ }
804
+
805
+ TEST(TestCuvsGpuIndexFlat, IDSelector_IP) {
806
+ testIDSelectorFlat(faiss::METRIC_INNER_PRODUCT);
807
+ }
808
+ #endif
809
+
764
810
  int main(int argc, char** argv) {
765
811
  testing::InitGoogleTest(&argc, argv);
766
812
 
@@ -28,6 +28,8 @@
28
28
  #include <faiss/gpu/StandardGpuResources.h>
29
29
  #include <faiss/gpu/test/TestUtils.h>
30
30
  #include <faiss/gpu/utils/DeviceUtils.h>
31
+ #include <faiss/impl/IDSelector.h>
32
+ #include <faiss/utils/distances.h>
31
33
  #include <gtest/gtest.h>
32
34
  #include <cmath>
33
35
  #include <sstream>
@@ -354,6 +356,11 @@ TEST(TestGpuIndexIVFFlat, Float32_Query_IP) {
354
356
  }
355
357
 
356
358
  TEST(TestGpuIndexIVFFlat, LargeBatch) {
359
+ // With low-dim vectors, CPU will use non-BLAS. Force the CPU to use
360
+ // the BLAS for consistent comparison.
361
+ int saved_threshold = faiss::distance_compute_blas_threshold;
362
+ faiss::distance_compute_blas_threshold = 1;
363
+
357
364
  Options opt;
358
365
  opt.dim = 3;
359
366
  opt.numQuery = 100000;
@@ -364,6 +371,8 @@ TEST(TestGpuIndexIVFFlat, LargeBatch) {
364
371
  opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
365
372
  queryTest(opt, faiss::METRIC_L2, false);
366
373
  #endif
374
+
375
+ faiss::distance_compute_blas_threshold = saved_threshold;
367
376
  }
368
377
 
369
378
  // float16 coarse quantizer
@@ -909,6 +918,53 @@ TEST(TestGpuIndexIVFFlat, Reconstruct_n) {
909
918
  EXPECT_EQ(gpuVals, cpuVals);
910
919
  }
911
920
 
921
+ void testIDSelectorIVFFlat(faiss::MetricType metricType) {
922
+ Options opt;
923
+
924
+ std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
925
+ std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
926
+
927
+ faiss::gpu::StandardGpuResources res;
928
+ res.noTempMemory();
929
+
930
+ faiss::gpu::GpuIndexIVFFlatConfig config;
931
+ config.device = opt.device;
932
+ config.indicesOptions = faiss::gpu::INDICES_64_BIT;
933
+ config.use_cuvs = true;
934
+
935
+ faiss::gpu::GpuIndexIVFFlat gpuIndex(
936
+ &res, opt.dim, opt.numCentroids, metricType, config);
937
+ gpuIndex.nprobe = opt.nprobe;
938
+
939
+ gpuIndex.train(opt.numTrain, trainVecs.data());
940
+ gpuIndex.add(opt.numAdd, addVecs.data());
941
+
942
+ auto queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
943
+ faiss::gpu::TestIDSelectorStruct selector_struct(opt.numAdd);
944
+ faiss::SearchParametersIVF search_params;
945
+ search_params.nprobe = opt.nprobe;
946
+ for (auto& [selectorName, selector] : selector_struct.selector_map) {
947
+ search_params.sel = selector.get();
948
+ faiss::gpu::testIDSelectorSearch(
949
+ &gpuIndex,
950
+ &search_params,
951
+ queryVecs,
952
+ opt.numQuery,
953
+ opt.k,
954
+ selectorName);
955
+ }
956
+ }
957
+
958
+ #if defined USE_NVIDIA_CUVS
959
+ TEST(TestCuvsGpuIndexIVFFlat, IDSelector_L2) {
960
+ testIDSelectorIVFFlat(faiss::METRIC_L2);
961
+ }
962
+
963
+ TEST(TestCuvsGpuIndexIVFFlat, IDSelector_IP) {
964
+ testIDSelectorIVFFlat(faiss::METRIC_INNER_PRODUCT);
965
+ }
966
+ #endif
967
+
912
968
  int main(int argc, char** argv) {
913
969
  testing::InitGoogleTest(&argc, argv);
914
970
 
@@ -6,11 +6,14 @@
6
6
  */
7
7
 
8
8
  #include <faiss/IndexFlat.h>
9
+ #include <faiss/IndexIVF.h>
9
10
  #include <faiss/IndexIVFPQ.h>
10
11
  #include <faiss/gpu/GpuIndexIVFPQ.h>
11
12
  #include <faiss/gpu/StandardGpuResources.h>
12
13
  #include <faiss/gpu/test/TestUtils.h>
13
14
  #include <faiss/gpu/utils/DeviceUtils.h>
15
+ #include <faiss/impl/IDSelector.h>
16
+ #include <faiss/utils/distances.h>
14
17
  #include <gtest/gtest.h>
15
18
  #include <cmath>
16
19
  #include <sstream>
@@ -101,7 +104,11 @@ struct Options {
101
104
  }
102
105
 
103
106
  float getCompareEpsilon() const {
104
- return 0.035f;
107
+ // With very low dimensionality (e.g., dim=4, codes=2 giving
108
+ // dimPerSubQuantizer=2), L2 distances can be very small
109
+ // (near-zero), causing relative error comparisons to be
110
+ // unstable despite tiny absolute differences.
111
+ return (dim <= 8) ? 0.15f : 0.035f;
105
112
  }
106
113
 
107
114
  float getPctMaxDiff1() const {
@@ -190,6 +197,11 @@ TEST(TestGpuIndexIVFPQ, Query_IP) {
190
197
 
191
198
  // Large batch sizes (>= 65536) should also work
192
199
  TEST(TestGpuIndexIVFPQ, LargeBatch) {
200
+ // With low-dim vectors, CPU will use non-BLAS. Force the CPU to use
201
+ // the BLAS for consistent comparison.
202
+ int saved_threshold = faiss::distance_compute_blas_threshold;
203
+ faiss::distance_compute_blas_threshold = 1;
204
+
193
205
  for (bool usePrecomputed : {false, true}) {
194
206
  Options opt;
195
207
 
@@ -202,6 +214,8 @@ TEST(TestGpuIndexIVFPQ, LargeBatch) {
202
214
 
203
215
  queryTest(opt, faiss::MetricType::METRIC_L2);
204
216
  }
217
+
218
+ faiss::distance_compute_blas_threshold = saved_threshold;
205
219
  }
206
220
 
207
221
  void testMMCodeDistance(faiss::MetricType mt) {
@@ -697,6 +711,11 @@ TEST(TestGpuIndexIVFPQ, Query_IP_Cuvs) {
697
711
 
698
712
  // Large batch sizes (>= 65536) should also work
699
713
  TEST(TestGpuIndexIVFPQ, LargeBatch_Cuvs) {
714
+ // See LargeBatch comment: force CPU BLAS path to match GPU GEMM
715
+ // decomposition for consistent L2 distance computation.
716
+ int saved_threshold = faiss::distance_compute_blas_threshold;
717
+ faiss::distance_compute_blas_threshold = 1;
718
+
700
719
  Options opt;
701
720
 
702
721
  // override for large sizes
@@ -711,6 +730,8 @@ TEST(TestGpuIndexIVFPQ, LargeBatch_Cuvs) {
711
730
  opt.bitsPerCode = 8;
712
731
 
713
732
  queryTest(opt, faiss::MetricType::METRIC_L2);
733
+
734
+ faiss::distance_compute_blas_threshold = saved_threshold;
714
735
  }
715
736
 
716
737
  TEST(TestGpuIndexIVFPQ, CopyFrom_Cuvs) {
@@ -877,6 +898,62 @@ TEST(TestGpuIndexIVFPQ, UnifiedMemory) {
877
898
  #endif
878
899
  }
879
900
 
901
+ void testIDSelectorIVFPQ(faiss::MetricType metricType) {
902
+ Options opt;
903
+
904
+ std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
905
+ std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
906
+
907
+ faiss::IndexFlatL2 quantizerL2(opt.dim);
908
+ faiss::IndexFlatIP quantizerIP(opt.dim);
909
+ faiss::Index* quantizer = metricType == faiss::METRIC_L2
910
+ ? (faiss::Index*)&quantizerL2
911
+ : (faiss::Index*)&quantizerIP;
912
+
913
+ faiss::IndexIVFPQ cpuIndex(
914
+ quantizer, opt.dim, opt.numCentroids, opt.codes, opt.bitsPerCode);
915
+ cpuIndex.metric_type = metricType;
916
+ cpuIndex.train(opt.numTrain, trainVecs.data());
917
+ cpuIndex.add(opt.numAdd, addVecs.data());
918
+
919
+ faiss::gpu::StandardGpuResources res;
920
+ res.noTempMemory();
921
+
922
+ faiss::gpu::GpuIndexIVFPQConfig config;
923
+ config.device = opt.device;
924
+ config.indicesOptions = faiss::gpu::INDICES_64_BIT;
925
+ config.interleavedLayout = true;
926
+ config.use_cuvs = true;
927
+
928
+ faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, &cpuIndex, config);
929
+ gpuIndex.nprobe = opt.nprobe;
930
+
931
+ auto queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
932
+ faiss::gpu::TestIDSelectorStruct selector_struct(opt.numAdd);
933
+ faiss::SearchParametersIVF search_params;
934
+ search_params.nprobe = opt.nprobe;
935
+ for (auto& [selectorName, selector] : selector_struct.selector_map) {
936
+ search_params.sel = selector.get();
937
+ faiss::gpu::testIDSelectorSearch(
938
+ &gpuIndex,
939
+ &search_params,
940
+ queryVecs,
941
+ opt.numQuery,
942
+ opt.k,
943
+ selectorName);
944
+ }
945
+ }
946
+
947
+ #if defined USE_NVIDIA_CUVS
948
+ TEST(TestCuvsGpuIndexIVFPQ, IDSelector_L2) {
949
+ testIDSelectorIVFPQ(faiss::METRIC_L2);
950
+ }
951
+
952
+ TEST(TestCuvsGpuIndexIVFPQ, IDSelector_IP) {
953
+ testIDSelectorIVFPQ(faiss::METRIC_INNER_PRODUCT);
954
+ }
955
+ #endif
956
+
880
957
  int main(int argc, char** argv) {
881
958
  testing::InitGoogleTest(&argc, argv);
882
959
 
@@ -6,6 +6,7 @@
6
6
  */
7
7
 
8
8
  #include <faiss/gpu/test/TestUtils.h>
9
+ #include <faiss/impl/IDSelector.h>
9
10
  #include <faiss/utils/random.h>
10
11
  #include <gtest/gtest.h>
11
12
  #include <time.h>
@@ -429,5 +430,76 @@ void compareLists(
429
430
  }
430
431
  }
431
432
 
433
+ TestIDSelectorStruct::TestIDSelectorStruct(int numAdd) {
434
+ // Range selector [20%, 80%] of the database
435
+ size_t min_id = numAdd / 5;
436
+ size_t max_id = numAdd * 4 / 5;
437
+ selector_map["Range"] =
438
+ std::make_unique<faiss::IDSelectorRange>(min_id, max_id);
439
+
440
+ // Array selector (every 3rd element)
441
+ array_ids.clear();
442
+ for (int i = 0; i < numAdd; i += 3) {
443
+ array_ids.push_back(i);
444
+ }
445
+ selector_map["Array"] = std::make_unique<faiss::IDSelectorArray>(
446
+ array_ids.size(), array_ids.data());
447
+
448
+ // Batch selector (every 5th element)
449
+ batch_ids.clear();
450
+ for (int i = 1; i < numAdd; i += 5) {
451
+ batch_ids.push_back(i);
452
+ }
453
+ selector_map["Batch"] = std::make_unique<faiss::IDSelectorBatch>(
454
+ batch_ids.size(), batch_ids.data());
455
+
456
+ // Bitmap selector (every 4th element selected)
457
+ size_t bitmap_size = (numAdd + 7) / 8;
458
+ bitmap.resize(bitmap_size, 0);
459
+ for (int i = 0; i < numAdd; i += 4) {
460
+ int byte_idx = i / 8;
461
+ int bit_idx = i % 8;
462
+ bitmap[byte_idx] |= (1 << bit_idx);
463
+ }
464
+ selector_map["Bitmap"] =
465
+ std::make_unique<faiss::IDSelectorBitmap>(numAdd, bitmap.data());
466
+
467
+ selector_map["Not"] =
468
+ std::make_unique<faiss::IDSelectorNot>(selector_map["Range"].get());
469
+ selector_map["And"] = std::make_unique<faiss::IDSelectorAnd>(
470
+ selector_map["Range"].get(), selector_map["Array"].get());
471
+ selector_map["Or"] = std::make_unique<faiss::IDSelectorOr>(
472
+ selector_map["Range"].get(), selector_map["Batch"].get());
473
+ selector_map["XOr"] = std::make_unique<faiss::IDSelectorXOr>(
474
+ selector_map["Not"].get(), selector_map["Array"].get());
475
+ }
476
+
477
+ void testIDSelectorSearch(
478
+ faiss::Index* index,
479
+ faiss::SearchParameters* search_params,
480
+ const std::vector<float>& queryVecs,
481
+ int numQuery,
482
+ int k,
483
+ const std::string& selectorName) {
484
+ FAISS_ASSERT(search_params && search_params->sel);
485
+ std::vector<float> distances(numQuery * k, 0);
486
+ std::vector<faiss::idx_t> labels(numQuery * k, -1);
487
+ index->search(
488
+ numQuery,
489
+ queryVecs.data(),
490
+ k,
491
+ distances.data(),
492
+ labels.data(),
493
+ search_params);
494
+ faiss::IDSelector* selector = search_params->sel;
495
+ for (int i = 0; i < numQuery * k; ++i) {
496
+ if (labels[i] >= 0) {
497
+ EXPECT_TRUE(selector->is_member(labels[i]))
498
+ << "Label " << labels[i] << " @ " << i << " not in "
499
+ << selectorName << " selector";
500
+ }
501
+ }
502
+ }
503
+
432
504
  } // namespace gpu
433
505
  } // namespace faiss
@@ -13,6 +13,7 @@
13
13
  #include <gtest/gtest.h>
14
14
  #include <cstring>
15
15
  #include <initializer_list>
16
+ #include <map>
16
17
  #include <memory>
17
18
  #include <string>
18
19
  #include <vector>
@@ -141,5 +142,27 @@ void testIVFEquality(A& cpuIndex, B& gpuIndex) {
141
142
  }
142
143
  }
143
144
 
145
+ /// Run search with the given search_params and verify all returned labels are
146
+ /// members of the selector. Works with any Index that supports search with
147
+ /// SearchParameters.
148
+ void testIDSelectorSearch(
149
+ faiss::Index* index,
150
+ faiss::SearchParameters* search_params,
151
+ const std::vector<float>& queryVecs,
152
+ int numQuery,
153
+ int k,
154
+ const std::string& selectorName);
155
+
156
+ // Structure to hold all IDSelector instances
157
+ struct TestIDSelectorStruct {
158
+ // Storage for selectors that need it
159
+ std::vector<faiss::idx_t> array_ids;
160
+ std::vector<faiss::idx_t> batch_ids;
161
+ std::vector<uint8_t> bitmap;
162
+ std::map<std::string, std::unique_ptr<faiss::IDSelector>> selector_map;
163
+
164
+ explicit TestIDSelectorStruct(int numAdd);
165
+ };
166
+
144
167
  } // namespace gpu
145
168
  } // namespace faiss
@@ -36,6 +36,6 @@ namespace faiss::gpu {
36
36
  void convert_to_bitset(
37
37
  faiss::gpu::GpuResources* res,
38
38
  const faiss::IDSelector& selector,
39
- cuvs::core::bitset_view<uint32_t, uint32_t> bitset,
39
+ cuvs::core::bitset_view<uint32_t, int64_t> bitset,
40
40
  int num_threads = 0);
41
41
  } // namespace faiss::gpu
@@ -27,32 +27,32 @@
27
27
  #include <faiss/gpu/GpuResources.h>
28
28
  #include <faiss/gpu/utils/Tensor.cuh>
29
29
 
30
- #include <cuvs/distance/distance.h>
30
+ #include <cuvs/distance/distance.hpp>
31
31
 
32
32
  #pragma GCC visibility push(default)
33
33
  namespace faiss {
34
34
  namespace gpu {
35
35
 
36
- inline cuvsDistanceType metricFaissToCuvs(
36
+ inline cuvs::distance::DistanceType metricFaissToCuvs(
37
37
  MetricType metric,
38
38
  bool exactDistance) {
39
39
  switch (metric) {
40
40
  case MetricType::METRIC_INNER_PRODUCT:
41
- return cuvsDistanceType::InnerProduct;
41
+ return cuvs::distance::DistanceType::InnerProduct;
42
42
  case MetricType::METRIC_L2:
43
- return cuvsDistanceType::L2Expanded;
43
+ return cuvs::distance::DistanceType::L2Expanded;
44
44
  case MetricType::METRIC_L1:
45
- return cuvsDistanceType::L1;
45
+ return cuvs::distance::DistanceType::L1;
46
46
  case MetricType::METRIC_Linf:
47
- return cuvsDistanceType::Linf;
47
+ return cuvs::distance::DistanceType::Linf;
48
48
  case MetricType::METRIC_Lp:
49
- return cuvsDistanceType::LpUnexpanded;
49
+ return cuvs::distance::DistanceType::LpUnexpanded;
50
50
  case MetricType::METRIC_Canberra:
51
- return cuvsDistanceType::Canberra;
51
+ return cuvs::distance::DistanceType::Canberra;
52
52
  case MetricType::METRIC_BrayCurtis:
53
- return cuvsDistanceType::BrayCurtis;
53
+ return cuvs::distance::DistanceType::BrayCurtis;
54
54
  case MetricType::METRIC_JensenShannon:
55
- return cuvsDistanceType::JensenShannon;
55
+ return cuvs::distance::DistanceType::JensenShannon;
56
56
  default:
57
57
  RAFT_FAIL("Distance type not supported");
58
58
  }
@@ -71,6 +71,17 @@ idx_t inplaceGatherFilteredRows(
71
71
  GpuResources* res,
72
72
  Tensor<float, 2, true>& vecs,
73
73
  Tensor<idx_t, 1, true>& indices);
74
+
75
+ /// Copy uint32_t indices to idx_t, replacing any index >= n with -1.
76
+ /// cuVS CAGRA returns sentinel values (e.g. INT32_MAX) for result slots
77
+ /// where filtered search couldn't find a valid neighbor.
78
+ void sanitizeCuvsIndices(
79
+ GpuResources* res,
80
+ uint32_t* src,
81
+ idx_t* dst,
82
+ size_t count,
83
+ idx_t n);
84
+
74
85
  } // namespace gpu
75
86
  } // namespace faiss
76
87
  #pragma GCC visibility pop
@@ -0,0 +1,22 @@
1
+ // @lint-ignore-every LICENSELINT
2
+ /**
3
+ * Copyright (c) Meta Platforms, Inc. and its affiliates.
4
+ *
5
+ * This source code is licensed under the MIT license found in the
6
+ * LICENSE file in the root directory of this source tree.
7
+ *
8
+ * Unified name for flat GPU index when Metal backend is built.
9
+ * Include this when using the Metal backend for API parity with
10
+ * faiss::gpu::GpuIndexFlat.
11
+ */
12
+
13
+ #pragma once
14
+
15
+ #include <faiss/gpu_metal/MetalIndexFlat.h>
16
+
17
+ namespace faiss {
18
+
19
+ /// When FAISS is built with Metal backend, GpuIndexFlat is MetalIndexFlat.
20
+ using GpuIndexFlat = gpu_metal::MetalIndexFlat;
21
+
22
+ } // namespace faiss
@@ -0,0 +1,35 @@
1
+ // @lint-ignore-every LICENSELINT
2
+ /**
3
+ * Copyright (c) Meta Platforms, Inc. and its affiliates.
4
+ *
5
+ * This source code is licensed under the MIT license found in the
6
+ * LICENSE file in the root directory of this source tree.
7
+ *
8
+ * Clone CPU <-> Metal GPU. Mirrors GpuCloner roles for Metal backend.
9
+ */
10
+
11
+ #pragma once
12
+
13
+ #include <faiss/Index.h>
14
+
15
+ namespace faiss {
16
+ namespace gpu_metal {
17
+
18
+ class StandardMetalResources;
19
+
20
+ /// Returns the number of Metal "devices" (1 if Metal is available, else 0).
21
+ int get_num_gpus();
22
+
23
+ /// Clone a CPU index to Metal GPU. Supports IndexFlat, IndexFlatL2,
24
+ /// IndexFlatIP. device must be 0. Caller owns the returned index.
25
+ faiss::Index* index_cpu_to_metal_gpu(
26
+ StandardMetalResources* res,
27
+ int device,
28
+ const faiss::Index* index);
29
+
30
+ /// Copy a Metal index back to CPU. Supports MetalIndexFlat -> IndexFlat.
31
+ /// Caller owns the returned index.
32
+ faiss::Index* index_metal_gpu_to_cpu(const faiss::Index* index);
33
+
34
+ } // namespace gpu_metal
35
+ } // namespace faiss