storage: Enable simsimd for usearch (#224)

Signed-off-by: Wish <breezewish@outlook.com>
Co-authored-by: JaySon <tshent@qq.com>

cmake/cpu_features.cmake

@@ -37,7 +37,8 @@ if (ARCH_AARCH64)
  option (TIFLASH_ENABLE_ASIMD_SUPPORT "Enable Advanced SIMD support." ON)
  option (TIFLASH_ENABLE_SVE_SUPPORT "Enable Scalable Vector Extension support." OFF)
  # TODO: default ON, to be changed after CI is updated
- option (NO_ARMV81_OR_HIGHER "Disable ARMv8.1 or higher on Aarch64 for maximum compatibility with older/embedded hardware." ON)
+ # Note: explicitly set to OFF for cse-branch
+ option (NO_ARMV81_OR_HIGHER "Disable ARMv8.1 or higher on Aarch64 for maximum compatibility with older/embedded hardware." OFF)
 
  if (NO_ARMV81_OR_HIGHER)
  # crc32 is optional in v8.0 and mandatory in v8.1. Enable it as __crc32()* is used in lot's of places and even very old ARM CPUs
@@ -95,7 +96,7 @@ elseif (ARCH_AMD64)
  # so we do not set the flags to avoid core dump in old machines
  option (TIFLASH_ENABLE_AVX_SUPPORT "Use AVX/AVX2 instructions on x86_64" ON)
  option (TIFLASH_ENABLE_AVX512_SUPPORT "Use AVX512 instructions on x86_64" ON)
- 
+
  # `haswell` was released since 2013 with cpu feature avx2, bmi2. It's a practical arch for optimizer
  option (TIFLASH_ENABLE_ARCH_HASWELL_SUPPORT "Use instructions based on architecture `haswell` on x86_64" ON)
 

dbms/src/Common/TiFlashBuildInfo.cpp

@@ -14,6 +14,8 @@
 
 #include <Common/config.h>
 #include <Common/config_version.h>
+#include <Storages/DeltaMerge/Index/VectorIndexHNSW/SIMDFeatures.h>
+#include <TiDB/Decode/Vector.h>
 #include <common/config_common.h>
 #include <fmt/core.h>
 #include <fmt/format.h>
@@ -128,6 +130,17 @@ std::string getEnabledFeatures()
  "fdo",
 #endif
  };
+ {
+ auto f = DB::DM::VectorIndexHNSWSIMDFeatures::get();
+ for (const auto & feature : f)
+ features.push_back(feature);
+ }
+ {
+ auto f = DB::VectorDistanceSIMDFeatures::get();
+ for (const auto & feature : f)
+ features.push_back(feature);
+ }
+
  return fmt::format("{}", fmt::join(features.begin(), features.end(), " "));
 }
 // clang-format on

dbms/src/Storages/DeltaMerge/Index/VectorIndexHNSW/Index.cpp

@@ -25,8 +25,6 @@
 
 #include <algorithm>
 #include <ext/scope_guard.h>
