faiss 0.6.2 → 0.6.3

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (144) hide show
  1. checksums.yaml +4 -4
  2. data/CHANGELOG.md +4 -0
  3. data/lib/faiss/version.rb +1 -1
  4. data/vendor/faiss/faiss/AutoTune.cpp +3 -1
  5. data/vendor/faiss/faiss/Clustering.cpp +9 -1
  6. data/vendor/faiss/faiss/IVFlib.cpp +14 -3
  7. data/vendor/faiss/faiss/Index.h +2 -2
  8. data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp +9 -10
  9. data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.cpp +2 -3
  10. data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +1 -2
  11. data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +4 -5
  12. data/vendor/faiss/faiss/IndexBinaryHash.cpp +5 -9
  13. data/vendor/faiss/faiss/IndexBinaryIVF.cpp +2 -4
  14. data/vendor/faiss/faiss/IndexEDEN.cpp +273 -0
  15. data/vendor/faiss/faiss/IndexEDEN.h +57 -0
  16. data/vendor/faiss/faiss/IndexFastScan.cpp +15 -4
  17. data/vendor/faiss/faiss/IndexFlat.cpp +13 -50
  18. data/vendor/faiss/faiss/IndexHNSW.cpp +10 -11
  19. data/vendor/faiss/faiss/IndexIDMap.cpp +16 -3
  20. data/vendor/faiss/faiss/IndexIDMap.h +2 -0
  21. data/vendor/faiss/faiss/IndexIVF.cpp +17 -6
  22. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp +1 -1
  23. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.cpp +3 -4
  24. data/vendor/faiss/faiss/IndexIVFEDEN.cpp +302 -0
  25. data/vendor/faiss/faiss/IndexIVFEDEN.h +70 -0
  26. data/vendor/faiss/faiss/IndexIVFFastScan.cpp +5 -6
  27. data/vendor/faiss/faiss/IndexIVFFlat.cpp +3 -4
  28. data/vendor/faiss/faiss/IndexIVFIndependentQuantizer.cpp +1 -1
  29. data/vendor/faiss/faiss/IndexIVFPQ.cpp +40 -22
  30. data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +0 -1
  31. data/vendor/faiss/faiss/IndexIVFRaBitQ.cpp +19 -49
  32. data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.cpp +180 -76
  33. data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.h +5 -4
  34. data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +8 -6
  35. data/vendor/faiss/faiss/IndexLSH.cpp +2 -3
  36. data/vendor/faiss/faiss/IndexLattice.cpp +5 -0
  37. data/vendor/faiss/faiss/IndexNNDescent.cpp +9 -2
  38. data/vendor/faiss/faiss/IndexNSG.cpp +7 -2
  39. data/vendor/faiss/faiss/IndexPQ.cpp +6 -8
  40. data/vendor/faiss/faiss/IndexPreTransform.cpp +15 -0
  41. data/vendor/faiss/faiss/IndexRaBitQ.cpp +2 -2
  42. data/vendor/faiss/faiss/IndexRaBitQFastScan.cpp +1 -2
  43. data/vendor/faiss/faiss/IndexRaBitQFastScan.h +5 -1
  44. data/vendor/faiss/faiss/IndexRefine.cpp +30 -1
  45. data/vendor/faiss/faiss/IndexReplicas.cpp +1 -2
  46. data/vendor/faiss/faiss/IndexShards.cpp +2 -2
  47. data/vendor/faiss/faiss/IndexShardsIVF.cpp +2 -2
  48. data/vendor/faiss/faiss/MetaIndexes.cpp +2 -4
  49. data/vendor/faiss/faiss/SuperKMeans.cpp +256 -240
  50. data/vendor/faiss/faiss/SuperKMeans.h +30 -0
  51. data/vendor/faiss/faiss/VectorTransform.cpp +33 -2
  52. data/vendor/faiss/faiss/clone_index.cpp +5 -0
  53. data/vendor/faiss/faiss/factory_tools.cpp +47 -4
  54. data/vendor/faiss/faiss/gpu/GpuCloner.cpp +11 -11
  55. data/vendor/faiss/faiss/gpu/GpuIndex.h +34 -11
  56. data/vendor/faiss/faiss/gpu/GpuIndexCagra.h +47 -0
  57. data/vendor/faiss/faiss/gpu/GpuIndexIVF.h +17 -0
  58. data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h +16 -0
  59. data/vendor/faiss/faiss/gpu/perf/PerfClustering.cpp +1 -1
  60. data/vendor/faiss/faiss/gpu/perf/PerfIVFPQAdd.cpp +2 -2
  61. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFScalarQuantizer.cpp +180 -0
  62. data/vendor/faiss/faiss/gpu_metal/MetalIndexIVFFlat.h +1 -5
  63. data/vendor/faiss/faiss/gpu_metal/MetalIndexIVFPQ.h +88 -0
  64. data/vendor/faiss/faiss/gpu_metal/impl/MetalIVFPQ.h +134 -0
  65. data/vendor/faiss/faiss/impl/ClusteringInitialization.cpp +2 -2
  66. data/vendor/faiss/faiss/impl/DistanceComputer.h +34 -0
  67. data/vendor/faiss/faiss/impl/EDENQuantizer.h +119 -0
  68. data/vendor/faiss/faiss/impl/HNSW.cpp +109 -152
  69. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +2 -2
  70. data/vendor/faiss/faiss/impl/NSG.cpp +3 -1
  71. data/vendor/faiss/faiss/impl/Panorama.h +9 -7
  72. data/vendor/faiss/faiss/impl/PolysemousTraining.cpp +152 -84
  73. data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +34 -22
  74. data/vendor/faiss/faiss/impl/RaBitQUtils.cpp +44 -36
  75. data/vendor/faiss/faiss/impl/RaBitQUtils.h +35 -0
  76. data/vendor/faiss/faiss/impl/RaBitQuantizer.cpp +168 -67
  77. data/vendor/faiss/faiss/impl/RaBitQuantizer.h +19 -0
  78. data/vendor/faiss/faiss/impl/RaBitQuantizerMultiBit.cpp +2 -11
  79. data/vendor/faiss/faiss/impl/ResultHandler.h +25 -31
  80. data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +258 -57
  81. data/vendor/faiss/faiss/impl/ScalarQuantizer.h +20 -0
  82. data/vendor/faiss/faiss/impl/ThreadedIndex-inl.h +2 -2
  83. data/vendor/faiss/faiss/impl/VisitedTable.cpp +22 -2
  84. data/vendor/faiss/faiss/impl/VisitedTable.h +20 -0
  85. data/vendor/faiss/faiss/impl/binary_hamming/IndexBinaryIVF_impl.h +1 -1
  86. data/vendor/faiss/faiss/impl/binary_hamming/avx2.cpp +4 -4
  87. data/vendor/faiss/faiss/impl/fast_scan/dispatching.h +35 -2
  88. data/vendor/faiss/faiss/impl/hnsw/LockVector.cpp +1 -1
  89. data/vendor/faiss/faiss/impl/index_read.cpp +376 -36
  90. data/vendor/faiss/faiss/impl/index_write.cpp +55 -4
  91. data/vendor/faiss/faiss/impl/lattice_Zn.cpp +8 -9
  92. data/vendor/faiss/faiss/impl/platform_macros.h +3 -1
  93. data/vendor/faiss/faiss/impl/polysemous_training/avx512.cpp +284 -0
  94. data/vendor/faiss/faiss/impl/polysemous_training/dispatch.h +115 -0
  95. data/vendor/faiss/faiss/impl/pq_code_distance/IVFPQ_QueryTables.cpp +0 -1
  96. data/vendor/faiss/faiss/impl/pq_code_distance/PQDistanceComputer_impl.h +26 -15
  97. data/vendor/faiss/faiss/impl/pq_code_distance/avx2.cpp +4 -4
  98. data/vendor/faiss/faiss/impl/result_handler/ResultHandler.cpp +195 -0
  99. data/vendor/faiss/faiss/impl/result_handler/avx2.cpp +133 -0
  100. data/vendor/faiss/faiss/impl/result_handler/avx512.cpp +281 -0
  101. data/vendor/faiss/faiss/impl/scalar_quantizer/EDENQuantizer-avx2.cpp +72 -0
  102. data/vendor/faiss/faiss/impl/scalar_quantizer/EDENQuantizer-avx512.cpp +228 -0
  103. data/vendor/faiss/faiss/impl/scalar_quantizer/EDENQuantizer.cpp +882 -0
  104. data/vendor/faiss/faiss/impl/scalar_quantizer/quantizers.h +9 -8
  105. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx2.cpp +85 -23
  106. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx512-impl.h +30 -30
  107. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx512-spr.cpp +4 -5
  108. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx512.cpp +101 -34
  109. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-dispatch.h +136 -0
  110. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-neon.cpp +16 -16
  111. data/vendor/faiss/faiss/impl/simd_dispatch.h +30 -9
  112. data/vendor/faiss/faiss/index_factory.cpp +32 -6
  113. data/vendor/faiss/faiss/invlists/DirectMap.cpp +1 -1
  114. data/vendor/faiss/faiss/invlists/InvertedLists.cpp +2 -2
  115. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +19 -4
  116. data/vendor/faiss/faiss/python/python_callbacks.cpp +3 -1
  117. data/vendor/faiss/faiss/svs/IndexSVSFaissUtils.h +60 -0
  118. data/vendor/faiss/faiss/svs/IndexSVSFlat.cpp +26 -1
  119. data/vendor/faiss/faiss/svs/IndexSVSFlat.h +13 -0
  120. data/vendor/faiss/faiss/svs/IndexSVSIVF.cpp +1 -1
  121. data/vendor/faiss/faiss/svs/IndexSVSIVFLeanVec.cpp +1 -1
  122. data/vendor/faiss/faiss/svs/IndexSVSVamana.cpp +31 -1
  123. data/vendor/faiss/faiss/svs/IndexSVSVamana.h +15 -2
  124. data/vendor/faiss/faiss/svs/IndexSVSVamanaLeanVec.cpp +1 -2
  125. data/vendor/faiss/faiss/utils/approx_topk_hamming/approx_topk_hamming.h +1 -1
  126. data/vendor/faiss/faiss/utils/distances.cpp +14 -2
  127. data/vendor/faiss/faiss/utils/distances_simd.cpp +4 -3
  128. data/vendor/faiss/faiss/utils/extra_distances.cpp +4 -14
  129. data/vendor/faiss/faiss/utils/extra_distances.h +1 -2
  130. data/vendor/faiss/faiss/utils/hamming.cpp +1 -1
  131. data/vendor/faiss/faiss/utils/quantize_lut.cpp +29 -8
  132. data/vendor/faiss/faiss/utils/rabitq_simd.h +202 -0
  133. data/vendor/faiss/faiss/utils/simd_impl/distances_avx2.cpp +0 -1
  134. data/vendor/faiss/faiss/utils/simd_impl/distances_avx512.cpp +263 -15
  135. data/vendor/faiss/faiss/utils/simd_impl/distances_rvv.cpp +160 -18
  136. data/vendor/faiss/faiss/utils/simd_impl/rabitq_avx2.cpp +245 -0
  137. data/vendor/faiss/faiss/utils/simd_impl/rabitq_avx512.cpp +273 -0
  138. data/vendor/faiss/faiss/utils/simd_impl/rabitq_avx512_spr.cpp +92 -0
  139. data/vendor/faiss/faiss/utils/simd_impl/rabitq_neon.cpp +11 -0
  140. data/vendor/faiss/faiss/utils/simd_impl/rabitq_rvv.cpp +143 -6
  141. data/vendor/faiss/faiss/utils/simd_levels.cpp +44 -0
  142. data/vendor/faiss/faiss/utils/simd_levels.h +14 -0
  143. data/vendor/faiss/faiss/utils/utils.cpp +9 -27
  144. metadata +16 -1
