storage: Update usearch and simsimd (pingcap#270)

Signed-off-by: Wish <breezewish@outlook.com>
Lloyd-Pottiger · Sep 20, 2024 · 34348e0 · 34348e0
1 parent c9ed0fd
commit 34348e0
Show file tree

Hide file tree

Showing 8 changed files with 40 additions and 21 deletions.
diff --git a/contrib/simsimd b/contrib/simsimd
diff --git a/contrib/usearch b/contrib/usearch
diff --git a/dbms/src/Storages/DeltaMerge/Index/VectorIndexHNSW/Index.cpp b/dbms/src/Storages/DeltaMerge/Index/VectorIndexHNSW/Index.cpp
@@ -84,10 +84,7 @@ void VectorIndexHNSWBuilder::addBlock(
 
     const auto * del_mark_data = (!del_mark) ? nullptr : &(del_mark->getData());
 
-    if (!index.reserve(unum::usearch::ceil2(index.size() + column.size())))
-    {
-        throw Exception(ErrorCodes::CANNOT_ALLOCATE_MEMORY, "Could not reserve memory for HNSW index");
-    }
+    index.reserve(unum::usearch::ceil2(index.size() + column.size()));
 
     Stopwatch w;
     SCOPE_EXIT({ total_duration += w.elapsedSeconds(); });
@@ -170,6 +167,7 @@ VectorIndexViewerPtr VectorIndexHNSWViewer::view(const dtpb::VectorIndexFileProp
     SCOPE_EXIT({ GET_METRIC(tiflash_vector_index_duration, type_view).Observe(w.elapsedSeconds()); });
 
     auto vi = std::make_shared<VectorIndexHNSWViewer>(file_props);
+
     vi->index = USearchImplType::make(
         unum::usearch::metric_punned_t( //
             file_props.dimensions(),
@@ -179,6 +177,10 @@ VectorIndexViewerPtr VectorIndexHNSWViewer::view(const dtpb::VectorIndexFileProp
             unum::usearch::default_expansion_add(),
             16 /* default is 64 */));
 
+    // Currently may have a lot of threads querying concurrently
+    auto limit = unum::usearch::index_limits_t(0, /* threads */ std::thread::hardware_concurrency() * 10);
+    vi->index.reserve(limit);
+
     auto result = vi->index.view(unum::usearch::memory_mapped_file_t(path.data()));
     RUNTIME_CHECK_MSG(result, "Failed to load vector index: {}", result.error.what());
 

diff --git a/dbms/src/Storages/DeltaMerge/tests/gtest_dm_delta_merge_store_vector_index.cpp b/dbms/src/Storages/DeltaMerge/tests/gtest_dm_delta_merge_store_vector_index.cpp
@@ -262,7 +262,7 @@ try
 
         auto filter = std::make_shared<PushDownFilter>(wrapWithANNQueryInfo(nullptr, ann_query_info));
 
-        read(left_segment_range, filter, createVecFloat32Column<Array>({})); // FIXME: should be 63.0
+        read(left_segment_range, filter, createVecFloat32Column<Array>({{63.0}}));
     }
 
     // merge segment

diff --git a/dbms/src/TiDB/Decode/Vector.cpp b/dbms/src/TiDB/Decode/Vector.cpp
@@ -16,7 +16,7 @@
 #include <IO/Buffer/WriteBufferFromString.h>
 #include <IO/WriteHelpers.h>
 #include <TiDB/Decode/Vector.h>
-#include <VectorSearch/simdsimd-internals.h>
+#include <VectorSearch/SimSIMD.h>
 
 #include <compare>
 

diff --git a/dbms/src/VectorSearch/DistanceSIMDFeatures.cpp b/dbms/src/VectorSearch/DistanceSIMDFeatures.cpp
@@ -16,7 +16,7 @@
 // polluting all compile units.
 
 #include <VectorSearch/DistanceSIMDFeatures.h>
-#include <VectorSearch/simdsimd-internals.h>
+#include <VectorSearch/SimSIMD.h>
 
 namespace simsimd_details
 {
@@ -52,21 +52,39 @@ std::vector<std::string> VectorDistanceSIMDFeatures::get()
     simsimd_capability_t cap_l2 = simsimd_details::actual_capability(simsimd_datatype_f32_k, simsimd_metric_l2sq_k);
     simsimd_capability_t cap_cos = simsimd_details::actual_capability(simsimd_datatype_f32_k, simsimd_metric_cos_k);
 
-    auto cap_to_string = [](simsimd_capability_t cap) -> std::string {
-        switch (cap)
+    auto cap_to_string = [](simsimd_capability_t isa_kind) -> std::string {
+        switch (isa_kind)
         {
-        case simsimd_cap_sve2_k:
-            return "sve2";
-        case simsimd_cap_sve_k:
-            return "sve";
+        case simsimd_cap_serial_k:
+            return "serial";
         case simsimd_cap_neon_k:
             return "neon";
-        case simsimd_cap_skylake_k:
-            return "skylake";
+        case simsimd_cap_neon_i8_k:
+            return "neon_i8";
+        case simsimd_cap_neon_f16_k:
+            return "neon_f16";
+        case simsimd_cap_neon_bf16_k:
+            return "neon_bf16";
+        case simsimd_cap_sve_k:
+            return "sve";
+        case simsimd_cap_sve_i8_k:
+            return "sve_i8";
+        case simsimd_cap_sve_f16_k:
+            return "sve_f16";
+        case simsimd_cap_sve_bf16_k:
+            return "sve_bf16";
         case simsimd_cap_haswell_k:
             return "haswell";
+        case simsimd_cap_skylake_k:
+            return "skylake";
+        case simsimd_cap_ice_k:
+            return "ice";
+        case simsimd_cap_genoa_k:
+            return "genoa";
+        case simsimd_cap_sapphire_k:
+            return "sapphire";
         default:
-            return "serial";
+            return "unknown";
         }
     };
 

diff --git a/dbms/src/VectorSearch/simdsimd-internals.h → dbms/src/VectorSearch/SimSIMD.h b/dbms/src/VectorSearch/simdsimd-internals.h → dbms/src/VectorSearch/SimSIMD.h
@@ -29,12 +29,11 @@
 #define SIMSIMD_TARGET_SVE 1
 #define SIMSIMD_TARGET_HASWELL 1
 #define SIMSIMD_TARGET_SKYLAKE 1
-#define SIMSIMD_TARGET_ICE 0
+#define SIMSIMD_TARGET_ICE 1
 #define SIMSIMD_TARGET_GENOA 0
 #define SIMSIMD_TARGET_SAPPHIRE 0
 #include <simsimd/simsimd.h>
 
-
 namespace simsimd_details
 {
 

diff --git a/dbms/src/VectorSearch/USearch.h b/dbms/src/VectorSearch/USearch.h
@@ -26,7 +26,7 @@
 #define SIMSIMD_TARGET_SVE 1
 #define SIMSIMD_TARGET_HASWELL 1
 #define SIMSIMD_TARGET_SKYLAKE 1
-#define SIMSIMD_TARGET_ICE 0
+#define SIMSIMD_TARGET_ICE 1
 #define SIMSIMD_TARGET_GENOA 0
 #define SIMSIMD_TARGET_SAPPHIRE 0