compute: Use SimSIMD for vectors (#221)

Signed-off-by: Wish <[email protected]>
pingcap · Aug 6, 2024 · 9a9daf8 · 9a9daf8
1 parent 86629a5
commit 9a9daf8
Show file tree

Hide file tree

Showing 12 changed files with 47 additions and 2,300 deletions.
diff --git a/.gitmodules b/.gitmodules
@@ -143,3 +143,6 @@
 [submodule "contrib/usearch"]
 	path = contrib/usearch
 	url = https://github.com/unum-cloud/usearch.git
+[submodule "contrib/simsimd"]
+	path = contrib/simsimd
+	url = https://github.com/ashvardanian/SimSIMD
diff --git a/contrib/CMakeLists.txt b/contrib/CMakeLists.txt
@@ -192,3 +192,5 @@ add_subdirectory(magic_enum)
 add_subdirectory(aws-cmake)
 
 add_subdirectory(usearch-cmake)
+
+add_subdirectory(simsimd-cmake)
diff --git a/contrib/simsimd b/contrib/simsimd
diff --git a/contrib/simsimd-cmake/CMakeLists.txt b/contrib/simsimd-cmake/CMakeLists.txt
@@ -0,0 +1,13 @@
+set(SIMSIMD_PROJECT_DIR "${TiFlash_SOURCE_DIR}/contrib/simsimd")
+set(SIMSIMD_SOURCE_DIR "${SIMSIMD_PROJECT_DIR}/include")
+
+add_library(_simsimd INTERFACE)
+
+if (NOT EXISTS "${SIMSIMD_SOURCE_DIR}/simsimd/simsimd.h")
+    message (FATAL_ERROR "submodule contrib/simsimd not found")
+endif()
+
+target_include_directories(_simsimd SYSTEM INTERFACE
+    ${SIMSIMD_SOURCE_DIR})
+
+add_library(tiflash_contrib::simsimd ALIAS _simsimd)
diff --git a/contrib/usearch b/contrib/usearch
diff --git a/contrib/usearch-cmake/CMakeLists.txt b/contrib/usearch-cmake/CMakeLists.txt
@@ -8,7 +8,7 @@ if (NOT EXISTS "${USEARCH_SOURCE_DIR}/usearch/index.hpp")
 endif()
 
 target_include_directories(_usearch SYSTEM INTERFACE
-    ${USEARCH_PROJECT_DIR}/simsimd/include
+    # ${USEARCH_PROJECT_DIR}/simsimd/include  # Use our simsimd
     ${USEARCH_PROJECT_DIR}/fp16/include
     ${USEARCH_SOURCE_DIR})
 

diff --git a/dbms/CMakeLists.txt b/dbms/CMakeLists.txt
@@ -195,6 +195,7 @@ target_link_libraries (dbms
     ${BTRIE_LIBRARIES}
     absl::synchronization
     tiflash_contrib::usearch
+    tiflash_contrib::simsimd
     tiflash_contrib::aws_s3
 
     etcdpb

diff --git a/dbms/src/Functions/tests/gtest_vector.cpp b/dbms/src/Functions/tests/gtest_vector.cpp
@@ -203,7 +203,15 @@ TEST_F(Vector, CosineDistance)
 try
 {
     ASSERT_COLUMN_EQ(
-        createColumn<Nullable<Float64>>({0.0, std::nullopt, 0.0, 1.0, 2.0, 0.0, 2.0, std::nullopt}),
+        createColumn<Nullable<Float64>>(
+            {0.004130363464355469,
+             1.0, // CosDistance to (0,0) cannot be calculated, clapped to 1.0
+             0.00572967529296875,
+             1.0,
+             1.9942703247070312,
+             0.00022123707458376884,
+             1.9997787475585938,
+             std::nullopt}),
         executeFunction(
             "vecCosineDistance",
             createColumn<Array>(

diff --git a/dbms/src/Storages/DeltaMerge/Index/VectorIndexHNSW/Index.cpp b/dbms/src/Storages/DeltaMerge/Index/VectorIndexHNSW/Index.cpp
@@ -187,15 +187,15 @@ std::vector<VectorIndexBuilder::Key> VectorIndexHNSWViewer::search(
     std::atomic<size_t> discarded_nodes = 0;
     std::atomic<bool> has_exception_in_search = false;
 
-    auto predicate = [&](typename USearchImplType::member_cref_t const & member) {
+    auto predicate = [&](const Key & key) {
         // Must catch exceptions in the predicate, because search runs on other threads.
         try
         {
             // Note: We don't increase the thread_local perf, because search runs on other threads.
             visited_nodes++;
-            if (!valid_rows[member.key])
+            if (!valid_rows[key])
                 discarded_nodes++;
-            return valid_rows[member.key];
+            return valid_rows[key];
         }
         catch (...)
         {
@@ -209,7 +209,7 @@ std::vector<VectorIndexBuilder::Key> VectorIndexHNSWViewer::search(
     SCOPE_EXIT({ GET_METRIC(tiflash_vector_index_duration, type_search).Observe(w.elapsedSeconds()); });
 
     // TODO: Support efSearch.
-    auto result = index.search( //
+    auto result = index.filtered_search( //
         reinterpret_cast<const Float32 *>(queryInfo->ref_vec_f32().data() + sizeof(UInt32)),
         queryInfo->top_k(),
         predicate);

diff --git a/dbms/src/Storages/DeltaMerge/Index/VectorIndexHNSW/Index.h b/dbms/src/Storages/DeltaMerge/Index/VectorIndexHNSW/Index.h
@@ -17,12 +17,7 @@
 #include <Storages/DeltaMerge/File/dtpb/dmfile.pb.h>
 #include <Storages/DeltaMerge/Index/VectorIndex.h>
 
-#if __clang__
-#pragma clang diagnostic push
-#pragma clang diagnostic ignored "-Wpass-failed"
-#include <Storages/DeltaMerge/Index/VectorIndexHNSW/usearch_index_dense.h>
-#pragma clang diagnostic pop
-#endif
+#include <usearch/index_dense.hpp>
 
 namespace DB::DM
 {