@@ -9,6 +9,35 @@
9
9
 
10
10
  #include <immintrin.h>
11
11
 
12
+ #include <memory>
13
+
14
+ #include <faiss/impl/AuxIndexStructures.h>
15
+ #include <faiss/impl/ResultHandler.h>
16
+ #include <faiss/utils/distances_fused/distances_fused.h>
17
+ #include <faiss/utils/simd_impl/exhaustive_L2sqr_blas_cmax.h>
18
+
19
+ #ifndef FINTEGER
20
+ #define FINTEGER long
21
+ #endif
22
+
23
+ extern "C" {
24
+
25
+ int sgemm_(
26
+ const char* transa,
27
+ const char* transb,
28
+ FINTEGER* m,
29
+ FINTEGER* n,
30
+ FINTEGER* k,
31
+ const float* alpha,
32
+ const float* a,
33
+ FINTEGER* lda,
34
+ const float* b,
35
+ FINTEGER* ldb,
36
+ float* beta,
37
+ float* c,
38
+ FINTEGER* ldc);
39
+ }
40
+
12
41
  #define THE_SIMD_LEVEL SIMDLevel::AVX512
13
42
  #include <faiss/utils/simd_impl/distances_autovec-inl.h>
14
43
  // NOLINTNEXTLINE(facebook-hte-InlineHeader)