-#include <usearch/index.hpp>
-#include <usearch/index_plugins.hpp>
 
 namespace DB::ErrorCodes
 {

dbms/src/Storages/DeltaMerge/Index/VectorIndexHNSW/Index.h

@@ -16,8 +16,7 @@
 
 #include <Storages/DeltaMerge/File/dtpb/dmfile.pb.h>
 #include <Storages/DeltaMerge/Index/VectorIndex.h>
-
-#include <usearch/index_dense.hpp>
+#include <Storages/DeltaMerge/Index/VectorIndexHNSW/USearch.h>
 
 namespace DB::DM
 {

dbms/src/Storages/DeltaMerge/Index/VectorIndexHNSW/SIMDFeatures.cpp

@@ -0,0 +1,33 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <Common/Logger.h>
+#include <Storages/DeltaMerge/Index/VectorIndexHNSW/SIMDFeatures.h>
+#include <Storages/DeltaMerge/Index/VectorIndexHNSW/USearch.h>
+#include <common/logger_useful.h>
+
+namespace DB::DM
+{
+
+std::vector<std::string> VectorIndexHNSWSIMDFeatures::get()
+{
+ auto m_l2 = unum::usearch::metric_punned_t(3, unum::usearch::metric_kind_t::l2sq_k);
+ auto m_cos = unum::usearch::metric_punned_t(3, unum::usearch::metric_kind_t::cos_k);
+ return {
+ fmt::format("hnsw.l2={}", m_l2.isa_name()),
+ fmt::format("hnsw.cosine={}", m_cos.isa_name()),
+ };
+}
+
+} // namespace DB::DM

dbms/src/Storages/DeltaMerge/Index/VectorIndexHNSW/SIMDFeatures.h

@@ -0,0 +1,28 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <string>
+
+namespace DB::DM
+{
+
+class VectorIndexHNSWSIMDFeatures
+{
+public:
+ static std::vector<std::string> get();
+};
+
+} // namespace DB::DM

dbms/src/Storages/DeltaMerge/Index/VectorIndexHNSW/USearch.h

@@ -0,0 +1,35 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+// USearch is header only. We don't use cmake to make these defines to avoid
+// polluting all compile units.
+
+#define USEARCH_USE_SIMSIMD 1
+#define SIMSIMD_NATIVE_F16 0
+#define SIMSIMD_NATIVE_BF16 0
+
+// Force enable all target features.
+#define SIMSIMD_TARGET_NEON 1
+#define SIMSIMD_TARGET_SVE 0 // Clang13's header does not support enableing SVE for region
+#define SIMSIMD_TARGET_HASWELL 1
+#define SIMSIMD_TARGET_SKYLAKE 0 // Clang13 does not support AVX512
+#define SIMSIMD_TARGET_ICE 0
+#define SIMSIMD_TARGET_GENOA 0
+#define SIMSIMD_TARGET_SAPPHIRE 0
+
+#include <usearch/index.hpp>
+#include <usearch/index_dense.hpp>
+#include <usearch/index_plugins.hpp>

dbms/src/TiDB/Decode/Vector.cpp

@@ -17,12 +17,28 @@
 #include <IO/WriteHelpers.h>
 #include <TiDB/Decode/Vector.h>
 
+#include <compare>
+
+// SIMSIMD is header only. We don't use cmake to make these defines to avoid
+// polluting all compile units.
+
+// Note: Be careful that usearch also includes simsimd with a customized config.
+// Don't include simsimd and usearch at the same time. Otherwise, the effective
+// config depends on the include order.
 #define SIMSIMD_NATIVE_F16 0
 #define SIMSIMD_NATIVE_BF16 0
+#define SIMSIMD_DYNAMIC_DISPATCH 0
+
+// Force enable all target features. We will do our own dynamic dispatch.
+#define SIMSIMD_TARGET_NEON 1
+#define SIMSIMD_TARGET_SVE 0 // Clang13's header does not support enableing SVE for region
+#define SIMSIMD_TARGET_HASWELL 1
+#define SIMSIMD_TARGET_SKYLAKE 0 // Clang13 does not support AVX512
+#define SIMSIMD_TARGET_ICE 0
+#define SIMSIMD_TARGET_GENOA 0
+#define SIMSIMD_TARGET_SAPPHIRE 0
 #include <simsimd/simsimd.h>
 
-#include <compare>
-
 namespace DB
 {
 
@@ -31,6 +47,36 @@ namespace ErrorCodes
 extern const int BAD_ARGUMENTS;
 } // namespace ErrorCodes
 
+namespace simsimd_details
+{
+
+simsimd_capability_t simd_capabilities()
+{
+ static simsimd_capability_t static_capabilities = simsimd_cap_any_k;
+ if (static_capabilities == simsimd_cap_any_k)
+ static_capabilities = simsimd_capabilities_implementation();
+ return static_capabilities;
+}
+
+} // namespace simsimd_details
+
+std::vector<std::string> VectorDistanceSIMDFeatures::get()
+{
+ simsimd_capability_t caps = simsimd_details::simd_capabilities();
+ std::vector<std::string> ret{};
+ if (caps & simsimd_cap_neon_k)
+ ret.push_back("vec_distance=neon");
+ if (caps & simsimd_cap_sve_k)
+ ret.push_back("vec_distance=sve");
+ if (caps & simsimd_cap_sve2_k)
+ ret.push_back("vec_distance=sve2");
+ if (caps & simsimd_cap_haswell_k)
+ ret.push_back("vec_distance=haswell");
+ if (caps & simsimd_cap_skylake_k)
+ ret.push_back("vec_distance=skylake");
+ return ret;
+}
+
 VectorFloat32Ref::VectorFloat32Ref(const Float32 * elements, size_t n)
  : elements(elements)
  , elements_n(n)
@@ -56,8 +102,23 @@ Float64 VectorFloat32Ref::l2SquaredDistance(VectorFloat32Ref b) const
 {
  checkDims(b);
 
+ static simsimd_metric_punned_t metric = nullptr;
+ if (metric == nullptr)
+ {
+ simsimd_capability_t used_capability;
+ simsimd_find_metric_punned(
+ simsimd_metric_l2sq_k,
+ simsimd_datatype_f32_k,
+ simsimd_details::simd_capabilities(),
+ simsimd_cap_any_k,
+ &metric,
+ &used_capability);
+ if (!metric)
+ return std::numeric_limits<double>::quiet_NaN();
+ }
+
  simsimd_distance_t distance;
- simsimd_l2sq_f32(elements, b.elements, elements_n, &distance);
+ metric(elements, b.elements, elements_n, &distance);
 
  return distance;
 }
@@ -66,8 +127,23 @@ Float64 VectorFloat32Ref::innerProduct(VectorFloat32Ref b) const
 {
  checkDims(b);
 
+ static simsimd_metric_punned_t metric = nullptr;
+ if (metric == nullptr)
+ {
+ simsimd_capability_t used_capability;
+ simsimd_find_metric_punned(
+ simsimd_metric_dot_k,
+ simsimd_datatype_f32_k,
+ simsimd_details::simd_capabilities(),
+ simsimd_cap_any_k,
+ &metric,
+ &used_capability);
+ if (!metric)
+ return std::numeric_limits<double>::quiet_NaN();
+ }
+
  simsimd_distance_t distance;
- simsimd_dot_f32(elements, b.elements, elements_n, &distance);
+ metric(elements, b.elements, elements_n, &distance);
 
  return distance;
 }
@@ -76,8 +152,23 @@ Float64 VectorFloat32Ref::cosineDistance(VectorFloat32Ref b) const
 {
  checkDims(b);
 
+ static simsimd_metric_punned_t metric = nullptr;
+ if (metric == nullptr)
+ {
+ simsimd_capability_t used_capability;
+ simsimd_find_metric_punned(
+ simsimd_metric_cos_k,
+ simsimd_datatype_f32_k,
+ simsimd_details::simd_capabilities(),
+ simsimd_cap_any_k,
+ &metric,
+ &used_capability);
+ if (!metric)
+ return std::numeric_limits<double>::quiet_NaN();
+ }
+
  simsimd_distance_t distance;
- simsimd_cos_f32(elements, b.elements, elements_n, &distance);
+ metric(elements, b.elements, elements_n, &distance);
 
  return distance;
 }

dbms/src/TiDB/Decode/Vector.h

@@ -23,6 +23,12 @@
 namespace DB
 {
 
+class VectorDistanceSIMDFeatures
+{
+public:
+ static std::vector<std::string> get();
+};
+
 class VectorFloat32Ref
 {
 public:

libs/libdaemon/src/BaseDaemon.cpp

@@ -1179,7 +1179,10 @@ void BaseDaemon::logRevision() const
  LOG_INFO(log, "Starting daemon with revision " + Poco::NumberFormatter::format(ClickHouseRevision::get()));
  std::stringstream ss;
  TiFlashBuildInfo::outputDetail(ss);
- LOG_INFO(log, "TiFlash build info: {}", ss.str());
+
+ std::string line;
+ while (std::getline(ss, line, '\n'))
+ LOG_INFO(log, "{}", line);
 }
 
 /// Used for exitOnTaskError()