@@ -954,21 +983,6 @@ size_t fvec_L2sqr_ny_nearest<SIMDLevel::AVX512>(
954
983
  &fvec_L2sqr_ny_nearest_D8<SIMDLevel::AVX512>);
955
984
  }
956
985
 
957
- template <>
958
- size_t fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::AVX512>(
959
- float* distances_tmp_buffer,
960
- const float* x,
961
- const float* y,
962
- const float* y_sqlen,
963
- size_t d,
964
- size_t d_offset,
965
- size_t ny) {
966
- return fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::NONE>(
967
- distances_tmp_buffer, x, y, y_sqlen, d, d_offset, ny);
968
- }
969
-
970
- // TODO: Following functions are not used in the current codebase. Check AVX2 ,
971
- // respective implementation has been used
972
986
  template <size_t DIM>
973
987
  size_t fvec_L2sqr_ny_nearest_y_transposed_D(
974
988
  float* /* distances_tmp_buffer */,
@@ -1082,6 +1096,33 @@ size_t fvec_L2sqr_ny_nearest_y_transposed_D(
1082
1096
  return current_min_index;
1083
1097
  }
1084
1098
 
1099
+ template <>
1100
+ size_t fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::AVX512>(
1101
+ float* distances_tmp_buffer,
1102
+ const float* x,
1103
+ const float* y,
1104
+ const float* y_sqlen,
1105
+ size_t d,
1106
+ size_t d_offset,
1107
+ size_t ny) {
1108
+ // optimized for a few special cases
1109
+ #define DISPATCH(dval) \
1110
+ case dval: \
1111
+ return fvec_L2sqr_ny_nearest_y_transposed_D<dval>( \
1112
+ distances_tmp_buffer, x, y, y_sqlen, d_offset, ny);
1113
+
1114
+ switch (d) {
1115
+ DISPATCH(1)
1116
+ DISPATCH(2)
1117
+ DISPATCH(4)
1118
+ DISPATCH(8)
1119
+ default:
1120
+ return fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::NONE>(
1121
+ distances_tmp_buffer, x, y, y_sqlen, d, d_offset, ny);
1122
+ }
1123
+ #undef DISPATCH
1124
+ }
1125
+
1085
1126
  template <>
1086
1127
  int fvec_madd_and_argmin<SIMDLevel::AVX512>(
1087
1128
  size_t n,
@@ -1092,4 +1133,211 @@ int fvec_madd_and_argmin<SIMDLevel::AVX512>(
1092
1133
  return fvec_madd_and_argmin_sse(n, a, bf, b, c);
1093
1134
  }
1094
1135
 
1136
+ template <>
1137
+ void exhaustive_L2sqr_blas_cmax<SIMDLevel::AVX512>(
1138
+ const float* x,
1139
+ const float* y,
1140
+ size_t d,
1141
+ size_t nx,
1142
+ size_t ny,
1143
+ Top1BlockResultHandler<CMax<float, int64_t>>& res,
1144
+ const float* y_norms) {
1145
+ // BLAS does not like empty matrices
1146
+ if (nx == 0 || ny == 0) {
1147
+ return;
1148
+ }
1149
+
1150
+ /* block sizes */
1151
+ const size_t bs_x = distance_compute_blas_query_bs;
1152
+ const size_t bs_y = distance_compute_blas_database_bs;
1153
+ std::unique_ptr<float[]> ip_block(new float[bs_x * bs_y]);
1154
+ std::unique_ptr<float[]> x_norms(new float[nx]);
1155
+ std::unique_ptr<float[]> del2;
1156
+
1157
+ fvec_norms_L2sqr(x_norms.get(), x, d, nx);
1158
+
1159
+ if (!y_norms) {
1160
+ float* y_norms2 = new float[ny];
1161
+ del2.reset(y_norms2);
1162
+ fvec_norms_L2sqr(y_norms2, y, d, ny);
1163
+ y_norms = y_norms2;
1164
+ }
1165
+
1166
+ for (size_t i0 = 0; i0 < nx; i0 += bs_x) {
1167
+ size_t i1 = i0 + bs_x;
1168
+ if (i1 > nx) {
1169
+ i1 = nx;
1170
+ }
1171
+
1172
+ res.begin_multiple(i0, i1);
1173
+
1174
+ for (size_t j0 = 0; j0 < ny; j0 += bs_y) {
1175
+ size_t j1 = j0 + bs_y;
1176
+ if (j1 > ny) {
1177
+ j1 = ny;
1178
+ }
1179
+ /* compute the actual dot products */
1180
+ {
1181
+ float one = 1, zero = 0;
1182
+ FINTEGER nyi = j1 - j0, nxi = i1 - i0, di = d;
1183
+ sgemm_("Transpose",
1184
+ "Not transpose",
1185
+ &nyi,
1186
+ &nxi,
1187
+ &di,
1188
+ &one,
1189
+ y + j0 * d,
1190
+ &di,
1191
+ x + i0 * d,
1192
+ &di,
1193
+ &zero,
1194
+ ip_block.get(),
1195
+ &nyi);
1196
+ }
1197
+ #pragma omp parallel for schedule(static) if ((i1 - i0) >= 16)
1198
+ for (int64_t i = static_cast<int64_t>(i0);
1199
+ i < static_cast<int64_t>(i1);
1200
+ i++) {
1201
+ float* ip_line = ip_block.get() + (i - i0) * (j1 - j0);
1202
+
1203
+ _mm_prefetch((const char*)ip_line, _MM_HINT_NTA);
1204
+ _mm_prefetch((const char*)(ip_line + 16), _MM_HINT_NTA);
1205
+
1206
+ // constant
1207
+ const __m512 mul_minus2 = _mm512_set1_ps(-2);
1208
+
1209
+ // Track 16 min distances + 16 min indices.
1210
+ // All the distances tracked do not take x_norms[i]
1211
+ // into account in order to get rid of extra
1212
+ // _mm512_add_ps(x_norms[i], ...) instructions
1213
+ // in distance computations.
1214
+ __m512 min_distances =
1215
+ _mm512_set1_ps(res.dis_tab[i] - x_norms[i]);
1216
+
1217
+ // these indices are local and are relative to j0.
1218
+ // so, value 0 means j0.
1219
+ __m512i min_indices = _mm512_set1_epi32(0);
1220
+
1221
+ __m512i current_indices = _mm512_setr_epi32(
1222
+ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);
1223
+ const __m512i indices_delta = _mm512_set1_epi32(16);
1224
+
1225
+ // current j index
1226
+ size_t idx_j = 0;
1227
+ size_t count = j1 - j0;
1228
+
1229
+ // process 32 elements per loop
1230
+ for (; idx_j < (count / 32) * 32; idx_j += 32, ip_line += 32) {
1231
+ _mm_prefetch((const char*)(ip_line + 32), _MM_HINT_NTA);
1232
+ _mm_prefetch((const char*)(ip_line + 48), _MM_HINT_NTA);
1233
+
1234
+ // load values for norms
1235
+ const __m512 y_norm_0 =
1236
+ _mm512_loadu_ps(y_norms + idx_j + j0 + 0);
1237
+ const __m512 y_norm_1 =
1238
+ _mm512_loadu_ps(y_norms + idx_j + j0 + 16);
1239
+
1240
+ // load values for dot products
1241
+ const __m512 ip_0 = _mm512_loadu_ps(ip_line + 0);
1242
+ const __m512 ip_1 = _mm512_loadu_ps(ip_line + 16);
1243
+
1244
+ // compute dis = y_norm[j] - 2 * dot(x_norm[i], y_norm[j]).
1245
+ // x_norm[i] was dropped off because it is a constant for a
1246
+ // given i.
1247
+ __m512 distances_0 =
1248
+ _mm512_fmadd_ps(ip_0, mul_minus2, y_norm_0);
1249
+ __m512 distances_1 =
1250
+ _mm512_fmadd_ps(ip_1, mul_minus2, y_norm_1);
1251
+
1252
+ // compare the new distances to the min distances
1253
+ // for each of the first group of 16 AVX512 components.
1254
+ const __mmask16 comparison_0 = _mm512_cmp_ps_mask(
1255
+ min_distances, distances_0, _CMP_LE_OS);
1256
+
1257
+ // update min distances and indices with closest vectors if
1258
+ // needed.
1259
+ min_distances = _mm512_mask_blend_ps(
1260
+ comparison_0, distances_0, min_distances);
1261
+ min_indices = _mm512_mask_blend_epi32(
1262
+ comparison_0, current_indices, min_indices);
1263
+ current_indices =
1264
+ _mm512_add_epi32(current_indices, indices_delta);
1265
+
1266
+ // compare the new distances to the min distances
1267
+ // for each of the second group of 16 AVX512 components.
1268
+ const __mmask16 comparison_1 = _mm512_cmp_ps_mask(
1269
+ min_distances, distances_1, _CMP_LE_OS);
1270
+
1271
+ // update min distances and indices with closest vectors if
1272
+ // needed.
1273
+ min_distances = _mm512_mask_blend_ps(
1274
+ comparison_1, distances_1, min_distances);
1275
+ min_indices = _mm512_mask_blend_epi32(
1276
+ comparison_1, current_indices, min_indices);
1277
+ current_indices =
1278
+ _mm512_add_epi32(current_indices, indices_delta);
1279
+ }
1280
+
1281
+ // dump values and find the minimum distance / minimum index
1282
+ float min_distances_scalar[16];
1283
+ uint32_t min_indices_scalar[16];
1284
+ _mm512_storeu_ps(min_distances_scalar, min_distances);
1285
+ _mm512_storeu_si512(
1286
+ (__m512i*)(min_indices_scalar), min_indices);
1287
+
1288
+ float current_min_distance = res.dis_tab[i];
1289
+ uint32_t current_min_index = res.ids_tab[i];
1290
+
1291
+ // if two indices are represented with equal distance values,
1292
+ // then the index with the min value is returned.
1293
+ for (size_t jv = 0; jv < 16; jv++) {
1294
+ // add missing x_norms[i]
1295
+ float distance_candidate =
1296
+ min_distances_scalar[jv] + x_norms[i];
1297
+
1298
+ // negative values can occur for identical vectors
1299
+ // due to roundoff errors.
1300
+ if (distance_candidate < 0) {
1301
+ distance_candidate = 0;
1302
+ }
1303
+
1304
+ int64_t index_candidate = min_indices_scalar[jv] + j0;
1305
+
1306
+ if (current_min_distance > distance_candidate) {
1307
+ current_min_distance = distance_candidate;
1308
+ current_min_index = index_candidate;
1309
+ } else if (
1310
+ current_min_distance == distance_candidate &&
1311
+ current_min_index > index_candidate) {
1312
+ current_min_index = index_candidate;
1313
+ }
1314
+ }
1315
+
1316
+ // process leftovers
1317
+ for (; idx_j < count; idx_j++, ip_line++) {
1318
+ float ip = *ip_line;
1319
+ float dis = x_norms[i] + y_norms[idx_j + j0] - 2 * ip;
1320
+ // negative values can occur for identical vectors
1321
+ // due to roundoff errors.
1322
+ if (dis < 0) {
1323
+ dis = 0;
1324
+ }
1325
+
1326
+ if (current_min_distance > dis) {
1327
+ current_min_distance = dis;
1328
+ current_min_index = idx_j + j0;
1329
+ }
1330
+ }
1331
+
1332
+ //
1333
+ res.add_result(i, current_min_distance, current_min_index);
1334
+ }
1335
+ }
1336
+ // Does nothing for SingleBestResultHandler, but
1337
+ // keeping the call for the consistency.
1338
+ res.end_multiple();
1339
+ InterruptCallback::check();
1340
+ }
1341
+ }
1342
+
1095
1343
  } // namespace faiss
@@ -12,12 +12,24 @@
12
12
  #ifdef COMPILE_SIMD_RISCV_RVV
13
13
 
14
14
  #include <faiss/utils/extra_distances.h>
15
+ #include <riscv_vector.h>
15
16
 
16
17
  namespace faiss {
17
18
 
18
19
  template <>
19
20
  float fvec_norm_L2sqr<SIMDLevel::RISCV_RVV>(const float* x, size_t d) {
20
- return fvec_norm_L2sqr<SIMDLevel::NONE>(x, d);
21
+ size_t vlmax = __riscv_vsetvlmax_e32m8();
22
+ vfloat32m8_t acc = __riscv_vfmv_v_f_f32m8(0.0f, vlmax);
23
+ size_t i = 0;
24
+ while (i < d) {
25
+ size_t vl = __riscv_vsetvl_e32m8(d - i);
26
+ vfloat32m8_t vx = __riscv_vle32_v_f32m8(x + i, vl);
27
+ acc = __riscv_vfmacc_vv_f32m8_tu(acc, vx, vx, vl);
28
+ i += vl;
29
+ }
30
+ vfloat32m1_t sum = __riscv_vfmv_s_f_f32m1(0.0f, 1);
31
+ sum = __riscv_vfredusum_vs_f32m8_f32m1(acc, sum, vlmax);
32
+ return __riscv_vfmv_f_s_f32m1_f32(sum);
21
33
  }
22
34
 
23
35
  template <>
@@ -25,7 +37,20 @@ float fvec_L2sqr<SIMDLevel::RISCV_RVV>(
25
37
  const float* x,
26
38
  const float* y,
27
39
  size_t d) {
28
- return fvec_L2sqr<SIMDLevel::NONE>(x, y, d);
40
+ size_t vlmax = __riscv_vsetvlmax_e32m8();
41
+ vfloat32m8_t acc = __riscv_vfmv_v_f_f32m8(0.0f, vlmax);
42
+ size_t i = 0;
43
+ while (i < d) {
44
+ size_t vl = __riscv_vsetvl_e32m8(d - i);
45
+ vfloat32m8_t vx = __riscv_vle32_v_f32m8(x + i, vl);
46
+ vfloat32m8_t vy = __riscv_vle32_v_f32m8(y + i, vl);
47
+ vx = __riscv_vfsub_vv_f32m8(vx, vy, vl);
48
+ acc = __riscv_vfmacc_vv_f32m8_tu(acc, vx, vx, vl);
49
+ i += vl;
50
+ }
51
+ vfloat32m1_t sum = __riscv_vfmv_s_f_f32m1(0.0f, 1);
52
+ sum = __riscv_vfredusum_vs_f32m8_f32m1(acc, sum, vlmax);
53
+ return __riscv_vfmv_f_s_f32m1_f32(sum);
29
54
  }
30
55
 
31
56
  template <>
@@ -33,12 +58,38 @@ float fvec_inner_product<SIMDLevel::RISCV_RVV>(
33
58
  const float* x,
34
59
  const float* y,
35
60
  size_t d) {
36
- return fvec_inner_product<SIMDLevel::NONE>(x, y, d);
61
+ size_t vlmax = __riscv_vsetvlmax_e32m8();
62
+ vfloat32m8_t acc = __riscv_vfmv_v_f_f32m8(0.0f, vlmax);
63
+ size_t i = 0;
64
+ while (i < d) {
65
+ size_t vl = __riscv_vsetvl_e32m8(d - i);
66
+ vfloat32m8_t vx = __riscv_vle32_v_f32m8(x + i, vl);
67
+ vfloat32m8_t vy = __riscv_vle32_v_f32m8(y + i, vl);
68
+ acc = __riscv_vfmacc_vv_f32m8_tu(acc, vx, vy, vl);
69
+ i += vl;
70
+ }
71
+ vfloat32m1_t sum = __riscv_vfmv_s_f_f32m1(0.0f, 1);
72
+ sum = __riscv_vfredusum_vs_f32m8_f32m1(acc, sum, vlmax);
73
+ return __riscv_vfmv_f_s_f32m1_f32(sum);
37
74
  }
38
75
 
39
76
  template <>
40
77
  float fvec_L1<SIMDLevel::RISCV_RVV>(const float* x, const float* y, size_t d) {
41
- return fvec_L1<SIMDLevel::NONE>(x, y, d);
78
+ size_t vlmax = __riscv_vsetvlmax_e32m8();
79
+ vfloat32m8_t acc = __riscv_vfmv_v_f_f32m8(0.0f, vlmax);
80
+ size_t i = 0;
81
+ while (i < d) {
82
+ size_t vl = __riscv_vsetvl_e32m8(d - i);
83
+ vfloat32m8_t vx = __riscv_vle32_v_f32m8(x + i, vl);
84
+ vfloat32m8_t vy = __riscv_vle32_v_f32m8(y + i, vl);
85
+ vx = __riscv_vfsub_vv_f32m8(vx, vy, vl);
86
+ vx = __riscv_vfsgnjx_vv_f32m8(vx, vx, vl);
87
+ acc = __riscv_vfadd_vv_f32m8_tu(acc, acc, vx, vl);
88
+ i += vl;
89
+ }
90
+ vfloat32m1_t sum = __riscv_vfmv_s_f_f32m1(0.0f, 1);
91
+ sum = __riscv_vfredusum_vs_f32m8_f32m1(acc, sum, vlmax);
92
+ return __riscv_vfmv_f_s_f32m1_f32(sum);
42
93
  }
43
94
 
44
95
  template <>
@@ -46,7 +97,21 @@ float fvec_Linf<SIMDLevel::RISCV_RVV>(
46
97
  const float* x,
47
98
  const float* y,
48
99
  size_t d) {
49
- return fvec_Linf<SIMDLevel::NONE>(x, y, d);
100
+ size_t vlmax = __riscv_vsetvlmax_e32m8();
101
+ vfloat32m8_t vmax = __riscv_vfmv_v_f_f32m8(0.0f, vlmax);
102
+ size_t i = 0;
103
+ while (i < d) {
104
+ size_t vl = __riscv_vsetvl_e32m8(d - i);
105
+ vfloat32m8_t vx = __riscv_vle32_v_f32m8(x + i, vl);
106
+ vfloat32m8_t vy = __riscv_vle32_v_f32m8(y + i, vl);
107
+ vx = __riscv_vfsub_vv_f32m8(vx, vy, vl);
108
+ vx = __riscv_vfsgnjx_vv_f32m8(vx, vx, vl);
109
+ vmax = __riscv_vfmax_vv_f32m8_tu(vmax, vmax, vx, vl);
110
+ i += vl;
111
+ }
112
+ vfloat32m1_t max = __riscv_vfmv_s_f_f32m1(0.0f, 1);
113
+ max = __riscv_vfredmax_vs_f32m8_f32m1(vmax, max, vlmax);
114
+ return __riscv_vfmv_f_s_f32m1_f32(max);
50
115
  }
51
116
 
52
117
  template <>
@@ -61,8 +126,10 @@ void fvec_inner_product_batch_4<SIMDLevel::RISCV_RVV>(
61
126
  float& dis1,
62
127
  float& dis2,
63
128
  float& dis3) {
64
- fvec_inner_product_batch_4<SIMDLevel::NONE>(
65
- x, y0, y1, y2, y3, d, dis0, dis1, dis2, dis3);
129
+ dis0 = fvec_inner_product<SIMDLevel::RISCV_RVV>(x, y0, d);
130
+ dis1 = fvec_inner_product<SIMDLevel::RISCV_RVV>(x, y1, d);
131
+ dis2 = fvec_inner_product<SIMDLevel::RISCV_RVV>(x, y2, d);
132
+ dis3 = fvec_inner_product<SIMDLevel::RISCV_RVV>(x, y3, d);
66
133
  }
67
134
 
68
135
  template <>
@@ -77,8 +144,10 @@ void fvec_L2sqr_batch_4<SIMDLevel::RISCV_RVV>(
77
144
  float& dis1,
78
145
  float& dis2,
79
146
  float& dis3) {
80
- fvec_L2sqr_batch_4<SIMDLevel::NONE>(
81
- x, y0, y1, y2, y3, d, dis0, dis1, dis2, dis3);
147
+ dis0 = fvec_L2sqr<SIMDLevel::RISCV_RVV>(x, y0, d);
148
+ dis1 = fvec_L2sqr<SIMDLevel::RISCV_RVV>(x, y1, d);
149
+ dis2 = fvec_L2sqr<SIMDLevel::RISCV_RVV>(x, y2, d);
150
+ dis3 = fvec_L2sqr<SIMDLevel::RISCV_RVV>(x, y3, d);
82
151
  }
83
152
 
84
153
  template <>
@@ -90,8 +159,33 @@ void fvec_L2sqr_ny_transposed<SIMDLevel::RISCV_RVV>(
90
159
  size_t d,
91
160
  size_t d_offset,
92
161
  size_t ny) {
93
- fvec_L2sqr_ny_transposed<SIMDLevel::NONE>(
94
- dis, x, y, y_sqlen, d, d_offset, ny);
162
+ size_t vlmax = __riscv_vsetvlmax_e32m8();
163
+ vfloat32m8_t acc = __riscv_vfmv_v_f_f32m8(0.0f, vlmax);
164
+ size_t i = 0;
165
+ while (i < d) {
166
+ size_t vl = __riscv_vsetvl_e32m8(d - i);
167
+ vfloat32m8_t vx = __riscv_vle32_v_f32m8(x + i, vl);
168
+ acc = __riscv_vfmacc_vv_f32m8_tu(acc, vx, vx, vl);
169
+ i += vl;
170
+ }
171
+ vfloat32m1_t sum = __riscv_vfmv_s_f_f32m1(0.0f, 1);
172
+ sum = __riscv_vfredusum_vs_f32m8_f32m1(acc, sum, vlmax);
173
+ float x_sqlen = __riscv_vfmv_f_s_f32m1_f32(sum);
174
+ i = 0;
175
+ while (i < ny) {
176
+ size_t vl = __riscv_vsetvl_e32m8(ny - i);
177
+ acc = __riscv_vfmv_v_f_f32m8(0.0f, vl);
178
+ for (size_t j = 0; j < d; j++) {
179
+ vfloat32m8_t vy = __riscv_vle32_v_f32m8(y + j * d_offset + i, vl);
180
+ acc = __riscv_vfmacc_vf_f32m8(acc, x[j], vy, vl);
181
+ }
182
+ vfloat32m8_t vres = __riscv_vle32_v_f32m8(y_sqlen + i, vl);
183
+ vres = __riscv_vfadd_vf_f32m8(vres, x_sqlen, vl);
184
+ acc = __riscv_vfmul_vf_f32m8(acc, 2.0f, vl);
185
+ vres = __riscv_vfsub_vv_f32m8(vres, acc, vl);
186
+ __riscv_vse32_v_f32m8(dis + i, vres, vl);
187
+ i += vl;
188
+ }
95
189
  }
96
190
 
97
191
  template <>
@@ -101,7 +195,10 @@ void fvec_inner_products_ny<SIMDLevel::RISCV_RVV>(
101
195
  const float* y,
102
196
  size_t d,
103
197
  size_t ny) {
104
- fvec_inner_products_ny<SIMDLevel::NONE>(ip, x, y, d, ny);
198
+ for (size_t i = 0; i < ny; i++) {
199
+ ip[i] = fvec_inner_product<SIMDLevel::RISCV_RVV>(x, y, d);
200
+ y += d;
201
+ }
105
202
  }
106
203
 
107
204
  template <>
@@ -111,7 +208,39 @@ void fvec_L2sqr_ny<SIMDLevel::RISCV_RVV>(
111
208
  const float* y,
112
209
  size_t d,
113
210
  size_t ny) {
114
- fvec_L2sqr_ny<SIMDLevel::NONE>(dis, x, y, d, ny);
211
+ for (size_t i = 0; i < ny; i++) {
212
+ dis[i] = fvec_L2sqr<SIMDLevel::RISCV_RVV>(x, y, d);
213
+ y += d;
214
+ }
215
+ }
216
+
217
+ // Index of the first element equal to the minimum of values[0..n), or n when
218
+ // there is none (e.g. n == 0). Shared by the *_nearest and madd_and_argmin
219
+ // kernels so the vfmin/vfredmin/vfirst sequence lives in one place.
220
+ static size_t rvv_argmin(const float* values, size_t n) {
221
+ size_t vlmax = __riscv_vsetvlmax_e32m8();
222
+ vfloat32m8_t vmin = __riscv_vfmv_v_f_f32m8(__builtin_inff(), vlmax);
223
+ size_t i = 0;
224
+ while (i < n) {
225
+ size_t vl = __riscv_vsetvl_e32m8(n - i);
226
+ vfloat32m8_t vd = __riscv_vle32_v_f32m8(values + i, vl);
227
+ vmin = __riscv_vfmin_vv_f32m8_tu(vmin, vmin, vd, vl);
228
+ i += vl;
229
+ }
230
+ vfloat32m1_t rmin = __riscv_vfmv_s_f_f32m1(__builtin_inff(), 1);
231
+ rmin = __riscv_vfredmin_vs_f32m8_f32m1(vmin, rmin, vlmax);
232
+ float min_val = __riscv_vfmv_f_s_f32m1_f32(rmin);
233
+ i = 0;
234
+ while (i < n) {
235
+ size_t vl = __riscv_vsetvl_e32m8(n - i);
236
+ vfloat32m8_t vd = __riscv_vle32_v_f32m8(values + i, vl);
237
+ long j = __riscv_vfirst_m_b4(
238
+ __riscv_vmfeq_vf_f32m8_b4(vd, min_val, vl), vl);
239
+ if (j >= 0)
240
+ return i + static_cast<size_t>(j);
241
+ i += vl;
242
+ }
243
+ return n;
115
244
  }
116
245
 
117
246
  template <>
@@ -121,8 +250,9 @@ size_t fvec_L2sqr_ny_nearest<SIMDLevel::RISCV_RVV>(
121
250
  const float* y,
122
251
  size_t d,
123
252
  size_t ny) {
124
- return fvec_L2sqr_ny_nearest<SIMDLevel::NONE>(
125
- distances_tmp_buffer, x, y, d, ny);
253
+ fvec_L2sqr_ny<SIMDLevel::RISCV_RVV>(distances_tmp_buffer, x, y, d, ny);
254
+ const size_t j = rvv_argmin(distances_tmp_buffer, ny);
255
+ return j < ny ? j : 0;
126
256
  }
127
257
 
128
258
  template <>
@@ -134,8 +264,10 @@ size_t fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::RISCV_RVV>(
134
264
  size_t d,
135
265
  size_t d_offset,
136
266
  size_t ny) {
137
- return fvec_L2sqr_ny_nearest_y_transposed<SIMDLevel::NONE>(
267
+ fvec_L2sqr_ny_transposed<SIMDLevel::RISCV_RVV>(
138
268
  distances_tmp_buffer, x, y, y_sqlen, d, d_offset, ny);
269
+ const size_t j = rvv_argmin(distances_tmp_buffer, ny);
270
+ return j < ny ? j : 0;
139
271
  }
140
272
 
141
273
  template <>
@@ -145,7 +277,15 @@ void fvec_madd<SIMDLevel::RISCV_RVV>(
145
277
  float bf,
146
278
  const float* b,
147
279
  float* c) {
148
- fvec_madd<SIMDLevel::NONE>(n, a, bf, b, c);
280
+ size_t i = 0;
281
+ while (i < n) {
282
+ size_t vl = __riscv_vsetvl_e32m8(n - i);
283
+ vfloat32m8_t va = __riscv_vle32_v_f32m8(a + i, vl);
284
+ vfloat32m8_t vb = __riscv_vle32_v_f32m8(b + i, vl);
285
+ va = __riscv_vfmacc_vf_f32m8(va, bf, vb, vl);
286
+ __riscv_vse32_v_f32m8(c + i, va, vl);
287
+ i += vl;
288
+ }
149
289
  }
150
290
 
151
291
  template <>
@@ -155,7 +295,9 @@ int fvec_madd_and_argmin<SIMDLevel::RISCV_RVV>(
155
295
  float bf,
156
296
  const float* b,
157
297
  float* c) {
158
- return fvec_madd_and_argmin<SIMDLevel::NONE>(n, a, bf, b, c);
298
+ fvec_madd<SIMDLevel::RISCV_RVV>(n, a, bf, b, c);
299
+ const size_t j = rvv_argmin(c, n);
300
+ return j < n ? static_cast<int>(j) : -1;
159
301
  }
160
302
 
161
303
  #define DEFINE_VECTOR_DISTANCE_RVV_FALLBACK(metric) \